From e23efe25c34f3c47f79bb912fd99507347c13f63 Mon Sep 17 00:00:00 2001 From: Cesar Rodas Date: Mon, 20 Jul 2026 17:54:08 -0300 Subject: [PATCH 1/3] Drop client queries abandoned during maintenance mode Maintenance mode parked a fully-buffered request on a bare await that ignored the client socket. If the client gave up and disconnected while the database was paused, its query was still dispatched to the backend once maintenance lifted, running work for a client that was already gone. Wait for maintenance to lift while watching the client socket, so a disconnect aborts the request before it reaches the query engine. Add a cancel-safe MessageBuffer::read_more that detects the disconnect without consuming any pipelined data, and cover the behavior with an integration test that asserts a write from a disconnected client never lands. --- .../tests/integration/maintenance_mode.rs | 66 ++++++++++++++- pgdog/src/frontend/client/mod.rs | 83 ++++++++++++++++--- pgdog/src/net/messages/buffer.rs | 29 +++++++ 3 files changed, 164 insertions(+), 14 deletions(-) diff --git a/integration/rust/tests/integration/maintenance_mode.rs b/integration/rust/tests/integration/maintenance_mode.rs index 80eff213f..a17fd3e49 100644 --- a/integration/rust/tests/integration/maintenance_mode.rs +++ b/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,69 @@ 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. + let query = tokio::spawn(async move { + let _ = client + .execute("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() { diff --git a/pgdog/src/frontend/client/mod.rs b/pgdog/src/frontend/client/mod.rs index 56f9004ed..038f74c7c 100644 --- a/pgdog/src/frontend/client/mod.rs +++ b/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}; @@ -495,12 +496,24 @@ impl Client { } buffer = self.buffer(client_state) => { - let event = buffer?; + let mut event = buffer?; // Only send requests to the backend if they are complete. if self.client_request.is_complete() && !self.client_request.messages.is_empty() { - self.client_messages(&mut query_engine).await?; + // Hold the request at the door while the database is + // in maintenance mode, before it reaches the query + // engine. If the client gives up and disconnects + // while we're paused, drop the request instead of + // running it once maintenance lifts. + match self.wait_for_maintenance(&mut query_engine).await? { + MaintenanceWait::Proceed => { + self.client_messages(&mut query_engine).await?; + } + MaintenanceWait::Disconnected => { + event = BufferEvent::DisconnectAbrupt; + } + } } match event { @@ -529,19 +542,54 @@ impl Client { Ok(()) } - /// 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); + /// Wait out maintenance mode before a buffered request reaches the backend. + /// + /// We only park between transactions; a transaction already in flight is + /// allowed to finish. While parked we keep watching the client socket, so a + /// client that gives up and disconnects during a long maintenance window + /// never has its query dispatched once maintenance lifts. + async fn wait_for_maintenance( + &mut self, + query_engine: &mut QueryEngine, + ) -> Result { + if self.in_transaction() || self.admin { + return Ok(MaintenanceWait::Proceed); } + let Some(waiter) = maintenance_mode::waiter(&self.database) else { + return Ok(MaintenanceWait::Proceed); + }; + + let saved_state = query_engine.get_state(); + query_engine.set_state(State::Waiting); + + // `Pin>` is `Unpin`, so `&mut waiter` polls the same future + // across loop iterations. + let mut waiter = waiter.into_future(); + let outcome = loop { + select! { + _ = &mut waiter => break MaintenanceWait::Proceed, + + // 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() { + break MaintenanceWait::Disconnected; + } + } + } + }; + + query_engine.set_state(saved_state); + + Ok(outcome) + } + + /// Handle client messages. + async fn client_messages(&mut self, query_engine: &mut QueryEngine) -> Result<(), Error> { // If client sent multiple requests, split them up and execute individually. let spliced = self.client_request.spliced()?; if spliced.is_empty() { @@ -714,3 +762,12 @@ enum BufferEvent { DisconnectAbrupt, HaveRequest, } + +/// Outcome of waiting for a database to leave maintenance mode. +#[derive(Copy, Clone, PartialEq, Debug)] +enum MaintenanceWait { + /// Maintenance isn't on (or has lifted): dispatch the request. + Proceed, + /// The client disconnected while paused: drop the request. + Disconnected, +} diff --git a/pgdog/src/net/messages/buffer.rs b/pgdog/src/net/messages/buffer.rs index 33bf4cc59..bf325f8fb 100644 --- a/pgdog/src/net/messages/buffer.rs +++ b/pgdog/src/net/messages/buffer.rs @@ -192,6 +192,35 @@ impl MessageBuffer { ) -> Result { 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)] From 092d3b2f65b5f78f8e1eff9c0bff2cb38afdb8e6 Mon Sep 17 00:00:00 2001 From: Cesar Rodas Date: Tue, 21 Jul 2026 00:41:41 -0300 Subject: [PATCH 2/3] Move the maintenance-mode gate into the client buffer loop The disconnect fix for maintenance mode added a separate wait_for_maintenance method, a small result enum, and matching branches in the client run loop. buffer() already owns the client socket and already turns an EOF into a disconnect, so it is the natural home for the wait and the extra scaffolding is redundant. Fold the gate into buffer(): once a request is fully buffered, park on the maintenance waiter while watching the socket, and on a disconnect clear the request and return the existing abrupt-disconnect event. This removes wait_for_maintenance, its result enum, and the run-loop plumbing. The check goes at the bottom of buffer, after the read loop, not at the top. A client sits idle inside buffer's read loop between requests; if maintenance turns on while it is parked there, a top-of-function check has already run and would let the next query through. Gating right before the buffered request is returned keeps the hold-at-dispatch behavior and is exercised by the existing idle-client maintenance test. The gate no longer sets the frontend Waiting state during a park, because buffer cannot borrow the query engine: the run loop already holds it to read the backend. A parked client keeps its prior state in SHOW CLIENTS, an accepted trade for the simpler shape. --- pgdog/src/frontend/client/mod.rs | 103 ++++++++++--------------------- 1 file changed, 34 insertions(+), 69 deletions(-) diff --git a/pgdog/src/frontend/client/mod.rs b/pgdog/src/frontend/client/mod.rs index 038f74c7c..e969963e4 100644 --- a/pgdog/src/frontend/client/mod.rs +++ b/pgdog/src/frontend/client/mod.rs @@ -496,24 +496,12 @@ impl Client { } buffer = self.buffer(client_state) => { - let mut event = buffer?; + let event = buffer?; // Only send requests to the backend if they are complete. if self.client_request.is_complete() && !self.client_request.messages.is_empty() { - // Hold the request at the door while the database is - // in maintenance mode, before it reaches the query - // engine. If the client gives up and disconnects - // while we're paused, drop the request instead of - // running it once maintenance lifts. - match self.wait_for_maintenance(&mut query_engine).await? { - MaintenanceWait::Proceed => { - self.client_messages(&mut query_engine).await?; - } - MaintenanceWait::Disconnected => { - event = BufferEvent::DisconnectAbrupt; - } - } + self.client_messages(&mut query_engine).await?; } match event { @@ -542,52 +530,6 @@ impl Client { Ok(()) } - /// Wait out maintenance mode before a buffered request reaches the backend. - /// - /// We only park between transactions; a transaction already in flight is - /// allowed to finish. While parked we keep watching the client socket, so a - /// client that gives up and disconnects during a long maintenance window - /// never has its query dispatched once maintenance lifts. - async fn wait_for_maintenance( - &mut self, - query_engine: &mut QueryEngine, - ) -> Result { - if self.in_transaction() || self.admin { - return Ok(MaintenanceWait::Proceed); - } - - let Some(waiter) = maintenance_mode::waiter(&self.database) else { - return Ok(MaintenanceWait::Proceed); - }; - - let saved_state = query_engine.get_state(); - query_engine.set_state(State::Waiting); - - // `Pin>` is `Unpin`, so `&mut waiter` polls the same future - // across loop iterations. - let mut waiter = waiter.into_future(); - let outcome = loop { - select! { - _ = &mut waiter => break MaintenanceWait::Proceed, - - // 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() { - break MaintenanceWait::Disconnected; - } - } - } - }; - - query_engine.set_state(saved_state); - - Ok(outcome) - } - /// Handle client messages. async fn client_messages(&mut self, query_engine: &mut QueryEngine) -> Result<(), Error> { // If client sent multiple requests, split them up and execute individually. @@ -704,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>` 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) } @@ -762,12 +736,3 @@ enum BufferEvent { DisconnectAbrupt, HaveRequest, } - -/// Outcome of waiting for a database to leave maintenance mode. -#[derive(Copy, Clone, PartialEq, Debug)] -enum MaintenanceWait { - /// Maintenance isn't on (or has lifted): dispatch the request. - Proceed, - /// The client disconnected while paused: drop the request. - Disconnected, -} From 8218af93ba40d72c10131a02e3238a471f263878 Mon Sep 17 00:00:00 2001 From: Cesar Rodas Date: Wed, 22 Jul 2026 19:11:16 -0300 Subject: [PATCH 3/3] Use simple_query in the maintenance disconnect test The disconnected-client test needs the INSERT to reach PgDog as a single complete request that gets buffered and parked on maintenance mode. With execute() the extended protocol parks on the statement-prepare round-trip instead, so no INSERT is ever buffered and the test cannot distinguish the fix from the bug. simple_query sends the INSERT as one Query message, which is what the test needs to exercise the dispatch-after-disconnect path. --- integration/rust/tests/integration/maintenance_mode.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/integration/rust/tests/integration/maintenance_mode.rs b/integration/rust/tests/integration/maintenance_mode.rs index a17fd3e49..6912994a5 100644 --- a/integration/rust/tests/integration/maintenance_mode.rs +++ b/integration/rust/tests/integration/maintenance_mode.rs @@ -83,9 +83,13 @@ async fn test_maintenance_mode_drops_disconnected_client_query() { }); // 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 - .execute("INSERT INTO maintenance_probe VALUES (1)", &[]) + .simple_query("INSERT INTO maintenance_probe VALUES (1)") .await; drop(client); });