diff --git a/synapse/logging/context.py b/synapse/logging/context.py index b6535be388..d65620a88d 100644 --- a/synapse/logging/context.py +++ b/synapse/logging/context.py @@ -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] diff --git a/synapse/types/__init__.py b/synapse/types/__init__.py index 02889795bb..5cc470829a 100644 --- a/synapse/types/__init__.py +++ b/synapse/types/__init__.py @@ -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: """ diff --git a/synapse/util/async_helpers.py b/synapse/util/async_helpers.py index 1c3cd48a9c..792343c18c 100644 --- a/synapse/util/async_helpers.py +++ b/synapse/util/async_helpers.py @@ -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. diff --git a/tests/util/test_native_async.py b/tests/util/test_native_async.py new file mode 100644 index 0000000000..3ed452c3b8 --- /dev/null +++ b/tests/util/test_native_async.py @@ -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: +# . +# + +"""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()