drain connections when more messages are expected

The needs_drain() condition prior to this commit checks the connection
state is some kind of idle or otherwise "in sync". However, even when
the connection state is Idle, there can be expected messages in the
protocol queue. If a client disconnects abruptly at just the right time,
needs_drain() would not return true, but the protocol queue is not empty
and expected messages from the server are not drained.

In this scenario, the connection is checked back in in a tainted state
and subsequent usage can produce unexpected results (ProtocolOutOfSync
or off-by-1 results mixups).

This updates needs_drain() to also account for
prepared_statements.has_more_messages() in addition to the "in sync"
check, thus triggering a drain when there are more messages expected
from the server (and tracked in the protocol queue).

Author: jbielick

pgdog/src/backend/pool/test/mod.rs

@@ -653,3 +653,59 @@ async fn test_lsn_monitor() {
 
     pool.shutdown();
 }
+
+/// If a client disconnects while the backend still has protocol messages
+/// queued for it, but the high-level stats state is Idle, the connection must
+/// still be treated as needing a drain. Otherwise the connection is checked
+/// back into the pool with a non-empty protocol queue and will be "tainted"
+/// for the next client.
+#[tokio::test]
+async fn test_abrupt_disconnect_with_pending_protocol_state_requires_drain() {
+    crate::logger();
+
+    let pool = pool();
+    let mut guard = pool.get(&Request::default()).await.unwrap();
+
+    guard
+        .send(&vec![ProtocolMessage::from(Query::new("SELECT 1"))].into())
+        .await
+        .unwrap();
+
+    // Consume only the RowDescription ('T'), leaving the remainder of the
+    // response (D, C, Z) unread. At this point the protocol state still
+    // expects a ReadyForQuery, so has_more_messages() is true.
+    let msg = guard.read().await.unwrap();
+    assert_eq!(msg.code(), 'T');
+
+    assert!(
+        guard.prepared_statements().state().has_more_messages(),
+        "protocol state should have pending messages after partial read"
+    );
+    guard.stats_mut().state(State::Idle);
+    assert_eq!(guard.stats().state, State::Idle);
+
+    // Client disconnects, guard is dropped, runs Guard::cleanup. needs_drain()
+    // must report that there are pending protocol messages it's expecting,
+    // otherwise the connection is checked back into the pool with a non-empty
+    // protocol queue.
+    drop(guard);
+
+    // Re-acquire the same connection from the pool. If this
+    // assertion ever fails, the connection has been returned to the pool
+    // in a tainted state.
+    let mut guard = pool.get(&Request::default()).await.unwrap();
+    assert!(
+        !guard.prepared_statements().state().has_more_messages(),
+        "connection should not have pending protocol messages after cleanup"
+    );
+
+    // Run a fresh query and assert we get a valid result back; if the
+    // connection had been returned to the pool without a drain, this is
+    // where stray CommandComplete/ReadyForQuery messages from the previous
+    // query would surface and break the protocol. Each query would receive the
+    // prior query's results (off by 1).
+    let rows: Vec<i32> = guard.fetch_all("SELECT 2").await.unwrap();
+    assert_eq!(rows, vec![2]);
+    let rows: Vec<i32> = guard.fetch_all("SELECT 3").await.unwrap();
+    assert_eq!(rows, vec![3]);
+}

pgdog/src/backend/server.rs

@@ -562,7 +562,7 @@ impl Server {
         )
     }
 
-    /// Server is done executing all queries and isz
+    /// Server is done executing all queries and is
     /// not inside a transaction.
     pub fn can_check_in(&self) -> bool {
         self.stats().state == State::Idle
@@ -602,7 +602,17 @@ impl Server {
 
     /// Connection was left with an unfinished query.
     pub fn needs_drain(&self) -> bool {
-        !self.in_sync()
+        // We need to drain whenever the protocol state indicates that there
+        // are still tracked messages expected from the server *or* the high‑
+        // level connection state is not one of the "in sync" states.
+        //
+        // This is intentionally more conservative than just checking
+        // `in_sync()`: if a client disappears after we've enqueued expected
+        // backend responses (via `PreparedStatements::handle`) but before we
+        // see a new ReadyForQuery, the stats state can still be Idle while the
+        // protocol queue has pending items. In that case we must treat the
+        // connection as needing a drain before it is safe to reuse.
+        self.has_more_messages() || !self.in_sync()
     }
 
     /// Close the connection, don't do any recovery.
@@ -1771,6 +1781,23 @@ pub mod test {
         }
     }
 
+    #[test]
+    fn test_needs_drain_when_protocol_queue_not_empty_even_if_idle() {
+        use crate::state::State;
+
+        let mut server = Server::default();
+        assert_eq!(server.stats().state, State::Idle);
+        assert!(server.in_sync());
+        assert!(!server.has_more_messages());
+        assert!(!server.needs_drain());
+
+        server.prepared_statements_mut().state_mut().add('Z'); // expect a ReadyForQuery at some point
+
+        assert_eq!(server.stats().state, State::Idle);
+        assert!(server.prepared_statements().state().has_more_messages());
+        assert!(server.needs_drain());
+    }
+
     #[tokio::test]
     async fn test_partial_state() -> Result<(), Box<dyn std::error::Error>> {
         crate::logger();