Drop client queries abandoned during maintenance mode

integration/rust/tests/integration/maintenance_mode.rs

@@ -2,8 +2,9 @@ use std::time::Duration;
 
 use crate::setup::{admin_sqlx, admin_tokio, connection_failover, connections_sqlx};
 use serial_test::serial;
-use sqlx::Executor;
+use sqlx::{Executor, Row};
 use tokio::time::sleep;
+use tokio_postgres::NoTls;
 
 #[tokio::test]
 #[serial]
@@ -50,6 +51,73 @@ async fn test_maintenance_mode_client_queries() {
     conn.close().await;
 }
 
+#[tokio::test]
+#[serial]
+async fn test_maintenance_mode_drops_disconnected_client_query() {
+    // A client whose query is held by maintenance mode, then disconnects
+    // before it lifts, must not have that query executed once maintenance is
+    // turned off. We assert on a write's side effect: the row must never land.
+    let admin = admin_tokio().await;
+    let probe = connection_failover().await;
+
+    admin.simple_query("MAINTENANCE OFF").await.unwrap();
+
+    probe
+        .execute("CREATE TABLE IF NOT EXISTS maintenance_probe (id INT)")
+        .await
+        .unwrap();
+    probe.execute("TRUNCATE maintenance_probe").await.unwrap();
+
+    // Pause the database.
+    admin.simple_query("MAINTENANCE ON").await.unwrap();
+
+    // A raw connection we can kill mid-flight.
+    let (client, connection) = tokio_postgres::connect(
+        "host=127.0.0.1 user=pgdog dbname=failover password=pgdog port=6432",
+        NoTls,
+    )
+    .await
+    .unwrap();
+    let connection = tokio::spawn(async move {
+        let _ = connection.await;
+    });
+
+    // Fire the write. It blocks inside PgDog, parked on maintenance mode.
+    // Use simple_query so the INSERT travels as a single Query message and is
+    // buffered as one complete request. execute() would use the extended
+    // protocol and park on the statement-prepare round-trip instead, leaving
+    // no INSERT buffered to dispatch, so it can't tell the fix from the bug.
+    let query = tokio::spawn(async move {
+        let _ = client
+            .simple_query("INSERT INTO maintenance_probe VALUES (1)")
+            .await;
+        drop(client);
+    });
+
+    // Let PgDog buffer the query and park.
+    sleep(Duration::from_millis(100)).await;
+
+    // Disconnect abruptly while still paused: kill the driver and the query.
+    connection.abort();
+    query.abort();
+    sleep(Duration::from_millis(50)).await;
+
+    // Lift maintenance. The old code would now dispatch the buffered write.
+    admin.simple_query("MAINTENANCE OFF").await.unwrap();
+    sleep(Duration::from_millis(200)).await;
+
+    // The write from the gone client must not have run.
+    let count: i64 = probe
+        .fetch_one("SELECT count(*) FROM maintenance_probe")
+        .await
+        .unwrap()
+        .get(0);
+    assert_eq!(count, 0, "query from a disconnected client was executed");
+
+    probe.execute("DROP TABLE maintenance_probe").await.unwrap();
+    probe.close().await;
+}
+
 #[tokio::test]
 #[serial]
 async fn test_maintenance_mode_admin_connections() {

pgdog/src/frontend/client/mod.rs

@@ -3,6 +3,7 @@
 //! Entrypoint for client/server interactions.
 //!
 
+use std::future::IntoFuture;
 use std::net::SocketAddr;
 use std::sync::Arc;
 use std::time::{Duration, Instant};
@@ -531,17 +532,6 @@ 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() {
@@ -656,6 +646,38 @@ impl Client {
             );
         }
 
+        // Hold a complete request at the door while the database is in
+        // maintenance mode. buffer() already owns the client socket, so we
+        // watch it for a disconnect and drop the request instead of
+        // dispatching it once maintenance lifts. Only gate between
+        // transactions; a transaction already in flight is allowed to finish.
+        // Admin connections are never gated.
+        if !self.in_transaction()
+            && !self.admin
+            && let Some(waiter) = maintenance_mode::waiter(&self.database)
+        {
+            // `Pin<Box<_>>` is `Unpin`, so `&mut waiter` polls the same
+            // future across iterations.
+            let mut waiter = waiter.into_future();
+            loop {
+                select! {
+                    _ = &mut waiter => break,
+
+                    // A client waiting for its response shouldn't be sending
+                    // anything, so a readable socket means either a disconnect
+                    // (EOF/error) or pipelined traffic. Pipelined bytes stay
+                    // buffered for the next request; only a disconnect ends the
+                    // wait early.
+                    result = self.stream_buffer.read_more(&mut self.stream) => {
+                        if result.is_err() {
+                            self.client_request.clear();
+                            return Ok(BufferEvent::DisconnectAbrupt);
+                        }
+                    }
+                }
+            }
+        }
+
         Ok(BufferEvent::HaveRequest)
     }
 

pgdog/src/net/messages/buffer.rs

@@ -192,6 +192,35 @@ impl MessageBuffer {
     ) -> Result<Message, Error> {
         self.read_internal(stream).await
     }
+
+    /// Read whatever is currently available on the stream into the buffer,
+    /// without framing a message. A closed stream (EOF) surfaces as
+    /// `Error::UnexpectedEof`.
+    ///
+    /// This is meant for watching an otherwise-idle client for a disconnect:
+    /// any bytes read stay in the buffer for the next [`read`](Self::read), so
+    /// the read is not observed by the caller and no data is lost.
+    ///
+    /// # Cancellation safety
+    ///
+    /// Cancel-safe: bytes already read remain buffered.
+    pub async fn read_more(
+        &mut self,
+        stream: &mut (impl Unpin + AsyncReadExt),
+    ) -> Result<(), Error> {
+        // Keep some headroom so we reclaim memory frequently, matching
+        // `read_internal`.
+        self.ensure_capacity(self.capacity / 4);
+
+        let read = eof(stream.read_buf(&mut self.buffer).await)?;
+        self.stats.bytes_used += read;
+
+        if read == 0 {
+            return Err(Error::UnexpectedEof);
+        }
+
+        Ok(())
+    }
 }
 
 #[cfg(test)]