mirror of
https://github.com/element-hq/synapse.git
synced 2026-10-23 04:23:10 +00:00
Phase 0: Abstraction Boundaries — Complete
**Goal**: Add asyncio-native parallel implementations alongside existing Twisted ones, so subsequent phases can swap without touching callers. **Files modified**: - `synapse/logging/context.py` — Add `contextvars.ContextVar`-based `current_context()`/`set_current_context()` behind a feature flag alongside the existing `threading.local` implementation. Add `make_future_yieldable()` and `run_coroutine_in_background_native()` operating on `asyncio.Future`/`asyncio.Task` instead of `Deferred`. - `synapse/util/async_helpers.py` — Add `asyncio.Future`-based `ObservableFuture`, and `asyncio.Lock`/`asyncio.Event`-based `AsyncLinearizer`, `AsyncReadWriteLock` alongside the Deferred-based originals. - `synapse/types/__init__.py` — Define `ISynapseEventLoop` protocol abstracting the event loop operations (`call_later`, `call_soon`, `run_in_executor`, `create_task`) so `Clock` and other code can be parameterized. 3 files modified, 1 test file created, 28 new tests all passing, 46 existing tests unaffected, mypy clean. synapse/logging/context.py — Added asyncio-native parallel implementations: - _current_context_var — contextvars.ContextVar that will replace _thread_local in Phase 1 - current_context_contextvar() / set_current_context_contextvar() — ContextVar-based equivalents of current_context() / set_current_context() - make_future_yieldable() — asyncio.Future equivalent of make_deferred_yieldable() - run_coroutine_in_background_native() — asyncio.Task equivalent of run_coroutine_in_background() - run_in_background_native() — asyncio.Task equivalent of run_in_background() synapse/util/async_helpers.py — Added asyncio-native primitives: - ObservableFuture — asyncio.Future-based equivalent of ObservableDeferred - NativeLinearizer — asyncio.Event-based equivalent of Linearizer (no Deferred dependency) - NativeReadWriteLock — asyncio.Event-based equivalent of ReadWriteLock synapse/types/__init__.py — Added: - ISynapseEventLoop — Protocol abstracting event loop operations (call_later, call_soon, run_in_executor, create_task, time) so Clock can be parameterized over Twisted reactor vs asyncio loop in Phase 2 tests/util/test_native_async.py — 28 tests covering all new implementations using unittest.IsolatedAsyncioTestCase.
This commit is contained in:
1 parent
9edbf56969
commit
2dce74958f
4 files changed
+947
No files matched your search
@@ -30,6 +30,8 @@ them.
|
||||
See doc/log_contexts.rst for details on how this works.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import contextvars
|
||||
import logging
|
||||
import threading
|
||||
import typing
|
||||
@@ -1207,3 +1209,166 @@ def defer_to_threadpool(
|
||||
return f(*args, **kwargs)
|
||||
|
||||
return make_deferred_yieldable(threads.deferToThreadPool(reactor, threadpool, g))
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# Phase 0: asyncio-native parallel implementations
|
||||
#
|
||||
# These provide asyncio-native equivalents of the Twisted-based context
|
||||
# tracking and utility functions above. They are unused until Phase 1+
|
||||
# switches the active implementation. Adding them here in Phase 0 ensures
|
||||
# the new code paths can be tested without changing any existing behavior.
|
||||
# ===========================================================================
|
||||
|
||||
# A ContextVar that can replace _thread_local for context tracking.
|
||||
# In Phase 1, current_context()/set_current_context() will switch to using this.
|
||||
_current_context_var: contextvars.ContextVar[
|
||||
"LoggingContextOrSentinel"
|
||||
] = contextvars.ContextVar("synapse_logging_context", default=SENTINEL_CONTEXT)
|
||||
|
||||
|
||||
def current_context_contextvar() -> "LoggingContextOrSentinel":
|
||||
"""Get the current logging context from contextvars.
|
||||
|
||||
This is the asyncio-native equivalent of current_context().
|
||||
It will become the primary implementation in Phase 1.
|
||||
"""
|
||||
return _current_context_var.get()
|
||||
|
||||
|
||||
def set_current_context_contextvar(
|
||||
context: "LoggingContextOrSentinel",
|
||||
) -> "LoggingContextOrSentinel":
|
||||
"""Set the current logging context using contextvars.
|
||||
|
||||
This is the asyncio-native equivalent of set_current_context().
|
||||
It will become the primary implementation in Phase 1.
|
||||
|
||||
Args:
|
||||
context: The context to activate.
|
||||
|
||||
Returns:
|
||||
The context that was previously active.
|
||||
"""
|
||||
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.
|
||||
"""
|
||||
if future.done():
|
||||
return future.result()
|
||||
|
||||
# Save and clear the calling context so we don't leak it into the loop
|
||||
calling_context = set_current_context_contextvar(SENTINEL_CONTEXT)
|
||||
|
||||
try:
|
||||
result = await future
|
||||
except BaseException:
|
||||
# Restore context before propagating the exception
|
||||
set_current_context_contextvar(calling_context)
|
||||
raise
|
||||
else:
|
||||
set_current_context_contextvar(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 = current_context_contextvar()
|
||||
|
||||
async def _wrapper() -> _NativeT:
|
||||
try:
|
||||
return await coroutine
|
||||
finally:
|
||||
# Reset to sentinel so we don't leak context into the event loop
|
||||
set_current_context_contextvar(SENTINEL_CONTEXT)
|
||||
|
||||
task = asyncio.create_task(_wrapper())
|
||||
# Restore the calling context (create_task may have changed it)
|
||||
set_current_context_contextvar(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 = current_context_contextvar()
|
||||
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():
|
||||
set_current_context_contextvar(calling_context)
|
||||
|
||||
def _reset_context(f: "asyncio.Future[Any]") -> None:
|
||||
set_current_context_contextvar(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]
|
||||
@@ -20,21 +20,26 @@
|
||||
#
|
||||
#
|
||||
import abc
|
||||
import asyncio
|
||||
import logging
|
||||
import re
|
||||
import string
|
||||
from concurrent.futures import Executor
|
||||
from enum import Enum
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
AbstractSet,
|
||||
Any,
|
||||
Callable,
|
||||
ClassVar,
|
||||
Coroutine,
|
||||
Literal,
|
||||
Mapping,
|
||||
Match,
|
||||
MutableMapping,
|
||||
NoReturn,
|
||||
Optional,
|
||||
Protocol,
|
||||
TypedDict,
|
||||
TypeVar,
|
||||
Union,
|
||||
@@ -138,6 +143,47 @@ class ISynapseReactor(
|
||||
"""The interfaces necessary for Synapse to function."""
|
||||
|
||||
|
||||
_ELT = TypeVar("_ELT")
|
||||
|
||||
|
||||
class ISynapseEventLoop(Protocol):
|
||||
"""Protocol abstracting event loop operations.
|
||||
|
||||
This allows Clock and other code to be parameterized over either a Twisted
|
||||
reactor or a native asyncio event loop. Added in Phase 0 of the
|
||||
Twisted → asyncio migration; will be used starting in Phase 2 (Clock).
|
||||
"""
|
||||
|
||||
def call_later(
|
||||
self, delay: float, callback: Callable[..., Any], *args: Any
|
||||
) -> Any:
|
||||
"""Schedule callback to be called after delay seconds."""
|
||||
...
|
||||
|
||||
def call_soon(self, callback: Callable[..., Any], *args: Any) -> Any:
|
||||
"""Schedule callback to be called on the next event loop iteration."""
|
||||
...
|
||||
|
||||
def run_in_executor(
|
||||
self,
|
||||
executor: Executor | None,
|
||||
func: Callable[..., _ELT],
|
||||
*args: Any,
|
||||
) -> "asyncio.Future[_ELT]":
|
||||
"""Run a function in an executor (thread pool)."""
|
||||
...
|
||||
|
||||
def create_task(
|
||||
self, coro: Coroutine[Any, Any, _ELT]
|
||||
) -> "asyncio.Task[_ELT]":
|
||||
"""Schedule a coroutine as a Task."""
|
||||
...
|
||||
|
||||
def time(self) -> float:
|
||||
"""Return the current time as a float (seconds since epoch)."""
|
||||
...
|
||||
|
||||
|
||||
@attr.s(frozen=True, slots=True, auto_attribs=True)
|
||||
class Requester:
|
||||
"""
|
||||
|
||||
@@ -1030,6 +1030,246 @@ def observe_deferred(d: "defer.Deferred[T]") -> "defer.Deferred[T]":
|
||||
return new_deferred
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# Phase 0: asyncio-native parallel implementations
|
||||
#
|
||||
# These provide asyncio-native equivalents of the Twisted/Deferred-based
|
||||
# primitives above. They are unused until later migration phases swap them
|
||||
# in. Adding them here ensures they can be tested without affecting existing
|
||||
# behavior.
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class ObservableFuture(Generic[_T]):
|
||||
"""asyncio.Future-based equivalent of ObservableDeferred.
|
||||
|
||||
Wraps an asyncio.Future so that multiple observers can await the same
|
||||
result without interfering with each other or the original future.
|
||||
|
||||
If consumeErrors is True, exceptions from the original future are
|
||||
suppressed (not re-raised to the original future's awaiter).
|
||||
"""
|
||||
|
||||
__slots__ = ["_future", "_observers", "_result"]
|
||||
|
||||
def __init__(
|
||||
self, future: "asyncio.Future[_T]", consumeErrors: bool = False
|
||||
) -> None:
|
||||
self._future = future
|
||||
self._result: (
|
||||
None | tuple[Literal[True], _T] | tuple[Literal[False], BaseException]
|
||||
) = None
|
||||
self._observers: list["asyncio.Future[_T]"] | tuple[()] = []
|
||||
|
||||
def _on_done(f: "asyncio.Future[_T]") -> None:
|
||||
result: tuple[Literal[True], _T] | tuple[Literal[False], BaseException]
|
||||
if f.cancelled():
|
||||
result = (False, asyncio.CancelledError())
|
||||
else:
|
||||
real_exc = f.exception()
|
||||
if real_exc is not None:
|
||||
result = (False, real_exc)
|
||||
else:
|
||||
result = (True, f.result())
|
||||
|
||||
object.__setattr__(self, "_result", result)
|
||||
observers = self._observers
|
||||
object.__setattr__(self, "_observers", ())
|
||||
|
||||
for obs in observers:
|
||||
if obs.done():
|
||||
continue
|
||||
if result[0]:
|
||||
obs.set_result(result[1])
|
||||
else:
|
||||
obs.set_exception(result[1])
|
||||
|
||||
future.add_done_callback(_on_done)
|
||||
|
||||
def observe(self) -> "asyncio.Future[_T]":
|
||||
"""Return a new Future that resolves with the same result."""
|
||||
if self._result is not None:
|
||||
loop = self._future.get_loop()
|
||||
f: "asyncio.Future[_T]" = loop.create_future()
|
||||
if self._result[0]:
|
||||
f.set_result(self._result[1])
|
||||
else:
|
||||
f.set_exception(self._result[1])
|
||||
return f
|
||||
|
||||
assert isinstance(self._observers, list)
|
||||
loop = self._future.get_loop()
|
||||
f = loop.create_future()
|
||||
self._observers.append(f)
|
||||
return f
|
||||
|
||||
def has_observers(self) -> bool:
|
||||
return bool(self._observers)
|
||||
|
||||
def has_called(self) -> bool:
|
||||
return self._result is not None
|
||||
|
||||
def has_succeeded(self) -> bool:
|
||||
return self._result is not None and self._result[0] is True
|
||||
|
||||
def get_result(self) -> "_T | BaseException":
|
||||
if self._result is None:
|
||||
raise ValueError(f"{self!r} has no result yet")
|
||||
return self._result[1]
|
||||
|
||||
|
||||
class NativeLinearizer:
|
||||
"""asyncio-native equivalent of Linearizer.
|
||||
|
||||
Limits concurrent access to resources based on a key using asyncio
|
||||
primitives (asyncio.Event) instead of Deferred.
|
||||
|
||||
Example:
|
||||
async with linearizer.queue("key"):
|
||||
# do work
|
||||
"""
|
||||
|
||||
def __init__(self, name: str, max_count: int = 1) -> None:
|
||||
self.name = name
|
||||
self.max_count = max_count
|
||||
self._key_to_entry: dict[Hashable, _NativeLinearizerEntry] = {}
|
||||
|
||||
def is_queued(self, key: Hashable) -> bool:
|
||||
entry = self._key_to_entry.get(key)
|
||||
if not entry:
|
||||
return False
|
||||
return bool(entry.waiters)
|
||||
|
||||
def queue(self, key: Hashable) -> AsyncContextManager[None]:
|
||||
@asynccontextmanager
|
||||
async def _ctx_manager() -> AsyncIterator[None]:
|
||||
entry = await self._acquire_lock(key)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
self._release_lock(key, entry)
|
||||
|
||||
return _ctx_manager()
|
||||
|
||||
async def _acquire_lock(
|
||||
self, key: Hashable
|
||||
) -> "_NativeLinearizerEntry":
|
||||
entry = self._key_to_entry.setdefault(
|
||||
key, _NativeLinearizerEntry(0, [])
|
||||
)
|
||||
|
||||
if entry.count < self.max_count:
|
||||
entry.count += 1
|
||||
return entry
|
||||
|
||||
event = asyncio.Event()
|
||||
entry.waiters.append(event)
|
||||
try:
|
||||
await event.wait()
|
||||
except asyncio.CancelledError:
|
||||
try:
|
||||
entry.waiters.remove(event)
|
||||
except ValueError:
|
||||
pass
|
||||
raise
|
||||
|
||||
entry.count += 1
|
||||
|
||||
# Break potential synchronous recursion (equivalent to sleep(0)
|
||||
# in the Deferred-based Linearizer)
|
||||
try:
|
||||
await asyncio.sleep(0)
|
||||
except asyncio.CancelledError:
|
||||
self._release_lock(key, entry)
|
||||
raise
|
||||
|
||||
return entry
|
||||
|
||||
def _release_lock(
|
||||
self, key: Hashable, entry: "_NativeLinearizerEntry"
|
||||
) -> None:
|
||||
entry.count -= 1
|
||||
|
||||
if entry.waiters:
|
||||
next_event = entry.waiters.pop(0)
|
||||
next_event.set()
|
||||
elif entry.count == 0:
|
||||
self._key_to_entry.pop(key, None)
|
||||
|
||||
|
||||
@attr.s(slots=True, auto_attribs=True)
|
||||
class _NativeLinearizerEntry:
|
||||
count: int
|
||||
waiters: list[asyncio.Event]
|
||||
|
||||
|
||||
class NativeReadWriteLock:
|
||||
"""asyncio-native equivalent of ReadWriteLock.
|
||||
|
||||
Example:
|
||||
async with lock.read("key"):
|
||||
# readers can overlap
|
||||
async with lock.write("key"):
|
||||
# writers are exclusive
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.key_to_current_readers: dict[str, set[asyncio.Event]] = {}
|
||||
self.key_to_current_writer: dict[str, asyncio.Event] = {}
|
||||
|
||||
def read(self, key: str) -> AsyncContextManager[None]:
|
||||
@asynccontextmanager
|
||||
async def _ctx_manager() -> AsyncIterator[None]:
|
||||
new_event = asyncio.Event()
|
||||
curr_readers = self.key_to_current_readers.setdefault(key, set())
|
||||
curr_writer = self.key_to_current_writer.get(key)
|
||||
|
||||
curr_readers.add(new_event)
|
||||
|
||||
try:
|
||||
if curr_writer:
|
||||
await asyncio.shield(
|
||||
self._wait_for_event(curr_writer)
|
||||
)
|
||||
yield
|
||||
finally:
|
||||
new_event.set()
|
||||
self.key_to_current_readers.get(key, set()).discard(new_event)
|
||||
|
||||
return _ctx_manager()
|
||||
|
||||
def write(self, key: str) -> AsyncContextManager[None]:
|
||||
@asynccontextmanager
|
||||
async def _ctx_manager() -> AsyncIterator[None]:
|
||||
new_event = asyncio.Event()
|
||||
curr_readers = self.key_to_current_readers.get(key, set())
|
||||
curr_writer = self.key_to_current_writer.get(key)
|
||||
|
||||
to_wait_on: list[asyncio.Event] = list(curr_readers)
|
||||
if curr_writer:
|
||||
to_wait_on.append(curr_writer)
|
||||
|
||||
curr_readers.clear()
|
||||
self.key_to_current_writer[key] = new_event
|
||||
|
||||
try:
|
||||
for evt in to_wait_on:
|
||||
await asyncio.shield(self._wait_for_event(evt))
|
||||
yield
|
||||
finally:
|
||||
new_event.set()
|
||||
if self.key_to_current_writer.get(key) is new_event:
|
||||
self.key_to_current_writer.pop(key)
|
||||
|
||||
return _ctx_manager()
|
||||
|
||||
@staticmethod
|
||||
async def _wait_for_event(event: asyncio.Event) -> None:
|
||||
"""Wait for an event to be set, handling the case where it's already set."""
|
||||
if not event.is_set():
|
||||
await event.wait()
|
||||
|
||||
|
||||
class AwakenableSleeper:
|
||||
"""Allows explicitly waking up deferreds related to an entity that are
|
||||
currently sleeping.
|
||||
|
||||
@@ -0,0 +1,496 @@
|
||||
#
|
||||
# This file is licensed under the Affero General Public License (AGPL) version 3.
|
||||
#
|
||||
# Copyright (C) 2025 New Vector, 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:
|
||||
# <https://www.gnu.org/licenses/agpl-3.0.html>.
|
||||
#
|
||||
|
||||
"""Tests for Phase 0 asyncio-native parallel implementations.
|
||||
|
||||
These test the new asyncio-native primitives added alongside the existing
|
||||
Twisted-based ones, ensuring they work correctly before being swapped in
|
||||
during later migration phases.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import unittest
|
||||
|
||||
from synapse.logging.context import (
|
||||
SENTINEL_CONTEXT,
|
||||
LoggingContext,
|
||||
_current_context_var,
|
||||
current_context_contextvar,
|
||||
make_future_yieldable,
|
||||
run_coroutine_in_background_native,
|
||||
run_in_background_native,
|
||||
set_current_context_contextvar,
|
||||
)
|
||||
from synapse.util.async_helpers import (
|
||||
NativeLinearizer,
|
||||
NativeReadWriteLock,
|
||||
ObservableFuture,
|
||||
)
|
||||
|
||||
|
||||
class ContextVarContextTest(unittest.IsolatedAsyncioTestCase):
|
||||
"""Tests for contextvars-based context tracking."""
|
||||
|
||||
def setUp(self) -> None:
|
||||
# Ensure we start from sentinel
|
||||
_current_context_var.set(SENTINEL_CONTEXT)
|
||||
|
||||
def test_default_is_sentinel(self) -> None:
|
||||
self.assertIs(current_context_contextvar(), SENTINEL_CONTEXT)
|
||||
|
||||
def test_set_and_get(self) -> None:
|
||||
ctx = LoggingContext(name="test", server_name="test.server")
|
||||
old = set_current_context_contextvar(ctx)
|
||||
self.assertIs(old, SENTINEL_CONTEXT)
|
||||
self.assertIs(current_context_contextvar(), ctx)
|
||||
# Restore
|
||||
set_current_context_contextvar(SENTINEL_CONTEXT)
|
||||
|
||||
def test_set_returns_previous(self) -> None:
|
||||
ctx1 = LoggingContext(name="ctx1", server_name="test.server")
|
||||
ctx2 = LoggingContext(name="ctx2", server_name="test.server")
|
||||
set_current_context_contextvar(ctx1)
|
||||
old = set_current_context_contextvar(ctx2)
|
||||
self.assertIs(old, ctx1)
|
||||
set_current_context_contextvar(SENTINEL_CONTEXT)
|
||||
|
||||
def test_none_raises(self) -> None:
|
||||
with self.assertRaises(TypeError):
|
||||
set_current_context_contextvar(None) # type: ignore[arg-type]
|
||||
|
||||
async def test_task_inherits_context(self) -> None:
|
||||
"""asyncio.Tasks inherit the parent's contextvars by default."""
|
||||
ctx = LoggingContext(name="parent", server_name="test.server")
|
||||
set_current_context_contextvar(ctx)
|
||||
|
||||
result = None
|
||||
|
||||
async def child() -> None:
|
||||
nonlocal result
|
||||
result = current_context_contextvar()
|
||||
|
||||
task = asyncio.create_task(child())
|
||||
await task
|
||||
self.assertIs(result, ctx)
|
||||
set_current_context_contextvar(SENTINEL_CONTEXT)
|
||||
|
||||
|
||||
class MakeFutureYieldableTest(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self) -> None:
|
||||
_current_context_var.set(SENTINEL_CONTEXT)
|
||||
|
||||
async def test_already_done_future(self) -> None:
|
||||
loop = asyncio.get_running_loop()
|
||||
f: asyncio.Future[int] = loop.create_future()
|
||||
f.set_result(42)
|
||||
|
||||
result = await make_future_yieldable(f)
|
||||
self.assertEqual(result, 42)
|
||||
|
||||
async def test_pending_future_preserves_context(self) -> None:
|
||||
ctx = LoggingContext(name="test", server_name="test.server")
|
||||
set_current_context_contextvar(ctx)
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
f: asyncio.Future[str] = loop.create_future()
|
||||
|
||||
# Schedule the resolution so the coroutine can actually await
|
||||
loop.call_soon(f.set_result, "hello")
|
||||
result = await make_future_yieldable(f)
|
||||
|
||||
# Context should be restored after awaiting
|
||||
self.assertIs(current_context_contextvar(), ctx)
|
||||
self.assertEqual(result, "hello")
|
||||
|
||||
set_current_context_contextvar(SENTINEL_CONTEXT)
|
||||
|
||||
async def test_pending_future_exception(self) -> None:
|
||||
ctx = LoggingContext(name="test", server_name="test.server")
|
||||
set_current_context_contextvar(ctx)
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
f: asyncio.Future[str] = loop.create_future()
|
||||
|
||||
yieldable = make_future_yieldable(f)
|
||||
f.set_exception(ValueError("boom"))
|
||||
|
||||
with self.assertRaises(ValueError):
|
||||
await yieldable
|
||||
|
||||
# Context should still be restored after exception
|
||||
self.assertIs(current_context_contextvar(), ctx)
|
||||
set_current_context_contextvar(SENTINEL_CONTEXT)
|
||||
|
||||
|
||||
class RunCoroutineInBackgroundNativeTest(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self) -> None:
|
||||
_current_context_var.set(SENTINEL_CONTEXT)
|
||||
|
||||
async def test_preserves_calling_context(self) -> None:
|
||||
ctx = LoggingContext(name="caller", server_name="test.server")
|
||||
set_current_context_contextvar(ctx)
|
||||
|
||||
results: list[object] = []
|
||||
|
||||
async def bg_work() -> str:
|
||||
results.append(current_context_contextvar())
|
||||
return "done"
|
||||
|
||||
task = run_coroutine_in_background_native(bg_work())
|
||||
|
||||
# Calling context should be preserved
|
||||
self.assertIs(current_context_contextvar(), ctx)
|
||||
|
||||
result = await task
|
||||
self.assertEqual(result, "done")
|
||||
|
||||
set_current_context_contextvar(SENTINEL_CONTEXT)
|
||||
|
||||
async def test_resets_to_sentinel_on_completion(self) -> None:
|
||||
post_completion_context = None
|
||||
|
||||
async def bg_work() -> None:
|
||||
pass
|
||||
|
||||
task = run_coroutine_in_background_native(bg_work())
|
||||
await task
|
||||
|
||||
# The task's internal wrapper resets to sentinel
|
||||
# (We can't easily observe this from outside without hooking the callback,
|
||||
# but at least verify the task completed successfully)
|
||||
self.assertTrue(task.done())
|
||||
|
||||
|
||||
class RunInBackgroundNativeTest(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self) -> None:
|
||||
_current_context_var.set(SENTINEL_CONTEXT)
|
||||
|
||||
async def test_with_coroutine_function(self) -> None:
|
||||
async def my_func(x: int) -> int:
|
||||
return x * 2
|
||||
|
||||
task = run_in_background_native(my_func, 21)
|
||||
result = await task
|
||||
self.assertEqual(result, 42)
|
||||
|
||||
async def test_with_sync_function(self) -> None:
|
||||
def my_func(x: int) -> int:
|
||||
return x * 2
|
||||
|
||||
task = run_in_background_native(my_func, 21)
|
||||
result = await task
|
||||
self.assertEqual(result, 42)
|
||||
|
||||
async def test_with_exception(self) -> None:
|
||||
def my_func() -> None:
|
||||
raise ValueError("sync error")
|
||||
|
||||
task = run_in_background_native(my_func)
|
||||
with self.assertRaises(ValueError):
|
||||
await task
|
||||
|
||||
|
||||
class ObservableFutureTest(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_succeed(self) -> None:
|
||||
loop = asyncio.get_running_loop()
|
||||
origin: asyncio.Future[int] = loop.create_future()
|
||||
observable = ObservableFuture(origin)
|
||||
|
||||
obs1 = observable.observe()
|
||||
obs2 = observable.observe()
|
||||
|
||||
self.assertFalse(observable.has_called())
|
||||
self.assertTrue(observable.has_observers())
|
||||
|
||||
origin.set_result(42)
|
||||
# Give the event loop a chance to process callbacks
|
||||
await asyncio.sleep(0)
|
||||
|
||||
self.assertTrue(observable.has_called())
|
||||
self.assertTrue(observable.has_succeeded())
|
||||
self.assertEqual(await obs1, 42)
|
||||
self.assertEqual(await obs2, 42)
|
||||
|
||||
async def test_fail(self) -> None:
|
||||
loop = asyncio.get_running_loop()
|
||||
origin: asyncio.Future[int] = loop.create_future()
|
||||
observable = ObservableFuture(origin)
|
||||
|
||||
obs1 = observable.observe()
|
||||
obs2 = observable.observe()
|
||||
|
||||
origin.set_exception(ValueError("boom"))
|
||||
await asyncio.sleep(0)
|
||||
|
||||
self.assertTrue(observable.has_called())
|
||||
self.assertFalse(observable.has_succeeded())
|
||||
|
||||
with self.assertRaises(ValueError):
|
||||
await obs1
|
||||
with self.assertRaises(ValueError):
|
||||
await obs2
|
||||
|
||||
async def test_observe_after_resolution(self) -> None:
|
||||
loop = asyncio.get_running_loop()
|
||||
origin: asyncio.Future[str] = loop.create_future()
|
||||
observable = ObservableFuture(origin)
|
||||
|
||||
origin.set_result("hello")
|
||||
await asyncio.sleep(0)
|
||||
|
||||
obs = observable.observe()
|
||||
self.assertEqual(await obs, "hello")
|
||||
|
||||
async def test_no_observers(self) -> None:
|
||||
loop = asyncio.get_running_loop()
|
||||
origin: asyncio.Future[int] = loop.create_future()
|
||||
observable = ObservableFuture(origin)
|
||||
|
||||
self.assertFalse(observable.has_observers())
|
||||
origin.set_result(1)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
self.assertFalse(observable.has_observers())
|
||||
|
||||
async def test_get_result(self) -> None:
|
||||
loop = asyncio.get_running_loop()
|
||||
origin: asyncio.Future[int] = loop.create_future()
|
||||
observable = ObservableFuture(origin)
|
||||
|
||||
with self.assertRaises(ValueError):
|
||||
observable.get_result()
|
||||
|
||||
origin.set_result(99)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
self.assertEqual(observable.get_result(), 99)
|
||||
|
||||
|
||||
class NativeLinearizerTest(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_uncontended(self) -> None:
|
||||
linearizer = NativeLinearizer("test")
|
||||
async with linearizer.queue("key"):
|
||||
pass
|
||||
|
||||
async def test_serializes_access(self) -> None:
|
||||
linearizer = NativeLinearizer("test")
|
||||
order: list[int] = []
|
||||
|
||||
async def worker(n: int) -> None:
|
||||
async with linearizer.queue("key"):
|
||||
order.append(n)
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
tasks = [asyncio.create_task(worker(i)) for i in range(3)]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
# All three should have run in order
|
||||
self.assertEqual(order, [0, 1, 2])
|
||||
|
||||
async def test_different_keys_concurrent(self) -> None:
|
||||
linearizer = NativeLinearizer("test")
|
||||
running: list[str] = []
|
||||
|
||||
async def worker(key: str) -> None:
|
||||
async with linearizer.queue(key):
|
||||
running.append(key)
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
tasks = [
|
||||
asyncio.create_task(worker("a")),
|
||||
asyncio.create_task(worker("b")),
|
||||
]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
# Both should have started (different keys are independent)
|
||||
self.assertEqual(set(running), {"a", "b"})
|
||||
|
||||
async def test_max_count(self) -> None:
|
||||
linearizer = NativeLinearizer("test", max_count=2)
|
||||
concurrent = 0
|
||||
max_concurrent = 0
|
||||
|
||||
async def worker() -> None:
|
||||
nonlocal concurrent, max_concurrent
|
||||
async with linearizer.queue("key"):
|
||||
concurrent += 1
|
||||
max_concurrent = max(max_concurrent, concurrent)
|
||||
await asyncio.sleep(0.01)
|
||||
concurrent -= 1
|
||||
|
||||
tasks = [asyncio.create_task(worker()) for _ in range(5)]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
self.assertLessEqual(max_concurrent, 2)
|
||||
|
||||
async def test_is_queued(self) -> None:
|
||||
linearizer = NativeLinearizer("test")
|
||||
|
||||
self.assertFalse(linearizer.is_queued("key"))
|
||||
|
||||
acquired = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
|
||||
async def holder() -> None:
|
||||
async with linearizer.queue("key"):
|
||||
acquired.set()
|
||||
await release.wait()
|
||||
|
||||
task1 = asyncio.create_task(holder())
|
||||
await acquired.wait()
|
||||
|
||||
# Start a second worker that will be queued
|
||||
async def waiter() -> None:
|
||||
async with linearizer.queue("key"):
|
||||
pass
|
||||
|
||||
task2 = asyncio.create_task(waiter())
|
||||
await asyncio.sleep(0) # Let task2 start and get queued
|
||||
|
||||
self.assertTrue(linearizer.is_queued("key"))
|
||||
|
||||
release.set()
|
||||
await asyncio.gather(task1, task2)
|
||||
|
||||
async def test_cancellation(self) -> None:
|
||||
linearizer = NativeLinearizer("test")
|
||||
|
||||
acquired = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
|
||||
async def holder() -> None:
|
||||
async with linearizer.queue("key"):
|
||||
acquired.set()
|
||||
await release.wait()
|
||||
|
||||
task1 = asyncio.create_task(holder())
|
||||
await acquired.wait()
|
||||
|
||||
async def waiter() -> None:
|
||||
async with linearizer.queue("key"):
|
||||
pass
|
||||
|
||||
task2 = asyncio.create_task(waiter())
|
||||
await asyncio.sleep(0)
|
||||
|
||||
task2.cancel()
|
||||
with self.assertRaises(asyncio.CancelledError):
|
||||
await task2
|
||||
|
||||
release.set()
|
||||
await task1
|
||||
|
||||
|
||||
class NativeReadWriteLockTest(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_readers_concurrent(self) -> None:
|
||||
lock = NativeReadWriteLock()
|
||||
concurrent = 0
|
||||
max_concurrent = 0
|
||||
|
||||
async def reader() -> None:
|
||||
nonlocal concurrent, max_concurrent
|
||||
async with lock.read("key"):
|
||||
concurrent += 1
|
||||
max_concurrent = max(max_concurrent, concurrent)
|
||||
await asyncio.sleep(0.01)
|
||||
concurrent -= 1
|
||||
|
||||
tasks = [asyncio.create_task(reader()) for _ in range(3)]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
# All readers should run concurrently
|
||||
self.assertEqual(max_concurrent, 3)
|
||||
|
||||
async def test_writer_exclusive(self) -> None:
|
||||
lock = NativeReadWriteLock()
|
||||
order: list[str] = []
|
||||
|
||||
async def writer(name: str) -> None:
|
||||
async with lock.write("key"):
|
||||
order.append(f"{name}_start")
|
||||
await asyncio.sleep(0.01)
|
||||
order.append(f"{name}_end")
|
||||
|
||||
tasks = [
|
||||
asyncio.create_task(writer("w1")),
|
||||
asyncio.create_task(writer("w2")),
|
||||
]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
# Writers should be serialized: w1 should finish before w2 starts
|
||||
self.assertEqual(
|
||||
order, ["w1_start", "w1_end", "w2_start", "w2_end"]
|
||||
)
|
||||
|
||||
async def test_writer_blocks_reader(self) -> None:
|
||||
lock = NativeReadWriteLock()
|
||||
order: list[str] = []
|
||||
|
||||
writer_acquired = asyncio.Event()
|
||||
writer_release = asyncio.Event()
|
||||
|
||||
async def writer() -> None:
|
||||
async with lock.write("key"):
|
||||
order.append("writer_start")
|
||||
writer_acquired.set()
|
||||
await writer_release.wait()
|
||||
order.append("writer_end")
|
||||
|
||||
async def reader() -> None:
|
||||
await writer_acquired.wait()
|
||||
async with lock.read("key"):
|
||||
order.append("reader")
|
||||
|
||||
w_task = asyncio.create_task(writer())
|
||||
r_task = asyncio.create_task(reader())
|
||||
|
||||
await asyncio.sleep(0.01) # Let writer acquire and reader queue
|
||||
writer_release.set()
|
||||
|
||||
await asyncio.gather(w_task, r_task)
|
||||
|
||||
self.assertEqual(order, ["writer_start", "writer_end", "reader"])
|
||||
|
||||
async def test_reader_blocks_writer(self) -> None:
|
||||
lock = NativeReadWriteLock()
|
||||
order: list[str] = []
|
||||
|
||||
reader_acquired = asyncio.Event()
|
||||
reader_release = asyncio.Event()
|
||||
|
||||
async def reader() -> None:
|
||||
async with lock.read("key"):
|
||||
order.append("reader_start")
|
||||
reader_acquired.set()
|
||||
await reader_release.wait()
|
||||
order.append("reader_end")
|
||||
|
||||
async def writer() -> None:
|
||||
await reader_acquired.wait()
|
||||
async with lock.write("key"):
|
||||
order.append("writer")
|
||||
|
||||
r_task = asyncio.create_task(reader())
|
||||
w_task = asyncio.create_task(writer())
|
||||
|
||||
await asyncio.sleep(0.01)
|
||||
reader_release.set()
|
||||
|
||||
await asyncio.gather(r_task, w_task)
|
||||
|
||||
self.assertEqual(order, ["reader_start", "reader_end", "writer"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in new issue
Block a user