fix: client disconnect inside transaction caused deadlock (#848)

If a client sent `Terminate` while inside a transaction (i.e. closing
connection), we would deadlock, leaving the client task stuck, blocking
the server connection inside the transaction as well.

This is a regression introduced in #820, in an attempt to fix session
mode clients not disconnecting properly.

Contributing factors include a special "test mode" for the query engine
/ client query handling code that required code that only ran in tests
but not in production, and vice versa. This was removed, and the client
query handling loop is now fully tested the way it works in prod.

Author: levkk

.config/nextest.toml

@@ -2,3 +2,6 @@
 fail-fast = true
 test-threads = 1
 default-filter = "not package(rust)"
+
+[profile.default]
+slow-timeout = "15s"

integration/python/test_idle_in_transaction.py

@@ -1,4 +1,5 @@
 import asyncio
+import asyncpg
 import pytest
 from sqlalchemy import text
 from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker
@@ -101,3 +102,46 @@ async def test_idle_in_transaction_timeout():
     under concurrent load."""
     tasks = [_run_idle_in_transaction() for _ in range(WORKERS)]
     await asyncio.gather(*tasks)
+
+
+async def _disconnect_mid_transaction():
+    """Single worker: begin a transaction, disconnect, verify rollback."""
+    conn = await asyncpg.connect(
+        user="pgdog",
+        password="pgdog",
+        database="pgdog",
+        host="127.0.0.1",
+        port=6432,
+    )
+
+    # Start a transaction and grab the backend pid
+    await conn.execute("BEGIN")
+    pid = await conn.fetchval("SELECT pg_backend_pid()")
+    await conn.execute("SELECT 1")
+
+    # Disconnect abruptly (without COMMIT/ROLLBACK)
+    await conn.close()
+
+    # Give pgdog a moment to clean up
+    await asyncio.sleep(0.5)
+
+    # Check via direct connection that the backend is idle (transaction was rolled back)
+    direct = direct_sync()
+    direct.autocommit = True
+    cur = direct.cursor()
+    cur.execute(
+        "SELECT state FROM pg_stat_activity WHERE pid = %s",
+        (pid,),
+    )
+    row = cur.fetchone()
+    direct.close()
+
+    assert row is not None, f"expected backend {pid} to still exist in the pool"
+    assert row[0] == "idle", f"expected backend {pid} to be idle (rolled back), got: {row[0]}"
+
+
+@pytest.mark.asyncio
+async def test_disconnect_mid_transaction_rolls_back():
+    """Disconnect in the middle of a transaction and verify pgdog rolls it back."""
+    tasks = [_disconnect_mid_transaction() for _ in range(WORKERS)]
+    await asyncio.gather(*tasks)

pgdog/src/backend/pool/connection/mirror/mod.rs

@@ -96,11 +96,6 @@ impl Mirror {
         let mut query_engine =
             QueryEngine::new(&params, &ClientComms::new(&BackendKeyData::new()), false)?;
 
-        // Mirror must read server responses to keep the connection synchronized,
-        // so disable test_mode which skips reading responses.
-        #[cfg(test)]
-        query_engine.set_test_mode(false);
-
         // Mirror traffic handler.
         let mut mirror = Self::new(&params, &config);
 

pgdog/src/frontend/client/mod.rs

@@ -47,7 +47,6 @@ pub struct Client {
     comms: ClientComms,
     admin: bool,
     streaming: bool,
-    shutdown: bool,
     prepared_statements: PreparedStatements,
     transaction: Option<TransactionType>,
     timeouts: Timeouts,
@@ -306,7 +305,6 @@ impl Client {
             timeouts: Timeouts::from_config(&config.config.general),
             client_request: ClientRequest::new(),
             stream_buffer: MessageBuffer::new(config.config.memory.message_buffer),
-            shutdown: false,
             sticky: Sticky::from_params(&params),
             connect_params: params,
         }))
@@ -338,7 +336,6 @@ impl Client {
             timeouts: Timeouts::from_config(&config().config.general),
             client_request: ClientRequest::new(),
             stream_buffer: MessageBuffer::new(4096),
-            shutdown: false,
             sticky: Sticky::from_params(&connect_params),
             params: connect_params,
         }
@@ -380,28 +377,28 @@ impl Client {
     /// Run the client.
     async fn run(&mut self) -> Result<(), Error> {
         let shutdown = self.comms.shutting_down();
-        let mut offline;
         let mut query_engine = QueryEngine::from_client(self)?;
-        let mut terminating = false;
 
         loop {
-            offline = (self.comms.offline() && !self.admin || self.shutdown)
-                && query_engine.can_disconnect();
-            if offline {
-                break;
-            }
-
-            if terminating && query_engine.can_disconnect() {
+            // Check if we should be shutting down.
+            let offline = self.comms.offline();
+            // Check that there are no active transactions.
+            let query_engine_done = query_engine.can_disconnect();
+
+            // If query engine is idle and we requested shutdown, we're done.
+            if query_engine_done && offline {
+                // Send shutdown notification to client.
+                self.stream
+                    .send_flush(&ErrorResponse::shutting_down())
+                    .await?;
                 break;
             }
 
             let client_state = query_engine.client_state();
 
             select! {
                 _ = shutdown.notified() => {
-                    if query_engine.can_disconnect() {
-                        continue; // Wake up task.
-                    }
+                    continue; // Wake up task.
                 }
 
                 // Async messages.
@@ -410,7 +407,7 @@ impl Client {
                     self.server_message(&mut query_engine, message).await?;
                 }
 
-                buffer = self.buffer(client_state), if !terminating => {
+                buffer = self.buffer(client_state) => {
                     let event = buffer?;
 
                     // Only send requests to the backend if they are complete.
@@ -420,25 +417,14 @@ impl Client {
                         }
 
                     match event {
-                        BufferEvent::DisconnectAbrupt => break,
-                        BufferEvent::DisconnectGraceful => {
-                            if query_engine.can_disconnect() {
-                                break;
-                            }
-                            terminating = true;
-                        }
+                        // Client disconnected, we're done.
+                        BufferEvent::DisconnectAbrupt | BufferEvent::DisconnectGraceful => break,
                         BufferEvent::HaveRequest => (),
                     }
                 }
             }
         }
 
-        if offline && !self.shutdown {
-            self.stream
-                .send_flush(&ErrorResponse::shutting_down())
-                .await?;
-        }
-
         Ok(())
     }
 
@@ -547,7 +533,6 @@ impl Client {
 
             // Terminate (B & F).
             if message.code() == 'X' {
-                self.shutdown = true;
                 return Ok(BufferEvent::DisconnectGraceful);
             } else {
                 let message = ProtocolMessage::from_bytes(message.to_bytes()?)?;

pgdog/src/frontend/client/query_engine/mod.rs

@@ -43,28 +43,6 @@ use notify_buffer::NotifyBuffer;
 pub use two_pc::phase::TwoPcPhase;
 use two_pc::TwoPc;
 
-#[derive(Debug)]
-pub struct TestMode {
-    pub enabled: bool,
-}
-
-impl Default for TestMode {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
-impl TestMode {
-    pub fn new() -> Self {
-        Self {
-            #[cfg(test)]
-            enabled: true,
-            #[cfg(not(test))]
-            enabled: false,
-        }
-    }
-}
-
 #[derive(Debug)]
 pub struct QueryEngine {
     begin_stmt: Option<BufferedQuery>,
@@ -73,7 +51,6 @@ pub struct QueryEngine {
     stats: Stats,
     backend: Connection,
     streaming: bool,
-    test_mode: TestMode,
     two_pc: TwoPc,
     notify_buffer: NotifyBuffer,
     pending_explain: Option<ExplainResponseState>,
@@ -92,7 +69,6 @@ impl QueryEngine {
             backend,
             comms: comms.clone(),
             hooks: QueryEngineHooks::new(),
-            test_mode: TestMode::new(),
             stats: Stats::default(),
             streaming: bool::default(),
             two_pc: TwoPc::default(),

pgdog/src/frontend/client/query_engine/query.rs

@@ -95,7 +95,6 @@ impl QueryEngine {
                 while self.backend.has_more_messages()
                     && !self.backend.in_copy_mode()
                     && !self.streaming
-                    && !self.test_mode.enabled
                 {
                     let message = self.read_server_message().await?;
                     self.process_server_message(context, message).await?;

pgdog/src/frontend/client/query_engine/test/close_parse.rs

@@ -0,0 +1,41 @@
+use crate::{
+    expect_message,
+    net::{
+        CloseComplete, CommandComplete, DataRow, Parameters, ParseComplete, ReadyForQuery,
+        RowDescription,
+    },
+};
+
+use super::prelude::*;
+
+/// Test closing a non-existent prepared statement, then a query,
+/// then close+parse+flush for the same name.
+#[tokio::test]
+async fn test_close_parse() {
+    let mut client = TestClient::new_replicas(Parameters::default()).await;
+
+    // Close non-existent statement + Sync
+    client.send(Close::named("test")).await;
+    client.send(Sync).await;
+    client.try_process().await.unwrap();
+
+    expect_message!(client.read().await, CloseComplete); // '3'
+    expect_message!(client.read().await, ReadyForQuery); // 'Z'
+
+    // Simple SELECT query
+    client.send_simple(Query::new("SELECT 1")).await;
+
+    expect_message!(client.read().await, RowDescription); // 'T'
+    expect_message!(client.read().await, DataRow); // 'D'
+    expect_message!(client.read().await, CommandComplete); // 'C'
+    expect_message!(client.read().await, ReadyForQuery); // 'Z'
+
+    // Close + Parse same name + Flush
+    client.send(Close::named("test1")).await;
+    client.send(Parse::named("test1", "SELECT $1")).await;
+    client.send(Flush).await;
+    client.try_process().await.unwrap();
+
+    expect_message!(client.read().await, CloseComplete); // '3'
+    expect_message!(client.read().await, ParseComplete); // '1'
+}

pgdog/src/frontend/client/query_engine/test/close_parse_global_cache.rs

@@ -0,0 +1,42 @@
+use crate::{
+    expect_message,
+    net::{CloseComplete, Parameters, ParseComplete},
+};
+
+use super::prelude::*;
+
+/// Test that Close + Parse for the same name correctly updates
+/// the global prepared statement cache.
+#[tokio::test]
+async fn test_close_parse_same_name_global_cache() {
+    let mut client = TestClient::new_replicas(Parameters::default()).await;
+
+    // Close + Parse same name + Flush
+    client.send(Close::named("test_stmt")).await;
+    client.send(Parse::named("test_stmt", "SELECT $1")).await;
+    client.send(Flush).await;
+    client.try_process().await.unwrap();
+
+    expect_message!(client.read().await, CloseComplete); // '3'
+    expect_message!(client.read().await, ParseComplete); // '1'
+
+    // Verify the statement is registered correctly in the global cache
+    let global_cache = client.client().prepared_statements.global.clone();
+    assert_eq!(global_cache.read().len(), 1);
+    let binding = global_cache.write();
+    let (_, cached_stmt) = binding.statements().iter().next().unwrap();
+    assert_eq!(cached_stmt.used, 1);
+
+    // Verify the SQL content in the global cache
+    let global_stmt_name = cached_stmt.name();
+    let cached_query = binding.query(&global_stmt_name).unwrap();
+    assert_eq!(cached_query, "SELECT $1");
+
+    // Verify the client's local cache
+    assert_eq!(client.client().prepared_statements.len_local(), 1);
+    assert!(client
+        .client()
+        .prepared_statements
+        .name("test_stmt")
+        .is_some());
+}

pgdog/src/frontend/client/query_engine/test/extended.rs

@@ -0,0 +1,43 @@
+use crate::{
+    expect_message,
+    net::{
+        BindComplete, CommandComplete, DataRow, ParameterDescription, Parameters, ParseComplete,
+        ReadyForQuery, RowDescription,
+    },
+};
+
+use super::prelude::*;
+
+/// Test extended protocol with named prepared statement, parameter binding, and Describe.
+#[tokio::test]
+async fn test_extended_with_params_and_describe() {
+    let mut client = TestClient::new_replicas(Parameters::default()).await;
+
+    client.send(Parse::named("test", "SELECT $1")).await;
+    client
+        .send(Bind::new_params(
+            "test",
+            &[Parameter {
+                len: 3,
+                data: "123".into(),
+            }],
+        ))
+        .await;
+    client.send(Describe::new_statement("test")).await;
+    client.send(Execute::new()).await;
+    client.send(Sync).await;
+    client.try_process().await.unwrap();
+
+    expect_message!(client.read().await, ParseComplete); // '1'
+    expect_message!(client.read().await, BindComplete); // '2'
+    expect_message!(client.read().await, ParameterDescription); // 't'
+    expect_message!(client.read().await, RowDescription); // 'T'
+    let row = expect_message!(client.read().await, DataRow); // 'D'
+    assert_eq!(
+        row.get::<String>(0, crate::net::Format::Text),
+        Some("123".into())
+    );
+    expect_message!(client.read().await, CommandComplete); // 'C'
+    let rfq = expect_message!(client.read().await, ReadyForQuery); // 'Z'
+    assert_eq!(rfq.status, 'I');
+}

pgdog/src/frontend/client/query_engine/test/graceful_disconnect.rs

@@ -0,0 +1,32 @@
+use std::time::Duration;
+
+use tokio::{io::AsyncReadExt, time::timeout};
+
+use super::prelude::*;
+
+/// Regression test: client sends BEGIN, gets response, then sends Terminate.
+/// The backend is idle-in-transaction, so it won't send any messages.
+/// The client event loop must exit promptly instead of deadlocking.
+#[tokio::test]
+async fn test_graceful_disconnect_idle_in_transaction() {
+    let mut client = SpawnedClient::new_default(Parameters::default()).await;
+
+    // Start a transaction.
+    client.send(Query::new("BEGIN")).await;
+    client.read_until('Z').await;
+
+    // Send Terminate while idle-in-transaction.
+    client.send(Terminate).await;
+
+    // The client task should exit and close the connection.
+    // With the deadlock bug, this times out because the event loop hangs.
+    let mut buf = [0u8; 1];
+    let result = timeout(Duration::from_secs(2), client.conn.read(&mut buf)).await;
+    assert!(
+        result.is_ok(),
+        "client.run() deadlocked after Terminate during idle-in-transaction"
+    );
+
+    // EOF means the connection was closed cleanly.
+    assert_eq!(result.unwrap().unwrap(), 0);
+}