diff --git a/integration/rust/tests/integration/maintenance_mode.rs b/integration/rust/tests/integration/maintenance_mode.rs index 80eff213f..6912994a5 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,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() { diff --git a/pgdog/src/frontend/client/mod.rs b/pgdog/src/frontend/client/mod.rs index 56f9004ed..e969963e4 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}; @@ -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>` 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) } 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)]