feat: handle more stream edge-cases

Author: brian14708

pyproject.toml

@@ -79,12 +79,12 @@ preview = true
 future-annotations = true
 
 [tool.ruff.lint.per-file-ignores]
-"tests/**/*.py" = ["S101", "S311", "PLC2701", "PLR2004", "RUF029"]
+"tests/**/*.py" = ["S311", "PLC2701", "PLR2004", "RUF029"]
 "examples/**/*.py" = ["T201", "S311"]
 
 [tool.ruff.lint.isort]
 extra-standard-library = ["typing_extensions"]
-split-on-trailing-comma=false
+split-on-trailing-comma = false
 
 [tool.ruff.lint.flake8-type-checking]
 runtime-evaluated-base-classes = ["typing_extensions.TypedDict"]

src/duron/_core/session.py

@@ -31,7 +31,7 @@
     create_op,
 )
 from duron._core.signal import Signal
-from duron._core.stream import OpWriter, Stream, StreamWriter, create_buffer_stream
+from duron._core.stream import Stream, StreamWriter, create_buffer_stream
 from duron._core.stream_manager import StreamManager
 from duron._core.task_manager import TaskError, TaskManager
 from duron._core.utils import decode_error, encode_error
@@ -754,7 +754,7 @@ async def open_stream(
             return self._streams.pop(name)
 
         sid = await self._stream_manager.wait_stream(name)
-        w: OpWriter[Any] = OpWriter(sid, self._loop)
+        w: StreamWriter[Any] = StreamWriter(sid, self._loop)
         return w
 
     def is_future_pending(self, future_id: str) -> bool:

src/duron/_core/signal.py

@@ -7,13 +7,12 @@
 from typing_extensions import Any, TypeVar, final, override
 
 from duron._core.ops import StreamCreate, create_op
-from duron._core.stream import OpWriter
+from duron._core.stream import StreamWriter
 
 if TYPE_CHECKING:
     from types import TracebackType
 
     from duron._core.ops import OpMetadata
-    from duron._core.stream import StreamWriter
     from duron.loop import EventLoop
     from duron.typing import TypeHint
 
@@ -120,5 +119,5 @@ async def create_signal(
     sid = await create_op(
         loop, StreamCreate(dtype=dtype, name=name, observer=s, metadata=metadata)
     )
-    w: OpWriter[_T] = OpWriter(sid, loop)
+    w: StreamWriter[_T] = StreamWriter(sid, loop)
     return (s, w)

src/duron/_core/stream.py

@@ -8,9 +8,8 @@
 from asyncio.exceptions import CancelledError
 from collections import deque
 from collections.abc import AsyncIterable
-from contextlib import AbstractAsyncContextManager
 from typing import TYPE_CHECKING, Concatenate, Generic, cast
-from typing_extensions import Any, ParamSpec, Protocol, TypeVar, final, override
+from typing_extensions import Any, ParamSpec, TypeVar, final, override
 
 from duron._core.ops import (
     Barrier,
@@ -21,7 +20,7 @@
     StreamEmit,
     create_op,
 )
-from duron.loop import EventLoop, wrap_future
+from duron.loop import EventLoop, LoopClosedError, wrap_future
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Awaitable, Callable, Sequence
@@ -63,32 +62,10 @@ def reason(self) -> Exception | None:
         return cast("Exception | None", self.__cause__)
 
 
-class StreamWriter(
-    AbstractAsyncContextManager["StreamWriter[_T_contra]"], Protocol, Generic[_T_contra]
-):
+@final
+class StreamWriter(Generic[_T_contra]):
     """Protocol for writing values to a stream."""
 
-    async def send(self, value: _T_contra, /) -> None:
-        """Send a value to the stream.
-
-        Args:
-            value: The value to send to stream consumers.
-
-        """
-        ...
-
-    async def close(self, error: Exception | None = None, /) -> None:
-        """Close the stream, optionally with an error.
-
-        Args:
-            error: Optional exception to signal an error condition to consumers.
-
-        """
-        ...
-
-
-@final
-class OpWriter(Generic[_T_contra]):
     __slots__ = ("_closed", "_loop", "_stream_id")
 
     def __init__(self, stream_id: str, loop: EventLoop) -> None:
@@ -97,11 +74,35 @@ def __init__(self, stream_id: str, loop: EventLoop) -> None:
         self._closed = False
 
     async def send(self, value: _T_contra, /) -> None:
+        """Send a value to the stream.
+
+        Raises:
+            RuntimeError: If the stream is already closed.
+
+        Args:
+            value: The value to send to stream consumers.
+
+        """
+        if self._closed:
+            msg = "Cannot send to a closed stream"
+            raise RuntimeError(msg)
         await wrap_future(
             create_op(self._loop, StreamEmit(stream_id=self._stream_id, value=value))
         )
 
     async def close(self, exception: Exception | None = None, /) -> None:
+        """Close the stream, optionally with an error.
+
+        Raises:
+            RuntimeError: If the stream is already closed.
+
+        Args:
+            exception: Optional exception to signal an error condition to consumers.
+
+        """
+        if self._closed:
+            msg = "Cannot send to a closed stream"
+            raise RuntimeError(msg)
         await wrap_future(
             create_op(
                 self._loop, StreamClose(stream_id=self._stream_id, exception=exception)
@@ -120,14 +121,15 @@ async def __aexit__(
     ) -> None:
         if self._closed:
             return
-        if not exc_value:
-            await self.close()
-        elif isinstance(exc_value, Exception):
-            await self.close(exc_value)
-        else:
-            await self.close(
-                Exception(f"StreamWriter exited with exception: {exc_value}")
-            )
+        with contextlib.suppress(LoopClosedError):
+            if not exc_value:
+                await self.close()
+            elif isinstance(exc_value, Exception):
+                await self.close(exc_value)
+            else:
+                await self.close(
+                    Exception(f"StreamWriter exited with exception: {exc_value}")
+                )
 
 
 class Stream(ABC, AsyncIterable[_T], Generic[_T]):
@@ -189,7 +191,7 @@ async def create_stream(
     sid = await create_op(
         loop, StreamCreate(dtype=dtype, observer=w, name=name, metadata=metadata)
     )
-    writer: OpWriter[_T] = OpWriter(sid, loop)
+    writer: StreamWriter[_T] = StreamWriter(sid, loop)
     return (s, writer)
 
 
@@ -206,7 +208,7 @@ def __init__(self) -> None:
         self._loop: asyncio.AbstractEventLoop | None = None
         self._event: asyncio.Event | None = None
         self._buffer: deque[tuple[int, _T | StreamClosed]] = deque()
-        self._cursor: int = -1
+        self._cursor: int = 0
 
     @final
     @override
@@ -216,34 +218,41 @@ async def next(self, *, block: bool) -> Sequence[_T]:
             self._event = asyncio.Event()
 
         if not block:
-            return await self._next_nowait()
-
-        while True:
-            _ = await self._event.wait()
             self._event.clear()
-            if it := await self._next_nowait():
-                return it
-
-    async def _next_nowait(self) -> Sequence[_T]:
-        if not self._loop:
-            self._loop = asyncio.get_running_loop()
+            begin, end = await self._next_cursor()
+            return self._pop(begin, end)
 
-        if isinstance(self._loop, EventLoop):
-
-            def cb(f: asyncio.Future[tuple[int, int]]) -> None:
-                if not f.cancelled():
-                    offset, _ = f.result()
-                    self._cursor = max(self._cursor, offset)
-
-            begin = self._cursor
-            op = create_op(self._loop, Barrier())
-            op.add_done_callback(cb)
-            end, _ = await asyncio.shield(op)
-            self._cursor = max(self._cursor, end)
-        else:
+        while True:
+            try:
+                _ = await self._event.wait()
+            finally:
+                self._event.clear()
+
+            begin, end = await self._next_cursor()
+            items = self._pop(begin, end)
+            if items:
+                return items
+
+    async def _next_cursor(self) -> tuple[int, int | None]:
+        if not isinstance(self._loop, EventLoop):
+            return (0, None)
+
+        def cb(f: asyncio.Future[tuple[int, int]]) -> None:
+            if not f.cancelled():
+                offset, _ = f.result()
+                self._cursor = max(self._cursor, offset)
+
+        begin = self._cursor
+        op = create_op(self._loop, Barrier())
+        op.add_done_callback(cb)
+        end, _ = await asyncio.shield(op)
+        self._cursor = max(self._cursor, end)
+        return (begin, end)
+
+    def _pop(self, begin: int, end: int | None) -> Sequence[_T]:
+        if end is None:
             if not self._buffer:
                 return ()
-            begin = 0
             end = self._buffer[-1][0] + 1
 
         result: list[_T] = []
@@ -328,7 +337,7 @@ async def run_stateful(
     stream: _StatefulStream[_U, _T] = _StatefulStream(
         reducer, fn, initial, *args, **kwargs
     )
-    sink: StreamWriter[_U] = OpWriter(
+    sink: StreamWriter[_U] = StreamWriter(
         await create_op(
             loop,
             StreamCreate(

src/duron/loop.py

@@ -354,7 +354,7 @@ async def shutdown_asyncgens(self) -> None:
         pass
 
     @override
-    async def shutdown_default_executor(self) -> None:
+    async def shutdown_default_executor(self, timeout: float | None = None) -> None:
         pass
 
     def _timer_handle_cancelled(self, _th: asyncio.TimerHandle) -> None:

tests/test_stream.py

@@ -10,14 +10,18 @@
 
 from duron import (
     Context,
+    Provided,
     Reducer,
     Session,
+    Signal,
     SignalInterrupt,
+    Stream,
     StreamClosed,
     durable,
     effect,
 )
 from duron.contrib.storage import MemoryLogStorage
+from duron.loop import LoopClosedError
 
 if TYPE_CHECKING:
     from duron import StreamWriter
@@ -254,7 +258,79 @@ async def trigger_signal() -> None:
         result2 = await (await t.start(activity)).result()
     b = len(await log.entries())
     for chunk1, chunk2 in zip(result1, result2, strict=True):
-        length = min(len(chunk1), len(chunk2))
-        assert len(chunk2) >= len(chunk1)
-        assert chunk1[:length] == chunk2[:length]
+        assert len(chunk2) == len(chunk1)
+        assert chunk1 == chunk2
+    assert a == b
+
+
+@pytest.mark.asyncio
+async def test_external_stream_signal_timing() -> None:
+    @durable()
+    async def activity(
+        _ctx: Context,
+        input_stream: Stream[int] = Provided,
+        interrupt_signal: Signal[str] = Provided,
+    ) -> list[list[list[int]]]:
+        results: list[list[list[int]]] = []
+
+        # Consume with signal context
+        for i in range(1000):
+            batch: list[list[int]] = []
+            try:
+                async with interrupt_signal:
+                    # Consume values until interrupted
+                    while True:
+                        if i % 3 == 0:
+                            batch.append(list(await input_stream.next(block=False)))
+                            await asyncio.sleep(0.01)
+                        else:
+                            batch.append(list(await input_stream.next(block=True)))
+            except StreamClosed:
+                results.append(batch)
+            except SignalInterrupt as e:
+                results.append(batch)
+                if e.value is True:
+                    break
+        return results
+
+    # First run
+    log = MemoryLogStorage()
+    async with Session(log) as sess:
+        run = await sess.start(activity)
+        input_stream = await run.open_stream("input_stream", "w")
+        interrupt_signal = await run.open_stream("interrupt_signal", "w")
+
+        # Emit values continuously
+        async def emitter() -> None:
+            async with input_stream as w:
+                for i in range(500):
+                    try:
+                        await w.send(i)
+                    except LoopClosedError:
+                        break
+                    await asyncio.sleep(random.random() * 0.001)
+
+        # Trigger signal at specific time
+        async def trigger_signal() -> None:
+            async with interrupt_signal as w:
+                for i in range(5):
+                    await asyncio.sleep(0.015)
+                    await w.send(i == 4)
+
+        result1, _, _ = await asyncio.gather(
+            run.result(),
+            asyncio.create_task(emitter()),
+            asyncio.create_task(trigger_signal()),
+        )
+
+    # Replay
+    a = len(await log.entries())
+    async with Session(log) as sess:
+        result2 = await (await sess.start(activity)).result()
+    b = len(await log.entries())
+
+    # Verify results match
+    for chunk1, chunk2 in zip(result1, result2, strict=True):
+        assert len(chunk2) == len(chunk1)
+        assert chunk1 == chunk2
     assert a == b