Present the _db_pool interface on RustConnectionPool

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) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01W3G4M92AmwSSZCbmtMJU3d

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

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)