LogicalTransactions in frontend Client

integration/go/go_pgx/load_balancer_test.go

@@ -231,10 +231,18 @@ outer:
 }
 
 func prewarm(t *testing.T, pool *pgxpool.Pool) {
+	ctx := context.Background()
 	for range 25 {
-		for _, q := range []string{"BEGIN", "SELECT 1", "COMMIT", "SELECT 1"} {
-			_, err := pool.Exec(context.Background(), q)
-			assert.NoError(t, err)
-		}
+		// transaction `BEGIN; SELECT 1; COMMIT;`
+		tx, err := pool.Begin(ctx)
+		assert.NoError(t, err)
+		_, err = tx.Exec(ctx, "SELECT 1")
+		assert.NoError(t, err)
+		err = tx.Commit(ctx)
+		assert.NoError(t, err)
+
+		// no-transaction `SELECT 1;`
+		_, err = pool.Exec(ctx, "SELECT 1")
+		assert.NoError(t, err)
 	}
 }

integration/load_balancer/run.sh

@@ -10,7 +10,7 @@ export PGHOST=127.0.0.1
 export PGDATABASE=postgres
 export PGPASSWORD=postgres
 
-docker-compose up -d
+docker compose up -d
 
 
 echo "Waiting for Postgres to be ready"
@@ -45,4 +45,4 @@ popd
 
 killall pgdog
 
-docker-compose down
+docker compose down

pgdog/Cargo.toml

@@ -59,6 +59,7 @@ indexmap = "2.9"
 lru = "0.16"
 hickory-resolver = "0.25.2"
 lazy_static = "1"
+smallvec = "1.15.1"
 
 [target.'cfg(not(target_env = "msvc"))'.dependencies]
 tikv-jemallocator = "0.6"

pgdog/src/backend/pool/connection/binding.rs

@@ -1,7 +1,7 @@
 //! Binding between frontend client and a connection on the backend.
 
 use crate::{
-    net::{parameter::Parameters, ProtocolMessage},
+    net::{parameter::Parameters, ProtocolMessage, Query},
     state::State,
 };
 
@@ -237,6 +237,13 @@ impl Binding {
         Ok(())
     }
 
+    /// Execute a BEGIN on all servers
+    /// TODO: Block mutli-shard BEGINs as transaction should not occur on multiple shards
+    pub async fn begin(&mut self) -> Result<(), Error> {
+        let query = Query::new("BEGIN");
+        self.execute(query.query()).await
+    }
+
     pub async fn link_client(&mut self, params: &Parameters) -> Result<usize, Error> {
         match self {
             Binding::Server(Some(ref mut server)) => server.link_client(params).await,

pgdog/src/frontend/client/engine/context.rs

@@ -1,5 +1,7 @@
 use crate::{
-    frontend::{client::Inner, Buffer, Client, PreparedStatements},
+    frontend::{
+        client::Inner, logical_transaction::LogicalTransaction, Buffer, Client, PreparedStatements,
+    },
     net::Parameters,
 };
 
@@ -13,7 +15,7 @@ pub struct EngineContext<'a> {
     /// Client parameters.
     pub(super) params: &'a Parameters,
     /// Is the client inside a transaction?
-    pub(super) in_transaction: bool,
+    pub(super) logical_transaction: &'a LogicalTransaction,
     /// Messages currently in client's buffer.
     pub(super) buffer: &'a Buffer,
 }
@@ -23,9 +25,13 @@ impl<'a> EngineContext<'a> {
         Self {
             prepared_statements: &mut client.prepared_statements,
             params: &client.params,
-            in_transaction: client.in_transaction,
+            logical_transaction: &client.logical_transaction,
             connected: inner.connected(),
             buffer: &client.request_buffer,
         }
     }
+
+    pub fn in_transaction(&self) -> bool {
+        self.logical_transaction.in_transaction()
+    }
 }

pgdog/src/frontend/client/engine/mod.rs

@@ -54,7 +54,8 @@ impl<'a> Engine<'a> {
                 'S' => {
                     if only_close || only_sync && !self.context.connected {
                         messages.push(
-                            ReadyForQuery::in_transaction(self.context.in_transaction).message()?,
+                            ReadyForQuery::in_transaction(self.context.in_transaction())
+                                .message()?,
                         )
                     }
                 }
@@ -73,7 +74,7 @@ impl<'a> Engine<'a> {
 #[cfg(test)]
 mod test {
     use crate::{
-        frontend::{Buffer, PreparedStatements},
+        frontend::{logical_transaction::LogicalTransaction, Buffer, PreparedStatements},
         net::{Parameters, Parse, Sync},
     };
 
@@ -93,11 +94,13 @@ mod test {
             Sync.into(),
         ]);
 
+        let logical_transaction = LogicalTransaction::new();
+
         let context = EngineContext {
             connected: false,
             prepared_statements: &mut prepared,
             params: &params,
-            in_transaction: false,
+            logical_transaction: &logical_transaction,
             buffer: &buf,
         };
 

pgdog/src/frontend/client/inner.rs

@@ -6,8 +6,8 @@ use crate::{
         Error as BackendError,
     },
     frontend::{
-        buffer::BufferedQuery, router::Error as RouterError, Buffer, Command, Comms,
-        PreparedStatements, Router, RouterContext, Stats,
+        logical_transaction::LogicalTransaction, router::Error as RouterError, Buffer, Command,
+        Comms, PreparedStatements, Router, RouterContext, Stats,
     },
     net::Parameters,
     state::State,
@@ -29,8 +29,6 @@ pub struct Inner {
     pub(super) router: Router,
     /// Client stats.
     pub(super) stats: Stats,
-    /// Start transaction statement, intercepted by the router.
-    pub(super) start_transaction: Option<BufferedQuery>,
     /// Client-wide comms.
     pub(super) comms: Comms,
 }
@@ -47,7 +45,6 @@ impl Inner {
             backend,
             router,
             stats: Stats::new(),
-            start_transaction: None,
             comms: client.comms.clone(),
         })
     }
@@ -58,7 +55,7 @@ impl Inner {
         buffer: &mut Buffer,
         prepared_statements: &mut PreparedStatements,
         params: &Parameters,
-        in_transaction: bool,
+        logical_transaction: &LogicalTransaction,
     ) -> Result<Option<&Command>, RouterError> {
         let command = self
             .backend
@@ -67,11 +64,11 @@ impl Inner {
             .map(|cluster| {
                 // Build router context.
                 let context = RouterContext::new(
-                    buffer,              // Query and parameters.
-                    cluster,             // Cluster configuration.
-                    prepared_statements, // Prepared statements.
-                    params,              // Client connection parameters.
-                    in_transaction,      // Client in explcitely started transaction.
+                    buffer,                               // Query and parameters.
+                    cluster,                              // Cluster configuration.
+                    prepared_statements,                  // Prepared statements.
+                    params,                               // Client connection parameters.
+                    logical_transaction.in_transaction(), // Client in explcitely started transaction.
                 )?;
                 self.router.query(context)
             })