Wire the native Rust Postgres backend behind a use_rust_driver flag

Turn the Rust backend on with `database.use_rust_driver: true` (alongside
`name: psycopg2`). Keeping `name: psycopg2` means the engine is still selected as
Postgres and all the `isinstance(engine, PostgresEngine)` dialect checks hold;
the flag only swaps the driver implementation.

  - `create_engine` returns `RustPostgresEngine` instead of `Psycopg2Engine` when
    the flag is set.
  - `make_pool` gains a `clock` argument and, for the Rust engine, builds a
    `RustConnectionPool` (via `_make_rust_pool`) instead of an adbapi pool:
    it derives the libpq DSN from the config `args`, sizes threads/connections
    to `cp_max`, passes the engine's synchronous_commit / statement_timeout,
    starts the pool, and registers a shutdown hook via the clock. Its declared
    return type stays `adbapi.ConnectionPool` — `RustConnectionPool` provides the
    `_db_pool` subset `DatabasePool` uses (runWithConnection / threadID /
    running / threadpool) — so callers that lean on adbapi-specific methods keep
    type-checking.
  - `RustConnectionPool` takes synchronous_commit / statement_timeout_ms and
    threads them into its pooled connections' session setup.
  - `DatabasePool` passes its clock to `make_pool`; the test harness's
    `make_fake_db_pool` accepts the new argument.

Tested: create_engine honours the flag (and the Rust engine is still a
PostgresEngine, not a Psycopg2Engine); make_pool builds a started
RustConnectionPool that serves a query as `_db_pool`. psycopg2 and sqlite
homeserver boots remain green. A full homeserver boot on the Rust backend isn't
covered here because the test harness patches `make_pool` with a synchronous
adbapi fake.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01W3G4M92AmwSSZCbmtMJU3d

synapse/storage/database.py

@@ -46,7 +46,6 @@ from prometheus_client import Counter, Histogram
 from typing_extensions import Concatenate, ParamSpec
 
 from twisted.enterprise import adbapi
-from twisted.internet.interfaces import IReactorCore
 
 from synapse.api.errors import StoreError
 from synapse.config.database import DatabaseConnectionConfig
@@ -67,6 +66,10 @@ from synapse.util.iterutils import batch_iter
 
 if TYPE_CHECKING:
     from synapse.server import HomeServer
+    from synapse.storage.engines.postgres_rust import RustPostgresEngine
+    from synapse.storage.rust_pool import RustConnectionPool
+    from synapse.types import ISynapseReactor
+    from synapse.util.clock import Clock
 
 # python 3 does not have a maximum int value
 MAX_TXN_ID = 2**63 - 1
@@ -133,13 +136,28 @@ class _PoolConnection(Connection):
 
 def make_pool(
     *,
-    reactor: IReactorCore,
+    reactor: "ISynapseReactor",
+    clock: "Clock",
     db_config: DatabaseConnectionConfig,
     engine: BaseDatabaseEngine,
     server_name: str,
 ) -> adbapi.ConnectionPool:
     """Get the connection pool for the database."""
 
+    from synapse.storage.engines.postgres_rust import RustPostgresEngine
+
+    if isinstance(engine, RustPostgresEngine):
+        # RustConnectionPool isn't an adbapi.ConnectionPool, but it provides the
+        # subset of the interface DatabasePool uses of `_db_pool`
+        # (runWithConnection / threadID / running / threadpool).
+        return _make_rust_pool(  # type: ignore[return-value]
+            reactor=reactor,
+            clock=clock,
+            db_config=db_config,
+            engine=engine,
+            server_name=server_name,
+        )
+
     # By default enable `cp_reconnect`. We need to fiddle with db_args in case
     # someone has explicitly set `cp_reconnect`.
     db_args = dict(db_config.config.get("args", {}))
@@ -174,6 +192,50 @@ def make_pool(
     return connection_pool
 
 
+def _make_rust_pool(
+    *,
+    reactor: "ISynapseReactor",
+    clock: "Clock",
+    db_config: DatabaseConnectionConfig,
+    engine: "RustPostgresEngine",
+    server_name: str,
+) -> "RustConnectionPool":
+    """Build a native-Rust-backed connection pool for the database.
+
+    The `_db_pool` counterpart to the adbapi pool built by `make_pool`, used when
+    the database is configured with `use_rust_driver`.
+    """
+    from synapse.storage import rust_dbapi
+    from synapse.storage.rust_pool import RustConnectionPool
+
+    db_args = db_config.config.get("args", {})
+    dsn = rust_dbapi.build_dsn(
+        {k: v for k, v in db_args.items() if not k.startswith("cp_")}
+    )
+    # Size the pool (threads and connections, 1:1) to the configured cp_max;
+    # Twisted's adbapi default is 5.
+    threads = db_args.get("cp_max", 5)
+
+    pool = RustConnectionPool(
+        reactor,
+        dsn,
+        name=f"database-{db_config.name}",
+        threads=threads,
+        synchronous_commit=engine.synchronous_commit,
+        statement_timeout_ms=engine.statement_timeout,
+    )
+    pool.start()
+    clock.add_system_event_trigger("during", "shutdown", pool.close)
+
+    register_threadpool(
+        name=f"database-{db_config.name}",
+        server_name=server_name,
+        threadpool=pool.threadpool,
+    )
+
+    return pool
+
+
 def make_conn(
     *,
     db_config: DatabaseConnectionConfig,
@@ -629,6 +691,7 @@ class DatabasePool:
         self._database_config = database_config
         self._db_pool = make_pool(
             reactor=hs.get_reactor(),
+            clock=self._clock,
             db_config=database_config,
             engine=engine,
             server_name=self.server_name,

synapse/storage/engines/__init__.py

@@ -60,6 +60,13 @@ def create_engine(database_config: Mapping[str, Any]) -> BaseDatabaseEngine:
         return Sqlite3Engine(database_config)
 
     if name == "psycopg2":
+        # Opt in to the native Rust Postgres driver (same wire protocol and SQL
+        # dialect; a different, psycopg2-free connection/cursor implementation).
+        if database_config.get("use_rust_driver", False):
+            from .postgres_rust import RustPostgresEngine
+
+            return RustPostgresEngine(database_config)
+
         return Psycopg2Engine(database_config)
 
     raise RuntimeError("Unsupported database engine '%s'" % (name,))

synapse/storage/rust_pool.py

@@ -72,6 +72,8 @@ class RustConnectionPool:
         *,
         name: str,
         threads: int = 10,
+        synchronous_commit: bool = True,
+        statement_timeout_ms: int | None = None,
     ) -> None:
         """
         Args:
@@ -82,6 +84,9 @@ class RustConnectionPool:
                 pool is sized to match, since each worker holds at most one
                 connection at a time — a 1:1 cap avoids both starvation and
                 idle connections.
+            synchronous_commit: passed to each pooled connection's session setup.
+            statement_timeout_ms: passed to each pooled connection's session
+                setup (statements running longer are aborted).
 
         The owner is responsible for the lifecycle: call :meth:`start` before
         use and :meth:`close` on shutdown (e.g. via the Synapse clock's
@@ -89,7 +94,12 @@ class RustConnectionPool:
         hook itself, so it needs no clock and stays trivially testable.
         """
         self._reactor = reactor
-        self._pool = postgres.ConnectionPool(dsn, threads)
+        self._pool = postgres.ConnectionPool(
+            dsn,
+            threads,
+            synchronous_commit=synchronous_commit,
+            statement_timeout_ms=statement_timeout_ms,
+        )
         self.threadpool = ThreadPool(minthreads=1, maxthreads=threads, name=name)
         self.running = False
 

tests/server.py

@@ -734,6 +734,7 @@ def validate_connector(connector: tcp.Connector, expected_ip: str) -> None:
 
 def make_fake_db_pool(
     reactor: ISynapseReactor,
+    clock: "Clock",
     db_config: DatabaseConnectionConfig,
     engine: BaseDatabaseEngine,
     server_name: str,
@@ -746,7 +747,11 @@ def make_fake_db_pool(
     pool.
     """
     pool = make_pool(
-        reactor=reactor, db_config=db_config, engine=engine, server_name=server_name
+        reactor=reactor,
+        clock=clock,
+        db_config=db_config,
+        engine=engine,
+        server_name=server_name,
     )
 
     def runWithConnection(

tests/storage/test_rust_engine.py

@@ -14,6 +14,7 @@
 
 from typing import Any
 
+from synapse.storage.engines import PostgresEngine, Psycopg2Engine, create_engine
 from synapse.storage.engines.postgres_rust import RustPostgresEngine
 from synapse.synapse_rust.database import postgres
 
@@ -84,6 +85,17 @@ class RustPostgresEngineTestCase(unittest.TestCase):
         with self.assertRaises(NotImplementedError):
             self.engine.attempt_to_set_isolation_level(object(), None)
 
+    def test_create_engine_selects_rust_only_when_opted_in(self) -> None:
+        # Default: the psycopg2 engine.
+        self.assertIsInstance(create_engine({"name": "psycopg2"}), Psycopg2Engine)
+
+        # With the opt-in flag: the Rust engine — still a PostgresEngine, so the
+        # storage layer's `isinstance(engine, PostgresEngine)` dialect checks hold.
+        engine = create_engine({"name": "psycopg2", "use_rust_driver": True})
+        self.assertIsInstance(engine, RustPostgresEngine)
+        self.assertIsInstance(engine, PostgresEngine)
+        self.assertNotIsInstance(engine, Psycopg2Engine)
+
 
 @skip_unless(
     bool(USE_POSTGRES_FOR_TESTS), "requires a Postgres server (set SYNAPSE_POSTGRES)"

tests/storage/test_rust_pool.py

@@ -20,13 +20,15 @@ unless the suite is configured to run against Postgres.
 """
 
 from typing import TYPE_CHECKING, Any, cast
+from unittest.mock import Mock
 
 # The reactor the trial runner spins up; real, so threads and callFromThread work.
 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.config.database import DatabaseConnectionConfig
+from synapse.storage.database import LoggingDatabaseConnection, make_pool
 from synapse.storage.engines.postgres_rust import RustPostgresEngine
 from synapse.storage.rust_pool import RustConnectionPool
 
@@ -42,6 +44,7 @@ from tests.utils import (
 
 if TYPE_CHECKING:
     from synapse.types import ISynapseReactor
+    from synapse.util.clock import Clock
 
 # `twisted.internet.reactor` is a module-level singleton that is the reactor
 # object; narrow it for the type checker.
@@ -169,6 +172,46 @@ class RustConnectionPoolTestCase(trial_unittest.TestCase):
         result = yield self.pool.runWithConnection(interaction)
         self.assertEqual(result, (5,))
 
+    @inlineCallbacks
+    def test_make_pool_builds_a_running_rust_pool(self) -> Any:
+        # `make_pool` returns a started RustConnectionPool for a use_rust_driver
+        # config, ready to serve as `_db_pool`.
+        args: dict = {"dbname": POSTGRES_BASE_DB}
+        if POSTGRES_USER is not None:
+            args["user"] = POSTGRES_USER
+        if POSTGRES_HOST is not None:
+            args["host"] = POSTGRES_HOST
+        if POSTGRES_PORT is not None:
+            args["port"] = POSTGRES_PORT
+        if POSTGRES_PASSWORD is not None:
+            args["password"] = POSTGRES_PASSWORD
+
+        db_config = DatabaseConnectionConfig(
+            "master", {"name": "psycopg2", "use_rust_driver": True, "args": args}
+        )
+        engine = RustPostgresEngine(db_config.config)
+
+        pool = make_pool(
+            reactor=reactor,
+            clock=cast("Clock", Mock()),
+            db_config=db_config,
+            engine=engine,
+            server_name="test",
+        )
+        self.addCleanup(pool.close)
+
+        self.assertIsInstance(pool, RustConnectionPool)
+        self.assertTrue(pool.running)
+
+        def txn(conn: Any) -> Any:
+            cursor = conn.cursor()
+            cursor.execute("SELECT 1")
+            row = cursor.fetchone()
+            conn.commit()
+            return row
+
+        self.assertEqual((yield pool.runWithConnection(txn)), (1,))
+
     def test_run_when_not_running_raises(self) -> None:
         self.pool.close()
         with self.assertRaises(RuntimeError):