feat: delta waitset

Author: brian14708

src/duron/_core/invoke.py

@@ -299,7 +299,6 @@ class _InvokeRun:
         "_loop",
         "_now",
         "_pending_msg",
-        "_prev_ops",
         "_running",
         "_stream_manager",
         "_task",
@@ -322,7 +321,6 @@ def __init__(
         self._running: bool = False
         self._lease: bytes | None = None
         self._pending_msg: list[Entry] = []
-        self._prev_ops: set[str] = set()
         self._now = 0
         self._stream_manager = StreamManager(watchers)
         self._tracer: Tracer | None = Tracer.current()
@@ -409,20 +407,12 @@ async def _step(self) -> WaitSet | None:
 
         while True:
             result = self._loop.poll_completion(self._task)
-            if result is None:
+            if result is None or not result.added:
                 return result
 
-            new_ops = False
-            for s in result.ops:
+            for s in result.added:
                 sid = s.id
-                if sid not in self._prev_ops:
-                    await self.enqueue_op(sid, s)
-                    new_ops = True
-            if new_ops or len(self._prev_ops) != len(result.ops):
-                self._prev_ops.clear()
-                self._prev_ops.update(s.id for s in result.ops)
-            if not new_ops:
-                return result
+                await self.enqueue_op(sid, s)
 
     async def _send_traces(self, *, flush: bool = False) -> None:
         if not self._tracer:

src/duron/_loop.py

@@ -17,7 +17,7 @@
 if TYPE_CHECKING:
     import sys
     from asyncio.futures import Future
-    from collections.abc import Callable, Coroutine, Generator
+    from collections.abc import Callable, Coroutine, Generator, Sequence
     from contextvars import Context
 
     _T = TypeVar("_T")
@@ -46,7 +46,7 @@ def __init__(
 
 @dataclass(slots=True)
 class WaitSet:
-    ops: list[OpFuture]
+    added: Sequence[OpFuture]
     timer: int | None
     event: asyncio.Event
 
@@ -74,6 +74,7 @@ class _TaskCtx:
 
 class EventLoop(asyncio.AbstractEventLoop):
     __slots__: tuple[str, ...] = (
+        "_added",
         "_closed",
         "_ctx",
         "_event",
@@ -98,6 +99,7 @@ def __init__(self, host: asyncio.AbstractEventLoop) -> None:
         self._closed: bool = False
         self._event: asyncio.Event = asyncio.Event()  # loop = _host
         self._timers: list[asyncio.TimerHandle] = []
+        self._added: list[OpFuture] = []
 
     def set_key(self, key: bytes) -> None:
         # Derive a fixed-length key from the provided key
@@ -223,12 +225,17 @@ def poll_completion(self, task: Future[_T]) -> WaitSet | None:
 
             if task.done():
                 return None
-            return WaitSet(ops=[*self._ops.values()], timer=deadline, event=self._event)
+
+            added, self._added = self._added, []
+            return WaitSet(added=added, timer=deadline, event=self._event)
         finally:
             events._set_running_loop(self._host)  # noqa: SLF001
             if prev_task:
                 tasks._enter_task(self._host, prev_task)  # noqa: SLF001
 
+    def pending_ops(self) -> Sequence[OpFuture]:
+        return tuple(self._ops.values())
+
     def create_op(self, params: object, *, external: bool = False) -> OpFuture:
         if external:
             id_ = random_id()
@@ -237,6 +244,7 @@ def create_op(self, params: object, *, external: bool = False) -> OpFuture:
             id_ = self.generate_op_id()
         op_fut = OpFuture(id_, params, self)
         self._ops[id_] = op_fut
+        self._added.append(op_fut)
         return op_fut
 
     @overload
@@ -276,7 +284,7 @@ def close(self) -> None:
                 p = self.poll_completion(op)
                 if p is not None:
                     logger.warning(
-                        "Event loop closed with pending operations: %r", p.ops
+                        "Event loop closed with pending operations: %r", self._ops
                     )
         while self._timers:
             th = heappop(self._timers)

tests/test_loop.py

@@ -57,10 +57,11 @@ def tick(n: int, expect: set[int] | None) -> list[str] | None:
         waitset = loop.poll_completion(tsk)
         if expect:
             assert waitset
-            assert len(waitset.ops) == n
-            assert {o.params for o in waitset.ops}.issubset(expect)
-            ids.update(o.id for o in waitset.ops)
-            return [o.id for o in waitset.ops]
+            pending = loop.pending_ops()
+            assert len(pending) == n
+            assert {o.params for o in pending}.issubset(expect)
+            ids.update(o.id for o in pending)
+            return [o.id for o in pending]
         assert waitset is None
         return None