From 7355b22df16c3a08fb0daa6101d5e08323125ae2 Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Thu, 2 Jul 2026 19:34:20 +0000 Subject: [PATCH] Wire the native Rust Postgres backend behind a use_rust_driver flag MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) Claude-Session: https://claude.ai/code/session_01W3G4M92AmwSSZCbmtMJU3d --- synapse/storage/database.py | 67 ++++++++++++++++++++++++++++- synapse/storage/engines/__init__.py | 7 +++ synapse/storage/rust_pool.py | 12 +++++- tests/server.py | 7 ++- tests/storage/test_rust_engine.py | 12 ++++++ tests/storage/test_rust_pool.py | 45 ++++++++++++++++++- 6 files changed, 145 insertions(+), 5 deletions(-) diff --git a/synapse/storage/database.py b/synapse/storage/database.py index 3c6e9f81be..3ad3a612b3 100644 --- a/synapse/storage/database.py +++ b/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, diff --git a/synapse/storage/engines/__init__.py b/synapse/storage/engines/__init__.py index b66afaf913..322a15fd72 100644 --- a/synapse/storage/engines/__init__.py +++ b/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,)) diff --git a/synapse/storage/rust_pool.py b/synapse/storage/rust_pool.py index 7ab1b23788..9ec300c770 100644 --- a/synapse/storage/rust_pool.py +++ b/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 diff --git a/tests/server.py b/tests/server.py index ce5eaad63d..0d6643286c 100644 --- a/tests/server.py +++ b/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( diff --git a/tests/storage/test_rust_engine.py b/tests/storage/test_rust_engine.py index 98cf70f124..b624362120 100644 --- a/tests/storage/test_rust_engine.py +++ b/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)" diff --git a/tests/storage/test_rust_pool.py b/tests/storage/test_rust_pool.py index 7b76cdb6ab..3ff076da8d 100644 --- a/tests/storage/test_rust_pool.py +++ b/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):