From 504caa3184258c7cdeda44ec8737d94ab6238ebb Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Thu, 2 Jul 2026 20:37:31 +0000 Subject: [PATCH] Support reconnect on the Rust-backed connection MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `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) Claude-Session: https://claude.ai/code/session_01W3G4M92AmwSSZCbmtMJU3d --- rust/src/database/postgres/connection.rs | 12 ++++++++ synapse/storage/rust_dbapi.py | 30 +++++++++++++++---- synapse/storage/rust_pool.py | 10 ++++--- synapse/synapse_rust/database/postgres.pyi | 5 ++++ tests/storage/test_rust_dbapi.py | 35 +++++++++++++++++++++- 5 files changed, 81 insertions(+), 11 deletions(-) diff --git a/rust/src/database/postgres/connection.rs b/rust/src/database/postgres/connection.rs index 07eddb6f44..1d039325e0 100644 --- a/rust/src/database/postgres/connection.rs +++ b/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 diff --git a/synapse/storage/rust_dbapi.py b/synapse/storage/rust_dbapi.py index 2766c41863..94c6038099 100644 --- a/synapse/storage/rust_dbapi.py +++ b/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) --------------------- diff --git a/synapse/storage/rust_pool.py b/synapse/storage/rust_pool.py index 352933c201..395edb42c4 100644 --- a/synapse/storage/rust_pool.py +++ b/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() diff --git a/synapse/synapse_rust/database/postgres.pyi b/synapse/synapse_rust/database/postgres.pyi index 24303b3416..7dab0d87aa 100644 --- a/synapse/synapse_rust/database/postgres.pyi +++ b/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: ... diff --git a/tests/storage/test_rust_dbapi.py b/tests/storage/test_rust_dbapi.py index ccf564fe1b..964a206062 100644 --- a/tests/storage/test_rust_dbapi.py +++ b/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)