Harden async orchestration dispatch cleanup

Close async generators and drain or cancel pending tasks at dispatch teardown, share control-flow exception classification, and add a Windows async-runtime CI job.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Author: andystaples

.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:

durabletask/internal/asyncio_loop.py

@@ -9,13 +9,15 @@
 import contextvars
 import math
 import logging
+import sys
 from collections import deque
 from collections.abc import Callable, Coroutine, Generator
 from datetime import datetime, timedelta
 from fractions import Fraction
 from typing import Any, TypeVar, cast
 
 from durabletask import task
+from durabletask.internal.exceptions import is_control_flow_exception
 
 T = TypeVar("T")
 
@@ -130,9 +132,11 @@ def __init__(
         self._observers: dict[int, _CompletionObserver] = {}
         self._completions: deque[Callable[[], None]] = deque()
         self._timers: dict[_DurableTimerHandle, _CompletionObserver] = {}
+        self._asyncgens: dict[int, Any] = {}
         self._closed = False
         self._closing = False
         self._fatal: BaseException | None = None
+        self._cleanup_failure: BaseException | None = None
         self._root: asyncio.Task[Any] | None = None
 
     def time(self) -> float:
@@ -237,6 +241,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()
@@ -283,7 +289,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():
@@ -297,33 +304,57 @@ def _has_durable_work(self) -> bool:
     def _drain(self) -> None:
         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
+            cleanup_limit = max(10000, 4 * (len(self._tasks) + len(self._asyncgens)))
             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._record_cleanup_error(RuntimeError(
+                        "Async cleanup did not quiesce; cleanup must cooperate with cancellation."))
+                    self._ready.clear()
+                    self._cancel_ready.clear()
+                    break
         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:
+        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):
@@ -332,26 +363,47 @@ def call_exception_handler(self, context: dict[str, Any]) -> None:
             error = failure
         self._fatal = error if isinstance(error, BaseException) else RuntimeError(context["message"])
 
+    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 self._cleanup_failure is None or is_control_flow_exception(error):
+                self._cleanup_failure = error
+
     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
+            for generator in tuple(self._asyncgens.values()):
+                self.create_task(generator.aclose())
+            self._drain()
+            for future in tuple(self._tasks):
+                if not future.done():
+                    future.cancel()
+            self._drain()
+            failure: BaseException | None = self._cleanup_failure
             for future in tuple(self._tasks):
                 if not future.done():
-                    future.get_coro().close()
+                    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:
@@ -365,6 +417,7 @@ def close(self) -> None:
             self._waits.clear()
             self._timers.clear()
             self._tasks.clear()
+            self._asyncgens.clear()
 
     def run_in_executor(self, executor: Any, func: Any, *args: Any) -> asyncio.Future[Any]:
         raise RuntimeError("Threads/executors are unsupported in orchestrators; use an 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__()

durabletask/worker.py

@@ -63,14 +63,6 @@
 _WorkItem = tuple[Callable[..., Any], Callable[..., Any], tuple[Any, ...], dict[str, Any]]
 
 
-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 ConcurrencyOptions:
     """Configuration options for controlling concurrency of different work item types and the thread pool size.
 
@@ -1671,6 +1663,7 @@ def __init__(self,
         self._exception_properties_provider = exception_properties_provider
         self._logger = logger
         self._async_loop: DurableEventLoop | None = None
+        self._async_cleanup = False
 
     def run(self, orchestrator: Generator[task.Task[Any], Any, Any] | Coroutine[Any, Any, Any]) -> None:
         if isinstance(orchestrator, Coroutine):
@@ -1683,6 +1676,7 @@ def run(self, orchestrator: Generator[task.Task[Any], Any, Any] | Coroutine[Any,
 
     def close(self) -> None:
         if self._async_loop is not None:
+            self._async_cleanup = True
             self._async_loop.close()
         elif self._generator is not None:
             self._generator.close()
@@ -1849,6 +1843,8 @@ def get_actions(self) -> list[pb.OrchestratorAction]:
         return current_actions
 
     def next_sequence_number(self) -> int:
+        if self._async_cleanup:
+            raise RuntimeError("Dispatch cleanup cannot schedule durable work; use normal orchestrator control flow.")
         self._sequence_number += 1
         return self._sequence_number
 
@@ -2493,7 +2489,7 @@ def execute(
                 ctx.set_failed(ex)
 
             except BaseException as ex:
-                if _is_control_flow_exception(ex):
+                if pe.is_control_flow_exception(ex):
                     raise
                 self._logger.error(
                     f"{instance_id}: Orchestrator {orchestration_name} exited unexpectedly.",
@@ -2527,7 +2523,7 @@ def execute(
             result = ExecutionResults(
                 actions=actions, encoded_custom_status=ctx._encoded_custom_status,  # pyright: ignore[reportPrivateUsage]
                 orchestration_trace_context=ctx._orchestration_trace_context,  # pyright: ignore[reportPrivateUsage]
-                # Snapshot before generator teardown can exit a critical section.
+                # Snapshot before orchestrator teardown can exit a critical section.
                 entity_lock_releases=list(ctx._entity_context.get_lock_release_messages()),  # pyright: ignore[reportPrivateUsage]
             )
             if unexpected_failure and result.entity_lock_releases:
@@ -2560,7 +2556,7 @@ def execute(
                     exc_info=True,
                 )
                 if not dispatch_exited:
-                    if _is_control_flow_exception(ex) or result is None:
+                    if pe.is_control_flow_exception(ex) or result is None:
                         raise
                     # Keep an existing failure rather than mask its cause.
                     if not any(