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``.