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,12 +9,14 @@
 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 typing import Any, TypeVar, cast
 
 from durabletask import task
+from durabletask.internal.exceptions import is_control_flow_exception
 
 T = TypeVar("T")
 
@@ -113,9 +115,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:
@@ -217,6 +221,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()
@@ -263,7 +269,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():
@@ -277,33 +284,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):
@@ -312,26 +343,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:
@@ -345,6 +397,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(