Support reconnect on the Rust-backed connection

`DatabasePool`'s `inner_func` calls `conn.reconnect()` when a connection is
found closed, or to recycle one that has exceeded the per-connection
transaction limit (`txn_limit`). The Rust DBAPI2 adapter had no `reconnect`, so
those paths (off by default, but real) would have raised.

Add `reconnect` to the adapter: it returns the current connection to the pool
(or discards it if unusable) and checks out a fresh one. The adapter now holds
the pool it was checked out of, with an `owns_pool` flag distinguishing a shared
pool (RustConnectionPool — reconnect from it, don't close it) from a bootstrap
pool-of-one (`rust_dbapi.connect` — closed together with the connection).

Tested: reconnect swaps in a working connection; runWithConnection and the
bootstrap/adapter paths are unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01W3G4M92AmwSSZCbmtMJU3d

rust/src/database/postgres/connection.rs

@@ -365,6 +365,18 @@ impl Connection {
         Ok(())
     }
 
+    /// Discard the connection instead of returning it to the pool: the
+    /// underlying socket is closed and the pool shrinks (growing back on
+    /// demand). For forced recycling — Synapse's per-connection transaction
+    /// limit calls `reconnect()` expecting a *fresh server session*, which
+    /// `close()` can't deliver (it hands the same live session back to the
+    /// pool, where the next checkout just picks it up again). Idempotent,
+    /// like `close`.
+    fn discard(&self) -> PyResult<()> {
+        self.lock()?.discard();
+        Ok(())
+    }
+
     /// Whether this connection has been closed (or discarded).
     ///
     /// Mirrors psycopg2's `connection.closed`: it is true once the connection

synapse/storage/rust_dbapi.py

@@ -183,7 +183,7 @@ def connect(
         synchronous_commit=synchronous_commit,
         statement_timeout_ms=statement_timeout_ms,
     )
-    return Connection(pool.connect(), pool=pool)
+    return Connection(pool.connect(), pool=pool, owns_pool=True)
 
 
 class Cursor:
@@ -274,12 +274,14 @@ class Connection:
     engine-facing methods delegate straight to the shim.
     """
 
-    def __init__(self, conn: Any, pool: Any = None) -> None:
+    def __init__(self, conn: Any, pool: Any = None, owns_pool: bool = False) -> None:
         self._conn = conn
-        # A pool this connection owns (a bootstrap "pool of one"), kept alive for
-        # the connection's lifetime and closed with it. `None` for connections
-        # borrowed from a shared pool.
+        # The pool this connection was checked out of, if any. Used to
+        # `reconnect` (check out a fresh connection). When `owns_pool` is set the
+        # pool belongs solely to this connection (a bootstrap "pool of one") and
+        # is closed together with it.
         self._pool = pool
+        self._owns_pool = owns_pool
 
     def cursor(self) -> Cursor:
         return Cursor(self._conn.cursor())
@@ -290,9 +292,25 @@ class Connection:
     def rollback(self) -> None:
         self._conn.rollback()
 
+    def reconnect(self) -> None:
+        """Replace the underlying connection with a fresh one from the pool.
+
+        Mirrors ``adbapi.Connection.reconnect`` — which closes the DBAPI
+        connection and opens a brand-new one: the transaction driver calls it
+        when a connection is found closed, or to recycle one that has hit the
+        per-connection transaction limit (``txn_limit``, which exists to bound
+        per-session server-side state). The current connection is therefore
+        *discarded*, not returned to the pool — returning it would just hand
+        the same live session back out on the next checkout.
+        """
+        if self._pool is None:
+            raise RuntimeError("cannot reconnect a connection with no pool")
+        self._conn.discard()
+        self._conn = self._pool.connect()
+
     def close(self) -> None:
         self._conn.close()
-        if self._pool is not None:
+        if self._owns_pool and self._pool is not None:
             self._pool.close()
 
     # -- engine-facing methods (see RustPostgresEngine) ---------------------

synapse/storage/rust_pool.py

@@ -24,9 +24,8 @@ It presents the slice of ``twisted.enterprise.adbapi.ConnectionPool`` that
 ``DatabasePool`` uses — ``runWithConnection``, ``threadID``, ``threadpool`` and
 ``running`` — so it can stand in for ``_db_pool`` (paired with
 :class:`~synapse.storage.engines.RustPostgresEngine`, which drives the wrapped
-connection). What remains before ``make_pool`` can return it: ``reconnect`` on
-the connection (only used on the transaction-limit / closed-connection paths)
-and the startup path (``make_conn`` / ``check_database``).
+connection). ``make_pool`` returns one of these when the database is configured
+with ``use_rust_driver``.
 """
 
 import logging
@@ -188,7 +187,10 @@ class RustConnectionPool:
         A checkout failure surfaces as the raised exception (→ errback) before
         there is any connection to release.
         """
-        conn = DBAPI2Connection(self._pool.connect())
+        # Pass the pool so `func` can `reconnect` (checking out a fresh
+        # connection); `owns_pool=False` since the pool is shared and outlives
+        # this checkout.
+        conn = DBAPI2Connection(self._pool.connect(), pool=self._pool)
         try:
             result = func(conn, *args, **kwargs)
             conn.commit()

synapse/synapse_rust/database/postgres.pyi

@@ -47,6 +47,11 @@ class Connection:
     def commit(self) -> None: ...
     def rollback(self) -> None: ...
     def close(self) -> None: ...
+    def discard(self) -> None:
+        """Close the underlying session instead of returning it to the pool
+        (which shrinks, growing back on demand). For forced recycling, where a
+        fresh server session is wanted."""
+
     def is_closed(self) -> bool: ...
     def in_transaction(self) -> bool: ...
     def set_autocommit(self, autocommit: bool) -> None: ...

tests/storage/test_rust_dbapi.py

@@ -178,7 +178,10 @@ class RustDBAPIAdapterTestCase(unittest.TestCase):
 
     def setUp(self) -> None:
         self._pool = postgres.ConnectionPool(_build_dsn())
-        self.conn = rust_dbapi.Connection(self._pool.connect())
+        # Pass the pool (owns_pool=False) so `reconnect` can check out a fresh
+        # connection; tearDown closes the pool itself.
+        self.conn = rust_dbapi.Connection(self._pool.connect(), pool=self._pool)
+        self.engine = RustPostgresEngine({})
 
     def tearDown(self) -> None:
         del self.conn
@@ -279,3 +282,33 @@ class RustDBAPIAdapterTestCase(unittest.TestCase):
         self.assertTrue(self.conn.autocommit)
         self.conn.set_autocommit(False)
         self.assertFalse(self.conn.autocommit)
+
+    def test_reconnect(self) -> None:
+        # `reconnect` swaps in a fresh pooled connection; the connection is still
+        # usable afterwards.
+        self.conn.reconnect()
+        cursor = self.conn.cursor()
+        cursor.execute("SELECT 1")
+        self.assertEqual(cursor.fetchone(), (1,))
+        self.conn.commit()
+
+    def test_reconnect_gets_a_fresh_server_session(self) -> None:
+        # `reconnect` exists to recycle a connection (txn_limit bounds
+        # per-session server state), so it must discard the old session rather
+        # than return it to the pool — where, like adbapi's close-and-reopen,
+        # the next checkout would just get the same session back.
+        def backend_pid(conn: rust_dbapi.Connection) -> int:
+            cursor = conn.cursor()
+            cursor.execute("SELECT pg_backend_pid()")
+            (pid,) = cursor.fetchone()  # type: ignore[misc]
+            conn.commit()
+            return pid
+
+        pool = postgres.ConnectionPool(_build_dsn(), 1)
+        self.addCleanup(pool.close)
+        conn = rust_dbapi.Connection(pool.connect(), pool=pool)
+        self.addCleanup(conn.close)
+
+        before = backend_pid(conn)
+        conn.reconnect()
+        self.assertNotEqual(backend_pid(conn), before)