From 67a87a9966372477a78080927cfb9ff76bca97cc Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Fri, 3 Jul 2026 11:19:56 +0000 Subject: [PATCH] Route Rust pool checkouts through a swappable connectionFactory `adbapi.ConnectionPool` obtains connections via a `connectionFactory` attribute, which the database-outage tests swap out to make every checkout fail (and restore afterwards). `RustConnectionPool` checked out connections inline, so it had no such hook and the outage tests errored with `AttributeError`. Give it a `connectionFactory` attribute (defaulting to a checkout from the native pool, wrapped in the DBAPI2 adapter) and route both `runWithConnection` and `connect` through it, so a test's replacement takes effect. Also make the native pool reopenable: `close` drops it and `start` opens a fresh one, so the tests' `close()`/`start()` outage cycle recovers rather than leaving the pool permanently shut. Fixes the tests.storage.databases.main.test_events_worker DatabaseOutage failures on the Rust backend (12/12 pass); psycopg2 and sqlite are unaffected. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01W3G4M92AmwSSZCbmtMJU3d --- synapse/storage/rust_pool.py | 59 +++++++++++++++++++++++++++--------- 1 file changed, 45 insertions(+), 14 deletions(-) diff --git a/synapse/storage/rust_pool.py b/synapse/storage/rust_pool.py index 0860c1b0eb..ed99505d07 100644 --- a/synapse/storage/rust_pool.py +++ b/synapse/storage/rust_pool.py @@ -94,22 +94,49 @@ class RustConnectionPool: hook itself, so it needs no clock and stays trivially testable. """ self._reactor = reactor - self._pool = postgres.ConnectionPool( - dsn, - threads, - synchronous_commit=synchronous_commit, - statement_timeout_ms=statement_timeout_ms, - ) + self._dsn = dsn + self._synchronous_commit = synchronous_commit + self._statement_timeout_ms = statement_timeout_ms + self._threads = threads + self._pool: Any = self._open_pool() self.threadpool = ThreadPool(minthreads=1, maxthreads=threads, name=name) self.running = False # Connections handed out by `connect()` for direct synchronous use, one # cached per thread (adbapi's model). Closed together with the pool. self._connections: dict[int, DBAPI2Connection] = {} + # How a connection is obtained from the pool (adbapi interface). Held as + # a swappable attribute so tests can replace it to simulate a database + # outage; called as `connectionFactory(self)`. + self.connectionFactory = self._default_connection_factory + + def _open_pool(self) -> Any: + """Open a fresh native Rust connection pool from the stored config.""" + return postgres.ConnectionPool( + self._dsn, + self._threads, + synchronous_commit=self._synchronous_commit, + statement_timeout_ms=self._statement_timeout_ms, + ) + + def _default_connection_factory( + self, _pool: "RustConnectionPool" + ) -> DBAPI2Connection: + # Check a connection out of the native pool and wrap it in the DBAPI2 + # adapter. `owns_pool=False`: the pool is shared and outlives the + # checkout. Mirrors `adbapi.ConnectionPool.connectionFactory`. + return DBAPI2Connection(self._pool.connect(), pool=self._pool) def start(self) -> None: - """Start the thread pool. Idempotent.""" + """Start the thread pool. Idempotent. + + Reopens the native connection pool if a previous :meth:`close` shut it + down, so a closed pool can be brought back up (as the database-outage + tests do with a `close()`/`start()` cycle). + """ if self.running: return + if self._pool is None: + self._pool = self._open_pool() self.threadpool.start() self.running = True @@ -119,7 +146,7 @@ class RustConnectionPool: Stops the thread pool first (waiting for in-flight work to finish, which returns its connection to the pool), then closes the Rust pool so its server connections are dropped promptly rather than lingering until - garbage collection. + garbage collection. The pool can be reopened by a later :meth:`start`. """ if not self.running: return @@ -128,7 +155,9 @@ class RustConnectionPool: for conn in self._connections.values(): conn.close() self._connections.clear() - self._pool.close() + if self._pool is not None: + self._pool.close() + self._pool = None def runWithConnection( # noqa: N802 (implements adbapi's interface) self, @@ -209,7 +238,7 @@ class RustConnectionPool: tid = self.threadID() conn = self._connections.get(tid) if conn is None or conn.is_closed(): - conn = DBAPI2Connection(self._pool.connect(), pool=self._pool) + conn = self.connectionFactory(self) self._connections[tid] = conn return conn @@ -265,10 +294,12 @@ class RustConnectionPool: A checkout failure surfaces as the raised exception (→ errback) before there is any connection to release. """ - # 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) + # Check a connection out via `connectionFactory` (see + # `_default_connection_factory`); the wrapper keeps the pool so `func` + # can `reconnect`. Routing through the factory lets tests swap it to + # simulate a checkout failure. A checkout failure surfaces as the raised + # exception (→ errback) before there is any connection to release. + conn = self.connectionFactory(self) try: result = func(conn, *args, **kwargs) conn.commit()