fix(linux-postgres): wait for the server to take a cancel request

linux/CHANGELOG.md

@@ -114,6 +114,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - A connection added, changed or deleted in one window did not show in the other windows' connection lists.
 - A connection name with & or < in it showed wrong in the connection list.
 - 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.

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

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
 

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,
 }

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;

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)

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,42 +11,132 @@ 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<dyn TransportStream>);
+pub(crate) struct TransportIo {
+    stream: Box<dyn TransportStream>,
+    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<dyn TransportStream>) -> 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<dyn TransportStream>) -> 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<io::Result<()>> {
-        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<io::Result<usize>> {
-        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<io::Result<()>> {
-        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<io::Result<()>> {
-        Pin::new(&mut *self.0).poll_shutdown(cx)
+    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        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(
         mut self: Pin<&mut Self>,
         cx: &mut Context<'_>,
         bufs: &[io::IoSlice<'_>],
     ) -> Poll<io::Result<usize>> {
-        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<io::Result<()>> {
+    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:?}");
     }
 }

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<Hold>) -> Arc<dyn Connection> {
+    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
+    );
+}

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;