Add a connection-checkout timeout to the native Rust Postgres pool

A checkout from the Rust connection pool (`ConnectionPool.connect`) previously
blocked indefinitely while waiting for a free slot or establishing a new
connection, so an unreachable or overloaded server could hang the worker
threads handling requests forever.

Bound the whole checkout with a configurable timeout: deadpool's `wait` and
`create` timeouts (enabled via its `rt_tokio_1` runtime feature, driven on our
own tokio runtime). Hitting it raises an OperationalError, which Synapse
already treats as a retryable operational error.

Configurable via the `pool_checkout_timeout` database option (milliseconds;
default 30s, `0` disables), read by the Rust engine and plumbed to both the
connection pool and the bootstrap (pool-of-one) connection.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01W3G4M92AmwSSZCbmtMJU3d

rust/Cargo.toml

@@ -69,7 +69,9 @@ tokio-postgres = "0.7"
 # Connection pooling for the native Rust Postgres backend. We use the generic
 # `deadpool::managed` pool with our own manager (rather than `deadpool-postgres`)
 # so it reuses our `connect`/default-host logic and connection-task spawning.
-deadpool = "0.12"
+# `rt_tokio_1` lets the pool enforce checkout timeouts (it sleeps via tokio time),
+# which we drive on our own tokio runtime.
+deadpool = { version = "0.12", features = ["rt_tokio_1"] }
 once_cell = "1.18.0"
 itertools = "0.14.0"
 postgres-protocol = "0.6.10"

rust/src/database/postgres/pool.rs

@@ -17,8 +17,12 @@
 
 use std::collections::HashMap;
 use std::sync::Mutex;
+use std::time::Duration;
 
-use deadpool::managed::{Manager, Metrics, Object, Pool, PoolError, RecycleError, RecycleResult};
+use deadpool::managed::{
+    Manager, Metrics, Object, Pool, PoolError, RecycleError, RecycleResult, TimeoutType, Timeouts,
+};
+use deadpool::Runtime;
 use log::warn;
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
@@ -243,26 +247,45 @@ pub type ConnectionPool = Pool<ConnectionManager>;
 pub type PooledConnection = Object<ConnectionManager>;
 
 /// Build a [`ConnectionPool`] from a libpq-style DSN, capped at `max_size`
-/// connections, using the default [`SessionConfig`].
+/// connections, using the default [`SessionConfig`] and no checkout timeout.
 pub fn create_pool(dsn: &str, max_size: usize) -> Result<ConnectionPool, anyhow::Error> {
     create_pool_with_session(
         dsn,
         max_size,
         SessionConfig::default(),
         &TlsParams::default(),
+        None,
     )
 }
 
 /// Build a [`ConnectionPool`] with explicit per-connection [`SessionConfig`] and
 /// TLS params.
+///
+/// If `checkout_timeout` is set, a checkout ([`PyConnectionPool::connect`] /
+/// [`Pool::get`]) blocks for at most that long before failing with a
+/// [`PoolError::Timeout`] rather than waiting indefinitely. We bound both the
+/// wait-for-a-free-slot phase and the establish-a-new-connection phase, so a
+/// checkout can never hang forever (e.g. against an unreachable server).
 pub fn create_pool_with_session(
     dsn: &str,
     max_size: usize,
     session: SessionConfig,
     tls_params: &TlsParams,
+    checkout_timeout: Option<Duration>,
 ) -> Result<ConnectionPool, anyhow::Error> {
     let manager = ConnectionManager::from_dsn(dsn, session, tls_params)?;
-    Ok(Pool::builder(manager).max_size(max_size).build()?)
+    let mut builder = Pool::builder(manager).max_size(max_size);
+    if let Some(timeout) = checkout_timeout {
+        // deadpool enforces timeouts by sleeping on a runtime (we drive `get()`
+        // on our own tokio runtime, so `Runtime::Tokio1` is correct). `recycle`
+        // is left unbounded: it's a cheap local `is_closed` check, not I/O.
+        builder = builder.runtime(Runtime::Tokio1).timeouts(Timeouts {
+            wait: Some(timeout),
+            create: Some(timeout),
+            recycle: None,
+        });
+    }
+    Ok(builder.build()?)
 }
 
 /// The Python-facing connection pool.
@@ -286,13 +309,14 @@ impl PyConnectionPool {
     /// (REPEATABLE READ) isolation level, plus `synchronous_commit` /
     /// `statement_timeout` if configured here.
     #[new]
-    #[pyo3(signature = (dsn, max_size = 10, *, synchronous_commit = true, statement_timeout_ms = None, sslmode = None, sslrootcert = None, sslcert = None, sslkey = None, sslpassword = None))]
+    #[pyo3(signature = (dsn, max_size = 10, *, synchronous_commit = true, statement_timeout_ms = None, checkout_timeout_ms = None, sslmode = None, sslrootcert = None, sslcert = None, sslkey = None, sslpassword = None))]
     #[allow(clippy::too_many_arguments)]
     fn new(
         dsn: &str,
         max_size: usize,
         synchronous_commit: bool,
         statement_timeout_ms: Option<i32>,
+        checkout_timeout_ms: Option<u64>,
         sslmode: Option<String>,
         sslrootcert: Option<String>,
         sslcert: Option<String>,
@@ -310,9 +334,15 @@ impl PyConnectionPool {
             sslkey,
             sslpassword,
         };
-        let pool = create_pool_with_session(dsn, max_size, session, &tls_params).map_err(|e| {
-            PyRuntimeError::new_err(format!("failed to build connection pool: {e}"))
-        })?;
+        // A `checkout_timeout_ms` of `None` or `0` means no timeout (block until a
+        // connection is available), matching the previous behaviour.
+        let checkout_timeout = checkout_timeout_ms
+            .filter(|&ms| ms > 0)
+            .map(Duration::from_millis);
+        let pool = create_pool_with_session(dsn, max_size, session, &tls_params, checkout_timeout)
+            .map_err(|e| {
+                PyRuntimeError::new_err(format!("failed to build connection pool: {e}"))
+            })?;
         Ok(Self { pool })
     }
 
@@ -344,9 +374,21 @@ fn pool_err_to_py(err: PoolError<tokio_postgres::Error>) -> PyErr {
         // The backend failed to establish the connection: reuse the exact
         // mapping (and `pgcode` tagging) a query error gets.
         PoolError::Backend(e) => pg_err_to_py(&e),
-        // Timed out waiting for a slot, pool closed, no runtime, or a
-        // post-create hook failure: all connection-level problems, which
-        // Synapse treats as retryable operational errors.
+        // Hit the configured checkout timeout. Report which phase so a slow
+        // server (Create) is distinguishable from pool exhaustion (Wait).
+        PoolError::Timeout(t) => {
+            let phase = match t {
+                TimeoutType::Wait => "waiting for a free connection slot",
+                TimeoutType::Create => "establishing a new connection",
+                TimeoutType::Recycle => "recycling a connection",
+            };
+            OperationalError::new_err(format!(
+                "timed out checking out a database connection ({phase})"
+            ))
+        }
+        // Pool closed, no runtime, or a post-create hook failure: all
+        // connection-level problems, which Synapse treats as retryable
+        // operational errors.
         other => OperationalError::new_err(format!("failed to acquire connection: {other}")),
     }
 }
@@ -436,7 +478,7 @@ mod tests {
             synchronous_commit: false,
             statement_timeout_ms: Some(1234),
         };
-        let pool = create_pool_with_session(&dsn, 1, session, &TlsParams::default()).unwrap();
+        let pool = create_pool_with_session(&dsn, 1, session, &TlsParams::default(), None).unwrap();
 
         runtime().block_on(async {
             let client = pool.get().await.unwrap();
@@ -461,4 +503,39 @@ mod tests {
             assert_eq!(show("statement_timeout").await, "1234ms");
         });
     }
+
+    #[test]
+    fn checkout_times_out_when_pool_is_exhausted() {
+        let Some(dsn) = test_dsn() else {
+            eprintln!("skipping: set SYNAPSE_TEST_POSTGRES_DSN to run");
+            return;
+        };
+
+        // A single-connection pool with a short checkout timeout.
+        let pool = create_pool_with_session(
+            &dsn,
+            1,
+            SessionConfig::default(),
+            &TlsParams::default(),
+            Some(Duration::from_millis(200)),
+        )
+        .unwrap();
+
+        runtime().block_on(async {
+            // Hold the only connection, so a second checkout has to wait.
+            let _held = pool.get().await.unwrap();
+
+            let started = std::time::Instant::now();
+            let result = pool.get().await;
+            assert!(
+                matches!(result, Err(PoolError::Timeout(TimeoutType::Wait))),
+                "expected a wait timeout when the pool is exhausted"
+            );
+            // It should fail promptly (around the 200ms budget), not hang.
+            assert!(
+                started.elapsed() < Duration::from_secs(5),
+                "checkout did not time out promptly"
+            );
+        });
+    }
 }