Bound async-generator cleanup discovery budget

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

Author: andystaples

CHANGELOG.md

@@ -13,6 +13,9 @@ FIXED
 when an intermediate async wrapper decorates a synchronous function, including
 callable instances and partials. Cyclic wrapper chains fail instead of hanging
 worker dispatch.
+- Async dispatch cleanup no longer hangs when async-generator finalizers
+continually start successor generators; exceeding the cleanup budget is reported
+without changing the snapshotted response.
 - Coroutine task-group cancellation now schedules shared-Future callbacks
 deterministically across replays, including when a dispatch resumes on another
 worker using the same SDK and Python runtime. Previously recorded cancellation

docs/async-orchestrators.md

@@ -242,7 +242,10 @@ dispatch boundary**, before workflow completion. 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.
+queue is bounded and reported, not allowed to run indefinitely. Each
+async-generator sweep has a step budget fixed at sweep entry; tasks or generators
+created during that sweep cannot increase it. This 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.

durabletask-azuremanaged/CHANGELOG.md

@@ -9,6 +9,9 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 FIXED
 
+- With the corresponding core SDK update, async dispatch cleanup no longer
+hangs when async-generator finalizers continually start successor generators.
+Exceeding the cleanup budget is reported without changing the snapshotted response.
 - With the corresponding core SDK update, coroutine task-group cancellation
 schedules shared-Future callbacks deterministically when replayed on another
 worker using the same SDK and Python runtime. See the

durabletask/internal/asyncio_loop.py

@@ -389,12 +389,13 @@ 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 = max(10000, 4 * (len(self._tasks) + len(self._asyncgens)))
         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."))

tests/durabletask/test_asyncio_orchestrators.py

@@ -1101,32 +1101,91 @@ async def branch(index: int) -> None:
     assert cleaned == [0, 1]
 
 
+@pytest.mark.parametrize("count", [2, 32, 6000])
 def test_cleanup_closes_async_generators_started_during_generator_cleanup(
-        caplog: pytest.LogCaptureFixture) -> None:
-    closed: list[str] = []
+        count: int, caplog: pytest.LogCaptureFixture) -> None:
+    closed: list[int] = []
 
     async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None:
-        async def child() -> AsyncIterator[int]:
+        async def resource(index: int) -> AsyncIterator[int]:
             try:
                 yield 1
             finally:
+                if index + 1 < count:
+                    await anext(resource(index + 1))
                 await asyncio.sleep(0)
-                closed.append("child")
-
-        async def outer() -> AsyncIterator[int]:
-            try:
-                yield 1
-            finally:
-                await anext(child())
-                closed.append("outer")
+                closed.append(index)
 
-        await anext(outer())
+        await anext(resource(0))
         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
+    assert closed == list(range(count))
+    if count == 6000:
+        assert "Async cleanup did not quiesce" in caplog.text
+    else:
+        assert "Async" not in caplog.text
+
+
+@pytest.mark.parametrize("control", [None, "system_exit", "keyboard_interrupt", "nested_group"])
+def test_infinite_asyncgen_discovery_has_bounded_cleanup(control: str | None) -> None:
+    code = """
+import asyncio
+import logging
+import sys
+from durabletask import task
+from tests.durabletask.test_asyncio_orchestrators import _cleanup_signal
+from tests.durabletask.test_coroutine_orchestrators import _execute, _start
+
+control = sys.argv[1]
+exit_error = None if control == "none" else _cleanup_signal(control)
+warnings = []
+
+class Capture(logging.Handler):
+    def emit(self, record):
+        warnings.append(logging.Formatter().format(record))
+
+logging.getLogger().addHandler(Capture())
+
+async def orchestrator(ctx: task.OrchestrationContext, _):
+    async def resource():
+        try:
+            yield 1
+        finally:
+            await anext(resource())
+
+    async def signal_on_forced_close():
+        try:
+            await ctx.wait_for_external_event("never")
+        finally:
+            try:
+                while True:
+                    try:
+                        await asyncio.Future()
+                    except asyncio.CancelledError:
+                        pass
+            finally:
+                raise exit_error
+
+    if exit_error is not None:
+        asyncio.create_task(signal_on_forced_close())
+    await anext(resource())
+    await ctx.call_activity("wait")
+
+try:
+    result = _execute(orchestrator, [], _start())
+except BaseException as error:
+    assert exit_error is not None and error is exit_error, repr(error)
+else:
+    assert exit_error is None
+    assert [action.scheduleTask.name for action in result.actions] == ["wait"]
+assert any("Async cleanup did not quiesce" in message for message in warnings), warnings
+"""
+    process = subprocess.run(
+        [sys.executable, "-c", code, control or "none"],
+        capture_output=True, text=True, timeout=30)
+    assert process.returncode == 0, process.stdout + process.stderr
 
 
 def test_exhausted_async_generators_are_not_retained_until_dispatch_end() -> None:
@@ -1146,7 +1205,10 @@ async def values() -> AsyncIterator[int]:
         assert sum(reference() is not None for reference in references) < 100
         await ctx.call_activity("wait")
 
-    _execute(orchestrator, [], _start())
+    actions = _execute(orchestrator, [], _start()).actions
+    assert len(actions) == 1
+    assert actions[0].HasField("scheduleTask")
+    assert actions[0].scheduleTask.name == "wait"
 
 
 @pytest.mark.parametrize("suspend", ["activity", "entity"])