[Async stack 2/6] asyncio composition

durabletask/internal/asyncio_loop.py

@@ -7,11 +7,12 @@
 
 import asyncio
 import contextvars
+import math
 import logging
 from collections import deque
 from collections.abc import Callable, Coroutine, Generator
-from datetime import datetime
-from typing import Any, TypeVar
+from datetime import datetime, timedelta
+from typing import Any, TypeVar, cast
 
 from durabletask import task
 
@@ -33,11 +34,67 @@ def detach(self) -> None:
         self._notify_parent()
 
 
+class _DurableAsyncTask(asyncio.Task[T]):
+    def __await__(self) -> Generator[Any, None, T]:
+        try:
+            return (yield from super().__await__())
+        finally:
+            self._retire_if_observed()
+
+    __iter__ = __await__
+
+    def result(self) -> T:
+        try:
+            return super().result()
+        finally:
+            self._retire_if_observed()
+
+    def exception(self) -> BaseException | None:
+        try:
+            return super().exception()
+        finally:
+            self._retire_if_observed()
+
+    def _retire_if_observed(self) -> None:
+        loop = self.get_loop()
+        assert isinstance(loop, DurableEventLoop)
+        loop.retire_task(self)
+
+    def cancel(self, msg: Any = None) -> bool:
+        loop = self.get_loop()
+        assert isinstance(loop, DurableEventLoop)
+        previous = loop.cancelling_task
+        loop.cancelling_task = self
+        try:
+            return super().cancel(msg)
+        finally:
+            loop.cancelling_task = previous
+
+
+class _DurableTimerHandle(asyncio.TimerHandle):
+    def __eq__(self, other: object) -> bool:
+        return self is other
+
+    __hash__ = object.__hash__
+
+    def __init__(
+        self, when: float, callback: Callable[..., Any], args: tuple[Any, ...],
+        loop: DurableEventLoop, context: contextvars.Context | None,
+    ) -> None:
+        super().__init__(when, callback, args, loop, context=context)
+        self.durable_timer: task.TimerTask | None = None
+
+    def cancel(self) -> None:
+        super().cancel()
+        if self.durable_timer is not None:
+            self.durable_timer.cancel()
+
+
 class DurableEventLoop(asyncio.AbstractEventLoop):
     """A FIFO microtask loop; it never polls a selector or advances wall time.
 
-    Only history events can complete durable waits. asyncio owns coroutine
-    stepping; this loop owns the ordering of ready callbacks.
+    Only history events can complete durable waits/timers. asyncio owns coroutine
+    stepping, gather and structured cancellation; this loop owns their ordering.
     """
 
     def __init__(
@@ -48,10 +105,14 @@ def __init__(
         self._is_complete = is_complete
         self._logger = logger or logging.getLogger(__name__)
         self._ready: deque[asyncio.Handle] = deque()
-        self._tasks: list[asyncio.Task[Any]] = []
+        self._cancel_ready: list[tuple[int, int, asyncio.Handle]] = []
+        self._tasks: dict[asyncio.Task[Any], int] = {}
+        self._task_sequence = 0
+        self.cancelling_task: asyncio.Task[Any] | None = None
         self._waits: dict[int, tuple[task.Task[Any], list[asyncio.Future[Any]]]] = {}
         self._observers: dict[int, _CompletionObserver] = {}
         self._completions: deque[Callable[[], None]] = deque()
+        self._timers: dict[_DurableTimerHandle, _CompletionObserver] = {}
         self._closed = False
         self._closing = False
         self._fatal: BaseException | None = None
@@ -86,18 +147,75 @@ def create_task(
         if kwargs:
             coro.close()
             raise TypeError("Custom/eager task execution is not supported in durable orchestrators.")
-        future = asyncio.Task(coro, loop=self, name=name, context=context)
-        self._tasks.append(future)
+        future = _DurableAsyncTask(coro, loop=self, name=name, context=context)
+        self._tasks[future] = self._task_sequence
+        self._task_sequence += 1
+        future.add_done_callback(self.retire_task)
         return future
 
+    def retire_task(self, future: asyncio.Task[Any]) -> None:
+        # A failed task may be observed after its done callbacks have run.
+        # Inspect without retrieving exceptions so unobserved failures still log.
+        if not self._closing and future.done() and not getattr(future, "_log_traceback", False):
+            self._tasks.pop(future, None)
+
     def call_soon(
         self, callback: Callable[..., Any], *args: Any,
         context: contextvars.Context | None = None,
     ) -> asyncio.Handle:
         handle = asyncio.Handle(callback, args, self, context=context)
-        self._ready.append(handle)
+        if self.cancelling_task is not None:
+            # TaskGroup cancels a set of tasks. Do not let object hashes
+            # determine the order in which their finally blocks execute.
+            owner = getattr(callback, "__self__", None)
+            index = self._tasks[self.cancelling_task]
+            if isinstance(owner, asyncio.Task):
+                index = self._tasks.get(cast(asyncio.Task[Any], owner), index)
+            self._cancel_ready.append((index, len(self._cancel_ready), handle))
+        else:
+            self._ready.append(handle)
         return handle
 
+    def call_later(
+        self, delay: float, callback: Callable[..., Any], *args: Any,
+        context: contextvars.Context | None = None,
+    ) -> asyncio.TimerHandle:
+        if not math.isfinite(delay):
+            raise ValueError("Durable delays must be finite.")
+        return self.call_at(self.time() + max(0, delay), callback, *args, context=context)
+
+    def call_at(
+        self, when: float, callback: Callable[..., Any], *args: Any,
+        context: contextvars.Context | None = None,
+    ) -> asyncio.TimerHandle:
+        if not math.isfinite(when):
+            raise ValueError("Durable deadlines must be finite.")
+        handle = _DurableTimerHandle(when, callback, args, self, context)
+        if when <= self.time():
+            self._ready.append(handle)
+        else:
+            if self._closing:
+                raise RuntimeError("Dispatch cleanup cannot schedule durable timers.")
+            deadline = datetime(1970, 1, 1) + timedelta(seconds=when)
+            if (deadline - datetime(1970, 1, 1)).total_seconds() < when:
+                deadline += timedelta(microseconds=1)
+            handle.durable_timer = self.ctx.create_timer(deadline)
+            self._timers[handle] = _CompletionObserver(
+                handle.durable_timer,
+                lambda: self._completions.append(lambda: self._complete_timer(handle)))
+        return handle
+
+    def _timer_handle_cancelled(self, handle: asyncio.TimerHandle) -> None:
+        if isinstance(handle, _DurableTimerHandle):
+            observer = self._timers.pop(handle, None)
+            if observer is not None:
+                observer.detach()
+
+    def _complete_timer(self, handle: _DurableTimerHandle) -> None:
+        self._timers.pop(handle, None)
+        if not handle.cancelled():
+            self._ready.append(handle)
+
     async def wait_durable(self, durable: task.Task[T]) -> T:
         if durable.is_complete:
             return durable.get_result()
@@ -154,7 +272,7 @@ def resume(self) -> None:
                 "Await durable activities, events or timers; external I/O is unsupported.")
 
     def _has_durable_work(self) -> bool:
-        return any(not durable.is_complete for durable, _ in self._waits.values())
+        return any(not durable.is_complete for durable, _ in self._waits.values()) or bool(self._timers)
 
     def _drain(self) -> None:
         previous_loop = asyncio._get_running_loop()
@@ -163,7 +281,9 @@ def _drain(self) -> None:
             asyncio.tasks._leave_task(previous_loop, previous_task)
         asyncio._set_running_loop(self)
         try:
-            while self._ready:
+            while self._ready or self._cancel_ready:
+                self._ready.extend(handle for _, _, handle in sorted(self._cancel_ready))
+                self._cancel_ready.clear()
                 handle = self._ready.popleft()
                 if not handle.cancelled():
                     handle._run()
@@ -173,11 +293,11 @@ def _drain(self) -> None:
                         raise error
                     self._logger.warning(
                         "%s: Async cleanup callback failed.", self.ctx.instance_id, exc_info=error)
+                if not self._closing and self._is_complete():
+                    break
+                if not self._closing and self._root is not None and self._root.done():
+                    break
                 if not self._closing:
-                    if self._is_complete():
-                        break
-                    if self._root is not None and self._root.done():
-                        break
                     self._deliver_completions()
         finally:
             asyncio._set_running_loop(previous_loop)
@@ -217,9 +337,11 @@ def close(self) -> None:
         finally:
             self._closed = True
             self._ready.clear()
-            for observer in self._observers.values():
+            self._cancel_ready.clear()
+            for observer in (*self._observers.values(), *self._timers.values()):
                 observer.detach()
             self._completions.clear()
             self._observers.clear()
             self._waits.clear()
+            self._timers.clear()
             self._tasks.clear()