From a7f41c80fd6dd840d828eb7fcd41b760ff690b36 Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Thu, 2 Jul 2026 16:58:51 +0000 Subject: [PATCH] Present the _db_pool interface on RustConnectionPool MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Make RustConnectionPool a drop-in for DatabasePool._db_pool so real runInteraction traffic can flow through it (once make_pool is switched): - it now hands each function a DBAPI2-adapter connection (rust_dbapi.Connection) wrapping the pooled shim, so LoggingDatabaseConnection.cursor() yields a working LoggingTransaction and the engine's in_transaction / is_closed / set_autocommit operate on it; - the entry point is named `runWithConnection` (matching twisted.enterprise.adbapi.ConnectionPool, which database.py calls by that name), alongside `threadID` for the transaction-limit path. With `threadpool` and `running` already present, the pool covers the slice of the adbapi interface DatabasePool uses. A new test drives a full transaction through the pool the way DatabasePool.runWithConnection's inner_func does — engine.in_transaction check, LoggingDatabaseConnection + LoggingTransaction, `?`→`$n` conversion, commit — off the reactor thread, and gets the result back via the Deferred. Still outstanding before make_pool can return it: `reconnect` on the connection (transaction-limit / closed-connection paths) and the startup path (make_conn / check_database). Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01W3G4M92AmwSSZCbmtMJU3d --- synapse/storage/rust_pool.py | 95 +++++++++++++++++++------ tests/storage/test_rust_pool.py | 119 +++++++++++++++++++++++++++++--- 2 files changed, 181 insertions(+), 33 deletions(-) diff --git a/synapse/storage/rust_pool.py b/synapse/storage/rust_pool.py index 4bec9db322..afcce1d581 100644 --- a/synapse/storage/rust_pool.py +++ b/synapse/storage/rust_pool.py @@ -15,26 +15,32 @@ Synapse's transaction functions are synchronous and expect a DBAPI2 connection. This adapter lets them run unchanged against the Rust ``Connection`` / ``Cursor`` shim: it owns a dedicated Twisted thread pool and, for each call, checks a -connection out of the native Rust ``ConnectionPool``, runs the caller's function -against it on a worker thread, then returns the connection to the pool and hands -the result back to the reactor as a ``Deferred``. +connection out of the native Rust ``ConnectionPool``, wraps it in the DBAPI2 +adapter (:mod:`synapse.storage.rust_dbapi`), runs the caller's function against +it on a worker thread, then returns the connection to the pool and hands the +result back to the reactor as a ``Deferred``. -This is deliberately a thin *execution bridge*. Slotting it into -``DatabasePool`` (so ``runInteraction`` flows through it) additionally needs -engine-level support for the shim connection — ``in_transaction``, -``is_connection_closed``, autocommit / isolation and ``reconnect`` — and is left -to a follow-up. +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``). """ import logging +import threading from types import TracebackType from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar from typing_extensions import Concatenate, ParamSpec +from twisted.internet import threads from twisted.python.threadpool import ThreadPool from synapse.logging.context import defer_to_threadpool +from synapse.storage.rust_dbapi import Connection as DBAPI2Connection from synapse.synapse_rust.database import postgres if TYPE_CHECKING: @@ -51,11 +57,12 @@ R = TypeVar("R") class RustConnectionPool: """Runs blocking database functions against pooled Rust connections. - Each :meth:`run_with_connection` call runs its function on a worker thread - with a connection checked out of the native Rust pool, and returns a - ``Deferred`` that fires on the reactor thread with the result (or an - errback if it raised). Log contexts are preserved across the hop, following - the same rules as :func:`synapse.logging.context.defer_to_threadpool`. + 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 ``Deferred`` that fires on the reactor thread with the result (or + an errback if it raised). Log contexts are preserved across the hop, + following the same rules as + :func:`synapse.logging.context.defer_to_threadpool`. """ def __init__( @@ -107,7 +114,7 @@ class RustConnectionPool: self.threadpool.stop() self._pool.close() - def run_with_connection( + def runWithConnection( # noqa: N802 (implements adbapi's interface) self, func: Callable[Concatenate[Any, P], R], *args: P.args, @@ -115,22 +122,42 @@ class RustConnectionPool: ) -> "Deferred[R]": """Run ``func(conn, *args, **kwargs)`` on a worker thread. - ``conn`` is a connection checked out of the Rust pool for the duration - of the call. The function is responsible for committing or rolling back - (as Synapse's ``new_transaction`` does); the connection is returned to - the pool afterwards regardless. + ``conn`` is a DBAPI2-adapter connection wrapping one checked out of the + Rust pool for the duration of the call. As with adbapi, any transaction + the function leaves open is committed on success and rolled back on + error; the connection is returned to the pool afterwards regardless. + (One carve-out: ``DatabasePool.runWithConnection`` callers passing + ``isolation_level`` have anything uncommitted rolled back by the + isolation-level reset *before* this commit runs — psycopg2-matching + semantics; such functions must commit their own work.) + + Named to match ``twisted.enterprise.adbapi.ConnectionPool`` so this can + stand in for ``DatabasePool._db_pool``. Returns: - A ``Deferred`` firing with ``func``'s result, following the Synapse - logcontext rules (``yield`` / ``await`` it). + A raw ``Deferred`` — like ``adbapi.ConnectionPool.runWithConnection`` + (and unlike a `make_deferred_yieldable`-wrapped one), because the + caller wraps it: ``DatabasePool.runWithConnection`` supplies the + single ``make_deferred_yieldable``, and the DB function sets up its + own ``LoggingContext``. Wrapping it here as well (double + ``make_deferred_yieldable``, plus a nested logcontext) breaks + logcontext handling when the awaiting request is cancelled. """ if not self.running: raise RuntimeError("connection pool is not running") - return defer_to_threadpool( + return threads.deferToThreadPool( self._reactor, self.threadpool, self._run, func, args, kwargs ) + def threadID(self) -> int: # noqa: N802 (implements adbapi's interface) + """Identify the current worker thread (adbapi interface). + + Used by ``DatabasePool`` only when a per-connection transaction limit is + configured, to count transactions per thread. + """ + return threading.get_ident() + def _run( self, func: Callable[..., R], @@ -139,12 +166,34 @@ class RustConnectionPool: ) -> R: """Worker-thread body: check out a connection, run ``func``, return it. + Matches adbapi's ``_runWithConnection`` transaction contract: commit on + success, roll back on exception. The commit makes callers that rely on + the pool's implicit commit (e.g. one-shot background-update runners + that execute DDL and return) durable, and is a no-op after functions + like ``new_transaction`` that commit themselves. The rollback returns + the connection to the pool *clean* so it is reused, rather than being + discarded as mid-transaction (destroying the TCP/TLS session) every + time a transaction function raises. + A checkout failure surfaces as the raised exception (→ errback) before there is any connection to release. """ - conn = self._pool.connect() + conn = DBAPI2Connection(self._pool.connect()) try: - return func(conn, *args, **kwargs) + result = func(conn, *args, **kwargs) + conn.commit() + return result + except BaseException: + # BaseException, matching adbapi: even for e.g. SystemExit the + # transaction is rolled back so the connection goes back to the + # pool clean (reusable) rather than mid-transaction (discarded). + try: + conn.rollback() + except Exception: + # Re-raise the original error, not the rollback failure; the + # shim has already discarded the connection as unusable. + logger.warning("Rollback failed on connection check-in", exc_info=True) + raise finally: # Return the connection to the pool. The Rust shim returns a clean # connection for reuse and discards one left mid-transaction or diff --git a/tests/storage/test_rust_pool.py b/tests/storage/test_rust_pool.py index f4053b1866..9165511c65 100644 --- a/tests/storage/test_rust_pool.py +++ b/tests/storage/test_rust_pool.py @@ -26,6 +26,8 @@ from twisted.internet import reactor as _reactor from twisted.internet.defer import gatherResults, inlineCallbacks from twisted.trial import unittest as trial_unittest +from synapse.storage.database import LoggingDatabaseConnection +from synapse.storage.engines.postgres_rust import RustPostgresEngine from synapse.storage.rust_pool import RustConnectionPool from tests.unittest import skip_unless @@ -81,11 +83,11 @@ class RustConnectionPoolTestCase(trial_unittest.TestCase): def txn(conn: Any) -> Any: cursor = conn.cursor() cursor.execute("SELECT 42::int") - row = cursor.fetch_one() + row = cursor.fetchone() conn.commit() return row - result = yield self.pool.run_with_connection(txn) + result = yield self.pool.runWithConnection(txn) self.assertEqual(result, (42,)) @inlineCallbacks @@ -93,7 +95,7 @@ class RustConnectionPoolTestCase(trial_unittest.TestCase): def txn(conn: Any, a: int, b: int, c: int = 0) -> int: return a + b + c - result = yield self.pool.run_with_connection(txn, 1, 2, c=3) + result = yield self.pool.runWithConnection(txn, 1, 2, c=3) self.assertEqual(result, 6) @inlineCallbacks @@ -106,7 +108,7 @@ class RustConnectionPoolTestCase(trial_unittest.TestCase): # The failure crosses the thread boundary and surfaces as an errback. failure = yield self.assertFailure( - self.pool.run_with_connection(txn), MarkerError + self.pool.runWithConnection(txn), MarkerError ) self.assertEqual(str(failure), "boom") @@ -118,12 +120,83 @@ class RustConnectionPoolTestCase(trial_unittest.TestCase): def one(conn: Any) -> Any: cursor = conn.cursor() cursor.execute("SELECT 1::int") - row = cursor.fetch_one() + row = cursor.fetchone() conn.commit() return row - self.assertEqual((yield self.pool.run_with_connection(one)), (1,)) - self.assertEqual((yield self.pool.run_with_connection(one)), (1,)) + self.assertEqual((yield self.pool.runWithConnection(one)), (1,)) + self.assertEqual((yield self.pool.runWithConnection(one)), (1,)) + + @inlineCallbacks + def test_commits_work_left_open_on_success(self) -> Any: + # Like adbapi, the pool commits any transaction the function leaves + # open. One-shot background-update runners rely on this: they execute + # DDL/DML and return without committing. + def create_and_insert(conn: Any) -> None: + cursor = conn.cursor() + cursor.execute("CREATE TABLE implicit_commit_test (x INT)") + cursor.execute("INSERT INTO implicit_commit_test VALUES (1)") + # No commit: the pool must supply it. + + def count(conn: Any) -> Any: + cursor = conn.cursor() + cursor.execute("SELECT COUNT(*) FROM implicit_commit_test") + return cursor.fetchone() + + def drop(conn: Any) -> None: + conn.cursor().execute("DROP TABLE IF EXISTS implicit_commit_test") + + self.addCleanup(lambda: self.pool.runWithConnection(drop)) + + yield self.pool.runWithConnection(create_and_insert) + self.assertEqual((yield self.pool.runWithConnection(count)), (1,)) + + @inlineCallbacks + def test_rolls_back_and_reuses_connection_on_exception(self) -> Any: + # Like adbapi, an exception rolls the open transaction back before the + # connection is returned, so the pool reuses the (now clean) connection + # rather than discarding it as mid-transaction — otherwise every + # NotFoundError/StoreError raised inside a transaction would destroy a + # TCP/TLS session. + pool = RustConnectionPool( + reactor, _build_dsn(), name="test-rust-db-single", threads=1 + ) + pool.start() + self.addCleanup(pool.close) + + class MarkerError(Exception): + pass + + pids: list[int] = [] + + def create(conn: Any) -> None: + conn.cursor().execute("CREATE TABLE rollback_reuse_test (x INT)") + + def drop(conn: Any) -> None: + conn.cursor().execute("DROP TABLE IF EXISTS rollback_reuse_test") + + def insert_and_raise(conn: Any) -> None: + cursor = conn.cursor() + cursor.execute("SELECT pg_backend_pid()") + pids.append(cursor.fetchone()[0]) + cursor.execute("INSERT INTO rollback_reuse_test VALUES (1)") + raise MarkerError("boom") + + def check(conn: Any) -> Any: + cursor = conn.cursor() + cursor.execute("SELECT pg_backend_pid()") + pids.append(cursor.fetchone()[0]) + cursor.execute("SELECT COUNT(*) FROM rollback_reuse_test") + return cursor.fetchone() + + self.addCleanup(lambda: pool.runWithConnection(drop)) + + yield pool.runWithConnection(create) + yield self.assertFailure(pool.runWithConnection(insert_and_raise), MarkerError) + # The insert was rolled back... + self.assertEqual((yield pool.runWithConnection(check)), (0,)) + # ...and the single pooled connection survived the failed call. + self.assertEqual(pids[0], pids[1]) @inlineCallbacks def test_concurrent_calls_are_serviced(self) -> Any: @@ -132,16 +205,42 @@ class RustConnectionPoolTestCase(trial_unittest.TestCase): def txn(conn: Any, n: int) -> Any: cursor = conn.cursor() cursor.execute("SELECT $1::int", [n]) - row = cursor.fetch_one() + row = cursor.fetchone() conn.commit() return row results = yield gatherResults( - [self.pool.run_with_connection(txn, n) for n in range(10)] + [self.pool.runWithConnection(txn, n) for n in range(10)] ) self.assertEqual(results, [(n,) for n in range(10)]) + @inlineCallbacks + def test_drives_a_transaction_as_db_pool(self) -> Any: + # Mirror what DatabasePool.runWithConnection's inner_func does: the pool + # hands a DBAPI2 connection the engine can inspect, wrapped in a + # LoggingDatabaseConnection whose cursor is a real LoggingTransaction. + engine = RustPostgresEngine({}) + + def interaction(conn: Any) -> Any: + # A freshly checked-out connection is not mid-transaction. + self.assertFalse(engine.in_transaction(conn)) + + db_conn = LoggingDatabaseConnection( + conn=conn, + engine=engine, + default_txn_name="test", + server_name="test", + ) + txn = db_conn.cursor(txn_name="test") + txn.execute("SELECT ?::int + ?::int", (2, 3)) + row = txn.fetchone() + db_conn.commit() + return row + + result = yield self.pool.runWithConnection(interaction) + self.assertEqual(result, (5,)) + def test_run_when_not_running_raises(self) -> None: self.pool.close() with self.assertRaises(RuntimeError): - self.pool.run_with_connection(lambda conn: None) + self.pool.runWithConnection(lambda conn: None)