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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01ULdieJyQoHNzeHq42H2M1p

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<usize>,
 }
 
 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<Self::Error> {
+    async fn recycle(
+        &self,
+        client: &mut PooledClient,
+        metrics: &Metrics,
+    ) -> RecycleResult<Self::Error> {
         // 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<Duration>,
+    /// 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<usize>,
 }
 
 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<ConnectionPool, anyhow::Error> {
-    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<i32>,
         idle_in_transaction_session_timeout_ms: Option<i32>,
         checkout_timeout_ms: Option<u64>,
+        max_uses: Option<usize>,
         sslmode: Option<String>,
         sslrootcert: Option<String>,
         sslcert: Option<String>,
@@ -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 {

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)
 

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,
         )
 

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