diff --git a/rust/src/database/postgres/pool.rs b/rust/src/database/postgres/pool.rs index 19fafb93eb..6a21246f3b 100644 --- a/rust/src/database/postgres/pool.rs +++ b/rust/src/database/postgres/pool.rs @@ -123,6 +123,16 @@ impl PyConnectionPool { let conn = self.pool.get().block_on(py).map_err(pool_err_to_py)?; Ok(Connection::new(conn)) } + + /// Close the pool, closing every idle connection. + /// + /// After this, [`connect`](Self::connect) fails; a connection still checked + /// out is closed when it is returned. Idempotent. This lets the owning + /// Python code drop the pool's server connections deterministically rather + /// than waiting for garbage collection. + fn close(&self) { + self.pool.close(); + } } /// Map a `deadpool` checkout failure onto the DBAPI2 exception hierarchy. diff --git a/synapse/storage/rust_pool.py b/synapse/storage/rust_pool.py new file mode 100644 index 0000000000..4bec9db322 --- /dev/null +++ b/synapse/storage/rust_pool.py @@ -0,0 +1,164 @@ +# This file is licensed under the Affero General Public License (AGPL) version 3. +# +# Copyright (C) 2026 Element Creations Ltd +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# See the GNU Affero General Public License for more details: +# . + +"""A Twisted connection-pool adapter backed by the native Rust Postgres pool. + +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``. + +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. +""" + +import logging +from types import TracebackType +from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar + +from typing_extensions import Concatenate, ParamSpec + +from twisted.python.threadpool import ThreadPool + +from synapse.logging.context import defer_to_threadpool +from synapse.synapse_rust.database import postgres + +if TYPE_CHECKING: + from twisted.internet.defer import Deferred + + from synapse.types import ISynapseReactor + +logger = logging.getLogger(__name__) + +P = ParamSpec("P") +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`. + """ + + def __init__( + self, + reactor: "ISynapseReactor", + dsn: str, + *, + name: str, + threads: int = 10, + ) -> None: + """ + Args: + reactor: the reactor in whose main thread Deferreds are fired. + dsn: a libpq-style connection string for the Rust pool. + name: a label for the thread pool (used in logs / metrics). + threads: the maximum number of worker threads. The Rust connection + 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. + + The owner is responsible for the lifecycle: call :meth:`start` before + use and :meth:`close` on shutdown (e.g. via the Synapse clock's + ``add_system_event_trigger``). The pool does not register a shutdown + hook itself, so it needs no clock and stays trivially testable. + """ + self._reactor = reactor + self._pool = postgres.ConnectionPool(dsn, threads) + self.threadpool = ThreadPool(minthreads=1, maxthreads=threads, name=name) + self.running = False + + def start(self) -> None: + """Start the thread pool. Idempotent.""" + if self.running: + return + self.threadpool.start() + self.running = True + + def close(self) -> None: + """Stop the thread pool and close the connection pool. Idempotent. + + Stops the thread pool first (waiting for in-flight work to finish, which + returns its connection to the pool), then closes the Rust pool so its + server connections are dropped promptly rather than lingering until + garbage collection. + """ + if not self.running: + return + self.running = False + self.threadpool.stop() + self._pool.close() + + def run_with_connection( + self, + func: Callable[Concatenate[Any, P], R], + *args: P.args, + **kwargs: P.kwargs, + ) -> "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. + + Returns: + A ``Deferred`` firing with ``func``'s result, following the Synapse + logcontext rules (``yield`` / ``await`` it). + """ + if not self.running: + raise RuntimeError("connection pool is not running") + + return defer_to_threadpool( + self._reactor, self.threadpool, self._run, func, args, kwargs + ) + + def _run( + self, + func: Callable[..., R], + args: tuple, + kwargs: dict, + ) -> R: + """Worker-thread body: check out a connection, run ``func``, return it. + + A checkout failure surfaces as the raised exception (→ errback) before + there is any connection to release. + """ + conn = self._pool.connect() + try: + return func(conn, *args, **kwargs) + finally: + # Return the connection to the pool. The Rust shim returns a clean + # connection for reuse and discards one left mid-transaction or + # otherwise unusable (see the shim's disposal docs). + conn.close() + + def __enter__(self) -> "RustConnectionPool": + self.start() + return self + + def __exit__( + self, + exc_type: Optional[type[BaseException]], + exc: Optional[BaseException], + tb: Optional[TracebackType], + ) -> None: + self.close() diff --git a/tests/storage/test_rust_pool.py b/tests/storage/test_rust_pool.py new file mode 100644 index 0000000000..f4053b1866 --- /dev/null +++ b/tests/storage/test_rust_pool.py @@ -0,0 +1,147 @@ +# This file is licensed under the Affero General Public License (AGPL) version 3. +# +# Copyright (C) 2026 Element Creations Ltd +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# See the GNU Affero General Public License for more details: +# . + +"""Tests for the Rust-backed Twisted connection pool adapter +(:mod:`synapse.storage.rust_pool`). + +These drive real worker threads talking to a real Postgres over the real +reactor, so they use a plain Twisted trial ``TestCase`` (Synapse's in-memory +test reactor deliberately mocks the database thread pool out) and are skipped +unless the suite is configured to run against Postgres. +""" + +from typing import TYPE_CHECKING, Any, cast + +# 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.rust_pool import RustConnectionPool + +from tests.unittest import skip_unless +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 + +# `twisted.internet.reactor` is a module-level singleton that is the reactor +# object; narrow it for the type checker. +reactor = cast("ISynapseReactor", _reactor) + + +def _build_dsn() -> str: + """Build a libpq keyword/value connection string from the test config.""" + + parts = [f"dbname={POSTGRES_BASE_DB}"] + if POSTGRES_USER is not None: + parts.append(f"user={POSTGRES_USER}") + if POSTGRES_HOST is not None: + parts.append(f"host={POSTGRES_HOST}") + if POSTGRES_PORT is not None: + parts.append(f"port={POSTGRES_PORT}") + if POSTGRES_PASSWORD is not None: + parts.append(f"password={POSTGRES_PASSWORD}") + return " ".join(parts) + + +@skip_unless( + bool(USE_POSTGRES_FOR_TESTS), "requires a Postgres server (set SYNAPSE_POSTGRES)" +) +class RustConnectionPoolTestCase(trial_unittest.TestCase): + """The execution bridge: run blocking DB functions off the reactor thread.""" + + def setUp(self) -> None: + self.pool = RustConnectionPool( + reactor, _build_dsn(), name="test-rust-db", threads=4 + ) + self.pool.start() + self.addCleanup(self.pool.close) + + @inlineCallbacks + def test_runs_function_against_a_live_connection(self) -> Any: + # The function gets a usable connection: it can open a cursor, run a + # query, commit, and its return value comes back through the Deferred. + def txn(conn: Any) -> Any: + cursor = conn.cursor() + cursor.execute("SELECT 42::int") + row = cursor.fetch_one() + conn.commit() + return row + + result = yield self.pool.run_with_connection(txn) + self.assertEqual(result, (42,)) + + @inlineCallbacks + def test_forwards_args_and_kwargs(self) -> Any: + 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) + self.assertEqual(result, 6) + + @inlineCallbacks + def test_exception_propagates_as_errback(self) -> Any: + class MarkerError(Exception): + pass + + def txn(conn: Any) -> None: + raise MarkerError("boom") + + # The failure crosses the thread boundary and surfaces as an errback. + failure = yield self.assertFailure( + self.pool.run_with_connection(txn), MarkerError + ) + self.assertEqual(str(failure), "boom") + + @inlineCallbacks + def test_connection_is_reusable_across_calls(self) -> Any: + # Connections are returned to the pool after each call, so a second call + # (which may reuse the same underlying connection) works and sees a + # clean session rather than a leftover transaction. + def one(conn: Any) -> Any: + cursor = conn.cursor() + cursor.execute("SELECT 1::int") + row = cursor.fetch_one() + conn.commit() + return row + + self.assertEqual((yield self.pool.run_with_connection(one)), (1,)) + self.assertEqual((yield self.pool.run_with_connection(one)), (1,)) + + @inlineCallbacks + def test_concurrent_calls_are_serviced(self) -> Any: + # Fire more calls than we have threads/connections at once; the pool and + # thread pool should service them all and return the right answers. + def txn(conn: Any, n: int) -> Any: + cursor = conn.cursor() + cursor.execute("SELECT $1::int", [n]) + row = cursor.fetch_one() + conn.commit() + return row + + results = yield gatherResults( + [self.pool.run_with_connection(txn, n) for n in range(10)] + ) + self.assertEqual(results, [(n,) for n in range(10)]) + + def test_run_when_not_running_raises(self) -> None: + self.pool.close() + with self.assertRaises(RuntimeError): + self.pool.run_with_connection(lambda conn: None)