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

Cargo.lock

@@ -242,6 +242,9 @@ name = "deadpool-runtime"
 version = "0.1.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "092966b41edc516079bdf31ec78a2e0588d1d0c08f78b91d8307215928642b2b"
+dependencies = [
+ "tokio",
+]
 
 [[package]]
 name = "der"

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"
+            );
+        });
+    }
 }

synapse/storage/database.py

@@ -224,6 +224,7 @@ def _make_rust_pool(
         threads=threads,
         synchronous_commit=engine.synchronous_commit,
         statement_timeout_ms=engine.statement_timeout,
+        checkout_timeout_ms=engine.pool_checkout_timeout_ms,
         ssl_params=ssl_params,
     )
     pool.start()
@@ -271,6 +272,7 @@ def make_conn(
             rust_dbapi.build_dsn(dsn_args),
             synchronous_commit=engine.synchronous_commit,
             statement_timeout_ms=engine.statement_timeout,
+            checkout_timeout_ms=engine.pool_checkout_timeout_ms,
             ssl_params=ssl_params,
         )
     else:

synapse/storage/engines/postgres_rust.py

@@ -80,6 +80,14 @@ class RustPostgresEngine(PostgresEngine[Connection, Cursor]):
         super().__init__(postgres, database_config)  # type: ignore[arg-type]
         self._version: int | None = None  # set by check_database
 
+        # How long a connection checkout may block before failing with an
+        # operational error, rather than waiting indefinitely for the pool (a
+        # free slot or a new connection). In milliseconds; `0` disables it.
+        # Rust-only: the psycopg2/adbapi path has no equivalent hook.
+        self.pool_checkout_timeout_ms: int = database_config.get(
+            "pool_checkout_timeout", 30000
+        )
+
     def convert_param_style(self, sql: str) -> str:
         # The shim binds positional `$1, $2, ...` placeholders (like libpq),
         # not psycopg2's `%s`. Rewrite `?` left-to-right, matching the Rust-side

synapse/storage/rust_dbapi.py

@@ -204,6 +204,7 @@ def connect(
     *,
     synchronous_commit: bool = True,
     statement_timeout_ms: int | None = None,
+    checkout_timeout_ms: int | None = None,
     ssl_params: Mapping[str, Any] | None = None,
 ) -> "Connection":
     """Open a single standalone connection for bootstrap/one-off use.
@@ -212,13 +213,15 @@ def connect(
     returned :class:`Connection` keeps that pool alive for its lifetime. Used by
     ``make_conn`` for the startup connection that runs schema preparation before
     the real pool exists. ``ssl_params`` are the libpq ``ssl*`` keys (see
-    :func:`split_ssl_params`).
+    :func:`split_ssl_params`); ``checkout_timeout_ms`` bounds how long opening the
+    connection may block (``None``/``0`` waits indefinitely).
     """
     pool = postgres.ConnectionPool(
         dsn,
         1,
         synchronous_commit=synchronous_commit,
         statement_timeout_ms=statement_timeout_ms,
+        checkout_timeout_ms=checkout_timeout_ms,
         **(ssl_params or {}),
     )
     return Connection(pool.connect(), pool=pool, owns_pool=True)

synapse/storage/rust_pool.py

@@ -74,6 +74,7 @@ class RustConnectionPool:
         threads: int = 10,
         synchronous_commit: bool = True,
         statement_timeout_ms: int | None = None,
+        checkout_timeout_ms: int | None = None,
         ssl_params: "dict[str, Any] | None" = None,
     ) -> None:
         """
@@ -88,6 +89,9 @@ class RustConnectionPool:
             synchronous_commit: passed to each pooled connection's session setup.
             statement_timeout_ms: passed to each pooled connection's session
                 setup (statements running longer are aborted).
+            checkout_timeout_ms: how long a connection checkout may block before
+                failing with an operational error rather than waiting for the
+                pool indefinitely. ``None`` or ``0`` disables the timeout.
             ssl_params: the libpq ``ssl*`` keys (see
                 :func:`synapse.storage.rust_dbapi.split_ssl_params`), passed to
                 each pooled connection's TLS setup.
@@ -101,6 +105,7 @@ class RustConnectionPool:
         self._dsn = dsn
         self._synchronous_commit = synchronous_commit
         self._statement_timeout_ms = statement_timeout_ms
+        self._checkout_timeout_ms = checkout_timeout_ms
         self._ssl_params = dict(ssl_params or {})
         self._threads = threads
         self._pool: Any = self._open_pool()
@@ -121,6 +126,7 @@ class RustConnectionPool:
             self._threads,
             synchronous_commit=self._synchronous_commit,
             statement_timeout_ms=self._statement_timeout_ms,
+            checkout_timeout_ms=self._checkout_timeout_ms,
             **self._ssl_params,
         )
 

synapse/synapse_rust/database/postgres.pyi

@@ -28,6 +28,7 @@ class ConnectionPool:
         *,
         synchronous_commit: bool = True,
         statement_timeout_ms: Optional[int] = None,
+        checkout_timeout_ms: Optional[int] = None,
         sslmode: Optional[str] = None,
         sslrootcert: Optional[str] = None,
         sslcert: Optional[str] = None,

tests/storage/test_rust_pool.py

@@ -19,6 +19,7 @@ test reactor deliberately mocks the database thread pool out) and are skipped
 unless the suite is configured to run against Postgres.
 """
 
+import time
 from typing import TYPE_CHECKING, Any, cast
 from unittest.mock import Mock
 
@@ -287,3 +288,26 @@ class RustConnectionPoolTestCase(trial_unittest.TestCase):
         self.pool.close()
         with self.assertRaises(RuntimeError):
             self.pool.runWithConnection(lambda conn: None)
+
+    def test_checkout_timeout_raises_when_pool_exhausted(self) -> None:
+        # With a one-connection pool and a short checkout timeout, holding the
+        # only connection makes a second checkout fail with an operational error
+        # instead of blocking indefinitely (the `checkout_timeout_ms` plumbing).
+        from synapse.synapse_rust.database import (
+            postgres,
+        )
+
+        pool = postgres.ConnectionPool(_build_dsn(), 1, checkout_timeout_ms=200)
+        self.addCleanup(pool.close)
+
+        held = pool.connect()  # take the only connection and keep it checked out
+        self.addCleanup(held.close)
+
+        started = time.monotonic()
+        with self.assertRaises(postgres.OperationalError) as ctx:
+            pool.connect()
+        elapsed = time.monotonic() - started
+
+        self.assertIn("timed out", str(ctx.exception).lower())
+        # It should fail promptly (around the 200ms budget), not hang.
+        self.assertLess(elapsed, 5.0)