Harden async dispatch cleanup ownership and signal precedence

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

Author: andystaples

docs/async-orchestrators.md

@@ -245,13 +245,20 @@ 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

@@ -25,6 +25,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
@@ -133,6 +140,7 @@ def __init__(
         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
@@ -301,7 +309,7 @@ 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) -> 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()
@@ -342,13 +350,19 @@ def _drain(self) -> None:
                     self._ready.clear()
                     self._cancel_ready.clear()
                     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:
@@ -366,9 +380,30 @@ def call_exception_handler(self, context: dict[str, Any]) -> None:
     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
             if self._cleanup_failure is None or is_control_flow_exception(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
+        while True:
+            pending = [
+                generator for key, generator in self._asyncgens.items() if closed.get(key) is not generator]
+            if not pending:
+                return
+            limit = max(10000, 4 * (len(self._tasks) + len(self._asyncgens)))
+            if steps >= limit:
+                self._record_cleanup_error(RuntimeError(
+                    "Async cleanup did not quiesce; cleanup must cooperate with cancellation."))
+                return
+            for generator in pending:
+                closed[id(generator)] = generator
+                self.create_task(generator.aclose())
+            steps += max(1, self._drain())
+
     def close(self) -> None:
         self._closing = True
         try:
@@ -379,13 +414,13 @@ def close(self) -> None:
                 if not future.done():
                     future.cancel()
             self._drain()
-            for generator in tuple(self._asyncgens.values()):
-                self.create_task(generator.aclose())
-            self._drain()
+            closed_asyncgens: dict[int, Any] = {}
+            self._close_asyncgens(closed_asyncgens)
             for future in tuple(self._tasks):
                 if not future.done():
                     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():

durabletask/worker.py

@@ -2234,6 +2234,10 @@ def _exit_critical_section(self) -> None:
             # _exit_critical_section both from the EntityLock and the continue_as_new logic. We must keep both calls in
             # case the user code crashes after calling continue_as_new but before the EntityLock object is exited.
             return
+        if self._async_cleanup:
+            # Lock disposal during local teardown is not a durable release; the
+            # pre-teardown response already carries any release this dispatch owns.
+            return
         for entity_unlock_message in self._entity_context.emit_lock_release_messages():
             task_id = self.next_sequence_number()
             action = pb.OrchestratorAction(id=task_id, sendEntityMessage=entity_unlock_message)

tests/durabletask/test_asyncio_orchestrators.py

@@ -9,6 +9,7 @@
 import os
 import subprocess
 import sys
+from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from contextvars import ContextVar
 from dataclasses import dataclass
@@ -20,7 +21,7 @@
 
 import pytest
 
-from durabletask import task
+from durabletask import entities, task
 from durabletask.internal import helpers
 from durabletask.internal import orchestrator_service_pb2 as pb
 from tests.durabletask.test_coroutine_orchestrators import (
@@ -1057,3 +1058,132 @@ async def branch(name: str, error: BaseException | None) -> None:
         _execute(orchestrator, [], _start())
     assert caught.value is exit_error
     assert cleaned == ["custom", "exit", "last"]
+
+
+def _cleanup_signal(kind: str) -> BaseException:
+    if kind == "system_exit":
+        return SystemExit("stop")
+    if kind == "keyboard_interrupt":
+        return KeyboardInterrupt("interrupt")
+    if kind == "nested_group":
+        return BaseExceptionGroup("outer", [BaseExceptionGroup("inner", [SystemExit("stop")])])
+    if kind == "generator_exit":
+        return GeneratorExit("closed")
+    return BaseExceptionGroup("cancelled", [asyncio.CancelledError()])
+
+
+@pytest.mark.parametrize("control_first", [True, False])
+@pytest.mark.parametrize("other", ["generator_exit", "cancelled_group"])
+@pytest.mark.parametrize("control", ["system_exit", "keyboard_interrupt", "nested_group"])
+def test_cleanup_keeps_process_control_over_other_control_flow(
+        control: str, other: str, control_first: bool) -> None:
+    control_error = _cleanup_signal(control)
+    errors = [control_error, _cleanup_signal(other)]
+    if not control_first:
+        errors.reverse()
+    cleaned: list[int] = []
+
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None:
+        async def branch(index: int) -> None:
+            try:
+                await ctx.wait_for_external_event(f"never-{index}")
+            finally:
+                cleaned.append(index)
+                raise errors[index]
+
+        asyncio.create_task(branch(0))
+        asyncio.create_task(branch(1))
+        await ctx.call_activity("wait")
+
+    with pytest.raises(type(control_error)) as caught:
+        _execute(orchestrator, [], _start())
+    assert caught.value is control_error
+    assert cleaned == [0, 1]
+
+
+def test_cleanup_closes_async_generators_started_during_generator_cleanup(
+        caplog: pytest.LogCaptureFixture) -> None:
+    closed: list[str] = []
+
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None:
+        async def child() -> AsyncIterator[int]:
+            try:
+                yield 1
+            finally:
+                await asyncio.sleep(0)
+                closed.append("child")
+
+        async def outer() -> AsyncIterator[int]:
+            try:
+                yield 1
+            finally:
+                await anext(child())
+                closed.append("outer")
+
+        await anext(outer())
+        await ctx.call_activity("wait")
+
+    actions = _execute(orchestrator, [], _start()).actions
+    assert [action.scheduleTask.name for action in actions] == ["wait"]
+    assert closed == ["outer", "child"]
+    assert "Async" not in caplog.text
+
+
+def test_exhausted_async_generators_are_not_retained_until_dispatch_end() -> None:
+    references: list[ReferenceType[Any]] = []
+
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None:
+        async def values() -> AsyncIterator[int]:
+            yield 1
+
+        for _ in range(200):
+            generator = values()
+            references.append(ref(generator))
+            async for _ in generator:
+                pass
+            del generator
+        gc.collect()
+        assert sum(reference() is not None for reference in references) < 100
+        await ctx.call_activity("wait")
+
+    _execute(orchestrator, [], _start())
+
+
+@pytest.mark.parametrize("suspend", ["activity", "entity"])
+def test_entity_lock_disposal_during_teardown_does_not_schedule_releases(
+        suspend: str, caplog: pytest.LogCaptureFixture) -> None:
+    first = entities.EntityInstanceId("counter", "a")
+    second = entities.EntityInstanceId("counter", "b")
+
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any:
+        with await ctx.lock_entities([first, second]):
+            if suspend == "activity":
+                return await ctx.call_activity("inside-lock")
+            return await ctx.call_entity(first, "get")
+
+    old = _start()
+    lock = _execute(orchestrator, [], old).actions[0].sendEntityMessage.entityLockRequested
+    old += [helpers.new_event_sent_event(1, str(first), json.dumps({"id": lock.criticalSectionId}))]
+    new = [helpers.new_event_raised_event(lock.criticalSectionId, None)]
+    actions = _execute(orchestrator, old, new).actions
+    assert len(actions) == 1
+    if suspend == "activity":
+        assert actions[0].scheduleTask.name == "inside-lock"
+        old += new + [helpers.new_task_scheduled_event(2, "inside-lock")]
+        new = [helpers.new_task_completed_event(2, json.dumps(42))]
+    else:
+        call = actions[0].sendEntityMessage.entityOperationCalled
+        assert call.targetInstanceId.value == str(first)
+        old += new + [helpers.new_event_sent_event(2, str(first), json.dumps({"id": call.requestId}))]
+        new = [helpers.new_event_raised_event(call.requestId, json.dumps({"result": "42"}))]
+    assert "cleanup" not in caplog.text
+
+    result = _execute(orchestrator, old, new)
+    actions = result.actions
+    releases = [action.sendEntityMessage.entityUnlockSent for action in actions[:-1]]
+    # Entity calls inside a critical section keep the legacy shared-list release order.
+    expected = [first, second] if suspend == "activity" else [second, first]
+    assert [release.targetInstanceId.value for release in releases] == [str(e) for e in expected]
+    assert [action.id for action in actions] == [3, 4, 5]
+    assert json.loads(_completion(result).result.value) == 42
+    assert "cleanup" not in caplog.text