feat: client_idle_in_transaction_timeout (#696) (#697)

levkk · web-flow · commit 3fea92446c61 · 2025-12-30T14:23:51.000-08:00
fix #696
diff --git a/integration/rust/src/setup.rs b/integration/rust/src/setup.rs
@@ -44,6 +44,14 @@ pub async fn connections_sqlx() -> Vec<Pool<Postgres>> {
     pools
 }
 
+pub async fn connection_sqlx_direct() -> Pool<Postgres> {
+    PgPoolOptions::new()
+        .max_connections(1)
+        .connect("postgres://pgdog:pgdog@127.0.0.1:5432/pgdog?application_name=sqlx_direct")
+        .await
+        .unwrap()
+}
+
 #[derive(Debug, PartialEq, Clone)]
 pub struct Backend {
     pub pid: i32,
diff --git a/integration/rust/tests/integration/idle_in_transaction.rs b/integration/rust/tests/integration/idle_in_transaction.rs
@@ -0,0 +1,48 @@
+use std::time::Duration;
+
+use rust::setup::{admin_sqlx, connection_sqlx_direct, connections_sqlx};
+use sqlx::{Executor, Row};
+use tokio::time::sleep;
+
+#[tokio::test]
+async fn test_idle_in_transaction_timeout() {
+    let admin = admin_sqlx().await;
+    admin
+        .execute("SET client_idle_in_transaction_timeout TO 500")
+        .await
+        .unwrap();
+
+    let conn_direct = connection_sqlx_direct().await;
+
+    for conn in connections_sqlx().await {
+        let mut conn = conn.acquire().await.unwrap();
+
+        conn.execute("BEGIN").await.unwrap();
+        let pid_before = conn
+            .fetch_one("SELECT pg_backend_pid()")
+            .await
+            .unwrap()
+            .get::<i32, _>(0);
+        sleep(Duration::from_millis(750)).await;
+        let err = conn.execute("SELECT 1").await.unwrap_err();
+        assert!(err.to_string().contains("idle in transaction"));
+
+        sleep(Duration::from_millis(500)).await;
+
+        let (pid_after, query): (i32, String) =
+            sqlx::query_as("SELECT pid, query FROM pg_stat_activity WHERE pid = $1")
+                .bind(pid_before)
+                .fetch_one(&conn_direct)
+                .await
+                .unwrap();
+
+        assert_eq!(
+            pid_before, pid_after,
+            "expexted pooler not to cycle connection"
+        );
+        assert_eq!(query, "ROLLBACK", "expected a rollback on the connection",);
+    }
+
+    // Reset settings.
+    admin.execute("RELOAD").await.unwrap();
+}
diff --git a/integration/rust/tests/integration/mod.rs b/integration/rust/tests/integration/mod.rs
@@ -6,6 +6,7 @@ pub mod cross_shard_disabled;
 pub mod distinct;
 pub mod explain;
 pub mod fake_transactions;
+pub mod idle_in_transaction;
 pub mod maintenance_mode;
 pub mod notify;
 pub mod per_stmt_routing;
diff --git a/pgdog-config/src/general.rs b/pgdog-config/src/general.rs
@@ -135,6 +135,9 @@ pub struct General {
     /// Client idle timeout.
     #[serde(default = "General::default_client_idle_timeout")]
     pub client_idle_timeout: u64,
+    /// Client idle in transaction timeout.
+    #[serde(default = "General::default_client_idle_in_transaction_timeout")]
+    pub client_idle_in_transaction_timeout: u64,
     /// Server lifetime.
     #[serde(default = "General::server_lifetime")]
     pub server_lifetime: u64,
@@ -236,6 +239,7 @@ impl Default for General {
             dry_run: Self::dry_run(),
             idle_timeout: Self::idle_timeout(),
             client_idle_timeout: Self::default_client_idle_timeout(),
+            client_idle_in_transaction_timeout: Self::default_client_idle_in_transaction_timeout(),
             mirror_queue: Self::mirror_queue(),
             mirror_exposure: Self::mirror_exposure(),
             auth_type: Self::auth_type(),
@@ -364,6 +368,13 @@ impl General {
         )
     }
 
+    fn default_client_idle_in_transaction_timeout() -> u64 {
+        Self::env_or_default(
+            "PGDOG_CLIENT_IDLE_IN_TRANSACTION_TIMEOUT",
+            Duration::MAX.as_millis() as u64,
+        )
+    }
+
     fn default_query_timeout() -> u64 {
         Self::env_or_default("PGDOG_QUERY_TIMEOUT", Duration::MAX.as_millis() as u64)
     }
@@ -384,6 +395,10 @@ impl General {
         Duration::from_millis(self.connect_attempt_delay)
     }
 
+    pub fn client_idle_in_transaction_timeout(&self) -> Duration {
+        Duration::from_millis(self.client_idle_in_transaction_timeout)
+    }
+
     fn load_balancing_strategy() -> LoadBalancingStrategy {
         Self::env_enum_or_default("PGDOG_LOAD_BALANCING_STRATEGY")
     }
diff --git a/pgdog/src/admin/set.rs b/pgdog/src/admin/set.rs
@@ -149,6 +149,10 @@ impl Command for Set {
                 config.config.general.query_parser = Self::from_json(&self.value)?;
             }
 
+            "client_idle_in_transaction_timeout" => {
+                config.config.general.client_idle_in_transaction_timeout = self.value.parse()?;
+            }
+
             _ => return Err(Error::Syntax),
         }
 
diff --git a/pgdog/src/frontend/client/mod.rs b/pgdog/src/frontend/client/mod.rs
@@ -521,7 +521,7 @@ impl Client {
                 match timeout(idle_timeout, self.stream_buffer.read(&mut self.stream)).await {
                     Err(_) => {
                         self.stream
-                            .fatal(ErrorResponse::client_idle_timeout(idle_timeout))
+                            .fatal(ErrorResponse::client_idle_timeout(idle_timeout, &state))
                             .await?;
                         return Ok(BufferEvent::DisconnectAbrupt);
                     }
diff --git a/pgdog/src/frontend/client/timeouts.rs b/pgdog/src/frontend/client/timeouts.rs
@@ -6,13 +6,15 @@ use crate::{config::General, frontend::ClientRequest, state::State};
 pub struct Timeouts {
     pub(super) query_timeout: Duration,
     pub(super) client_idle_timeout: Duration,
+    pub(super) idle_in_transaction_timeout: Duration,
 }
 
 impl Default for Timeouts {
     fn default() -> Self {
         Self {
             query_timeout: Duration::MAX,
             client_idle_timeout: Duration::MAX,
+            idle_in_transaction_timeout: Duration::MAX,
         }
     }
 }
@@ -22,6 +24,7 @@ impl Timeouts {
         Self {
             query_timeout: general.query_timeout(),
             client_idle_timeout: general.client_idle_timeout(),
+            idle_in_transaction_timeout: general.client_idle_in_transaction_timeout(),
         }
     }
 
@@ -48,7 +51,40 @@ impl Timeouts {
                     Duration::MAX
                 }
             }
+            State::IdleInTransaction => {
+                // Client is sending the request, don't fire.
+                if !client_request.messages.is_empty() {
+                    Duration::MAX
+                } else {
+                    self.idle_in_transaction_timeout
+                }
+            }
+
             _ => Duration::MAX,
         }
     }
 }
+
+#[cfg(test)]
+mod test {
+
+    use crate::{config::config, net::Query};
+
+    use super::*;
+
+    #[test]
+    fn test_idle_in_transaction_timeout() {
+        let config = config(); // Will be default.
+        let timeout = Timeouts::from_config(&config.config.general);
+
+        let actual = timeout.client_idle_timeout(&State::IdleInTransaction, &ClientRequest::new());
+        assert_eq!(actual, timeout.idle_in_transaction_timeout);
+        assert_eq!(actual.as_millis(), u64::MAX.into());
+
+        let actual = timeout.client_idle_timeout(
+            &State::IdleInTransaction,
+            &ClientRequest::from(vec![Query::new("SELECT 1").into()]),
+        );
+        assert_eq!(actual, Duration::MAX);
+    }
+}
diff --git a/pgdog/src/net/messages/error_response.rs b/pgdog/src/net/messages/error_response.rs
@@ -4,7 +4,10 @@ use std::fmt::Display;
 use std::time::Duration;
 
 use super::prelude::*;
-use crate::net::{c_string_buf, code};
+use crate::{
+    net::{c_string_buf, code},
+    state::State,
+};
 
 /// ErrorResponse (B) message.
 #[derive(Debug, Clone)]
@@ -50,7 +53,7 @@ impl ErrorResponse {
     }
 
     pub fn client_login_timeout(timeout: Duration) -> ErrorResponse {
-        let mut error = Self::client_idle_timeout(timeout);
+        let mut error = Self::client_idle_timeout(timeout, &State::Active);
         error.message = "client login timeout".into();
         error.detail = Some(format!(
             "client_login_timeout of {}ms expired",
@@ -89,13 +92,25 @@ impl ErrorResponse {
         }
     }
 
-    pub fn client_idle_timeout(duration: Duration) -> ErrorResponse {
+    pub fn client_idle_timeout(duration: Duration, state: &State) -> ErrorResponse {
         ErrorResponse {
             severity: "FATAL".into(),
             code: "57P05".into(),
-            message: "disconnecting idle client".into(),
+            message: format!(
+                "disconnecting {} client",
+                if state == &State::IdleInTransaction {
+                    "idle in transaction"
+                } else {
+                    "idle"
+                }
+            ),
             detail: Some(format!(
-                "client_idle_timeout of {}ms expired",
+                "{} of {}ms expired",
+                if state == &State::IdleInTransaction {
+                    "client_idle_in_transaction_timeout"
+                } else {
+                    "client_idle_timeout"
+                },
                 duration.as_millis()
             )),
             context: None,

Original file line number	Diff line number	Diff line change
`@@ -149,6 +149,10 @@ impl Command for Set {`
`149`	`149`	`config.config.general.query_parser = Self::from_json(&self.value)?;`
`150`	`150`	`}`
`151`	`151`
	`152`	`+ "client_idle_in_transaction_timeout" => {`
	`153`	`+ config.config.general.client_idle_in_transaction_timeout = self.value.parse()?;`
	`154`	`+ }`
	`155`	`+`
`152`	`156`	`_ => return Err(Error::Syntax),`
`153`	`157`	`}`
`154`	`158`
Original file line number	Diff line number	Diff line change
`@@ -521,7 +521,7 @@ impl Client {`
`521`	`521`	`match timeout(idle_timeout, self.stream_buffer.read(&mut self.stream)).await {`
`522`	`522`	`Err(_) => {`
`523`	`523`	`self.stream`
`524`		`- .fatal(ErrorResponse::client_idle_timeout(idle_timeout))`
	`524`	`+ .fatal(ErrorResponse::client_idle_timeout(idle_timeout, &state))`
`525`	`525`	`.await?;`
`526`	`526`	`return Ok(BufferEvent::DisconnectAbrupt);`
`527`	`527`	`}`