feat: unify signal and stream writer

Author: brian14708

src/duron/__init__.py

@@ -4,7 +4,6 @@
 from duron._core.invoke import invoke as invoke
 from duron._core.signal import Signal as Signal
 from duron._core.signal import SignalInterrupt as SignalInterrupt
-from duron._core.signal import SignalWriter as SignalWriter
 from duron._core.stream import Stream as Stream
 from duron._core.stream import StreamClosed as StreamClosed
 from duron._core.stream import StreamOp as StreamOp

src/duron/_core/context.py

@@ -31,7 +31,7 @@
     from collections.abc import Awaitable, Callable, Coroutine, Mapping
     from contextlib import AbstractAsyncContextManager
 
-    from duron._core.signal import Signal, SignalWriter
+    from duron._core.signal import Signal
     from duron._core.stream import Stream, StreamWriter
     from duron._loop import EventLoop
     from duron.typing import TypeHint
@@ -189,7 +189,7 @@ async def create_signal(
         *,
         name: str | None = None,
         labels: Mapping[str, str] | None = None,
-    ) -> tuple[Signal[_T], SignalWriter[_T]]:
+    ) -> tuple[Signal[_T], StreamWriter[_T]]:
         """Create a new signal within the context.
 
         Args:

src/duron/_core/signal.py

@@ -2,10 +2,9 @@
 
 import asyncio
 import sys
-from asyncio.exceptions import CancelledError
 from collections import deque
-from typing import TYPE_CHECKING, Generic, cast
-from typing_extensions import Any, Protocol, TypeVar, final
+from typing import TYPE_CHECKING, Final, Generic, cast
+from typing_extensions import Any, TypeVar, final, override
 
 from duron._core.ops import Barrier, StreamClose, StreamCreate, StreamEmit, create_op
 from duron._loop import wrap_future
@@ -14,10 +13,11 @@
     from types import TracebackType
 
     from duron._core.ops import OpAnnotations
+    from duron._core.stream import StreamWriter
     from duron._loop import EventLoop
     from duron.typing._hint import TypeHint
 
-_In = TypeVar("_In", contravariant=True)  # noqa: PLC0105
+_InT = TypeVar("_InT", contravariant=True)  # noqa: PLC0105
 
 
 class SignalInterrupt(Exception):  # noqa: N818
@@ -27,53 +27,47 @@ class SignalInterrupt(Exception):  # noqa: N818
         value: The value passed to the signal trigger that caused the interrupt.
     """
 
-    def __init__(self, *args: object, value: object) -> None:
-        super().__init__(*args)
-        self.value: object = value
+    def __init__(self, value: object) -> None:
+        super().__init__()
+        self.value = value
 
-
-class SignalWriter(Protocol, Generic[_In]):
-    """Protocol for writing values to a signal to interrupt operations."""
-
-    async def trigger(self, value: _In, /) -> None:
-        """Trigger the signal with a value, interrupting active operations.
-
-        Args:
-            value: The value to send with the interrupt.
-        """
-        ...
-
-    async def close(self, /) -> None:
-        """Close the signal stream, preventing further triggers."""
-        ...
+    @override
+    def __repr__(self) -> str:
+        return f"SignalInterrupt(value={self.value!r})"
 
 
 @final
-class _Writer(Generic[_In]):
+class SignalWriter(Generic[_InT]):
+    """Object for writing values to a signal to interrupt operations."""
+
     __slots__ = ("_loop", "_stream_id")
 
     def __init__(self, stream_id: str, loop: EventLoop) -> None:
         self._stream_id = stream_id
         self._loop = loop
 
-    async def trigger(self, value: _In, /) -> None:
+    async def send(self, value: _InT) -> None:
+        """Trigger the signal with a value, interrupting active operations.
+
+        Args:
+            value: The value to send with the interrupt.
+        """
         await wrap_future(
             create_op(self._loop, StreamEmit(stream_id=self._stream_id, value=value))
         )
 
-    async def close(self, /) -> None:
+    async def close(self, exc: Exception | None = None) -> None:
+        """Close the signal stream, preventing further triggers."""
         await wrap_future(
-            create_op(
-                self._loop, StreamClose(stream_id=self._stream_id, exception=None)
-            )
+            create_op(self._loop, StreamClose(stream_id=self._stream_id, exception=exc))
         )
 
 
-_SENTINAL = object()
+_SIGNAL_TRIGGER: Final = object()
 
 
 @final
-class Signal(Generic[_In]):
+class Signal(Generic[_InT]):
     """Signal context manager for interruptible operations.
 
     Signal provides a mechanism for interrupting in-progress operations. When used
@@ -91,9 +85,9 @@ class Signal(Generic[_In]):
 
     def __init__(self, loop: EventLoop) -> None:
         self._loop = loop
-        # task -> [offset, refcnt]
+        # task -> [offset, stack depth]
         self._tasks: dict[asyncio.Task[Any], tuple[int, int]] = {}
-        self._trigger: deque[tuple[int, _In]] = deque()
+        self._trigger: deque[tuple[int, _InT]] = deque()
 
     async def __aenter__(self) -> None:
         task = asyncio.current_task()
@@ -104,8 +98,8 @@ async def __aenter__(self) -> None:
         for toffset, value in self._trigger:
             if toffset > offset:
                 raise SignalInterrupt(value=value)
-        _, refcnt = self._tasks.get(task, (0, 0))
-        self._tasks[task] = (offset, refcnt + 1)
+        _, depth = self._tasks.get(task, (0, -1))
+        self._tasks[task] = (offset, depth + 1)
         self._flush()
 
     async def __aexit__(
@@ -117,43 +111,49 @@ async def __aexit__(
         task = asyncio.current_task()
         if task is None:
             return
-        offset_start, refcnt = self._tasks.pop(task)
         offset_end = await create_op(self._loop, Barrier())
-        if refcnt > 1:
-            self._tasks[task] = (offset_end, refcnt - 1)
+
+        offset_start, depth = self._tasks.pop(task)
+        if depth > 0:
+            self._tasks[task] = (offset_end, depth - 1)
         for toffset, value in self._trigger:
-            if offset_start < toffset < offset_end:
-                if sys.version_info >= (3, 11) and exc_type is CancelledError:
-                    assert exc_value  # noqa: S101
-                    assert exc_value.args[0] is _SENTINAL  # noqa: S101
+            if (
+                offset_start < toffset < offset_end
+                and exc_type is asyncio.CancelledError
+                and (args := cast("asyncio.CancelledError", exc_value).args)
+                and args[0] is _SIGNAL_TRIGGER
+            ):
+                if sys.version_info >= (3, 11):
                     _ = task.uncancel()
+                self._flush()
                 raise SignalInterrupt(value=value)
 
-    def on_next(self, offset: int, value: _In) -> None:
+    def on_next(self, offset: int, value: _InT) -> None:
         self._trigger.append((offset, value))
-        for t, (toffset, _refcnt) in self._tasks.items():
+        for t, (toffset, _depth) in self._tasks.items():
             if toffset < offset:
-                _ = self._loop.call_soon(t.cancel, _SENTINAL)
+                _ = self._loop.call_soon(t.cancel, _SIGNAL_TRIGGER)
 
     def on_close(self, _offset: int, _exc: Exception | None) -> None:
         pass
 
     def _flush(self) -> None:
-        assert len(self._tasks) > 0  # noqa: S101
-        min_offset = min(offset for offset, _ in self._tasks.values())
+        if not self._tasks:
+            self._trigger.clear()
+            return
+        min_offset = min((offset for offset, _ in self._tasks.values()))
         while self._trigger and self._trigger[0][0] < min_offset:
             _ = self._trigger.popleft()
 
 
 async def create_signal(
-    loop: EventLoop, dtype: TypeHint[_In], annotations: OpAnnotations
-) -> tuple[Signal[_In], SignalWriter[_In]]:
-    assert asyncio.get_running_loop() is loop  # noqa: S101
-    s: Signal[_In] = Signal(loop)
+    loop: EventLoop, dtype: TypeHint[_InT], annotations: OpAnnotations
+) -> tuple[Signal[_InT], StreamWriter[_InT]]:
+    s: Signal[_InT] = Signal(loop)
     sid = await create_op(
         loop,
         StreamCreate(
             dtype=dtype, observer=cast("Signal[object]", s), annotations=annotations
         ),
     )
-    return (s, _Writer(sid, loop))
+    return (s, SignalWriter(sid, loop))

tests/test_signal.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 from typing import cast
 
 import pytest
 
-from duron import Context, SignalInterrupt, durable, invoke
+from duron import Context, Signal, SignalInterrupt, durable, invoke
 from duron.contrib.storage import MemoryLogStorage
+from duron.typing._hint import Provided
 
 
 @pytest.mark.asyncio
@@ -15,13 +17,13 @@ async def test_signal() -> None:
     async def activity(ctx: Context) -> list[int]:
         signal, handle = await ctx.create_signal(int)
 
-        await handle.trigger(1)
+        await handle.send(1)
 
         async def trigger() -> None:
             await asyncio.sleep(0.1)
-            await handle.trigger(2)
+            await handle.send(2)
             await asyncio.sleep(0.1)
-            await handle.trigger(3)
+            await handle.send(3)
 
         t = asyncio.create_task(trigger())
 
@@ -49,3 +51,47 @@ async def trigger() -> None:
     async with invoke(activity, log) as t:
         await t.resume()
         assert await t.wait() == [2, 3]
+
+
+@pytest.mark.asyncio
+async def test_signal_timing() -> None:
+    @durable()
+    async def activity(ctx: Context, s: Signal[int] = Provided) -> list[list[int]]:
+        async def tracker(signal: Signal[int], t: float) -> list[int]:
+            values: list[int] = []
+            while True:
+                await asyncio.sleep(t)
+                try:
+                    async with signal:
+                        await asyncio.sleep(9999)
+                except SignalInterrupt as e:
+                    values.append(cast("int", e.value))
+                    if len(values) > 10:
+                        return values
+
+        rnd = ctx.random()
+        return await asyncio.gather(*[
+            (tracker(s, rnd.random() * 0.01)) for _ in range(4)
+        ])
+
+    log = MemoryLogStorage()
+    async with invoke(activity, log) as t:
+        signal = t.open_stream("s", "w")
+        await t.start()
+
+        async def push() -> None:
+            i = 0
+            while True:
+                i += 1
+                await asyncio.sleep(0.001)
+                await signal.send(i)
+
+        pusher = asyncio.create_task(push())
+        a = await t.wait()
+        pusher.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await pusher
+
+    async with invoke(activity, log) as t:
+        await t.resume()
+        assert await t.wait() == a