diff --git a/Cargo.lock b/Cargo.lock index 7916865497..efe0edd5ac 100644 --- a/Cargo.lock +++ b/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" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 0b04e60180..d96a5e5820 100644 --- a/rust/Cargo.toml +++ b/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" diff --git a/rust/src/database/postgres/pool.rs b/rust/src/database/postgres/pool.rs index c1f29699d5..3ea3569722 100644 --- a/rust/src/database/postgres/pool.rs +++ b/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; pub type PooledConnection = Object; /// 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 { 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, ) -> Result { 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, + checkout_timeout_ms: Option, sslmode: Option, sslrootcert: Option, sslcert: Option, @@ -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) -> 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" + ); + }); + } } diff --git a/synapse/storage/database.py b/synapse/storage/database.py index 464fde9ccc..a01b678768 100644 --- a/synapse/storage/database.py +++ b/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: diff --git a/synapse/storage/engines/postgres_rust.py b/synapse/storage/engines/postgres_rust.py index dc67651732..a3b6e70954 100644 --- a/synapse/storage/engines/postgres_rust.py +++ b/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 diff --git a/synapse/storage/rust_dbapi.py b/synapse/storage/rust_dbapi.py index 585d8099a2..90b24d7caf 100644 --- a/synapse/storage/rust_dbapi.py +++ b/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) diff --git a/synapse/storage/rust_pool.py b/synapse/storage/rust_pool.py index 2ff98546fa..1ed636e380 100644 --- a/synapse/storage/rust_pool.py +++ b/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, ) diff --git a/synapse/synapse_rust/database/postgres.pyi b/synapse/synapse_rust/database/postgres.pyi index d9e56e52e5..4523e9ab43 100644 --- a/synapse/synapse_rust/database/postgres.pyi +++ b/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, diff --git a/tests/storage/test_rust_pool.py b/tests/storage/test_rust_pool.py index 9c17936eb3..b2003be40b 100644 --- a/tests/storage/test_rust_pool.py +++ b/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)