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:
Matthew Hodgson committed 2026-03-20 21:11:46 +00:00
1 parent 9edbf56969
commit 2dce74958f
4 files changed
+947

No files matched your search

+165
View File
@@ -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]
+46
View File
@@ -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:
"""
+240
View File
@@ -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.
+496
View File
@@ -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()