feat: support cancel of OpFuture

Author: brian14708

src/duron/event_loop.py

@@ -5,7 +5,6 @@
 import contextvars
 import heapq
 import logging
-import threading
 from asyncio import AbstractEventLoop, Handle, Task, TimerHandle, events
 from collections import deque
 from dataclasses import dataclass
@@ -86,7 +85,6 @@ def __init__(self, seed: bytes) -> None:
         self._closed: bool = False
         self._event: asyncio.Event = asyncio.Event()
 
-        self._lock: threading.Lock = threading.Lock()
         self._timers: list[TimerHandle] = []
 
     def _generate_id(self) -> bytes:
@@ -116,19 +114,9 @@ def call_soon_threadsafe(
         callback: Callable[[Unpack[_Ts]], object],
         *args: Unpack[_Ts],
         context: Context | None = None,
-        task_id: bytes | None = None,
     ) -> Handle:
-        self._event.set()
-        h = TimerHandle(
-            0,
-            callback,
-            args,
-            self,
-            context=self._context_with_task_id(context, task_id=task_id),
-        )
-        with self._lock:
-            heapq.heappush(self._timers, h)
-        return h
+        # the loop is not thread-safe, so we just call call_soon
+        return self.call_soon(callback, *args, context=context)
 
     @override
     def call_at(
@@ -143,10 +131,9 @@ def call_at(
             callback,
             args,
             loop=self,
-            context=self._context_with_task_id(context),
+            context=context,
         )
-        with self._lock:
-            heapq.heappush(self._timers, th)
+        heapq.heappush(self._timers, th)
         return th
 
     @override
@@ -179,7 +166,7 @@ def create_task(
         context: Context | None = None,
         **kwargs: Any,
     ) -> Task[_T]:
-        ctx = self._context_with_task_id(context)
+        ctx = self._context_new_task(context, self._generate_id())
         return ctx.run(
             cast("type[Task[_T]]", Task), coro, name=name, loop=self, **kwargs
         )
@@ -193,17 +180,16 @@ def poll_completion(self, task: Future[_T]) -> WaitSet | None:
             deadline: float | None = None
             while True:
                 deadline = None
-                with self._lock:
-                    while self._timers:
-                        ht = self._timers[0]
-                        if ht.cancelled():
-                            _ = heapq.heappop(self._timers)
-                        elif ht.when() <= now:
-                            _ = heapq.heappop(self._timers)
-                            self._ready.append(ht)
-                        else:
-                            deadline = ht.when()
-                            break
+                while self._timers:
+                    ht = self._timers[0]
+                    if ht.cancelled():
+                        _ = heapq.heappop(self._timers)
+                    elif ht.when() <= now:
+                        _ = heapq.heappop(self._timers)
+                        self._ready.append(ht)
+                    else:
+                        deadline = ht.when()
+                        break
 
                 if not self._ready:
                     break
@@ -213,7 +199,7 @@ def poll_completion(self, task: Future[_T]) -> WaitSet | None:
                         continue
                     try:
                         h._run()
-                    except BaseException as exc:
+                    except Exception as exc:
                         self.call_exception_handler({
                             "message": "exception in callback",
                             "exception": exc,
@@ -237,20 +223,20 @@ def create_op(self, params: object) -> OpFuture:
         return s
 
     @overload
-    def post_completion_threadsafe(
+    def post_completion(
         self,
         id: bytes,
         *,
         result: object,
     ) -> None: ...
     @overload
-    def post_completion_threadsafe(
+    def post_completion(
         self,
         id: bytes,
         *,
         exception: BaseException,
     ) -> None: ...
-    def post_completion_threadsafe(
+    def post_completion(
         self,
         id: bytes,
         *,
@@ -259,10 +245,15 @@ def post_completion_threadsafe(
     ) -> None:
         if op := self._ops.pop(id, None):
             tid = _mix_id(op.id, -1)
-            if exception is not None:
-                _ = self.call_soon_threadsafe(op.set_exception, exception, task_id=tid)
-            else:
-                _ = self.call_soon_threadsafe(op.set_result, result, task_id=tid)
+            self._event.set()
+            self._ready.append(
+                Handle(
+                    op.set_result if exception is None else op.set_exception,
+                    [(result if exception is None else exception)],
+                    self,
+                    context=self._context_new_task(None, tid),
+                )
+            )
 
     @override
     def is_closed(self) -> bool:
@@ -319,12 +310,8 @@ async def shutdown_default_executor(self):
     def _timer_handle_cancelled(self, _th: TimerHandle) -> None:
         pass
 
-    def _context_with_task_id(
-        self, context: Context | None, task_id: bytes | None = None
-    ) -> Context:
+    def _context_new_task(self, context: Context | None, task_id: bytes) -> Context:
         base = context.copy() if context is not None else contextvars.copy_context()
-        if task_id is None:
-            task_id = self._generate_id()
         _ = base.run(_task_ctx.set, _TaskCtx(parent_id=task_id))
         return base
 

src/duron/task.py

@@ -27,7 +27,7 @@
     from types import TracebackType
 
     from duron.codec import Codec
-    from duron.event_loop import WaitSet
+    from duron.event_loop import OpFuture, WaitSet
     from duron.fn import DurableFn
     from duron.log import (
         Entry,
@@ -36,7 +36,6 @@
         LogStorage,
         PromiseCompleteEntry,
     )
-    from duron.ops import Op
 
 
 _T = TypeVar("_T")
@@ -138,16 +137,17 @@ def __init__(
         self._codec = codec
         self._running: object | None = None
         self._pending_msg: list[Entry] = []
-        self._pending_task: dict[str, Callable[[], Coroutine[Any, Any, object]]] = {}
+        self._pending_task: dict[str, Callable[[], Coroutine[Any, Any, None]]] = {}
         self._pending_ops: set[bytes] = set()
         self._now = 0
         self._offset: object | None = None
+        self._tasks: dict[str, asyncio.Task[None]] = {}
 
     def now(self) -> int:
         if self._running:
             t = time.time_ns()
             t -= t % 1_000
-            self._now = max(self._now, t)
+            self._now = max(self._now + 1_000, t)
         return self._now
 
     async def resume(self) -> None:
@@ -177,8 +177,8 @@ async def run(self) -> object:
             for msg in self._pending_msg:
                 await self.enqueue_log(msg)
             self._pending_msg.clear()
-            for task in self._pending_task.values():
-                _ = asyncio.create_task(task())
+            for key, task in self._pending_task.items():
+                self._tasks[key] = asyncio.create_task(task())
             self._pending_task.clear()
 
             t1 = asyncio.create_task(self._follow_log())
@@ -193,6 +193,11 @@ async def run(self) -> object:
                 with contextlib.suppress(asyncio.CancelledError):
                     await t
 
+            for t in self._tasks.values():
+                _ = t.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await t
+
             for t in done:
                 if exc := t.exception():
                     raise exc
@@ -223,20 +228,20 @@ async def _step(self) -> WaitSet | None:
                 sid = s.id
                 if sid not in self._pending_ops:
                     self._pending_ops.add(sid)
-                    await self.enqueue_op(sid, s.params)
+                    await self.enqueue_op(sid, s)
 
     async def handle_message(self, e: Entry) -> None:
         if e["type"] == "promise/complete":
             _ = self._pending_task.pop(e["promise_id"], None)
             id = _decode_id(e["promise_id"])
             if "error" in e:
-                self._loop.post_completion_threadsafe(
+                self._loop.post_completion(
                     id,
                     exception=_decode_error(e["error"]),
                 )
                 self._pending_ops.discard(id)
             elif "result" in e:
-                self._loop.post_completion_threadsafe(
+                self._loop.post_completion(
                     id,
                     result=self._codec.decode_json(e["result"]),
                 )
@@ -251,11 +256,12 @@ async def enqueue_log(self, entry: Entry, flush: bool = False) -> None:
             self._pending_msg.append(entry)
         else:
             await self._log.append(self._running, entry)
-            await self.handle_message(entry)
             if flush:
                 await self._log.flush(self._running)
+            await self.handle_message(entry)
 
-    async def enqueue_op(self, id: bytes, op: Op | object) -> None:
+    async def enqueue_op(self, id: bytes, fut: OpFuture) -> None:
+        op = fut.params
         match op:
             case FnCall():
                 await self.enqueue_log({
@@ -278,12 +284,22 @@ async def cb() -> None:
                         entry["result"] = self._codec.encode_json(result)
                     except BaseException as e:
                         entry["error"] = _encode_error(e)
-                    await self.enqueue_log(entry)
-
+                    finally:
+                        await self.enqueue_log(entry)
+
+                def done(f: OpFuture) -> None:
+                    if f.cancelled():
+                        sid = _encode_id(f.id)
+                        _ = self._pending_task.pop(sid, None)
+                        if tsk := self._tasks.get(sid, None):
+                            _ = tsk.cancel()
+
+                fut.add_done_callback(done)
+                sid = _encode_id(id)
                 if self._running:
-                    _ = asyncio.create_task(cb())
+                    self._tasks[sid] = asyncio.create_task(cb())
                 else:
-                    self._pending_task[_encode_id(id)] = cb
+                    self._pending_task[sid] = cb
 
             case TaskRun():
                 await self.enqueue_log({
@@ -304,8 +320,8 @@ async def cb() -> None:
                         entry["result"] = self._codec.encode_json(result)
                     except BaseException as e:
                         entry["error"] = _encode_error(e)
-
-                    await self.enqueue_log(entry, True)
+                    finally:
+                        await self.enqueue_log(entry, True)
 
                 _ = self._loop.create_task(cb())
 
@@ -336,7 +352,7 @@ def _decode_timestamp(ts: int) -> int:
 def _encode_error(error: BaseException) -> ErrorInfo:
     return {
         "code": -1,
-        "message": str(error),
+        "message": repr(error),
     }
 
 

tests/test_loop.py

@@ -68,9 +68,9 @@ def tick(n: int, expect: set[int] | None) -> list[bytes] | None:
 
     for i in range(5, 0, -1):
         ws = tick(i, {1, 2, 3, 4, 5})
-        loop.post_completion_threadsafe(ws[random.randint(0, i - 1)], result=i - 1)
+        loop.post_completion(ws[random.randint(0, i - 1)], result=i - 1)
     ws = tick(1, {6})
-    loop.post_completion_threadsafe(ws[0], result=6)
+    loop.post_completion(ws[0], result=6)
     tick(0, None)
 
     return ids

tests/test_task.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import random
 import uuid
 from dataclasses import dataclass
@@ -108,6 +109,26 @@ async def activity(ctx: Context, s: str) -> str:
     assert x == "hello"
 
 
+@pytest.mark.asyncio
+async def test_cancel():
+    @durable()
+    async def activity(ctx: Context, s: str) -> str:
+        with contextlib.suppress(BaseException):
+            _ = await asyncio.wait_for(ctx.run(lambda: asyncio.sleep(9999)), 0.1)
+        _ = await asyncio.wait_for(ctx.run(lambda: asyncio.sleep(9999)), 0.1)
+        return s
+
+    log = MemoryLogStorage()
+    async with activity(log) as t:
+        await t.start("hello")
+        try:
+            _ = await t.wait()
+        except Exception as e:
+            assert "Timeout" in repr(e)
+
+    assert len(await log.entries()) == 8
+
+
 @dataclass
 class CustomPoint:
     x: int