Skip to content

Commit 5ca6f3c

Browse files
andystaplesCopilot
andcommitted
Merge andystaples-async-stack-2-asyncio-composition into andystaples-async-stack-3-dispatch-cleanup
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
2 parents f6a45e8 + ba1476e commit 5ca6f3c

7 files changed

Lines changed: 233 additions & 41 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,10 +9,6 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
99

1010
FIXED
1111

12-
- Async orchestrators wrapped with `functools.wraps` are now recognized even
13-
when an intermediate async wrapper decorates a synchronous function, including
14-
callable instances and partials. Cyclic wrapper chains fail instead of hanging
15-
worker dispatch.
1612
- Async dispatch cleanup no longer hangs when async-generator finalizers
1713
continually start successor generators; exceeding the cleanup budget is reported
1814
without changing the snapshotted response.

‎durabletask/internal/asyncio_loop.py‎

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,22 @@
11
# Copyright (c) Microsoft Corporation.
22
# Licensed under the MIT License.
33

4-
"""Dispatch-local asyncio scheduling on top of the shared durable history engine."""
4+
"""Dispatch-local asyncio scheduling on top of the shared durable history engine.
5+
6+
Private asyncio dependencies (verified against the supported Python versions):
7+
8+
- ``asyncio._get_running_loop`` / ``asyncio._set_running_loop``: install this
9+
loop as the running loop only while a dispatch drains ready callbacks, and
10+
restore any host loop afterwards. ``Task.__await__`` uses the getter to
11+
detect durable, foreign, or absent loops without raising.
12+
- ``asyncio.tasks._enter_task`` / ``asyncio.tasks._leave_task``: suspend and
13+
restore the host's current task while draining, so asyncio does not see two
14+
tasks running on one thread.
15+
- ``asyncio.Handle._run``: run a ready callback with asyncio's own context and
16+
exception-handler semantics.
17+
- ``asyncio.Task._log_traceback``: report unobserved local-task failures during
18+
dispatch cleanup without marking them as retrieved.
19+
"""
520

621
from __future__ import annotations
722

@@ -12,7 +27,7 @@
1227
import sys
1328
from collections import deque
1429
from collections.abc import Callable, Coroutine, Generator
15-
from datetime import datetime, timedelta
30+
from datetime import UTC, datetime, timedelta
1631
from fractions import Fraction
1732
from typing import Any, TypeVar, cast
1833

@@ -149,7 +164,10 @@ def __init__(
149164

150165
def time(self) -> float:
151166
# A fixed UTC epoch avoids local-time and platform timestamp differences.
152-
return (self.ctx.current_utc_datetime - _EPOCH).total_seconds()
167+
now = self.ctx.current_utc_datetime
168+
if now.tzinfo is not None:
169+
now = now.astimezone(UTC).replace(tzinfo=None)
170+
return (now - _EPOCH).total_seconds()
153171

154172
def get_debug(self) -> bool:
155173
return False

‎durabletask/internal/shared.py‎

Lines changed: 27 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -75,28 +75,40 @@ def from_json(json_str: str | bytes | bytearray, expected_type: type | None = No
7575

7676

7777
def is_async_callable(fn: Any) -> bool:
78-
"""Recognize async functions, wrapped functions, partials and callable objects."""
78+
"""Classify an orchestrator callable as async (``True``) or synchronous.
79+
80+
Each callable is checked directly before following ``functools.wraps``
81+
metadata: coroutine functions are async and generator functions are
82+
synchronous, even when they wrap a function of the other kind. Partials and
83+
callable instances are resolved through ``func`` and ``__call__``. Only a
84+
callable that is neither falls back to the function it wraps.
85+
"""
7986
# Keep nodes alive so temporary bound-method IDs cannot be reused.
8087
seen: dict[int, Any] = {}
88+
wrapped_fallbacks: list[Any] = []
89+
node: Any = fn
8190
while True:
82-
identity = id(fn)
91+
if node is None:
92+
if not wrapped_fallbacks:
93+
return False
94+
node = wrapped_fallbacks.pop()
95+
identity = id(node)
8396
if identity in seen:
8497
raise ValueError("wrapper loop when classifying async callable")
85-
seen[identity] = fn
86-
if inspect.iscoroutinefunction(fn):
98+
seen[identity] = node
99+
if inspect.iscoroutinefunction(node):
87100
return True
88-
if hasattr(fn, "__wrapped__"):
89-
fn = fn.__wrapped__
90-
continue
91-
if isinstance(fn, partial):
92-
fn = fn.func
93-
continue
94-
if (inspect.isfunction(fn) or inspect.ismethod(fn)
95-
or inspect.isbuiltin(fn) or inspect.ismethodwrapper(fn)):
96-
return False
97-
fn = getattr(fn, "__call__", None)
98-
if fn is None:
101+
if inspect.isgeneratorfunction(node) or inspect.isasyncgenfunction(node):
99102
return False
103+
if hasattr(node, "__wrapped__"):
104+
wrapped_fallbacks.append(node.__wrapped__)
105+
if isinstance(node, partial):
106+
node = node.func
107+
elif (inspect.isfunction(node) or inspect.ismethod(node)
108+
or inspect.isbuiltin(node) or inspect.ismethodwrapper(node)):
109+
node = None
110+
else:
111+
node = getattr(node, "__call__", None)
100112

101113

102114
def _get_legacy_logging_warning_stacklevel() -> int:

‎durabletask/task.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,9 @@
44
# See https://peps.python.org/pep-0563/
55
from __future__ import annotations
66

7+
import asyncio
78
import logging
89
import math
9-
import asyncio
1010
from abc import ABC, abstractmethod
1111
from collections.abc import Callable, Coroutine, Generator, Sequence
1212
from dataclasses import dataclass

‎durabletask/worker.py‎

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -350,6 +350,7 @@ class _Registry:
350350

351351
def __init__(self):
352352
self.orchestrators = {}
353+
self._orchestrator_kinds: dict[str, tuple[task.Orchestrator[Any, Any], bool]] = {}
353354
self.activities = {}
354355
self.entities = {}
355356

@@ -367,11 +368,20 @@ def add_named_orchestrator(self, name: str, fn: task.Orchestrator[TInput, TOutpu
367368
if name in self.orchestrators:
368369
raise ValueError(f"A '{name}' orchestrator already exists.")
369370

371+
# Classify once so wrapper problems surface at registration, not dispatch.
372+
is_async = shared.is_async_callable(fn)
370373
self.orchestrators[name] = fn
374+
self._orchestrator_kinds[name] = (fn, is_async)
371375

372376
def get_orchestrator(self, name: str) -> task.Orchestrator[Any, Any] | None:
373377
return self.orchestrators.get(name)
374378

379+
def is_async_orchestrator(self, name: str, fn: task.Orchestrator[Any, Any]) -> bool:
380+
kind = self._orchestrator_kinds.get(name)
381+
if kind is not None and kind[0] is fn:
382+
return kind[1]
383+
return shared.is_async_callable(fn)
384+
375385
def add_activity(self, fn: task.Activity[TInput, TOutput]) -> str:
376386
if fn is None:
377387
raise ValueError("An activity function argument is required.")
@@ -2745,7 +2755,7 @@ def process_event(
27452755
event.executionStarted.input.value, input_type)
27462756

27472757
# Layer 4 passes async orchestrators a cold context instead.
2748-
is_async = shared.is_async_callable(fn)
2758+
is_async = self._registry.is_async_orchestrator(event.executionStarted.name, fn)
27492759
result = fn(ctx, input)
27502760

27512761
if isinstance(result, CoroutineType) and not is_async:
@@ -2757,7 +2767,11 @@ def process_event(
27572767
if is_async and not isinstance(result, CoroutineType):
27582768
if isinstance(result, GeneratorType):
27592769
result.close()
2760-
raise TypeError("An async orchestrator must return a coroutine.")
2770+
raise TypeError(
2771+
"An async orchestrator must return a coroutine. A synchronous wrapper "
2772+
"decorated with functools.wraps(async_fn) is classified as async; declare "
2773+
"the wrapper with 'async def', or do not copy the async function's metadata."
2774+
)
27612775
if isinstance(result, (GeneratorType, CoroutineType)):
27622776
# Both protocols suspend on durable tasks and finish with StopIteration.
27632777
ctx.run(cast(Generator[task.Task[Any], Any, Any] | Coroutine[Any, Any, Any], result))

‎tests/durabletask/test_coroutine_orchestrators.py‎

Lines changed: 65 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -10,17 +10,21 @@
1010
import subprocess
1111
import sys
1212
from collections.abc import AsyncGenerator, Generator
13-
from datetime import datetime, timedelta
13+
from datetime import UTC, datetime, timedelta, timezone
1414
from functools import partial, update_wrapper, wraps
15-
from typing import Any
15+
from types import SimpleNamespace
16+
from typing import Any, cast
1617

1718
import pytest
1819

1920
from durabletask import entities, task, worker
2021
from durabletask.internal import helpers, shared
2122
from durabletask.internal import orchestrator_service_pb2 as pb
23+
from durabletask.internal.asyncio_loop import DurableEventLoop
2224
from durabletask.serialization import JsonDataConverter
2325

26+
pytestmark = pytest.mark.filterwarnings("error::RuntimeWarning")
27+
2428

2529
INSTANCE_ID = "coroutine-test"
2630
NOW = datetime(2026, 1, 1)
@@ -211,14 +215,55 @@ def wrapper(*args: Any, **kwargs: Any) -> Generator[int, None, None]:
211215
assert inspect.getgeneratorstate(created[0]) == inspect.GEN_CLOSED
212216

213217

214-
def test_cyclic_async_wrapper_fails_instead_of_hanging() -> None:
218+
@pytest.mark.parametrize("shape", ["function", "partial", "callable"])
219+
def test_generator_wrapper_of_async_function_is_synchronous(shape: str) -> None:
220+
async def async_orchestrator(ctx: task.OrchestrationContext, value: int) -> int:
221+
return value
222+
223+
@wraps(async_orchestrator)
224+
def generator(ctx: task.OrchestrationContext, value: int) -> Generator[task.Task[Any], Any, int]:
225+
return (yield ctx.call_activity("work", input=value, return_type=int))
226+
227+
class GeneratorCallable:
228+
def __call__(self, ctx: task.OrchestrationContext, value: int) -> Generator[task.Task[Any], Any, int]:
229+
return (yield ctx.call_activity("work", input=value, return_type=int))
230+
231+
update_wrapper(GeneratorCallable.__call__, async_orchestrator)
232+
variants: dict[str, Any] = {
233+
"function": generator,
234+
"partial": partial(generator),
235+
"callable": GeneratorCallable(),
236+
}
237+
fn = variants[shape]
238+
assert not shared.is_async_callable(fn)
239+
actions = _execute(fn, [], _start()).actions
240+
assert [a.scheduleTask.name for a in actions] == ["work"]
241+
242+
243+
def test_orchestrator_is_classified_once_at_registration(monkeypatch: pytest.MonkeyPatch) -> None:
244+
async def orchestrator(ctx: task.OrchestrationContext, value: int) -> int:
245+
return await ctx.call_activity("work", input=value, return_type=int)
246+
247+
registry = worker._Registry() # pyright: ignore[reportPrivateUsage]
248+
registry.add_named_orchestrator("test", orchestrator)
249+
calls: list[Any] = []
250+
original = shared.is_async_callable
251+
monkeypatch.setattr(shared, "is_async_callable", lambda fn: calls.append(fn) or original(fn))
252+
executor = worker._OrchestrationExecutor( # pyright: ignore[reportPrivateUsage]
253+
registry, logging.getLogger(__name__), JsonDataConverter())
254+
old = _start() + [helpers.new_task_scheduled_event(1, "work")]
255+
executor.execute(INSTANCE_ID, [], _start())
256+
executor.execute(INSTANCE_ID, old, [helpers.new_task_completed_event(1, "4")])
257+
assert calls == []
258+
259+
260+
def test_cyclic_async_wrapper_is_rejected_at_registration() -> None:
215261
def wrapper(ctx: task.OrchestrationContext, value: int) -> int:
216262
return value
217263

218264
wrapper.__wrapped__ = wrapper # type: ignore[attr-defined]
219-
complete = _completion(_execute(wrapper, [], _start()))
220-
assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED
221-
assert "wrapper loop" in complete.failureDetails.errorMessage
265+
with pytest.raises(ValueError, match="wrapper loop"):
266+
worker._Registry().add_named_orchestrator("test", wrapper) # pyright: ignore[reportPrivateUsage]
222267

223268

224269
def _assert_mixed_wrapper_cycle_fails(shape: str) -> None:
@@ -235,21 +280,16 @@ def __call__(self, ctx: task.OrchestrationContext, value: int) -> int:
235280
wrapper.__wrapped__ = CallableWrapper() # type: ignore[attr-defined]
236281
CallableWrapper.__call__.__wrapped__ = partial(wrapper) # type: ignore[attr-defined]
237282

238-
# No input: dispatch must detect the cycle without relying on type discovery.
239-
history = [
240-
helpers.new_orchestrator_started_event(NOW),
241-
helpers.new_execution_started_event("test", INSTANCE_ID),
242-
]
243-
complete = _completion(_execute(wrapper, [], history))
244-
assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED
245-
assert complete.failureDetails.errorType == "builtins.ValueError"
246-
assert "wrapper loop" in complete.failureDetails.errorMessage
283+
registry = worker._Registry() # pyright: ignore[reportPrivateUsage]
284+
with pytest.raises(ValueError, match="wrapper loop"):
285+
registry.add_named_orchestrator("test", wrapper)
286+
assert registry.get_orchestrator("test") is None
247287
with pytest.raises(ValueError, match="wrapper loop"):
248288
shared.is_async_callable(wrapper)
249289

250290

251291
@pytest.mark.parametrize("shape", ["partial", "callable-partial"])
252-
def test_mixed_wrapper_cycle_fails_within_bounded_worker_dispatch(shape: str) -> None:
292+
def test_mixed_wrapper_cycle_is_rejected_within_bounded_registration(shape: str) -> None:
253293
result = subprocess.run(
254294
[
255295
sys.executable, "-c",
@@ -555,3 +595,12 @@ async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> AsyncGenerator
555595
complete = _completion(_execute(orchestrator, [], _start()))
556596
assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED
557597
assert "Async generator orchestrators are not supported" in complete.failureDetails.errorMessage
598+
599+
600+
@pytest.mark.parametrize(("now", "offset"), [
601+
(NOW, 0), (NOW.replace(tzinfo=UTC), 0), (NOW.replace(tzinfo=timezone(timedelta(hours=2))), 7200),
602+
])
603+
def test_loop_time_accepts_naive_and_aware_orchestration_clocks(now: datetime, offset: int) -> None:
604+
ctx = cast(task.OrchestrationContext, SimpleNamespace(current_utc_datetime=now))
605+
loop = DurableEventLoop(ctx, is_complete=lambda: False)
606+
assert loop.time() == (NOW - datetime(1970, 1, 1)).total_seconds() - offset

0 commit comments

Comments
 (0)