From d04ec15bfca16142fd2992b181225cf7dae0aa36 Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Mon, 12 Oct 2026 18:43:26 +0000 Subject: [PATCH] Enforce txn_limit in the Rust pool rather than per thread `DatabasePool` enforces `txn_limit` by counting transactions per worker thread and reconnecting the thread's connection past the limit. The Rust pool does not pin connections to threads, so that count recycled whichever connection the thread happened to check out. Give the native pool a `max_uses` instead: deadpool's recycle hook retires a connection once it has served that many checkouts, which is what `txn_limit` means to bound. `DatabasePool` skips its own counting when the pool reports that it handles the limit. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01ULdieJyQoHNzeHq42H2M1p --- rust/src/database/postgres/pool.rs | 73 +++++++++++++++++++++- synapse/storage/database.py | 6 ++ synapse/storage/rust_pool.py | 15 +++++ synapse/synapse_rust/database/postgres.pyi | 5 +- 4 files changed, 95 insertions(+), 4 deletions(-) diff --git a/rust/src/database/postgres/pool.rs b/rust/src/database/postgres/pool.rs index fc2d56cdd2..f13d8cd0c7 100644 --- a/rust/src/database/postgres/pool.rs +++ b/rust/src/database/postgres/pool.rs @@ -109,6 +109,9 @@ pub struct ConnectionManager { session: SessionConfig, /// How to (or whether to) negotiate TLS for each connection. tls: TlsConnector, + /// If set, a connection is closed rather than reused once it has been + /// checked out this many times. See [`PoolConfig::max_uses`]. + max_uses: Option, } impl ConnectionManager { @@ -132,6 +135,7 @@ impl ConnectionManager { handle, session, tls: connector, + max_uses: None, }) } } @@ -246,7 +250,11 @@ impl Manager for ConnectionManager { }) } - async fn recycle(&self, client: &mut PooledClient, _: &Metrics) -> RecycleResult { + async fn recycle( + &self, + client: &mut PooledClient, + metrics: &Metrics, + ) -> RecycleResult { // Cheap liveness check before handing a pooled connection back out: if // the connection task has ended (server closed the socket, fatal error, // …) the client reports closed, so tell deadpool to drop it and create @@ -254,6 +262,15 @@ impl Manager for ConnectionManager { if client.is_closed() { return Err(RecycleError::message("connection closed")); } + // Retire a connection that has served its share of checkouts (the + // `txn_limit` config), bounding the server-side state a session + // accumulates. `recycle_count` is the number of checkouts after the + // first, so this runs when the `max_uses`th checkout would reuse it. + if let Some(max_uses) = self.max_uses { + if metrics.recycle_count + 1 >= max_uses { + return Err(RecycleError::message("connection reached its use limit")); + } + } Ok(()) } } @@ -282,6 +299,12 @@ pub struct PoolConfig { /// phase, so a checkout can never hang forever (e.g. against an /// unreachable server). pub checkout_timeout: Option, + /// If set, a connection is closed and replaced once it has been checked + /// out this many times, rather than reused forever. This is the pool-side + /// form of Synapse's `txn_limit`, which bounds the server-side state + /// (prepared statements, cached plans, memory) a long-lived session + /// accumulates. + pub max_uses: Option, } impl Default for PoolConfig { @@ -291,6 +314,7 @@ impl Default for PoolConfig { session: SessionConfig::default(), tls: TlsParams::default(), checkout_timeout: None, + max_uses: None, } } } @@ -302,7 +326,8 @@ pub fn create_pool( handle: Handle, config: PoolConfig, ) -> Result { - let manager = ConnectionManager::from_dsn(dsn, handle, config.session, &config.tls)?; + let mut manager = ConnectionManager::from_dsn(dsn, handle, config.session, &config.tls)?; + manager.max_uses = config.max_uses; let mut builder = Pool::builder(manager).max_size(config.max_size); if let Some(timeout) = config.checkout_timeout { // deadpool enforces timeouts by sleeping on a runtime (we drive `get()` @@ -341,7 +366,7 @@ impl PyConnectionPool { /// `statement_timeout` and `idle_in_transaction_session_timeout` if /// configured here. #[new] - #[pyo3(signature = (runtime, dsn, max_size = 10, *, synchronous_commit = true, statement_timeout_ms = None, idle_in_transaction_session_timeout_ms = None, checkout_timeout_ms = None, sslmode = None, sslrootcert = None, sslcert = None, sslkey = None, sslpassword = None))] + #[pyo3(signature = (runtime, dsn, max_size = 10, *, synchronous_commit = true, statement_timeout_ms = None, idle_in_transaction_session_timeout_ms = None, checkout_timeout_ms = None, max_uses = None, sslmode = None, sslrootcert = None, sslcert = None, sslkey = None, sslpassword = None))] #[allow(clippy::too_many_arguments)] fn new( runtime: &RustRuntime, @@ -351,6 +376,7 @@ impl PyConnectionPool { statement_timeout_ms: Option, idle_in_transaction_session_timeout_ms: Option, checkout_timeout_ms: Option, + max_uses: Option, sslmode: Option, sslrootcert: Option, sslcert: Option, @@ -377,6 +403,8 @@ impl PyConnectionPool { checkout_timeout: checkout_timeout_ms .filter(|&ms| ms > 0) .map(Duration::from_millis), + // `0` means no limit, like `txn_limit`. + max_uses: max_uses.filter(|&n| n > 0), }; let pool = create_pool(dsn, handle.clone(), config).map_err(|e| { PyRuntimeError::new_err(format!("failed to build connection pool: {e}")) @@ -602,6 +630,45 @@ mod tests { }); } + #[test] + fn connection_is_replaced_after_max_uses() { + let Some(dsn) = test_dsn() else { + eprintln!("skipping: set SYNAPSE_TEST_POSTGRES_DSN to run"); + return; + }; + + let rt = test_runtime(); + let pool = create_pool( + &dsn, + rt.handle().clone(), + PoolConfig { + max_size: 1, + max_uses: Some(2), + ..PoolConfig::default() + }, + ) + .unwrap(); + + async fn pid(pool: &ConnectionPool) -> i32 { + let client = pool.get().await.unwrap(); + client + .query_one("SELECT pg_backend_pid()", &[]) + .await + .unwrap() + .get::<_, i32>(0) + } + + rt.block_on(async { + // The first two checkouts share a server session; the third gets a + // fresh one because the second checkout used the connection up. + let first = pid(&pool).await; + let second = pid(&pool).await; + let third = pid(&pool).await; + assert_eq!(first, second); + assert_ne!(second, third); + }); + } + #[test] fn checkout_times_out_when_pool_is_exhausted() { let Some(dsn) = test_dsn() else { diff --git a/synapse/storage/database.py b/synapse/storage/database.py index d465e8c755..e8435fe9d3 100644 --- a/synapse/storage/database.py +++ b/synapse/storage/database.py @@ -241,6 +241,7 @@ def _make_rust_pool( statement_timeout_ms=engine.statement_timeout, idle_in_transaction_session_timeout_ms=engine.idle_in_transaction_session_timeout, checkout_timeout_ms=engine.pool_checkout_timeout_ms, + txn_limit=db_config.config.get("txn_limit", 0), ssl_params=ssl_params, ) pool.start() @@ -748,6 +749,11 @@ class DatabasePool: server_name=self.server_name, runtime=hs.get_rust_runtime(), ) + if getattr(self._db_pool, "handles_txn_limit", False): + # The pool retires connections itself (see RustConnectionPool), so + # the per-thread transaction counting below would only recycle a + # random connection on top of that. + self._txn_limit = 0 self.updates = BackgroundUpdater(hs, self) diff --git a/synapse/storage/rust_pool.py b/synapse/storage/rust_pool.py index 2a7d74397d..e7d67db7ab 100644 --- a/synapse/storage/rust_pool.py +++ b/synapse/storage/rust_pool.py @@ -84,6 +84,12 @@ rust_pool_checkout_wait_seconds = Histogram( class RustConnectionPool: """Runs blocking database functions against pooled Rust connections. + `txn_limit` is enforced here, by the native pool retiring a connection + after that many checkouts, rather than by `DatabasePool` counting + transactions per thread: connections are not pinned to threads, so a + per-thread count would recycle whichever connection the thread happened to + hold. `DatabasePool` reads `handles_txn_limit` and skips its own counting. + Each :meth:`runWithConnection` call runs its function on a worker thread with a (DBAPI2-adapter) connection checked out of the native Rust pool, and returns a *raw* ``Deferred`` that fires on the reactor thread with the @@ -93,6 +99,9 @@ class RustConnectionPool: ``make_deferred_yieldable``; see :meth:`runWithConnection`). """ + # See the class docstring: `DatabasePool` leaves `txn_limit` to this pool. + handles_txn_limit = True + def __init__( self, reactor: "ISynapseReactor", @@ -106,6 +115,7 @@ class RustConnectionPool: statement_timeout_ms: int | None = None, idle_in_transaction_session_timeout_ms: int | None = None, checkout_timeout_ms: int | None = None, + txn_limit: int = 0, ssl_params: "dict[str, Any] | None" = None, ) -> None: """ @@ -129,6 +139,9 @@ class RustConnectionPool: 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. + txn_limit: close and replace a connection after this many + checkouts (each ``runWithConnection`` is one); ``0`` keeps + connections for as long as they stay healthy. ssl_params: the libpq ``ssl*`` keys (see :func:`synapse.storage.rust_dbapi.split_ssl_params`), passed to each pooled connection's TLS setup. @@ -147,6 +160,7 @@ class RustConnectionPool: idle_in_transaction_session_timeout_ms ) self._checkout_timeout_ms = checkout_timeout_ms + self._txn_limit = txn_limit self._ssl_params = dict(ssl_params or {}) self._threads = threads self._pool: Any = self._open_pool() @@ -192,6 +206,7 @@ class RustConnectionPool: statement_timeout_ms=self._statement_timeout_ms, idle_in_transaction_session_timeout_ms=self._idle_in_transaction_session_timeout_ms, checkout_timeout_ms=self._checkout_timeout_ms, + max_uses=self._txn_limit, **self._ssl_params, ) diff --git a/synapse/synapse_rust/database/postgres.pyi b/synapse/synapse_rust/database/postgres.pyi index 2b93a2eb7c..166b569051 100644 --- a/synapse/synapse_rust/database/postgres.pyi +++ b/synapse/synapse_rust/database/postgres.pyi @@ -32,6 +32,7 @@ class ConnectionPool: statement_timeout_ms: int | None = None, idle_in_transaction_session_timeout_ms: int | None = None, checkout_timeout_ms: int | None = None, + max_uses: int | None = None, sslmode: str | None = None, sslrootcert: str | None = None, sslcert: str | None = None, @@ -48,7 +49,9 @@ class ConnectionPool: ``statement_timeout`` and ``idle_in_transaction_session_timeout`` (both in milliseconds) if configured. ``checkout_timeout_ms`` bounds how long ``connect`` may wait for a - connection; ``None`` or ``0`` waits indefinitely. + connection; ``None`` or ``0`` waits indefinitely. ``max_uses`` closes + and replaces a connection after that many checkouts; ``None`` or ``0`` + keeps connections for as long as they stay healthy. The ``ssl*`` keyword arguments are the libpq TLS settings; see ``rust/src/database/postgres/tls.rs``.