fix: don't buffer client request before checking maintenance mode

pgdog/src/frontend/client/mod.rs

@@ -531,43 +531,46 @@ impl Client {
 
     /// Handle client messages.
     async fn client_messages(&mut self, query_engine: &mut QueryEngine) -> Result<(), Error> {
-        // Check maintenance mode.
-        if !self.in_transaction()
-            && !self.admin
-            && let Some(waiter) = maintenance_mode::waiter(&self.database)
-        {
-            let state = query_engine.get_state();
-            query_engine.set_state(State::Waiting);
-            waiter.await;
-            query_engine.set_state(state);
-        }
-
         // If client sent multiple requests, split them up and execute individually.
         let spliced = self.client_request.spliced()?;
         if spliced.is_empty() {
+            // Execute the original request as-is.
             let mut context = QueryEngineContext::new(self);
             query_engine.handle(&mut context).await?;
             self.transaction = context.transaction();
         } else {
+            // Execute the spliced requests one at a time.
             let total = spliced.len();
             let mut reqs = spliced.into_iter().enumerate();
+
+            // Pipeliend requests are executed inside the same Postgres transaction.
             self.transaction.get_or_insert(TransactionType::Implicit);
+
             while let Some((num, mut req)) = reqs.next() {
                 debug!("processing spliced request {}/{}", num + 1, total);
+
                 let mut context = QueryEngineContext::new(self).spliced(&mut req, reqs.len());
                 query_engine.handle(&mut context).await?;
+
+                // Save transaction state after each request.
                 self.transaction = context.transaction();
 
                 // If pipeline is aborted due to error, skip to Sync to complete the pipeline.
-                // Postgres ignores all commands after an error until it receives Sync.
+                // Postgres ignores all commands after an error until it receives a Sync.
                 if query_engine.out_of_sync() && !req.is_sync_only() {
                     debug!("pipeline aborted, skipping to Sync");
+
                     for (_, mut next_req) in reqs.by_ref() {
+                        // The splicer (called above) leaves
+                        // the Sync in its own request. This is by design
+                        // to handle this exact scenario.
                         if next_req.is_sync_only() {
                             debug!("processing Sync to complete aborted pipeline");
+
                             let mut ctx = QueryEngineContext::new(self).spliced(&mut next_req, 0);
                             query_engine.handle(&mut ctx).await?;
                             self.transaction = ctx.transaction();
+
                             break;
                         }
                     }
@@ -582,13 +585,29 @@ impl Client {
         Ok(())
     }
 
-    /// Buffer extended protocol messages until client requests a sync.
+    /// Buffer protocol messages until client requests a sync or a flush.
     ///
     /// This ensures we don't check out a connection from the pool until the client
-    /// sent a complete request.
+    /// sends a complete request.
+    ///
     async fn buffer(&mut self, state: State) -> Result<BufferEvent, Error> {
         self.client_request.clear();
 
+        // Check maintenance mode.
+        //
+        // This will pause the client until the maintenance mode is off
+        // without checking any timeouts or reading data off of the socket
+        // so if the client disconnects while we wait, we won't send
+        // its pending query to the database.
+        if !self.in_transaction()
+            && !self.admin
+            && let Some(waiter) = maintenance_mode::waiter(&self.database)
+        {
+            self.comms.set_state(State::Waiting);
+            waiter.await;
+            self.comms.set_state(state);
+        }
+
         // Only start timer once we receive the first message.
         let mut timer = None;
 

pgdog/src/frontend/comms.rs

@@ -15,6 +15,7 @@ use tokio_util::task::TaskTracker;
 
 use crate::net::Parameters;
 use crate::net::messages::{BackendKeyData, FrontendPid};
+use crate::state::State;
 
 use super::{ConnectedClient, Stats};
 
@@ -155,23 +156,51 @@ impl Deref for ClientComms {
 }
 
 impl ClientComms {
-    pub fn disconnect(&self) {
+    /// Client has disconnected, remove it from the global state.
+    pub(crate) fn disconnect(&self) {
         self.comms.disconnect(self.id);
     }
 
-    pub fn update_stats(&self, stats: Stats) {
+    /// Update global client stats with the latest stats the caller knows about.
+    /// The caller must maintain up-to-date stats, or this will overwrite with bad
+    /// info.
+    pub(crate) fn update_stats(&self, stats: Stats) {
         self.comms.update_stats(self.id, stats);
     }
 
-    pub fn new(id: FrontendPid) -> Self {
+    /// Update the client state quickly.
+    ///
+    /// Use only when you're not updating other stats. If updating other stats,
+    /// use [`Self::update_stats`] instead.
+    ///
+    pub(crate) fn set_state(&self, state: State) {
+        if let Some(mut client) = self.comms.global.clients.get_mut(&self.id) {
+            client.stats.state = state;
+        }
+    }
+
+    /// Create new comms reference for a new client.
+    ///
+    /// Makes it easy to update globals without constantly reaching for the key.
+    pub(crate) fn new(id: FrontendPid) -> Self {
         Self { id, comms: comms() }
     }
 
-    pub fn connect(&self, key: BackendKeyData, addr: SocketAddr, params: &Parameters) {
+    /// Mark the client as connected to the proxy
+    /// and save some of its identifying information.
+    ///
+    /// # Parameters
+    ///
+    /// - `key`: The query cancellation key we gave to the client.
+    /// - `addr`: The client's IP address.
+    /// - `params`: The client's startup parameters.
+    ///
+    pub(crate) fn connect(&self, key: BackendKeyData, addr: SocketAddr, params: &Parameters) {
         self.comms.connect(key, addr, params)
     }
 
-    pub fn update_params(&self, params: &Parameters) {
+    /// Update the params the client has in its session state.
+    pub(crate) fn update_params(&self, params: &Parameters) {
         self.comms.update_params(self.id, params.clone());
     }
 }