[Async stack 3/6] Dispatch cleanup for coroutine orchestrators

.github/workflows/durabletask.yml

@@ -14,6 +14,30 @@ permissions:
   contents: read
 
 jobs:
+  async-runtime-windows:
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.11", "3.12", "3.13", "3.14", "3.15"]
+    runs-on: windows-latest
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+      - uses: actions/setup-python@6d4fd22b20465ccc0f7354476588e8df7c9ca0e0 # v7.1.0
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install runtime test dependencies
+        run: |
+          python -m pip install --upgrade "pip>=25.1"
+          python -m pip install --group test -e .
+      - name: Check asyncio, replay, and dispatch cleanup
+        run: >-
+          python -m pytest
+          tests\durabletask\test_asyncio_orchestrators.py
+          tests\durabletask\test_coroutine_orchestrators.py
+          tests\durabletask\test_orchestration_generator_cleanup.py
+          tests\durabletask\test_orchestration_dispatch_recovery.py
+          -q
+
   lint-and-unit-tests:
     runs-on: ubuntu-latest
     steps:

docs/async-orchestrators.md

@@ -302,22 +302,38 @@ continue as new, encounter a history failure, or reject a version.
 Local cleanup can await ready work such as `sleep(0)`; it must not wait for new
 durable or external work.
 Scheduling/awaiting durable work during async teardown raises an actionable
-error. Teardown must cooperate with cancellation; a non-quiescent cleanup
-queue is bounded and reported, not allowed to run indefinitely.
+error. Teardown must cooperate with cancellation.
+
+> [!WARNING]
+> Async teardown is capped at about 10,000 event-loop steps per cleanup phase
+> (more for dispatches with many tasks or async generators). Cleanup that is
+> still running when the cap is reached is truncated: remaining callbacks are
+> dropped, unfinished coroutines are closed, and a warning is logged. The
+> response decided before teardown is unchanged. Each async-generator sweep's
+> budget is fixed at sweep entry; tasks or generators created during that sweep
+> cannot increase it. The cap bounds cooperative scheduling steps, not
+> wall-clock time or code that never yields.
 
 Ordinary teardown failures (including `CancelledError`) are logged without
 changing the decided response or skipping other cooperative finalizers.
 Unexpected non-cancellation `BaseException`
 exits retain the existing terminal-failure/recovery policy; an existing primary
 failure is preserved. Teardown status/actions do not replace the response.
 Held entity-release obligations are captured before teardown for terminal
-payload/failure handling.
+payload/failure handling. Exiting an entity critical section during teardown
+does not schedule additional release messages, and a `SystemExit` or
+`KeyboardInterrupt` raised by any finalizer takes precedence over other
+cleanup failures.
 
 Keep workflow cleanup that must be durable in **normal orchestration control
 flow**, before return or continue-as-new, not solely in a dispatch finalizer.
 No cleanup is guaranteed after process crash or forced termination.
 Async-generator *helpers* are closed; async-generator *orchestrators*
-(`async def` containing `yield`) are rejected.
+(`async def` containing `yield`) are rejected. The dispatch owns every async
+generator it starts until teardown, including generators first iterated by
+another generator's cleanup. To run a generator's cleanup earlier, close it
+explicitly with `await gen.aclose()` or iterate it inside
+`contextlib.aclosing(...)`.
 
 ## Replay and framework adapters
 

durabletask/internal/asyncio_loop.py

@@ -31,6 +31,7 @@
 import contextvars
 import logging
 import math
+import sys
 from collections import deque
 from collections.abc import Callable, Coroutine, Generator
 from datetime import UTC, datetime, timedelta
@@ -45,6 +46,13 @@
 _MICROSECOND = timedelta(microseconds=1)
 
 
+def _is_process_control(error: BaseException | None) -> bool:
+    if isinstance(error, BaseExceptionGroup):
+        group = cast(BaseExceptionGroup[BaseException], error)
+        return group.subgroup((KeyboardInterrupt, SystemExit)) is not None
+    return isinstance(error, (KeyboardInterrupt, SystemExit))
+
+
 class _CompletionObserver(task.CompositeTask[None]):
     def __init__(self, durable: task.Task[Any], callback: Callable[[], None]) -> None:
         self._callback = callback
@@ -154,9 +162,13 @@ def __init__(
         self._observers: dict[int, _CompletionObserver] = {}
         self._completions: deque[Callable[[], None]] = deque()
         self._timers: dict[_DurableTimerHandle, _CompletionObserver] = {}
+        self._asyncgens: dict[int, Any] = {}
+        self._asyncgen_prune_at = 64
         self._closed = False
         self._closing = False
         self._fatal: BaseException | None = None
+        self._cleanup_failure: BaseException | None = None
+        self._cleanup_truncated = False
         self._root: asyncio.Task[Any] | None = None
 
     def _now(self) -> datetime:
@@ -268,6 +280,8 @@ def _complete_timer(self, handle: _DurableTimerHandle) -> None:
             self._ready.append(handle)
 
     async def wait_durable(self, durable: task.Task[T]) -> T:
+        if self._closing:
+            raise RuntimeError("Dispatch cleanup cannot await durable work.")
         if durable.is_complete:
             return durable.get_result()
         future = self.create_future()
@@ -317,7 +331,8 @@ def resume(self) -> None:
                 raise task.TaskCancelledError(
                     "The orchestrator coroutine was cancelled. Catch local cancellation "
                     "inside the orchestrator; terminate the durable instance through the client.")
-            # result() propagates the orchestrator's failure to the executor.
+            # result() deliberately propagates control-flow BaseExceptions to the
+            # existing executor recovery boundary, not a success-shaped fallback.
             result = self._root.result()
             raise StopIteration(result)
         if not self._has_durable_work():
@@ -328,36 +343,64 @@ def resume(self) -> None:
     def _has_durable_work(self) -> bool:
         return any(not durable.is_complete for durable, _ in self._waits.values()) or bool(self._timers)
 
-    def _drain(self) -> None:
+    def _drain(self, cleanup_limit: int | None = None) -> int:
         previous_loop = asyncio._get_running_loop()
         previous_task = asyncio.current_task(previous_loop) if previous_loop is not None else None
+        previous_hooks = sys.get_asyncgen_hooks()
         if previous_task is not None and previous_loop is not None:
             asyncio.tasks._leave_task(previous_loop, previous_task)
         asyncio._set_running_loop(self)
+        sys.set_asyncgen_hooks(firstiter=self._asyncgen_firstiter, finalizer=self._asyncgen_finalizer)
         try:
+            steps = 0
+            if cleanup_limit is None:
+                cleanup_limit = self._cleanup_budget()
             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()
+                    try:
+                        handle._run()
+                    except BaseException as error:
+                        if not self._closing:
+                            raise
+                        # Defer all exits until other cleanup tasks have drained.
+                        self._record_cleanup_error(error)
                 if self._fatal is not None:
                     error, self._fatal = self._fatal, None
                     if not self._closing:
                         raise error
-                    self._logger.warning(
-                        "%s: Async cleanup callback failed.", self.ctx.instance_id, exc_info=error)
+                    self._record_cleanup_error(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()
+                steps += 1
+                if self._closing and steps >= cleanup_limit:
+                    self._report_unquiesced_cleanup()
+                    break
+            return steps
         finally:
+            sys.set_asyncgen_hooks(*previous_hooks)
             asyncio._set_running_loop(previous_loop)
             if previous_task is not None and previous_loop is not None:
                 asyncio.tasks._enter_task(previous_loop, previous_task)
 
+    def _asyncgen_firstiter(self, generator: Any) -> None:
+        if len(self._asyncgens) >= self._asyncgen_prune_at:
+            # Exhausted generators have no cleanup left; amortized pruning bounds retention.
+            self._asyncgens = {
+                key: value for key, value in self._asyncgens.items() if value.ag_frame is not None}
+            self._asyncgen_prune_at = max(64, 2 * len(self._asyncgens))
+        self._asyncgens[id(generator)] = generator
+
+    def _asyncgen_finalizer(self, generator: Any) -> None:
+        if not self._closed:
+            self.create_task(generator.aclose())
+
     def call_exception_handler(self, context: dict[str, Any]) -> None:
         error = context.get("exception")
         if isinstance(error, StopIteration):
@@ -366,26 +409,84 @@ def call_exception_handler(self, context: dict[str, Any]) -> None:
             error = failure
         self._fatal = error if isinstance(error, BaseException) else RuntimeError(context["message"])
 
+    def _cleanup_budget(self) -> int:
+        # Cooperative cleanup is capped in scheduling steps, not wall-clock time.
+        return max(10000, 4 * (len(self._tasks) + len(self._asyncgens)))
+
+    def _report_unquiesced_cleanup(self) -> None:
+        self._ready.clear()
+        self._cancel_ready.clear()
+        if not self._cleanup_truncated:
+            self._cleanup_truncated = True
+            self._record_cleanup_error(RuntimeError(
+                "Async cleanup did not quiesce; cleanup must cooperate with cancellation."))
+
+    def _record_cleanup_error(self, error: BaseException, source: str = "cleanup callback") -> None:
+        self._logger.warning("%s: Async %s failed.", self.ctx.instance_id, source, exc_info=error)
+        if not isinstance(error, (Exception, asyncio.CancelledError)):
+            if _is_process_control(self._cleanup_failure):
+                return
+            # Keep the first failure unless a process-control exit supersedes it.
+            if self._cleanup_failure is None or _is_process_control(error):
+                self._cleanup_failure = error
+
+    def _close_asyncgens(self, closed: dict[int, Any]) -> None:
+        # Closing one generator may first-iterate another; sweep until no new generators appear.
+        # Closed generators stay referenced so their ids cannot be reused during cleanup.
+        steps = 0
+        # Cleanup-created tasks and generators must not replenish the sweep's budget.
+        limit = self._cleanup_budget()
+        while True:
+            pending = [
+                generator for key, generator in self._asyncgens.items() if closed.get(key) is not generator]
+            if not pending:
+                return
+            if steps >= limit:
+                self._report_unquiesced_cleanup()
+                return
+            for generator in pending:
+                closed[id(generator)] = generator
+                self.create_task(generator.aclose())
+            # Each drain spends from this sweep's remaining budget.
+            steps += max(1, self._drain(limit - steps))
+
     def close(self) -> None:
-        """Cancel unfinished coroutines so their cleanup runs before the dispatch ends."""
         if self._closed:
             return
         self._closing = True
         try:
+            completed_at_close = {future for future in self._tasks if future.done()}
+            unobserved_at_close = {
+                future for future in completed_at_close if getattr(future, "_log_traceback", False)}
             for future in tuple(self._tasks):
                 if not future.done():
                     future.cancel()
             self._drain()
-            failure: BaseException | None = None
+            closed_asyncgens: dict[int, Any] = {}
+            self._close_asyncgens(closed_asyncgens)
             for future in tuple(self._tasks):
                 if not future.done():
-                    future.get_coro().close()
+                    future.cancel()
+            self._drain()
+            self._close_asyncgens(closed_asyncgens)
+            failure: BaseException | None = self._cleanup_failure
+            for future in tuple(self._tasks):
+                if not future.done():
+                    try:
+                        future.get_coro().close()
+                    except BaseException as error:
+                        # Close siblings before propagating process-control exits.
+                        self._record_cleanup_error(error, "coroutine cleanup")
+                        failure = failure or error
                     failure = failure or RuntimeError(
                         "Async cleanup must finish without waiting for durable or external work.")
-                elif not future.cancelled() and getattr(future, "_log_traceback", False):
-                    self._logger.warning(
-                        "%s: Async background task failed.", self.ctx.instance_id,
-                        exc_info=future.exception())
+                elif not future.cancelled():
+                    error = future.exception()
+                    if error is not None and (future not in completed_at_close or future in unobserved_at_close):
+                        self._record_cleanup_error(
+                            error, "background task" if future in completed_at_close else "task cleanup")
+            if self._cleanup_failure is not None:
+                raise self._cleanup_failure
             if failure is not None:
                 raise failure
         finally:
@@ -399,6 +500,7 @@ def close(self) -> None:
             self._waits.clear()
             self._timers.clear()
             self._tasks.clear()
+            self._asyncgens.clear()
 
 
 _IO_HINT = "External I/O is unsupported in orchestrators; move it to an ordinary activity."

durabletask/internal/exceptions.py

@@ -1,9 +1,20 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
+import asyncio
+from typing import cast
+
 import durabletask.internal.orchestrator_service_pb2 as pb
 
 
+def is_control_flow_exception(error: BaseException) -> bool:
+    control_flow_types = (asyncio.CancelledError, GeneratorExit, KeyboardInterrupt, SystemExit)
+    if isinstance(error, BaseExceptionGroup):
+        group = cast(BaseExceptionGroup[BaseException], error)
+        return group.subgroup(control_flow_types) is not None
+    return isinstance(error, control_flow_types)
+
+
 class VersionFailureException(Exception):
     def __init__(self, error_details: pb.TaskFailureDetails) -> None:
         super().__init__()