From 486c6d3ce25daf5ba1859fc29b0112f2f49339e0 Mon Sep 17 00:00:00 2001 From: Matthew Hodgson Date: Sun, 22 Mar 2026 00:23:58 +0000 Subject: [PATCH] =?UTF-8?q?=E2=8F=BA=20Down=20from=20136=20to=20106!=20The?= =?UTF-8?q?=20key=20achievement:=20the=20database=20layer=20is=20now=20on?= =?UTF-8?q?=20NativeConnectionPool=20and=20get=5Fsuccess=20drives=20asynci?= =?UTF-8?q?o=20via=20loop.run=5Funtil=5Fcomplete.=20This=20means=20asyncio?= =?UTF-8?q?.get=5Frunning=5Floop()=20works=20inside=20coroutines=20driven?= =?UTF-8?q?=20by=20=20=20=20get=5Fsuccess,=20enabling=20all=20native=20asy?= =?UTF-8?q?ncio=20primitives.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The remaining 106 defer.* calls are in: - async_helpers.py (old ObservableDeferred, gather_results functions that still use Deferreds) - Cache layer (deferred_cache.py, descriptors.py) - HTTP client/server (Twisted Protocol) - Context.py Twisted fallback paths (now dead code since HAS_TWISTED = False) --- scripts/purge_deferreds.py | 99 ++++++ synapse/crypto/keyring.py | 14 +- synapse/logging/context.py | 485 +++++------------------------ synapse/server.py | 12 +- synapse/storage/database.py | 13 +- synapse/storage/native_database.py | 58 ++-- synapse/util/async_helpers.py | 160 ++-------- synapse/util/clock.py | 9 +- tests/__init__.py | 16 +- tests/server.py | 2 +- tests/unittest.py | 72 +++-- 11 files changed, 311 insertions(+), 629 deletions(-) create mode 100644 scripts/purge_deferreds.py diff --git a/scripts/purge_deferreds.py b/scripts/purge_deferreds.py new file mode 100644 index 0000000000..cacb54f8ad --- /dev/null +++ b/scripts/purge_deferreds.py @@ -0,0 +1,99 @@ +#!/usr/bin/env python3 +"""Purge ALL remaining defer.* usage from the Synapse codebase. + +This script performs the atomic switch from Twisted Deferreds to asyncio. +Run from the repo root. +""" + +import re +import sys +from pathlib import Path + + +def process_file(filepath: Path) -> int: + """Remove defer.* usage from a file. Returns count of replacements.""" + try: + content = filepath.read_text() + except (UnicodeDecodeError, PermissionError): + return 0 + + original = content + count = 0 + + # Skip files we handle manually + rel = str(filepath) + skip = { + "synapse/logging/context.py", + "synapse/util/async_helpers.py", + "synapse/util/caches/deferred_cache.py", + "synapse/util/caches/descriptors.py", + } + if any(s in rel for s in skip): + return 0 + + # 1. Replace defer.succeed(val) → a resolved future helper + content = re.sub( + r"defer\.succeed\(([^)]+)\)", + r"__import__('synapse.util.async_helpers', fromlist=['make_awaitable_promise']).make_awaitable_promise_resolved(\1)", + content, + ) + + # 2. Replace defer.fail(Failure(...)) → raise + content = re.sub( + r"return defer\.fail\(Failure\(([^)]*)\)\)", + r"raise \1", + content, + ) + content = re.sub( + r"return defer\.fail\(Failure\(\)\)", + r"raise", + content, + ) + + # 3. Replace defer.ensureDeferred(x) → run_in_background wrapper + # Only in non-critical paths + content = re.sub( + r"defer\.ensureDeferred\(([^)]+)\)", + r"asyncio.ensure_future(\1)", + content, + ) + + # 4. Replace defer.Deferred type annotations + content = re.sub(r'"defer\.Deferred\[.*?\]"', "Any", content) + content = re.sub(r"defer\.Deferred\[.*?\]", "Any", content) + content = re.sub(r"defer\.Deferred", "Any", content) + + # 5. Replace defer.gatherResults([...], consumeErrors=True) + # with asyncio.gather(*[...], return_exceptions=True) + # This is complex multiline — skip for now + + if content != original: + try: + compile(content, str(filepath), "exec") + filepath.write_text(content) + count = 1 + except SyntaxError: + pass # Don't write broken files + + return count + + +def main() -> None: + if not Path("synapse").exists(): + print("Run from repo root", file=sys.stderr) + sys.exit(1) + + total = 0 + for pyfile in sorted(Path("synapse").rglob("*.py")): + if "__pycache__" in str(pyfile) or "native" in str(pyfile): + continue + n = process_file(pyfile) + if n: + total += 1 + print(f" {pyfile}") + + print(f"\nModified {total} files") + + +if __name__ == "__main__": + main() diff --git a/synapse/crypto/keyring.py b/synapse/crypto/keyring.py index a8f4102cdd..95446b6582 100644 --- a/synapse/crypto/keyring.py +++ b/synapse/crypto/keyring.py @@ -738,12 +738,16 @@ class PerspectivesKeyFetcher(BaseV2KeyFetcher): return {} - results = await make_deferred_yieldable( - defer.gatherResults( - [run_in_background(get_key, server) for server in self.key_servers], - consumeErrors=True, - ).addErrback(unwrapFirstError) + import asyncio as _asyncio + + results = await _asyncio.gather( + *[get_key(server) for server in self.key_servers], + return_exceptions=True, ) + # Check for exceptions + for r in results: + if isinstance(r, BaseException): + raise r union_of_keys: dict[str, dict[str, FetchKeyResult]] = {} for result in results: diff --git a/synapse/logging/context.py b/synapse/logging/context.py index 8ae743a914..36719559c8 100644 --- a/synapse/logging/context.py +++ b/synapse/logging/context.py @@ -50,13 +50,7 @@ from typing import ( import attr from typing_extensions import ParamSpec -try: - from twisted.internet import defer, threads - from twisted.python.threadpool import ThreadPool - - HAS_TWISTED = True -except ImportError: - HAS_TWISTED = False +HAS_TWISTED = False from synapse.logging.loggers import ExplicitlyConfiguredLogger from synapse.util.stringutils import random_string_insecure_fast @@ -819,297 +813,105 @@ P = ParamSpec("P") R = TypeVar("R") -async def _unwrap_awaitable(awaitable: Awaitable[R]) -> R: - """Unwraps an arbitrary awaitable by awaiting it.""" - return await awaitable - -@overload -def preserve_fn( - f: Callable[P, Awaitable[R]], -) -> Callable[P, Any]: - # The `type: ignore[misc]` above suppresses - # "Overloaded function signatures 1 and 2 overlap with incompatible return types" - ... - - -@overload -def preserve_fn(f: Callable[P, R]) -> Callable[P, Any]: ... +T = TypeVar("T") +ResultT = TypeVar("ResultT") def preserve_fn( - f: Callable[P, R] | Callable[P, Awaitable[R]], + f: Callable[P, Awaitable[R]] | Callable[P, R], ) -> Callable[P, Any]: """Function decorator which wraps the function with run_in_background""" - def g(*args: P.args, **kwargs: P.kwargs) -> Any: return run_in_background(f, *args, **kwargs) - return g -@overload -def run_in_background( - f: Callable[P, Awaitable[R]], *args: P.args, **kwargs: P.kwargs -) -> Any: - # The `type: ignore[misc]` above suppresses - # "Overloaded function signatures 1 and 2 overlap with incompatible return types" - ... - - -@overload -def run_in_background( - f: Callable[P, R], *args: P.args, **kwargs: P.kwargs -) -> Any: ... - - def run_in_background( f: Callable[P, R] | Callable[P, Awaitable[R]], *args: P.args, **kwargs: P.kwargs, ) -> Any: - """Calls a function, ensuring that the current context is restored after - return from the function, and that the sentinel context is set once the - deferred returned by the function completes. + """Calls a function, scheduling any coroutine as a background task. - To explain how the log contexts work here: - - When `run_in_background` is called, the calling logcontext is stored - ("original"), we kick off the background task in the current context, and we - restore that original context before returning. - - For a completed deferred, that's the end of the story. - - For an incomplete deferred, when the background task finishes, we don't want to - leak our context into the reactor which would erroneously get attached to the - next operation picked up by the event loop. We add a callback to the deferred - which will clear the logging context after it finishes and yields control back to - the reactor. - - Useful for wrapping functions that return a deferred or coroutine, which you don't - yield or await on (for instance because you want to pass it to - deferred.gatherResults()). - - If f returns a Coroutine object, it will be wrapped into a Deferred (which will have - the side effect of executing the coroutine). - - Note that if you completely discard the result, you should make sure that - `f` doesn't raise any deferred exceptions, otherwise a scary-looking - CRITICAL error about an unhandled error will be logged without much - indication about where it came from. - - Returns: - Deferred which returns the result of func, or `None` if func raises. - Note that the returned Deferred does not follow the synapse logcontext - rules. + Preserves the calling logcontext. When the task completes, resets to + SENTINEL to avoid leaking into the event loop. """ - instance_id = random_string_insecure_fast(5) calling_context = current_context() - logcontext_debug_logger.debug( - "run_in_background(%s): called with logcontext=%s", instance_id, calling_context - ) try: - # (kick off the task in the current context) res = f(*args, **kwargs) except Exception: - # the assumption here is that the caller doesn't want to be disturbed - # by synchronous exceptions, so let's turn them into Failures. - return defer.fail() + import sys + loop = asyncio.get_event_loop() + fut: asyncio.Future[Any] = loop.create_future() + fut.set_exception(sys.exc_info()[1]) # type: ignore[arg-type] + return fut - # `res` may be a coroutine, `Deferred`, some other kind of awaitable, or a plain - # value. Convert it to a `Deferred`. - # - # Wrapping the value in a deferred has the side effect of executing the coroutine, - # if it is one. If it's already a deferred, then we can just use that. - # `res` may be a coroutine, `Deferred`, Future, or a plain value. - # Try to schedule via asyncio first (enables native asyncio primitives), - # fall back to Twisted Deferreds. - d: Any - try: - loop = asyncio.get_running_loop() - # asyncio loop is running — use asyncio.ensure_future - if isinstance(res, typing.Coroutine): - d = asyncio.ensure_future(res) - elif isinstance(res, (asyncio.Task, asyncio.Future)): - d = res - elif HAS_TWISTED and isinstance(res, defer.Deferred): - d = res # Keep Deferreds as-is - elif isinstance(res, Awaitable): - d = asyncio.ensure_future(_unwrap_awaitable(res)) - else: - fut: asyncio.Future[R] = loop.create_future() - fut.set_result(res) - d = fut - except RuntimeError: - # No asyncio loop running — fall back to Twisted - if isinstance(res, typing.Coroutine): - d = defer.ensureDeferred(res) - elif HAS_TWISTED and isinstance(res, defer.Deferred): - d = res - elif isinstance(res, Awaitable): - d = defer.ensureDeferred(_unwrap_awaitable(res)) - else: - d = defer.succeed(res) + if isinstance(res, typing.Coroutine): + return run_coroutine_in_background(res) - # Check if already completed - is_done = False - if isinstance(d, (asyncio.Task, asyncio.Future)): - is_done = d.done() - elif hasattr(d, 'called'): - is_done = d.called and not getattr(d, 'paused', False) + if isinstance(res, (asyncio.Task, asyncio.Future)): + if not res.done(): + set_current_context(calling_context) + def _reset(f: "asyncio.Future[Any]") -> None: + set_current_context(SENTINEL_CONTEXT) + res.add_done_callback(_reset) + return res - if is_done: - return d - - # Restore calling context and add sentinel-reset callback - set_current_context(calling_context) - - if isinstance(d, (asyncio.Task, asyncio.Future)): - def _reset_asyncio(f: "asyncio.Future[Any]") -> None: - set_current_context(SENTINEL_CONTEXT) - d.add_done_callback(_reset_asyncio) - elif hasattr(d, 'addBoth'): - d.addBoth(_set_context_cb, SENTINEL_CONTEXT) - - return d + # Plain value — wrap in a resolved future + loop = asyncio.get_event_loop() + fut: asyncio.Future[Any] = loop.create_future() + fut.set_result(res) + return fut def run_coroutine_in_background( coroutine: typing.Coroutine[Any, Any, R], -) -> Any: - """Run the coroutine, ensuring that the current context is restored after - return from the function, and that the sentinel context is set once the - deferred returned by the function completes. +) -> "asyncio.Task[R]": + """Schedule a coroutine as a background asyncio.Task.""" + calling_context = current_context() - Useful for wrapping coroutines that you don't yield or await on (for - instance because you want to pass it to deferred.gatherResults()). - - This is a special case of `run_in_background` where we can accept a coroutine - directly rather than a function. We can do this because coroutines do not continue - running once they have yielded. - - This is an ergonomic helper so we can do this: - ```python - run_coroutine_in_background(func1(arg1)) - ``` - Rather than having to do this: - ```python - run_in_background(lambda: func1(arg1)) - ``` - """ - return run_in_background(lambda: coroutine) - - -T = TypeVar("T") - - -ResultT = TypeVar("ResultT") - - -def _set_context_cb(result: ResultT, context: LoggingContextOrSentinel) -> ResultT: - """A callback function which just sets the logging context""" - set_current_context(context) - return result - - -def make_deferred_yieldable(deferred: Any) -> Any: - """Make a Deferred or awaitable follow the Synapse logcontext rules. - - For Twisted Deferreds: adds callbacks to save/restore logcontext. - For native awaitables: returns an async wrapper that preserves logcontext. - The returned value is always awaitable. - """ - # Handle Twisted Deferreds - if HAS_TWISTED and isinstance(deferred, defer.Deferred): - if deferred.called and not deferred.paused: - return deferred - calling_context = set_current_context(SENTINEL_CONTEXT) - deferred.addBoth(_set_context_cb, calling_context) - return deferred - - # For native awaitables, wrap in an async function - async def _wrap() -> Any: - calling_context = set_current_context(SENTINEL_CONTEXT) + async def _wrapper() -> R: try: - return await deferred + return await coroutine finally: - set_current_context(calling_context) + set_current_context(SENTINEL_CONTEXT) - return _wrap() + task = asyncio.ensure_future(_wrapper()) + set_current_context(calling_context) + return task -def defer_to_thread( - reactor: "ISynapseReactor", f: Callable[P, R], *args: P.args, **kwargs: P.kwargs -) -> Any: +async def make_deferred_yieldable(awaitable: Any) -> Any: + """Await an awaitable while preserving the logging context. + + Clears the logcontext before awaiting (so it doesn't leak into the + event loop) and restores it after completion. """ - Calls the function `f` using a thread from the reactor's default threadpool and - returns the result as a Deferred. + if awaitable is None or not hasattr(awaitable, '__await__') and not asyncio.isfuture(awaitable) and not asyncio.iscoroutine(awaitable): + # Plain value — return directly + return awaitable - Creates a new logcontext for `f`, which is created as a child of the current - logcontext (so its CPU usage metrics will get attributed to the current - logcontext). `f` should preserve the logcontext it is given. - - The result deferred follows the Synapse logcontext rules: you should `yield` - on it. - - Args: - reactor: The reactor in whose main thread the Deferred will be invoked, - and whose threadpool we should use for the function. - - Normally this will be hs.get_reactor(). - - f: The function to call. - - args: positional arguments to pass to f. - - kwargs: keyword arguments to pass to f. - - Returns: - A Deferred which fires a callback with the result of `f`, or an - errback if `f` throws an exception. - """ - return defer_to_threadpool(reactor, reactor.getThreadPool(), f, *args, **kwargs) + calling_context = set_current_context(SENTINEL_CONTEXT) + try: + return await awaitable + finally: + set_current_context(calling_context) -def defer_to_threadpool( - reactor: "ISynapseReactor", - threadpool: ThreadPool, - f: Callable[P, R], - *args: P.args, - **kwargs: P.kwargs, -) -> Any: - """ - A wrapper for twisted.internet.threads.deferToThreadpool, which handles - logcontexts correctly. +async def defer_to_thread( + reactor: Any = None, f: Callable[P, R] = None, *args: P.args, **kwargs: P.kwargs +) -> R: + """Run a function in a thread pool executor, preserving logcontext.""" + if f is None: + # Called as defer_to_thread(f, *args) without reactor + f = reactor + reactor = None - Calls the function `f` using a thread from the given threadpool and returns - the result as a Deferred. - - Creates a new logcontext for `f`, which is created as a child of the current - logcontext (so its CPU usage metrics will get attributed to the current - logcontext). `f` should preserve the logcontext it is given. - - The result deferred follows the Synapse logcontext rules: you should `yield` - on it. - - Args: - reactor: The reactor in whose main thread the Deferred will be invoked. - Normally this will be hs.get_reactor(). - - threadpool: The threadpool to use for running `f`. Normally this will be - hs.get_reactor().getThreadPool(). - - f: The function to call. - - args: positional arguments to pass to f. - - kwargs: keyword arguments to pass to f. - - Returns: - A Deferred which fires a callback with the result of `f`, or an - errback if `f` throws an exception. - """ curr_context = current_context() if not curr_context: logger.warning( - "Calling defer_to_threadpool from sentinel context: metrics will be lost" + "Calling defer_to_thread from sentinel context: metrics will be lost" ) parent_context = None server_name = "unknown_server_from_sentinel_context" @@ -1126,157 +928,24 @@ def defer_to_threadpool( ): return f(*args, **kwargs) - return make_deferred_yieldable(threads.deferToThreadPool(reactor, threadpool, g)) - - -# =========================================================================== -# asyncio-native utility functions -# -# These provide asyncio-native equivalents of the Twisted/Deferred-based -# utility functions above. They operate on _current_context_var directly -# (NOT threading.local) because they are designed for pure asyncio code -# where ContextVar propagation into child Tasks is the desired behavior. -# -# These functions should only be used from asyncio-native code paths -# (running inside asyncio.Task), not from Twisted Deferred chains. -# =========================================================================== - - -def _native_current_context() -> "LoggingContextOrSentinel": - """Read context from ContextVar (for asyncio-native code paths only).""" - return _current_context_var.get() - - -def _native_set_current_context( - context: "LoggingContextOrSentinel", -) -> "LoggingContextOrSentinel": - """Set context in ContextVar (for asyncio-native code paths only). - - Unlike set_current_context(), this does NOT write to threading.local, - since asyncio-native code runs in Tasks with their own ContextVar copy. - """ - if context is None: - raise TypeError("'context' argument may not be None") - - current = _current_context_var.get() - - if current is not context: - rusage = get_thread_resource_usage() - current.stop(rusage) - _current_context_var.set(context) - context.start(rusage) - - return current - - -_NativeT = TypeVar("_NativeT") - - -async def make_future_yieldable( - future: "asyncio.Future[_NativeT]", -) -> _NativeT: - """Given an asyncio.Future, make it follow the Synapse logcontext rules. - - This is the asyncio-native equivalent of make_deferred_yieldable(). - - - If the future has completed, awaits it directly (logcontext unchanged). - - If the future has not yet completed, resets the logcontext to SENTINEL - before awaiting, and restores the calling logcontext when the future - completes. - - The returned coroutine can be awaited without leaking the current logcontext - into the event loop. - - NOTE: This uses ContextVar directly and should only be called from - asyncio-native code paths (inside asyncio.Task), not from Twisted code. - """ - if future.done(): - return future.result() - - # Save and clear the calling context so we don't leak it into the loop - calling_context = _native_set_current_context(SENTINEL_CONTEXT) - - try: - result = await future - except BaseException: - # Restore context before propagating the exception - _native_set_current_context(calling_context) - raise - else: - _native_set_current_context(calling_context) - return result - - -def run_coroutine_in_background_native( - coroutine: "typing.Coroutine[Any, Any, _NativeT]", -) -> "asyncio.Task[_NativeT]": - """Schedule a coroutine as a background asyncio.Task, preserving logcontext. - - This is the asyncio-native equivalent of run_coroutine_in_background(). - - The calling logcontext is restored after the task is created. When the - background task completes, the logcontext is reset to SENTINEL to avoid - leaking into the event loop. - - Returns: - The asyncio.Task running the coroutine (does NOT follow logcontext rules; - callers should use make_future_yieldable if they want to await it). - """ - calling_context = _native_current_context() - - async def _wrapper() -> _NativeT: - try: - return await coroutine - finally: - # Reset to sentinel so we don't leak context into the event loop - _native_set_current_context(SENTINEL_CONTEXT) - - task = asyncio.create_task(_wrapper()) - # Restore the calling context (create_task may have changed it) - _native_set_current_context(calling_context) - return task - - -def run_in_background_native( - f: "Callable[P, Awaitable[_NativeT]] | Callable[P, _NativeT]", - *args: "P.args", - **kwargs: "P.kwargs", -) -> "asyncio.Task[_NativeT]": - """Call a function and schedule any resulting coroutine as a background task. - - This is the asyncio-native equivalent of run_in_background(). - - Preserves the calling logcontext. When the background task completes, - resets to SENTINEL context. - """ - calling_context = _native_current_context() - try: - res = f(*args, **kwargs) - except Exception: - # Return a future that contains the exception - loop = asyncio.get_running_loop() - future: "asyncio.Future[_NativeT]" = loop.create_future() - import sys - - future.set_exception(sys.exc_info()[1]) # type: ignore[arg-type] - return future # type: ignore[return-value] - - if isinstance(res, typing.Coroutine): - return run_coroutine_in_background_native(res) - - if isinstance(res, asyncio.Task) or isinstance(res, asyncio.Future): - # Already scheduled; add sentinel-reset callback if not done - if not res.done(): - _native_set_current_context(calling_context) - - def _reset_context(f: "asyncio.Future[Any]") -> None: - _native_set_current_context(SENTINEL_CONTEXT) - - res.add_done_callback(_reset_context) - return res # type: ignore[return-value] - - # Plain value — wrap in a completed future loop = asyncio.get_running_loop() - future2: "asyncio.Future[_NativeT]" = loop.create_future() - future2.set_result(res) # type: ignore[arg-type] - return future2 # type: ignore[return-value] + return await loop.run_in_executor(None, g) + + +async def defer_to_threadpool( + reactor: Any, + threadpool: Any, + f: Callable[P, R], + *args: P.args, + **kwargs: P.kwargs, +) -> R: + """Run a function in a threadpool, preserving logcontext.""" + return await defer_to_thread(f, *args, **kwargs) + + +# Legacy aliases +_native_current_context = current_context +_native_set_current_context = set_current_context +make_future_yieldable = make_deferred_yieldable +run_coroutine_in_background_native = run_coroutine_in_background +run_in_background_native = run_in_background diff --git a/synapse/server.py b/synapse/server.py index c8d6a44d00..2e66534016 100644 --- a/synapse/server.py +++ b/synapse/server.py @@ -421,15 +421,15 @@ class HomeServer(metaclass=abc.ABCMeta): deferred = run_as_background_process(desc, self.hostname, func, *args, **kwargs) # type: ignore[untracked-background-process] self._background_processes.add(deferred) - def on_done(res: R) -> R: + def on_done(fut: Any) -> None: try: - self._background_processes.remove(deferred) - except KeyError: - # If the background process isn't being tracked anymore we can just move on. + self._background_processes.discard(deferred) + except Exception: pass - return res - deferred.addBoth(on_done) + if hasattr(deferred, 'add_done_callback'): + deferred.add_done_callback(on_done) + return deferred async def shutdown(self) -> None: diff --git a/synapse/storage/database.py b/synapse/storage/database.py index ed7ee05229..6b1e8d8e69 100644 --- a/synapse/storage/database.py +++ b/synapse/storage/database.py @@ -601,11 +601,20 @@ class DatabasePool: self._clock = hs.get_clock() self._txn_limit = database_config.config.get("txn_limit", 0) self._database_config = database_config - self._db_pool = make_pool( - reactor=hs.get_reactor(), + from synapse.storage.native_database import NativeConnectionPool + + # Check for a pre-prepared connection (used in tests with in-memory SQLite) + prepped_conn = database_config.config.get("_TEST_PREPPED_CONN") + if prepped_conn and hasattr(prepped_conn, 'conn'): + initial_conn = prepped_conn.conn # Extract raw connection from LoggingDatabaseConnection + else: + initial_conn = None + + self._db_pool = NativeConnectionPool( db_config=database_config, engine=engine, server_name=self.server_name, + initial_connection=initial_conn, ) self.updates = BackgroundUpdater(hs, self) diff --git a/synapse/storage/native_database.py b/synapse/storage/native_database.py index 697279fb27..a7362e3b21 100644 --- a/synapse/storage/native_database.py +++ b/synapse/storage/native_database.py @@ -56,6 +56,7 @@ class NativeConnectionPool: engine: BaseDatabaseEngine, server_name: str, max_workers: int = 5, + initial_connection: Connection | None = None, ) -> None: self._db_config = db_config self._engine = engine @@ -68,6 +69,12 @@ class NativeConnectionPool: if not k.startswith("cp_") } + # For in-memory SQLite, use single worker and allow cross-thread access + db_path = self._db_args.get("database", "") + if db_path == ":memory:" or db_path == "": + max_workers = 1 + self._db_args["check_same_thread"] = False + self._executor = ThreadPoolExecutor( max_workers=max_workers, thread_name_prefix=f"synapse-db-{db_config.name}", @@ -76,15 +83,25 @@ class NativeConnectionPool: # Thread-local storage for per-thread connections self._thread_local = threading.local() + # For in-memory SQLite or when an initial connection is provided, + # use a shared connection (not thread-local) + self._shared_conn: Connection | None = initial_connection + if db_path == ":memory:" or db_path == "": + self._use_shared_conn = True + else: + self._use_shared_conn = initial_connection is not None + self._closed = False def _get_connection(self) -> Connection: """Get or create a connection for the current thread. - Each thread in the pool maintains its own persistent connection. - If the connection is closed or doesn't exist, a new one is created - and initialized via the engine's on_new_connection callback. + For shared mode (in-memory SQLite), uses a single shared connection. + Otherwise, each thread gets its own persistent connection. """ + if self._use_shared_conn and self._shared_conn is not None: + return self._shared_conn + conn = getattr(self._thread_local, "conn", None) if conn is not None and not self._engine.is_connection_closed(conn): @@ -105,7 +122,10 @@ class NativeConnectionPool: ) self._engine.on_new_connection(db_conn) - self._thread_local.conn = raw_conn + if self._use_shared_conn: + self._shared_conn = raw_conn + else: + self._thread_local.conn = raw_conn return raw_conn def _reconnect(self) -> Connection: @@ -145,18 +165,9 @@ class NativeConnectionPool: conn = self._get_connection() return func(conn, *args, **kwargs) - # Submit to thread pool. Use asyncio if available, Twisted otherwise. - try: - loop = asyncio.get_running_loop() - return await loop.run_in_executor(self._executor, _inner) - except RuntimeError: - # No running asyncio loop — use Twisted's thread pool - try: - from twisted.internet import threads - return await threads.deferToThread(_inner) - except ImportError: - # Last resort: blocking call - return self._executor.submit(_inner).result() + # Run in thread pool via asyncio + loop = asyncio.get_event_loop() + return await loop.run_in_executor(self._executor, _inner) async def runInteraction( self, @@ -190,18 +201,9 @@ class NativeConnectionPool: conn.rollback() raise - # Submit to thread pool. Use asyncio if available, Twisted otherwise. - try: - loop = asyncio.get_running_loop() - return await loop.run_in_executor(self._executor, _inner) - except RuntimeError: - # No running asyncio loop — use Twisted's thread pool - try: - from twisted.internet import threads - return await threads.deferToThread(_inner) - except ImportError: - # Last resort: blocking call - return self._executor.submit(_inner).result() + # Run in thread pool via asyncio + loop = asyncio.get_event_loop() + return await loop.run_in_executor(self._executor, _inner) def close(self) -> None: """Shut down the connection pool. diff --git a/synapse/util/async_helpers.py b/synapse/util/async_helpers.py index 9bfda9f3e2..8a4dc50085 100644 --- a/synapse/util/async_helpers.py +++ b/synapse/util/async_helpers.py @@ -583,99 +583,27 @@ async def gather_optional_coroutines( ) -def timeout_deferred( +async def timeout_deferred( *, - deferred: "defer.Deferred[_T]", + deferred: Any = None, timeout: float, cancel_on_shutdown: bool = True, - clock: Clock, -) -> "defer.Deferred[_T]": - """The in built twisted `Deferred.addTimeout` fails to time out deferreds - that have a canceller that throws exceptions. This method creates a new - deferred that wraps and times out the given deferred, correctly handling - the case where the given deferred's canceller throws. + clock: Any = None, +) -> Any: + """Await an awaitable with a timeout. - (See https://twistedmatrix.com/trac/ticket/9534) - - NOTE: Unlike `Deferred.addTimeout`, this function returns a new deferred. - - NOTE: the TimeoutError raised by the resultant deferred is - twisted.internet.asyncio.TimeoutError, which is *different* to the built-in - TimeoutError, as well as various other TimeoutErrors you might have imported. + Raises asyncio.TimeoutError if the timeout expires. Args: - deferred: The Deferred to potentially timeout. - timeout: Timeout in seconds - cancel_on_shutdown: Whether this call should be tracked for cleanup during - shutdown. In general, all calls should be tracked. There may be a use case - not to track calls with a `timeout` of 0 (or similarly short) since tracking - them may result in rapid insertions and removals of tracked calls - unnecessarily. But unless a specific instance of tracking proves to be an - issue, we can just track all delayed calls. - clock: The `Clock` instance used to track delayed calls. - + deferred: The awaitable to potentially timeout. + timeout: Timeout in seconds. + cancel_on_shutdown: Ignored (kept for API compatibility). + clock: Ignored (kept for API compatibility). Returns: - A new Deferred, which will errback with asyncio.TimeoutError on timeout. + The result of the awaitable. """ - new_d: "defer.Deferred[_T]" = defer.Deferred() - - timed_out = [False] - - def time_it_out() -> None: - timed_out[0] = True - - try: - with PreserveLoggingContext(): - deferred.cancel() - except Exception: # if we throw any exception it'll break time outs - logger.exception("Canceller failed during timeout") - - # the cancel() call should have set off a chain of errbacks which - # will have errbacked new_d, but in case it hasn't, errback it now. - - if not new_d.called: - with PreserveLoggingContext(): - new_d.errback(asyncio.TimeoutError("Timed out after %gs" % (timeout,))) - - # We don't track these calls since they are short. - delayed_call = clock.call_later( - Duration(seconds=timeout), - time_it_out, - call_later_cancel_on_shutdown=cancel_on_shutdown, - ) - - def convert_cancelled(value: Failure) -> Failure: - # if the original deferred was cancelled, and our timeout has fired, then - # the reason it was cancelled was due to our timeout. Turn the CancelledError - # into a TimeoutError. - if timed_out[0] and value.check(CancelledError): - raise asyncio.TimeoutError("Timed out after %gs" % (timeout,)) - return value - - deferred.addErrback(convert_cancelled) - - def cancel_timeout(result: _T) -> _T: - # stop the pending call to cancel the deferred if it's been fired - if delayed_call.active(): - delayed_call.cancel() - return result - - deferred.addBoth(cancel_timeout) - - def success_cb(val: _T) -> None: - if not new_d.called: - with PreserveLoggingContext(): - new_d.callback(val) - - def failure_cb(val: Failure) -> None: - if not new_d.called: - with PreserveLoggingContext(): - new_d.errback(val) - - deferred.addCallbacks(success_cb, failure_cb) - - return new_d + return await asyncio.wait_for(deferred, timeout=timeout) @attr.s(slots=True, frozen=True, auto_attribs=True) @@ -700,68 +628,20 @@ def maybe_awaitable(value: Awaitable[R] | R) -> Awaitable[R]: return DoneAwaitable(value) -@overload -def delay_cancellation(awaitable: "defer.Deferred[T]") -> "defer.Deferred[T]": ... - - -@overload -def delay_cancellation(awaitable: Coroutine[Any, Any, T]) -> "defer.Deferred[T]": ... - - -@overload -def delay_cancellation(awaitable: Awaitable[T]) -> Awaitable[T]: ... - - def delay_cancellation(awaitable: Awaitable[T]) -> Awaitable[T]: - """Delay cancellation of a coroutine or `Deferred` awaitable until it resolves. + """Shield an awaitable from cancellation. - Has the same effect as `stop_cancellation`, but the returned `Deferred` will not - resolve with a `CancelledError` until the original awaitable resolves. - - Args: - deferred: The coroutine or `Deferred` to protect against cancellation. May - optionally follow the Synapse logcontext rules. - - Returns: - A new `Deferred`, which will contain the result of the original coroutine or - `Deferred`. The new `Deferred` will not propagate cancellation through to the - original coroutine or `Deferred`. - - When cancelled, the new `Deferred` will wait until the original coroutine or - `Deferred` resolves before failing with a `CancelledError`. - - The new `Deferred` will follow the Synapse logcontext rules if `awaitable` - follows the Synapse logcontext rules. Otherwise the new `Deferred` should be - wrapped with `make_deferred_yieldable`. + The returned awaitable will not propagate cancellation to the original. + Uses asyncio.shield for native asyncio support. """ - - # First, convert the awaitable into a `Deferred`. - if isinstance(awaitable, defer.Deferred): - deferred = awaitable - elif asyncio.iscoroutine(awaitable): - # Ideally we'd use `Deferred.fromCoroutine()` here, to save on redundant - # type-checking, but we'd need Twisted >= 21.2. - deferred = defer.ensureDeferred(awaitable) + if asyncio.iscoroutine(awaitable): + return asyncio.shield(asyncio.ensure_future(awaitable)) + elif asyncio.isfuture(awaitable): + return asyncio.shield(awaitable) else: - # We have no idea what to do with this awaitable. - # We assume it's already resolved, such as `DoneAwaitable`s or `Future`s from - # `make_awaitable`, and let the caller `await` it normally. + # Already resolved or unknown type — return as-is return awaitable - def handle_cancel(new_deferred: "defer.Deferred[T]") -> None: - # before the new deferred is cancelled, we `pause` it to stop the cancellation - # propagating. we then `unpause` it once the wrapped deferred completes, to - # propagate the exception. - new_deferred.pause() - with PreserveLoggingContext(): - new_deferred.errback(Failure(CancelledError())) - - deferred.addBoth(lambda _: new_deferred.unpause()) - - new_deferred: "defer.Deferred[T]" = defer.Deferred(handle_cancel) - deferred.chainDeferred(new_deferred) - return new_deferred - # =========================================================================== # Phase 0: asyncio-native parallel implementations diff --git a/synapse/util/clock.py b/synapse/util/clock.py index d22923487c..0b1a5636be 100644 --- a/synapse/util/clock.py +++ b/synapse/util/clock.py @@ -260,7 +260,14 @@ class Clock: # current context into the reactor after the function finishes. with context.PreserveLoggingContext(): d = call.start(duration.as_secs(), now=now) - d.addErrback(log_failure, "Looping call died", consumeErrors=False) + if hasattr(d, 'addErrback'): + d.addErrback(log_failure, "Looping call died", consumeErrors=False) + elif hasattr(d, 'add_done_callback'): + def _log_err(f: Any) -> None: + exc = f.exception() if hasattr(f, 'exception') else None + if exc: + logger.exception("Looping call died", exc_info=exc) + d.add_done_callback(_log_err) self._looping_calls.add(call) clock_debug_logger.debug( diff --git a/tests/__init__.py b/tests/__init__.py index b4ab59035f..3d20f1f0bb 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -19,16 +19,12 @@ # # -# Install the asyncio reactor BEFORE any other Twisted imports. -# This ensures asyncio.get_running_loop() works inside Twisted-driven code, -# enabling native asyncio primitives (Future, Event, create_task, etc.) +# Set up an asyncio event loop for tests. import asyncio as _asyncio +import nest_asyncio as _nest_asyncio -try: - from twisted.internet import asyncioreactor +_test_asyncio_loop = _asyncio.new_event_loop() +_asyncio.set_event_loop(_test_asyncio_loop) - _test_asyncio_loop = _asyncio.new_event_loop() - _asyncio.set_event_loop(_test_asyncio_loop) - asyncioreactor.install(_test_asyncio_loop) -except Exception: - pass # Already installed or Twisted not available +# Allow nested event loop calls (run_until_complete inside run_until_complete) +_nest_asyncio.apply(_test_asyncio_loop) diff --git a/tests/server.py b/tests/server.py index fa62fc3cd2..1d4a004ebd 100644 --- a/tests/server.py +++ b/tests/server.py @@ -1169,7 +1169,7 @@ def setup_test_homeserver( if PREPPED_SQLITE_DB_CONN is None: temp_engine = create_engine(database_config) PREPPED_SQLITE_DB_CONN = LoggingDatabaseConnection( - conn=sqlite3.connect(":memory:"), + conn=sqlite3.connect(":memory:", check_same_thread=False), engine=temp_engine, default_txn_name="PREPPED_CONN", server_name=server_name, diff --git a/tests/unittest.py b/tests/unittest.py index c0d43715aa..2e4433a62b 100644 --- a/tests/unittest.py +++ b/tests/unittest.py @@ -869,47 +869,63 @@ class HomeserverTestCase(TestCase): def pump(self, by: float = 0.0) -> None: """ - Pump the reactor enough that Deferreds will fire. + Pump both the test reactor and the asyncio event loop. """ + import asyncio + + # Advance Twisted's fake clock self.reactor.pump([by] * 100) + # Drive the asyncio event loop + try: + loop = asyncio.get_event_loop() + if not loop.is_closed() and not loop.is_running(): + loop.run_until_complete(asyncio.sleep(0)) + except RuntimeError: + pass + def get_success(self, d: Awaitable[TV], by: float = 0.0) -> TV: - deferred: Deferred[TV] = ensureDeferred(d) # type: ignore[arg-type] - self.pump(by=by) - return self.successResultOf(deferred) + import asyncio + + # Advance Twisted's fake clock first (for any time-dependent setup) + if by > 0: + self.reactor.pump([by] * 100) + + # Drive the coroutine to completion on the global event loop + loop = asyncio.get_event_loop() + return loop.run_until_complete(d) # type: ignore[arg-type] def get_failure( self, d: Awaitable[Any], exc: type[_ExcType], by: float = 0.0 - ) -> _TypedFailure[_ExcType]: + ) -> Any: """ - Run an awaitable and get a Failure from it. The failure must be of the type `exc`. + Run an awaitable and get a Failure from it. """ - deferred: Deferred[Any] = ensureDeferred(d) # type: ignore[arg-type] + import asyncio + + future = asyncio.ensure_future(d) # type: ignore[arg-type] + + error_holder: list[BaseException] = [] + + def _on_done(f: asyncio.Future) -> None: # type: ignore[type-arg] + try: + f.result() + except BaseException as e: + error_holder.append(e) + + future.add_done_callback(_on_done) self.pump(by) - return self.failureResultOf(deferred, exc) + + if error_holder and isinstance(error_holder[0], exc): + return Failure(error_holder[0]) + elif error_holder: + self.fail(f"Expected {exc}, got {type(error_holder[0])}: {error_holder[0]}") + else: + self.fail("Expected failure, but awaitable succeeded") def get_success_or_raise(self, d: Awaitable[TV], by: float = 0.0) -> TV: """Drive awaitable to completion and return result or raise exception.""" - deferred: Deferred[TV] = ensureDeferred(d) # type: ignore[arg-type] - - results: list = [] - deferred.addBoth(results.append) - - self.pump(by=by) - - if not results: - self.fail( - "Success result expected on {!r}, found no result instead".format( - deferred - ) - ) - - result = results[0] - - if isinstance(result, Failure): - result.raiseException() - - return result + return self.get_success(d, by=by) def register_user( self,