From 8c0a54707965378472eaf2bc0171d178903b41bd Mon Sep 17 00:00:00 2001 From: Ngo Quoc Dat Date: Wed, 14 Oct 2026 23:54:31 +0700 Subject: [PATCH] fix(linux-postgres): wait for the server to take a cancel request tokio-postgres returns from a cancel once the request is written. A read's statement_timeout could end the statement first, and the server then applied the late request to the next statement on that backend. Wait for the server to close the cancel connection, as libpq does. --- linux/CHANGELOG.md | 1 + linux/crates/core/src/engine_stop.rs | 4 +- linux/crates/drivers/postgres/Cargo.toml | 2 +- linux/crates/drivers/postgres/src/cancel.rs | 2 + .../drivers/postgres/src/editor/stop_watch.rs | 6 +- linux/crates/drivers/postgres/src/physical.rs | 7 +- .../drivers/postgres/src/transport_io.rs | 110 ++++++++++++++++-- .../tests/integration/cancellation.rs | 107 ++++++++++++++++- .../postgres/tests/integration/main.rs | 1 + .../postgres/tests/integration/slow_route.rs | 109 +++++++++++++++++ 10 files changed, 332 insertions(+), 17 deletions(-) create mode 100644 linux/crates/drivers/postgres/tests/integration/slow_route.rs diff --git a/linux/CHANGELOG.md b/linux/CHANGELOG.md index 9c6ad6e5d3..8299c39a7f 100644 --- a/linux/CHANGELOG.md +++ b/linux/CHANGELOG.md @@ -68,6 +68,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - A query saved just before quitting could be lost. - The sidebar's table search button had no name for screen readers. - PostgreSQL statement timeouts showed as a plain server error when the server's messages were not in English. +- A PostgreSQL query run right after a timeout or Stop could fail as canceled by that earlier stop. - MySQL and MariaDB connections that verify the certificate did not trust certificate authorities installed on the system. - SQL Server money and smallmoney in the SQL editor and query results showed as approximate floats. - Saving a SQL Server row with a char or varchar key locked the whole table until the save ended. diff --git a/linux/crates/core/src/engine_stop.rs b/linux/crates/core/src/engine_stop.rs index 112ef86e76..d431b7c209 100644 --- a/linux/crates/core/src/engine_stop.rs +++ b/linux/crates/core/src/engine_stop.rs @@ -34,7 +34,9 @@ pub trait EngineStop: Send + Sync + 'static { fn delivery(&self) -> StopDelivery; /// Send the stop. Called at most once, and only for - /// `OutOfBand` delivery. + /// `OutOfBand` delivery. A `Connection` stop returns only once the + /// server has acted on it or it failed to reach the server, because + /// its connection can be handed on as soon as it returns. async fn stop(&self) -> Result<(), DriverError>; /// Whether this error is the engine acknowledging the stop rather diff --git a/linux/crates/drivers/postgres/Cargo.toml b/linux/crates/drivers/postgres/Cargo.toml index 7bed9123ae..9fd02cb8b9 100644 --- a/linux/crates/drivers/postgres/Cargo.toml +++ b/linux/crates/drivers/postgres/Cargo.toml @@ -37,7 +37,7 @@ chrono.workspace = true futures-util.workspace = true rustls.workspace = true secrecy.workspace = true -tokio = { workspace = true, features = ["macros", "rt-multi-thread", "test-util"] } +tokio = { workspace = true, features = ["io-util", "macros", "rt-multi-thread", "test-util"] } tokio-util = { workspace = true, features = ["rt"] } uuid.workspace = true diff --git a/linux/crates/drivers/postgres/src/cancel.rs b/linux/crates/drivers/postgres/src/cancel.rs index df2d70ffc7..ec7f6ed77b 100644 --- a/linux/crates/drivers/postgres/src/cancel.rs +++ b/linux/crates/drivers/postgres/src/cancel.rs @@ -19,6 +19,8 @@ pub(crate) enum StopProgress { /// The request was still out when the call returned: the transport /// did not open, or the server did not answer it. InFlight, + /// The server took the request and closed its stream, which it does + /// only once it has signalled the backend. Delivered, Failed, } diff --git a/linux/crates/drivers/postgres/src/editor/stop_watch.rs b/linux/crates/drivers/postgres/src/editor/stop_watch.rs index 5765241eee..7876565a9a 100644 --- a/linux/crates/drivers/postgres/src/editor/stop_watch.rs +++ b/linux/crates/drivers/postgres/src/editor/stop_watch.rs @@ -117,8 +117,10 @@ impl StopWatch { } /// Wait for an outstanding cancel request before anything else goes - /// to the server. One that has not settled by the end of the grace is - /// dropped: it could not reach the server, so it cannot land late. + /// to the server. The request settles once the server has signalled + /// the backend, so nothing sent after it can be what it cancels. One + /// the server has not taken by the end of the grace is let go, and + /// can still cancel a later statement. pub(crate) async fn settle(&mut self) { let Some(mut request) = self.request.take() else { return; diff --git a/linux/crates/drivers/postgres/src/physical.rs b/linux/crates/drivers/postgres/src/physical.rs index 18348f9dbc..0a03010a40 100644 --- a/linux/crates/drivers/postgres/src/physical.rs +++ b/linux/crates/drivers/postgres/src/physical.rs @@ -94,8 +94,13 @@ impl PgEndpoint { /// Ask the server to stop whatever the connection behind `token` is /// running. The request goes over a stream of its own, because the /// connection it names is busy. + /// + /// It returns once the server has acted on the request. The request + /// names the backend, not the statement, so one the server reads + /// after the statement ended cancels the next statement sent there. + /// A backend that is idle when the signal arrives ignores it. pub(crate) async fn cancel_request(&self, token: &CancelToken) -> Result<(), DriverError> { - let stream = TransportIo::new(self.transport.open().await.map_err(transport_error)?); + let stream = TransportIo::awaiting_close(self.transport.open().await.map_err(transport_error)?); let (connector, _) = self.tls.connector(); token .cancel_query_raw(stream, connector) diff --git a/linux/crates/drivers/postgres/src/transport_io.rs b/linux/crates/drivers/postgres/src/transport_io.rs index c48a19f0ce..0ac4f27678 100644 --- a/linux/crates/drivers/postgres/src/transport_io.rs +++ b/linux/crates/drivers/postgres/src/transport_io.rs @@ -1,6 +1,6 @@ use std::io; use std::pin::Pin; -use std::task::{Context, Poll}; +use std::task::{Context, Poll, ready}; use tablepro_core::TransportStream; use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; @@ -11,31 +11,65 @@ use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; /// trait object there carries an elided lifetime the compiler then asks /// every async caller to prove for all lifetimes. A named `'static` type /// has none to prove. -pub(crate) struct TransportIo(Box); +pub(crate) struct TransportIo { + stream: Box, + ending: Ending, +} + +/// What a shutdown waits for. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Ending { + WriteHalf, + PeerClose { write_closed: bool }, +} impl TransportIo { pub(crate) fn new(stream: Box) -> Self { - Self(stream) + Self { + stream, + ending: Ending::WriteHalf, + } + } + + /// A stream whose shutdown returns only once the server has closed + /// its side too, which the server does for a cancel request only + /// after it has signalled the backend. libpq's cancel waits the same + /// way; tokio-postgres's returns once the request is written. + pub(crate) fn awaiting_close(stream: Box) -> Self { + Self { + stream, + ending: Ending::PeerClose { write_closed: false }, + } } } impl AsyncRead for TransportIo { fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { - Pin::new(&mut *self.0).poll_read(cx, buf) + Pin::new(&mut *self.stream).poll_read(cx, buf) } } impl AsyncWrite for TransportIo { fn poll_write(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll> { - Pin::new(&mut *self.0).poll_write(cx, buf) + Pin::new(&mut *self.stream).poll_write(cx, buf) } fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - Pin::new(&mut *self.0).poll_flush(cx) + Pin::new(&mut *self.stream).poll_flush(cx) } - fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - Pin::new(&mut *self.0).poll_shutdown(cx) + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + let Ending::PeerClose { write_closed } = &mut this.ending else { + return Pin::new(&mut *this.stream).poll_shutdown(cx); + }; + if !*write_closed { + // A write half that fails to close still leaves the server to + // close its side, and that close is what is waited for. + let _ = ready!(Pin::new(&mut *this.stream).poll_shutdown(cx)); + *write_closed = true; + } + poll_peer_closed(&mut *this.stream, cx) } fn poll_write_vectored( @@ -43,10 +77,66 @@ impl AsyncWrite for TransportIo { cx: &mut Context<'_>, bufs: &[io::IoSlice<'_>], ) -> Poll> { - Pin::new(&mut *self.0).poll_write_vectored(cx, bufs) + Pin::new(&mut *self.stream).poll_write_vectored(cx, bufs) } fn is_write_vectored(&self) -> bool { - self.0.is_write_vectored() + self.stream.is_write_vectored() + } +} + +/// Read and drop whatever arrives until the server closes. A reset ends +/// the wait as a close does, as it does for libpq. +fn poll_peer_closed(stream: &mut dyn TransportStream, cx: &mut Context<'_>) -> Poll> { + let mut scratch = [0_u8; 64]; + loop { + let mut buf = ReadBuf::new(&mut scratch); + match ready!(Pin::new(&mut *stream).poll_read(cx, &mut buf)) { + Ok(()) if buf.filled().is_empty() => return Poll::Ready(Ok(())), + Ok(()) => {} + Err(_) => return Poll::Ready(Ok(())), + } + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + use super::*; + + #[tokio::test(start_paused = true)] + async fn a_shutdown_awaiting_close_returns_once_the_server_closes() { + let (client, mut server) = tokio::io::duplex(64); + let mut stream = TransportIo::awaiting_close(Box::new(client)); + let closing = tokio::spawn(async move { stream.shutdown().await }); + + let mut request = Vec::new(); + server + .read_to_end(&mut request) + .await + .expect("the client's end of stream"); + server.write_all(b"ignored").await.expect("a write to the client"); + tokio::time::sleep(Duration::from_secs(1)).await; + assert!( + !closing.is_finished(), + "the shutdown returned while the server was still open" + ); + drop(server); + + let closed = tokio::time::timeout(Duration::from_secs(1), closing).await; + assert!(matches!(closed, Ok(Ok(Ok(())))), "{closed:?}"); + } + + #[tokio::test(start_paused = true)] + async fn a_plain_shutdown_closes_only_the_write_half() { + let (client, _server) = tokio::io::duplex(64); + let mut stream = TransportIo::new(Box::new(client)); + + let closed = tokio::time::timeout(Duration::from_secs(1), stream.shutdown()).await; + + assert!(matches!(closed, Ok(Ok(()))), "{closed:?}"); } } diff --git a/linux/crates/drivers/postgres/tests/integration/cancellation.rs b/linux/crates/drivers/postgres/tests/integration/cancellation.rs index b29f0caad8..3a185ddb6a 100644 --- a/linux/crates/drivers/postgres/tests/integration/cancellation.rs +++ b/linux/crates/drivers/postgres/tests/integration/cancellation.rs @@ -1,13 +1,28 @@ //! A stop reaches the server: the backend goes idle rather than finishing //! work nobody waits for. +use std::sync::Arc; use std::time::Duration; -use tablepro_core::{CallOptions, Connection, DriverError, RowLimit, TimeoutPhase, Value}; +use tablepro_core::{ + CallOptions, Connection, DriverError, RowLimit, RunEnd, RunOptions, TimeoutPhase, TlsConfig, Value, + WorkspaceSettings, +}; +use tablepro_test_fixtures::PostgresFixture; +use tablepro_test_support::conformance::test_liveness; use tokio::time::Instant; use tokio_util::sync::CancellationToken; -use crate::support::{connect, scalar, start, statement, within}; +use crate::slow_route::{CANCEL_REQUEST, Hold, SlowRoute}; +use crate::support::{ + CALL, connect, connect_with, direct, options_over, password, run, run_with, scalar, session, start, statement, + within, +}; + +/// How long a cancel request takes to reach the server on a slow route. +/// It is inside the test cancel grace, so a driver that waits for the +/// server to take the request has the time to. +const CANCEL_TRANSIT: Duration = Duration::from_millis(700); const SLEEPERS: &str = "SELECT count(*) FROM pg_stat_activity \ WHERE state = 'active' AND query LIKE '%pg_sleep(60)%' AND pid <> pg_backend_pid()"; @@ -68,3 +83,91 @@ async fn deadline_1s_timeout_server_cancelled() { ); assert_eq!(backends_sleeping(connection.as_ref()).await, Value::Int(0)); } + +async fn connect_over(fixture: &PostgresFixture, holds: Vec) -> Arc { + let liveness = test_liveness(); + let route = SlowRoute::new(direct(fixture.host(), fixture.port(), &liveness), holds); + connect_with(options_over( + Arc::new(route), + password(&fixture.password_credentials()), + fixture.database(), + TlsConfig::disabled(), + WorkspaceSettings::workspace(false), + liveness, + )) + .await +} + +/// The read's own statement timeout ends it while its cancel request is +/// still on the way, and the server acts on that request whenever it +/// arrives, on whatever the backend runs then. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[ignore = "requires docker"] +async fn a_late_cancel_after_a_timeout_never_stops_the_next_read() { + let fixture = start().await; + // The statement reaches the server a little after the call started, + // so the server's timeout fires just after the deadline sent the + // cancel request. + let connection = connect_over( + &fixture, + vec![ + Hold::new(b"pg_sleep(60)", Duration::from_millis(50)), + Hold::new(&CANCEL_REQUEST, CANCEL_TRANSIT), + ], + ) + .await; + + let timed_out = connection + .query_bound( + &statement("SELECT pg_sleep(60)"), + RowLimit::EDITOR_DEFAULT, + within(Duration::from_secs(1)), + ) + .await; + let next = connection + .query_bound( + &statement("SELECT 1 FROM pg_sleep(1)"), + RowLimit::EDITOR_DEFAULT, + within(CALL), + ) + .await; + + assert!( + matches!( + timed_out, + Err(DriverError::Timeout { + phase: TimeoutPhase::Statement, + server_cancelled: true + }) + ), + "{timed_out:?}" + ); + assert!(next.is_ok(), "the read after the timeout gave {next:?}"); +} + +/// A statement that ends on its own after Stop leaves the cancel request +/// on its way to the backend the session keeps. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[ignore = "requires docker"] +async fn a_late_cancel_after_stop_never_stops_the_next_editor_run() { + let fixture = start().await; + let connection = connect_over(&fixture, vec![Hold::new(&CANCEL_REQUEST, CANCEL_TRANSIT)]).await; + let editor = session(connection.as_ref()).await; + let cancel = CancellationToken::new(); + let stopper = cancel.clone(); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(500)).await; + stopper.cancel(); + }); + + let (stopped, _) = run_with(editor.as_ref(), "SELECT pg_sleep(0.8)", RunOptions::new(cancel)).await; + let (next, _) = run(editor.as_ref(), "SELECT 1 FROM pg_sleep(1)").await; + + assert_eq!(stopped.end, RunEnd::Cancelled); + assert_eq!( + next.end, + RunEnd::Completed, + "the run after the stop ended {:?}", + next.end + ); +} diff --git a/linux/crates/drivers/postgres/tests/integration/main.rs b/linux/crates/drivers/postgres/tests/integration/main.rs index 08de5bf75a..91981dcf07 100644 --- a/linux/crates/drivers/postgres/tests/integration/main.rs +++ b/linux/crates/drivers/postgres/tests/integration/main.rs @@ -17,6 +17,7 @@ mod liveness; mod metadata; mod pgbouncer; mod read_only; +mod slow_route; mod support; mod tls; mod write_batch; diff --git a/linux/crates/drivers/postgres/tests/integration/slow_route.rs b/linux/crates/drivers/postgres/tests/integration/slow_route.rs new file mode 100644 index 0000000000..2dea1b4705 --- /dev/null +++ b/linux/crates/drivers/postgres/tests/integration/slow_route.rs @@ -0,0 +1,109 @@ +//! A route on which chosen requests reach the server late, as over a +//! congested link: the client's write and shutdown complete at once, +//! the way they do into a socket buffer, and the bytes arrive later. + +use std::sync::Arc; +use std::time::Duration; + +use async_trait::async_trait; +use tablepro_core::{NetworkEndpoint, Transport, TransportError, TransportKind, TransportRoute, TransportStream}; +use tokio::io::{AsyncReadExt, AsyncWriteExt, DuplexStream}; + +/// How a CancelRequest packet starts: its length, 16, then its code, +/// 80877102. +pub const CANCEL_REQUEST: [u8; 8] = [0, 0, 0, 16, 0x04, 0xd2, 0x16, 0x2e]; + +/// Bytes the route holds back, and for how long. +#[derive(Debug, Clone)] +pub struct Hold { + pub pattern: Vec, + pub delay: Duration, +} + +impl Hold { + pub fn new(pattern: &[u8], delay: Duration) -> Self { + Self { + pattern: pattern.to_vec(), + delay, + } + } +} + +#[derive(Debug)] +pub struct SlowRoute { + inner: Arc, + holds: Vec, +} + +impl SlowRoute { + /// A stream that writes one of the patterns has everything from that + /// write on delivered `delay` later. The patterns match plaintext. + pub fn new(inner: Arc, holds: Vec) -> Self { + Self { inner, holds } + } +} + +#[async_trait] +impl Transport for SlowRoute { + fn service_endpoint(&self) -> &NetworkEndpoint { + self.inner.service_endpoint() + } + + fn kind(&self) -> TransportKind { + self.inner.kind() + } + + fn route(&self) -> TransportRoute { + self.inner.route() + } + + async fn open(&self) -> Result, TransportError> { + let server = self.inner.open().await?; + let (client, relay) = tokio::io::duplex(64 * 1024); + tokio::spawn(relay_through(relay, server, self.holds.clone())); + Ok(Box::new(client)) + } + + async fn reroute(&self, endpoint: NetworkEndpoint) -> Result, TransportError> { + self.inner.reroute(endpoint).await + } + + async fn explain_closed_stream(&self) -> Option { + None + } +} + +/// Relay one stream. The relay outlives the client's end, so bytes the +/// client wrote before it closed still reach the server. +async fn relay_through(relay: DuplexStream, server: Box, holds: Vec) { + let (mut from_client, mut to_client) = tokio::io::split(relay); + let (mut from_server, mut to_server) = tokio::io::split(server); + let upstream = async move { + let mut chunk = vec![0_u8; 16 * 1024]; + loop { + let read = from_client.read(&mut chunk).await.unwrap_or(0); + let Some(bytes) = chunk.get(..read).filter(|bytes| !bytes.is_empty()) else { + break; + }; + if let Some(delay) = held_for(&holds, bytes) { + tokio::time::sleep(delay).await; + } + if to_server.write_all(bytes).await.is_err() { + return; + } + } + let _ = to_server.shutdown().await; + }; + let downstream = async move { + let _ = tokio::io::copy(&mut from_server, &mut to_client).await; + let _ = to_client.shutdown().await; + }; + tokio::join!(upstream, downstream); +} + +fn held_for(holds: &[Hold], bytes: &[u8]) -> Option { + holds + .iter() + .find(|hold| !hold.pattern.is_empty() && bytes.windows(hold.pattern.len()).any(|window| window == hold.pattern)) + .map(|hold| hold.delay) +}