feat: deterministic time steps

Author: brian14708

src/duron/_core/session.py

@@ -333,10 +333,15 @@ def __init__(
 
     async def _resume(self) -> bool:
         recvd_msgs: set[str] = set()
+        ws: WaitSet | None = None
         async for o, entry in self._log.stream():
             ts = entry["ts"]
+            while ws and ws.timer and ws.timer < ts:
+                self._now_us = max(self._now_us, ws.timer)
+                ws = await self._step()
+
             self._now_us = max(self._now_us, ts)
-            _ = await self._step()
+            ws = await self._step()
             if is_entry(entry):
                 if entry["source"] == "task":
                     if not self._handle_message(o, entry):
@@ -345,7 +350,7 @@ async def _resume(self) -> bool:
                     recvd_msgs.add(entry["id"])
                 else:
                     _ = self._handle_message(o, entry)
-                _ = await self._step()
+                ws = await self._step()
             while self._pending_msg:
                 id_ = self._pending_msg[-1]["id"]
                 if id_ not in recvd_msgs:
@@ -402,8 +407,12 @@ async def _run(self) -> _T_co:
                     await self._send_traces()
                 else:
                     await waitset.block(self._now_us)
-                _ = await self._step()
-                self._now_us = max(self._now_us + 1, time.time_ns() // 1_000)
+                waitset = await self._step()
+                now = time.time_ns() // 1_000
+                if waitset and waitset.timer and waitset.timer < now:
+                    self._now_us = max(self._now_us, waitset.timer)
+                else:
+                    self._now_us = max(self._now_us, now)
 
             # cleanup
             self._loop.close()

src/duron/_core/signal.py

@@ -68,6 +68,7 @@ async def __aenter__(self) -> None:
         if task is None:
             return
         assert task.get_loop() == self._loop
+
         if task not in self._tasks:
             val = _SignalState(depth=0, triggered=False)
             self._tasks[task] = val

src/duron/loop.py

@@ -147,9 +147,8 @@ def call_at(
         *args: Unpack[_Ts],
         context: Context | None = None,
     ) -> asyncio.TimerHandle:
-        th = asyncio.TimerHandle(
-            int(when * 1e6), callback, args, loop=self, context=context
-        )
+        t = int(when * 1e6 - self._now_us) + self._now_us
+        th = asyncio.TimerHandle(t, callback, args, loop=self, context=context)
         heappush(self._timers, th)
         if asyncio.get_running_loop() is self._host:
             self._event.set()
@@ -164,11 +163,7 @@ def call_later(
         context: Context | None = None,
     ) -> asyncio.TimerHandle:
         th = asyncio.TimerHandle(
-            self.time_us() + int(delay * 1e6),
-            callback,
-            args,
-            loop=self,
-            context=context,
+            self._now_us + int(delay * 1e6), callback, args, loop=self, context=context
         )
         heappush(self._timers, th)
         if asyncio.get_running_loop() is self._host:
@@ -244,15 +239,14 @@ def _poll(self, now: int) -> int | None:
 
     def poll_completion(self, task: Future[_T]) -> WaitSet | None:
         assert asyncio.get_running_loop() is self._host
-        now = self.time_us()
         self._event.clear()
 
         # hot path - inline task context switch
         if prev_task := tasks.current_task():
             tasks._leave_task(self._host, prev_task)  # noqa: SLF001
         events._set_running_loop(self)  # noqa: SLF001
         try:
-            next_deadline = self._poll(now)
+            next_deadline = self._poll(self._now_us)
             if task.done():
                 return None
             added, self._added = self._added, []
@@ -316,12 +310,12 @@ def close(self) -> None:
             tasks._leave_task(self._host, prev_task)  # noqa: SLF001
         events._set_running_loop(self)  # noqa: SLF001
         try:
-            _ = self._poll(self.time_us())
+            _ = self._poll(self._now_us)
             to_cancel = (*tasks.all_tasks(), *self._ops.values())
             for t in to_cancel:
                 t.cancel()
             _ = asyncio.gather(*to_cancel, return_exceptions=True)
-            _ = self._poll(self.time_us())
+            _ = self._poll(self._now_us)
             for task in to_cancel:
                 if task.cancelled() or not task.done():
                     continue

tests/test_signal.py

@@ -60,12 +60,16 @@ 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:
+                cnt = 0
                 await asyncio.sleep(t)
                 try:
                     async with signal:
-                        await asyncio.sleep(9999)
+                        while True:
+                            cnt += 1
+                            await asyncio.sleep(0.0001)
+                            await ctx.run(asyncio.sleep, 0.001)
                 except SignalInterrupt as e:
-                    values.append(cast("int", e.value))
+                    values.append(cast("int", e.value) + cnt)
                     if len(values) > 10:
                         return values
 
@@ -110,7 +114,6 @@ async def tracker() -> int:
                     if sys.version_info >= (3, 11):
                         async with asyncio.timeout(0.01):
                             await ctx.run(work)
-                            await ctx.time_ns()
                             values += 1
             return values