diff --git a/synapse/storage/rust_pool.py b/synapse/storage/rust_pool.py index 4bec9db322..7ab1b23788 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,38 @@ 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. The function is responsible for + committing or rolling back (as Synapse's ``new_transaction`` does); the + connection is returned to the pool afterwards regardless. + + 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], @@ -142,7 +165,7 @@ class RustConnectionPool: 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) finally: diff --git a/tests/storage/test_rust_pool.py b/tests/storage/test_rust_pool.py index f4053b1866..7b76cdb6ab 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,12 @@ 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_concurrent_calls_are_serviced(self) -> Any: @@ -132,16 +134,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)