diff --git a/synapse/storage/database.py b/synapse/storage/database.py index 8986282983..dad22ae38d 100644 --- a/synapse/storage/database.py +++ b/synapse/storage/database.py @@ -68,6 +68,7 @@ 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.synapse_rust.runtime import RustRuntime from synapse.types import ISynapseReactor from synapse.util.clock import Clock @@ -142,17 +143,27 @@ def make_pool( db_config: DatabaseConnectionConfig, engine: BaseDatabaseEngine, server_name: str, + runtime: "RustRuntime | None" = None, ) -> adbapi.ConnectionPool: - """Get the connection pool for the database.""" + """Get the connection pool for the database. + + Args: + runtime: The homeserver's `RustRuntime`, whose tokio runtime drives the + native connection pool. Required for the Rust backend, unused by + the others. + """ from synapse.storage.engines.postgres_rust import RustPostgresEngine if isinstance(engine, RustPostgresEngine): + if runtime is None: + raise ValueError("the Rust Postgres backend needs a RustRuntime") # 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, + runtime=runtime, clock=clock, db_config=db_config, engine=engine, @@ -196,6 +207,7 @@ def make_pool( def _make_rust_pool( *, reactor: "ISynapseReactor", + runtime: "RustRuntime", clock: "Clock", db_config: DatabaseConnectionConfig, engine: "RustPostgresEngine", @@ -219,6 +231,7 @@ def _make_rust_pool( pool = RustConnectionPool( reactor, + runtime, dsn, name=f"database-{db_config.name}", threads=threads, @@ -243,9 +256,16 @@ def make_conn( engine: BaseDatabaseEngine, default_txn_name: str, server_name: str, + runtime: "RustRuntime | None" = None, ) -> "LoggingDatabaseConnection": """Make a new connection to the database and return it. + Args: + runtime: The `RustRuntime` whose tokio runtime drives the connection. + Only used by the Rust backend; defaults to a standalone runtime on + the global reactor, for the scripts and tests that have no + homeserver to take it from. + Returns: Connection """ @@ -264,7 +284,11 @@ def make_conn( # connection from the same libpq args, with the engine's session settings. from synapse.storage import rust_dbapi + if runtime is None: + runtime = rust_dbapi.standalone_runtime() + native_db_conn = rust_dbapi.connect( + runtime, rust_dbapi.build_dsn(db_params), synchronous_commit=engine.synchronous_commit, statement_timeout_ms=engine.statement_timeout, @@ -696,6 +720,7 @@ class DatabasePool: db_config=database_config, engine=engine, server_name=self.server_name, + runtime=hs.get_rust_runtime(), ) self.updates = BackgroundUpdater(hs, self) diff --git a/synapse/storage/databases/__init__.py b/synapse/storage/databases/__init__.py index b44b84b913..30a9dec28f 100644 --- a/synapse/storage/databases/__init__.py +++ b/synapse/storage/databases/__init__.py @@ -88,6 +88,7 @@ class Databases(Generic[DataStoreT]): engine=engine, default_txn_name="startup", server_name=server_name, + runtime=hs.get_rust_runtime(), ) as db_conn: logger.info("[database config %r]: Checking database server", db_name) engine.check_database(db_conn) diff --git a/synapse/storage/rust_dbapi.py b/synapse/storage/rust_dbapi.py index 78486cea39..b9c76e1437 100644 --- a/synapse/storage/rust_dbapi.py +++ b/synapse/storage/rust_dbapi.py @@ -36,6 +36,7 @@ import logging from typing import TYPE_CHECKING, Any, Iterator, Mapping, Sequence from synapse.synapse_rust.database import postgres +from synapse.synapse_rust.runtime import RustRuntime if TYPE_CHECKING: from synapse.storage.types import SQLQueryParameters @@ -164,7 +165,50 @@ def build_dsn(params: Mapping[str, Any]) -> str: return " ".join(parts) +class _StandaloneReactor: + """What ``RustRuntime`` asks of the reactor: somewhere to register its + Twisted-dispatch reader. A standalone runtime only ever drives blocking + database calls, which never dispatch back to Twisted, so the reader is + never serviced and nothing is registered on the real reactor (which, under + trial, would otherwise fail the dirty-reactor check).""" + + def addReader(self, reader: Any) -> None: # noqa: N802 + pass + + def removeReader(self, reader: Any) -> None: # noqa: N802 + pass + + +class _StandaloneHomeserver: + """The two things ``RustRuntime`` needs from a homeserver, for callers + that have none (maintenance scripts, and tests that drive the backend + directly). The shutdown hook is dropped, so the runtime lives for the + process.""" + + def get_reactor(self) -> Any: + return _StandaloneReactor() + + def register_sync_shutdown_handler(self, **kwargs: Any) -> None: + pass + + +_standalone_runtime: RustRuntime | None = None + + +def standalone_runtime() -> RustRuntime: + """A process-wide ``RustRuntime`` for code with no homeserver to take one + from. Prefer ``hs.get_rust_runtime()`` wherever a homeserver exists.""" + global _standalone_runtime + if _standalone_runtime is None: + _standalone_runtime = RustRuntime( + hs=_StandaloneHomeserver(), # type: ignore[arg-type] + worker_threads=2, + ) + return _standalone_runtime + + def connect( + runtime: RustRuntime, dsn: str, *, synchronous_commit: bool = True, @@ -176,8 +220,12 @@ def connect( returned :class:`Connection` keeps that pool alive for its lifetime. Used by ``make_conn`` for the startup connection that runs schema preparation before the real pool exists. + + ``runtime`` is the ``RustRuntime`` whose tokio runtime drives the + connection; it is started on first use, so this works during startup. """ pool = postgres.ConnectionPool( + runtime, dsn, 1, synchronous_commit=synchronous_commit, diff --git a/tests/server.py b/tests/server.py index 33e313f1bb..3bb4ded1b2 100644 --- a/tests/server.py +++ b/tests/server.py @@ -101,7 +101,7 @@ from synapse.storage import DataStore from synapse.storage.database import LoggingDatabaseConnection, make_pool from synapse.storage.engines import BaseDatabaseEngine, create_engine from synapse.storage.prepare_database import prepare_database -from synapse.synapse_rust.runtime import set_virtual_time_msec +from synapse.synapse_rust.runtime import RustRuntime, set_virtual_time_msec from synapse.types import ISynapseReactor, JsonDict from synapse.util.clock import Clock from synapse.util.duration import Duration @@ -860,6 +860,7 @@ def make_fake_db_pool( db_config: DatabaseConnectionConfig, engine: BaseDatabaseEngine, server_name: str, + runtime: "RustRuntime | None" = None, ) -> adbapi.ConnectionPool: """Wrapper for `make_pool` which builds a pool which runs db queries synchronously. @@ -874,6 +875,7 @@ def make_fake_db_pool( db_config=db_config, engine=engine, server_name=server_name, + runtime=runtime, ) def runWithConnection( diff --git a/tests/storage/test_rust_dbapi.py b/tests/storage/test_rust_dbapi.py index fc72455aa5..c651af3a94 100644 --- a/tests/storage/test_rust_dbapi.py +++ b/tests/storage/test_rust_dbapi.py @@ -18,9 +18,9 @@ from synapse.config.database import DatabaseConnectionConfig from synapse.storage import rust_dbapi from synapse.storage.database import LoggingDatabaseConnection, make_conn from synapse.storage.engines.postgres_rust import RustPostgresEngine -from synapse.synapse_rust.database import postgres from tests import unittest +from tests.synapse_rust.test_database_postgres import _make_pool from tests.unittest import skip_unless from tests.utils import ( POSTGRES_BASE_DB, @@ -124,7 +124,7 @@ class BuildDsnTestCase(unittest.TestCase): value = samples.get(key, "1" if key.startswith("keepalives") else "x") if key in ("port", "tcp_user_timeout"): value = "5432" - pool = postgres.ConnectionPool(f"host=h {key}={value}") + pool = _make_pool(f"host=h {key}={value}") pool.close() @@ -177,7 +177,7 @@ class RustDBAPIAdapterTestCase(unittest.TestCase): """The adapter presents the DBAPI2 shape over the shim.""" def setUp(self) -> None: - self._pool = postgres.ConnectionPool(_build_dsn()) + self._pool = _make_pool(_build_dsn()) self.conn = rust_dbapi.Connection(self._pool.connect()) def tearDown(self) -> None: diff --git a/tests/storage/test_rust_engine.py b/tests/storage/test_rust_engine.py index b624362120..4c305e8a75 100644 --- a/tests/storage/test_rust_engine.py +++ b/tests/storage/test_rust_engine.py @@ -19,6 +19,7 @@ from synapse.storage.engines.postgres_rust import RustPostgresEngine from synapse.synapse_rust.database import postgres from tests import unittest +from tests.synapse_rust.test_database_postgres import _make_pool from tests.unittest import skip_unless from tests.utils import ( POSTGRES_BASE_DB, @@ -105,7 +106,7 @@ class RustPostgresEngineConnectionTestCase(unittest.TestCase): def setUp(self) -> None: self.engine = RustPostgresEngine({}) - self.pool = postgres.ConnectionPool(_build_dsn()) + self.pool = _make_pool(_build_dsn()) self.conn = self.pool.connect() def tearDown(self) -> None: diff --git a/tests/storage/test_rust_pool.py b/tests/storage/test_rust_pool.py index 8e2d1e66d9..2b3e51e0f9 100644 --- a/tests/storage/test_rust_pool.py +++ b/tests/storage/test_rust_pool.py @@ -34,7 +34,14 @@ from synapse.storage.rust_pool import RustConnectionPool from tests.synapse_rust.test_database_postgres import _build_dsn, _test_runtime from tests.unittest import skip_unless -from tests.utils import USE_POSTGRES_FOR_TESTS +from tests.utils import ( + POSTGRES_BASE_DB, + POSTGRES_HOST, + POSTGRES_PASSWORD, + POSTGRES_PORT, + POSTGRES_USER, + USE_POSTGRES_FOR_TESTS, +) if TYPE_CHECKING: from synapse.types import ISynapseReactor @@ -141,7 +148,11 @@ class RustConnectionPoolTestCase(trial_unittest.TestCase): # NotFoundError/StoreError raised inside a transaction would destroy a # TCP/TLS session. pool = RustConnectionPool( - reactor, _build_dsn(), name="test-rust-db-single", threads=1 + reactor, + _test_runtime(), + _build_dsn(), + name="test-rust-db-single", + threads=1, ) pool.start() self.addCleanup(pool.close) @@ -247,6 +258,7 @@ class RustConnectionPoolTestCase(trial_unittest.TestCase): db_config=db_config, engine=engine, server_name="test", + runtime=_test_runtime(), ) self.addCleanup(pool.close)