diff --git a/.github/workflows/durabletask.yml b/.github/workflows/durabletask.yml index 8b70cff5..1780df61 100644 --- a/.github/workflows/durabletask.yml +++ b/.github/workflows/durabletask.yml @@ -14,6 +14,28 @@ 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 -r requirements.txt -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: diff --git a/CHANGELOG.md b/CHANGELOG.md index 09b6148d..8aebacfe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,23 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## Unreleased +ADDED + +- Added native `async def` orchestrators and nested async helpers alongside +existing generator orchestrators. Await durable operations and compose them +with `asyncio.gather`, local tasks, and structured task groups. Async sleeps and +timeouts use durable timers and the orchestration clock; `sleep(0)` yields +locally without a timer. See the [supported contract](docs/async-orchestrators.md) +for cancellation, per-dispatch cleanup, replay/migration rules, and unsupported +APIs. External I/O still belongs in ordinary activities. +Sequential async batches release discarded results and observed-failure +tracebacks during replay. +- Added `OrchestrationContext.as_async()` for cold native coroutine operations +and owned timer/event waits, including event receive timeouts that withdraw +abandoned subscriptions. Added `Task.wait()` for independently cancellable +borrowers of existing durable handles. Existing eager and generator APIs are +unchanged. + ## v1.13.0 ADDED diff --git a/README.md b/README.md index fb947a98..5bc00318 100644 --- a/README.md +++ b/README.md @@ -32,6 +32,53 @@ The Azure Functions Durable provider requires Python 3.13 or later. ## Optional Features +### Async Orchestrators + +All three packages support `async def` orchestrators and nested async helpers. +Existing generator orchestrators continue to work. + +```python +import asyncio + +from durabletask import task + + +async def sequence(ctx: task.OrchestrationContext, value: int) -> list[int]: + ops = ctx.as_async() + first = await ops.call_activity("plus_one", input=value, return_type=int) + async with asyncio.TaskGroup() as group: + left = group.create_task(ops.call_activity("plus_one", input=first, return_type=int)) + right = group.create_task(ops.call_activity("plus_one", input=value, return_type=int)) + return [left.result(), right.result()] +``` + +Register the function with `TaskHubGrpcWorker`, `DurableTaskSchedulerWorker`, +or the Azure Functions `orchestration_trigger` decorator as usual. Functions +also supports single-argument `async def` orchestrators using the compatibility +context and `await context.task_all(...)`. +Use `ctx.as_async()` for cold native coroutines accepted by `create_task` and +task groups. Existing context methods still start eager, reusable durable +handles; `handle.wait()` creates an independently cancellable borrower. + +> [!WARNING] +> This is replay-driven await syntax, not general-purpose asyncio execution. +> Compose durable work with `asyncio.gather`, local tasks/task groups, and +> durable `asyncio.sleep` / timeouts within the +> [supported contract](docs/async-orchestrators.md). +> Perform network, storage, and other nondeterministic I/O in activities. +> Completed foreign awaitables and synchronous side effects cannot reliably be +> detected by the replay engine. +> Python `finally` blocks and async context-manager exits can run at ordinary +> dispatch suspension; they are not workflow-lifetime compensation scopes. + +An async helper is not a durable sub-orchestration: use +`await ctx.call_sub_orchestrator(...)` for a separate orchestration instance. +An `async def` function containing `yield` is an async generator, not a supported +orchestrator. Converting an existing orchestrator must preserve durable action +order and payloads for replay compatibility. +See the contract for cancellation, per-dispatch async cleanup, unsupported +ordering-sensitive APIs and framework adapters. + ### Task Hub Query Filters Set `OrchestrationQuery.task_hub_names` to filter orchestration queries by one or diff --git a/azure-functions-durable/CHANGELOG.md b/azure-functions-durable/CHANGELOG.md index 4b808272..63beae82 100644 --- a/azure-functions-durable/CHANGELOG.md +++ b/azure-functions-durable/CHANGELOG.md @@ -7,6 +7,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## Unreleased +ADDED + +- With the corresponding core SDK, native and v1-style orchestrators can use +`async def`, await durable tasks, and use deterministic asyncio composition and +durable sleeps/timeouts. See the [async contract](../docs/async-orchestrators.md). +The `orchestrator_generator_wrapper` helper accepts direct-await coroutines; +these work in synchronous and async tests. Asyncio scheduling requires +replay/backend tests. Sequential async batches release discarded results and +observed-failure tracebacks during replay. +- Added `as_async()` to native and compatibility orchestration contexts for +cold coroutine operations and owned timer/event waits. Event receive timeouts +withdraw abandoned subscriptions; existing eager APIs remain unchanged. + ## v2.0.0rc4 ADDED diff --git a/azure-functions-durable/README.md b/azure-functions-durable/README.md index 9bf0d77b..a3ab45ff 100644 --- a/azure-functions-durable/README.md +++ b/azure-functions-durable/README.md @@ -40,6 +40,28 @@ or more task hubs. Filtering support depends on the backend. so they can be used as dictionary keys or set elements. Construct a new ID instead of assigning or deleting attributes on an existing ID. +## Async orchestrators + +Both authoring styles accept `async def` and `await` without an adapter loop. +Await durable operations and compose branches with `asyncio.gather` or task +groups; `asyncio.sleep` and timeouts use durable timers. Synchronous generator +orchestrators remain supported, including mixed-style sub-orchestrations. + +Both native and compatibility contexts expose `as_async()` for cold native +coroutines accepted by `create_task` and task groups. This facade uses the +native SDK's keyword arguments (`input=`, `return_type=`, and `data_type=`), +not compatibility positional arguments. Its `receive_event()` owns its +registration and optional durable timeout, so an abandoned receive cannot +consume the next event. Existing compatibility methods still return eager, +reusable durable handles; use `handle.wait()` to create a borrowed coroutine. + +The [shared async contract](../docs/async-orchestrators.md) defines the supported +API subset, cancellation, per-dispatch `finally`/async-generator cleanup, +replay/migration rules, and framework adapter responsibilities. Arbitrary I/O, +threads, subprocesses, `asyncio.wait`/`as_completed`, and automatic recording of +ordinary async helpers are not supported. Put nondeterministic work in ordinary +activities, not orchestrators. + ## Unicode JSON serialization By default, durable payloads preserve the shared `azure-functions` codec's diff --git a/azure-functions-durable/azure/durable_functions/internal/compat/orchestration_context.py b/azure-functions-durable/azure/durable_functions/internal/compat/orchestration_context.py index f93176dd..83849252 100644 --- a/azure-functions-durable/azure/durable_functions/internal/compat/orchestration_context.py +++ b/azure-functions-durable/azure/durable_functions/internal/compat/orchestration_context.py @@ -41,6 +41,15 @@ def __init__(self, # an empty bag -- matching the common v1 runtime case. self._function_context = FunctionContext() + def as_async(self) -> task.AsyncOrchestrationContext: + """Get cold coroutine operations using the native SDK's keyword arguments. + + Existing compatibility methods keep their eager, reusable task handles. + The async surface owns its event receives and timers, not remote activity + or sub-orchestration execution. + """ + return self._ctx.as_async() + # -- input --------------------------------------------------------------- def get_input(self, expected_type: Optional[type] = None) -> Any: """Get the orchestration input. diff --git a/azure-functions-durable/azure/durable_functions/testing/orchestrator_generator_wrapper.py b/azure-functions-durable/azure/durable_functions/testing/orchestrator_generator_wrapper.py index 43daf0c5..491b10c0 100644 --- a/azure-functions-durable/azure/durable_functions/testing/orchestrator_generator_wrapper.py +++ b/azure-functions-durable/azure/durable_functions/testing/orchestrator_generator_wrapper.py @@ -1,16 +1,18 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -"""Replay simulation helper for unit-testing orchestrator generators.""" +"""Replay simulation helper for unit-testing generator or coroutine orchestrators.""" -from typing import Any, Generator +import asyncio +from collections.abc import Coroutine, Generator +from typing import Any def orchestrator_generator_wrapper( - generator: Generator[Any, Any, Any]) -> Generator[Any, None, None]: - """Drive a user orchestrator generator the way the Durable replay engine does. + generator: Generator[Any, Any, Any] | Coroutine[Any, Any, Any]) -> Generator[Any, None, None]: + """Drive a user orchestrator the way the Durable replay engine does. - A Durable Functions orchestrator is a generator that yields task objects and + A Durable Functions orchestrator yields or awaits durable task objects and is resumed with each task's result until it returns the orchestration output. In a real execution the Durable Task runtime performs this drive loop during replay. This wrapper reproduces that loop in-process so an @@ -20,10 +22,14 @@ def orchestrator_generator_wrapper( back into the orchestrator (mirroring how failed activities, timers, or other tasks surface during replay). + This mock-based helper handles direct durable awaits in synchronous or + async tests. It does not simulate the runtime's asyncio scheduler; use + replay/backend tests for gather, task groups, sleeps and timeouts. + Parameters ---------- - generator: Generator[Any, Any, Any] - The user orchestrator generator. It is expected to yield task objects + generator: Generator[Any, Any, Any] | Coroutine[Any, Any, Any] + The user orchestrator generator or coroutine. It must suspend on tasks that expose a ``result`` attribute and to be resumed with those results until it returns the orchestration output. @@ -33,14 +39,24 @@ def orchestrator_generator_wrapper( Each task object yielded by the orchestrator, followed by the orchestrator's final return value. """ + def step(value: Any = None, *, error: Exception | None = None) -> Any: + previous_loop = asyncio._get_running_loop() + try: + if isinstance(generator, Coroutine): + # Direct durable awaits belong to this mock driver, not the + # caller's asyncio loop (or an enclosing durable dispatch). + asyncio._set_running_loop(None) + return generator.throw(error) if error is not None else generator.send(value) + finally: + asyncio._set_running_loop(previous_loop) + try: - previous = next(generator) + previous = step() except StopIteration as e: # An orchestrator that returns immediately without yielding any task # completes right away (mirroring the runtime, which finishes the # orchestration on the first return). Surface its output as the final - # value and stop -- ``next`` raising here otherwise leaks as an - # unhandled ``StopIteration``. + # value and stop -- letting StopIteration escape would become RuntimeError. yield e.value return yield previous @@ -52,9 +68,9 @@ def orchestrator_generator_wrapper( # Simulated activity exceptions, timer-interrupted exceptions, # or any other case where reading a task result would throw: # push the exception back into the orchestrator as replay does. - previous = generator.throw(e) + previous = step(error=e) else: - previous = generator.send(previous_result) + previous = step(previous_result) yield previous except StopIteration as e: # The orchestrator returned; surface its output as the final value. diff --git a/docs/async-orchestrators.md b/docs/async-orchestrators.md new file mode 100644 index 00000000..553f92a1 --- /dev/null +++ b/docs/async-orchestrators.md @@ -0,0 +1,246 @@ +# Async orchestrators + +All three packages accept native `async def` orchestrators alongside existing +synchronous generator orchestrators. Both styles use the same history, action, +retry, payload, and entity engine. Register the function normally with +`TaskHubGrpcWorker`, `DurableTaskSchedulerWorker`, or the Functions +`orchestration_trigger` decorator. Functions supports both native two-argument +and compatibility single-argument orchestrators. + +```python +import asyncio +from datetime import timedelta + +from durabletask import task + + +async def workflow(ctx: task.OrchestrationContext, values: list[int]) -> list[int]: + ops = ctx.as_async() + + async def branch(value: int) -> int: + first = await ops.call_activity("increment", input=value, return_type=int) + await ops.create_timer(timedelta(seconds=1)) + return await ops.call_activity("increment", input=first, return_type=int) + + async with asyncio.TaskGroup() as group: + branches = [group.create_task(branch(value)) for value in values] + return [branch.result() for branch in branches] +``` + +Each branch resumes as its own durable operation completes, including when +multiple completions arrive in one dispatch. It does not wait for an entire +batch before scheduling its next operation. + +> [!WARNING] +> Python scopes are not persisted workflow resources. `finally` blocks and async +> context-manager exits can run when an ordinary dispatch suspends, not only +> when the workflow finishes. Keep durable compensation in normal orchestration +> control flow; see [dispatch teardown](#dispatch-teardown-is-not-workflow-cancellation). + +## Cold operations and reusable handles + +`ctx.as_async()` returns an explicit coroutine-native operation facade. It does +not replace the registered context, change generator behavior, or create a +second runtime. Functions compatibility contexts expose the same facade using +native keyword arguments such as `input=` and `data_type=`. + +| Expression | Meaning | +|---|---| +| `ctx.call_activity(...)` | Schedule eagerly and return a reusable durable handle. | +| `ops.call_activity(...)` | Create a single-use native coroutine; schedule when it starts. | +| `handle.wait()` | Create a native coroutine borrowing an existing handle. | +| `ops.receive_event(...)` | Own an event registration, withdrawing it on cancellation or timeout. | +| `ops.create_timer(...)` | Own a timer, cancelling its pending local work on cancellation. | + +Use the cold facade with `asyncio.create_task`, `TaskGroup.create_task`, and +ordered `gather`. Constructing a coroutine does not allocate action IDs, +serialize input, or consume buffered events. Cancelling its local task before +its first step starts no durable operation. Mutable inputs are read when the +coroutine executes, and relative timer delays start from that execution point. +Keep or copy inputs accordingly. + +Cold coroutines are single-use, like ordinary Python coroutines. For an +explicitly started operation with multiple independent waiters, retain its +durable handle: + +```python +operation = ctx.call_activity("quote", input="order-1", return_type=int) +first = asyncio.create_task(operation.wait()) +second = asyncio.create_task(operation.wait()) +first.cancel() +value = await second +``` + +Cancelling either borrower does not cancel the underlying operation. Direct +`await operation` remains supported, but `create_task(operation)` is not valid: +an awaitable durable handle is not a coroutine. Do not retry that expression +assuming the eager operation was rolled back. + +## Supported authoring contract + +Await activities, timers, external events, sub-orchestrations, durable +`task.when_all` / `task.when_any`, and entity calls/locks where supported by the +provider. Functions compatibility operations such as `context.task_all()` +remain awaitable. Generator and coroutine parents/children may call one another +through durable sub-orchestrations. + +The dispatch-local loop supports FIFO ready callbacks, local futures and tasks, +`asyncio.create_task`, `asyncio.gather`, `asyncio.sleep`, `asyncio.TaskGroup`, +and `asyncio.timeout` / `timeout_at`. +`asyncio.wait_for` uses the same task/future scheduling and durable timers. +Third-party async libraries are not part of this compatibility contract, even +when individual code paths happen to use supported asyncio primitives. + +Create branches from ordered inputs. Ready work follows FIFO scheduling; +history completions are delivered in history order, draining ready work after +each completion. Task-group cancellation wakeups are ordered by local task +creation, not set/object-hash iteration. Local tasks are dispatch resources, +not new durable instances. + +`loop.time()` uses seconds from a fixed UTC epoch and the current orchestration +clock, never machine monotonic time. `loop.call_at` and `call_later` schedule +durable timers for future finite deadlines. Timer conversion rounds up when +necessary so a durable timer does not precede its floating-point deadline. +Core/Functions long timers reuse the existing chunking; Azure Managed retains +its native long timers. Past deadlines are ready callbacks. `sleep(0)` yields +to ready coroutines without a durable timer. Clock advancement and timers +require history events, not elapsed time in the worker. + +> [!WARNING] +> This is not a general-purpose asyncio event loop. External network/storage +> I/O, threads, executors, subprocesses, custom/eager task factories, and +> cross-thread callbacks are outside the orchestrator contract. Perform +> nondeterministic work in ordinary activities. Detected unsupported operations +> raise errors, but the SDK cannot detect every unsafe operation, especially +> synchronous side effects or already-completed foreign awaitables. + +`asyncio.wait` and `as_completed` are **not supported deterministic composition +APIs**. `wait` returns unordered sets; `as_completed` converts inputs to a set +and can create tasks in hash-dependent order. Pre-completed/tied tasks and +differences between Python implementations can change ordering even when a +particular example executes successfully. Do not iterate sets to schedule +durable actions. Use ordered `gather`, task groups, or durable `when_any` and +explicitly handle the selected durable task. Other asyncio APIs are not +implicitly supported because they happen to execute. + +## Failures, timeouts, and cancellation + +Durable failures raise the existing SDK exceptions at the await, preserving +failure details, typed results, and retry policies. `task.when_all` still waits +for every durable child before surfacing failure. In contrast, +`asyncio.gather` with default options reports the first failure without +cancelling siblings. Retain and await siblings if their results matter. +`return_exceptions=True` collects their outcomes in input order. +Unobserved local background-task failures are reported during dispatch cleanup; +ordinary exceptions are logged, not silently discarded or automatically joined. +Discarded local-task results and observed-failure tracebacks are released during +replay rather than retained until dispatch cleanup. + +Task groups cancel and drain local sibling tasks on failure and use asyncio's +normal exception grouping rules. Python-version-specific exception grouping +and cancel-count behavior remain those of asyncio. + +Cancelling a borrower of an existing durable handle, including `handle.wait()`, +detaches that waiter; it does **not** cancel the activity, sub-orchestration, +entity operation, or external-event subscription. Another borrower can still +await the same handle. Keep the handle if you intend to await it again. +Use explicit SDK cancellation when abandoning a cancellable shared timer/event +wait and deliberately affecting every borrower. + +The cold facade owns its private timer and event registrations. For an event +with a deadline, use the owned receive instead of timing out a borrowed handle: + +```python +ops = ctx.as_async() +try: + approval = await ops.receive_event( + "approval", timeout=timedelta(seconds=30), data_type=bool) +except TimeoutError: + approval = await ops.receive_event("approval", data_type=bool) +``` + +The timed-out registration is withdrawn before the replacement is registered, +so it cannot consume the next approval. The receive also cancels its own +deadline timer when the event wins. A nonpositive timeout polls an already +buffered event and otherwise raises `TimeoutError` without creating a timer. +This ownership does not extend to an +activity, sub-orchestration, or entity call: cancellation cannot stop already +issued remote work. A task-group sibling failure can therefore leave remote +operations running. Retain and join their handles when completion is required +before compensation or releasing application-level capacity. + +Cancelling a timer callback removes an uncommitted timer when possible, but +cannot undo a timer already committed to history. Neither local nor durable +cancellation implies external side effects were undone. + +An unhandled local root-coroutine cancellation fails the orchestration with +`TaskCancelledError`; it does not request scheduler termination or repeatedly +abandon the same deterministic replay. Use the client to terminate the durable +instance. Worker/dispatch cancellation and process-control exceptions continue +to follow the existing worker recovery policy. + +## Dispatch teardown is not workflow cancellation + +Every dispatch reconstructs the coroutine and its local tasks from history. +After the response and entity-release obligations are snapshotted, teardown +cancels remaining local tasks, drains cooperative cleanup, closes tracked async +generators, and disposes the loop. This also happens on suspension, termination, +continue-as-new, history failure, and version rejection. + +Nested `finally` blocks and async context-manager exits may run at **every +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. + +Ordinary teardown failures (including `CancelledError`) are logged without +changing the decided response or skipping other cooperative finalizers. +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. + +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. + +## Replay and framework adapters + +Ordinary async helpers rerun during replay. Their return values are not +automatically checkpointed. There are no local/virtual activities and no +automatic recording or interception of arbitrary await expressions. +`is_replaying` is not a substitute for recording results. Activities remain +at-least-once work: side effects require appropriate idempotency/deduplication. + +Framework adapters can compose supported durable operations using async +helpers. Adapters remain responsible for registering ordinary model/tool +activities, encoding/decoding framework payloads, and supplying replay-stable +IDs, clocks, and message metadata. These responsibilities are not handled by +SDK-wide interception. Libraries that require unsupported asyncio APIs or +nondeterministic operations need an appropriate adapter or activity boundary. + +Changing an in-flight generator to equivalent coroutine code is safe **only +if its durable history contract is preserved**: action IDs/order, names, +payloads, retries, timer deadlines/chunks, event consumption, and entity +release order. Introducing concurrency, a different join, or a new timer may +change that contract. Test recorded histories and version/drain affected +instances; basic sequential parity tests are not blanket migration safety. +Switching from eager handles to cold operations can also change action order: +construction order is not execution order. + +Supporting several Python versions does not guarantee recorded-history +portability between them. Pin the runtime and dependencies for active +instances and validate histories before upgrades. In particular, Python 3.11 +and 3.12+ schedule coroutine operands of `asyncio.wait_for()` differently, which +can change durable action order inside concurrent branches. The cold facade +does not normalize arbitrary standard-library scheduling. + +The Functions `orchestrator_generator_wrapper` helper supports direct durable +awaits in synchronous and async mock-based unit tests; it does not simulate the +asyncio scheduler. +Use the real replay executor/provider or the in-memory backend to test task +groups, gather, and durable timeouts. diff --git a/docs/supported-patterns.md b/docs/supported-patterns.md index b67df626..1cbcfd29 100644 --- a/docs/supported-patterns.md +++ b/docs/supported-patterns.md @@ -2,6 +2,13 @@ The following orchestration patterns are currently supported. +Generator and native coroutine orchestrators share these durable patterns. +See the [async orchestrator contract](async-orchestrators.md) for supported +asyncio composition, durable clocks/timeouts, and dispatch-local cleanup. +Use `ctx.as_async()` for cold coroutine operations and owned event/timeouts; +existing eager handles remain available for explicit sharing through +`handle.wait()`. + ## Function chaining An orchestration can chain a sequence of function calls using the following syntax: @@ -81,6 +88,10 @@ to nested exception groups containing these signals on Python 3.11 and later. The existing warning-only handling of direct `asyncio.CancelledError` during forced teardown is unchanged. Other cancellation or process-control exits propagate after the worker attempts to abandon the accepted dispatch. +Native coroutine authoring distinguishes a cancelled local root task: an +unhandled local cancellation becomes `TaskCancelledError` and fails the +orchestration instead of repeatedly abandoning replay. It does not terminate +the durable instance; see the [async cancellation contract](async-orchestrators.md#failures-timeouts-and-cancellation). If a completion or abandonment RPC raises, its server-side outcome may be unknown. The worker logs the uncertainty and closes the dispatch's originating diff --git a/durabletask-azuremanaged/CHANGELOG.md b/durabletask-azuremanaged/CHANGELOG.md index 261117be..de9a55b9 100644 --- a/durabletask-azuremanaged/CHANGELOG.md +++ b/durabletask-azuremanaged/CHANGELOG.md @@ -7,6 +7,18 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## Unreleased +ADDED + +- With the corresponding core SDK, `DurableTaskSchedulerWorker` +accepts `async def` orchestrators that await durable tasks alongside generator +orchestrators, including deterministic asyncio composition and durable sleeps/ +timeouts. See the [async contract](../docs/async-orchestrators.md). +Sequential async batches release discarded results and observed-failure +tracebacks during replay. +- Added inherited `OrchestrationContext.as_async()` cold coroutine operations +and owned timer/event waits, plus `Task.wait()` for borrowing existing handles. +Remote operations are not cancelled when a local waiter is cancelled. + ## v1.13.0 ADDED diff --git a/durabletask/async_context.py b/durabletask/async_context.py new file mode 100644 index 00000000..9551c0a1 --- /dev/null +++ b/durabletask/async_context.py @@ -0,0 +1,223 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Cold coroutine operations over an existing durable orchestration context.""" + +from __future__ import annotations + +import math +from datetime import datetime, timedelta +from typing import TYPE_CHECKING, Any, TypeVar, overload + +from durabletask.entities import EntityInstanceId, EntityLock + +if TYPE_CHECKING: + from durabletask import task + +T = TypeVar("T") +TInput = TypeVar("TInput") +TOutput = TypeVar("TOutput") + + +class AsyncOrchestrationContext: + """A coroutine-native operation surface returned by ``ctx.as_async()``. + + Every operation returns a single-use native coroutine. Constructing it + schedules nothing: validation, input serialization, event consumption, + and relative deadline evaluation occur only when execution starts. Inputs + are not copied at construction; mutations before execution are observed. + Use the original context for properties, synchronous commands, and eager, + reusable durable handles; use ``handle.wait()`` for borrowed coroutine waits. + + Cancelling activity, sub-orchestration, entity, or lock waiters does not + cancel remote work. Timers and event receives own their private local + registrations and withdraw them when their coroutine exits. This cannot + undo actions already committed to history or external side effects. + """ + + def __init__(self, context: task.OrchestrationContext) -> None: + self._context = context + + @overload + async def call_activity( + self, activity: task.Activity[TInput, TOutput] | str, *, + input: TInput | None = ..., + retry_policy: task.RetryPolicy | None = ..., + tags: dict[str, str] | None = ..., + return_type: type[T], + ) -> T: + ... + + @overload + async def call_activity( + self, activity: task.Activity[TInput, TOutput], *, + input: TInput | None = ..., + retry_policy: task.RetryPolicy | None = ..., + tags: dict[str, str] | None = ..., + return_type: None = ..., + ) -> TOutput: + ... + + @overload + async def call_activity( + self, activity: str, *, input: Any = ..., + retry_policy: task.RetryPolicy | None = ..., + tags: dict[str, str] | None = ..., + return_type: None = ..., + ) -> Any: + ... + + async def call_activity( + self, activity: task.Activity[TInput, TOutput] | str, *, + input: TInput | None = None, + retry_policy: task.RetryPolicy | None = None, + tags: dict[str, str] | None = None, + return_type: type | None = None, + ) -> Any: + """Start and await an activity when this coroutine executes.""" + return await self._context.call_activity( + activity, input=input, retry_policy=retry_policy, tags=tags, + return_type=return_type) + + @overload + async def call_sub_orchestrator( + self, orchestrator: task.Orchestrator[TInput, TOutput] | str, *, + input: TInput | None = ..., + instance_id: str | None = ..., + retry_policy: task.RetryPolicy | None = ..., + version: str | None = ..., + return_type: type[T], + ) -> T: + ... + + @overload + async def call_sub_orchestrator( + self, orchestrator: task.Orchestrator[TInput, TOutput], *, + input: TInput | None = ..., + instance_id: str | None = ..., + retry_policy: task.RetryPolicy | None = ..., + version: str | None = ..., + return_type: None = ..., + ) -> TOutput: + ... + + @overload + async def call_sub_orchestrator( + self, orchestrator: str, *, input: Any = ..., + instance_id: str | None = ..., + retry_policy: task.RetryPolicy | None = ..., + version: str | None = ..., + return_type: None = ..., + ) -> Any: + ... + + async def call_sub_orchestrator( + self, orchestrator: task.Orchestrator[TInput, TOutput] | str, *, + input: TInput | None = None, + instance_id: str | None = None, + retry_policy: task.RetryPolicy | None = None, + version: str | None = None, + return_type: type | None = None, + ) -> Any: + """Start and await a sub-orchestration when this coroutine executes.""" + return await self._context.call_sub_orchestrator( + orchestrator, input=input, instance_id=instance_id, + retry_policy=retry_policy, version=version, return_type=return_type) + + @overload + async def call_entity( + self, entity: EntityInstanceId, operation: str, input: Any = ..., *, + return_type: type[T], + ) -> T: + ... + + @overload + async def call_entity( + self, entity: EntityInstanceId, operation: str, input: Any = ..., *, + return_type: None = ..., + ) -> Any: + ... + + async def call_entity( + self, entity: EntityInstanceId, operation: str, input: Any = None, *, + return_type: type | None = None, + ) -> Any: + """Start and await an entity operation when this coroutine executes.""" + return await self._context.call_entity( + entity, operation, input, return_type=return_type) + + async def lock_entities(self, entities: list[EntityInstanceId]) -> EntityLock: + """Acquire locks, for use as ``with await ops.lock_entities(ids):``. + + The returned synchronous context manager releases the locks. This + preserves the original context's entity locking and replay rules; + cancellation does not withdraw an already-issued lock request. + """ + return await self._context.lock_entities(entities) + + async def create_timer(self, fire_at: datetime | timedelta) -> None: + """Await a private durable timer, cancelling it when this wait exits. + + Relative deadlines are evaluated when execution starts, not when this + coroutine is constructed. Cancellation removes uncommitted timer work + where possible; it cannot retract a timer already present in history. + """ + timer = self._context.create_timer(fire_at) + try: + await timer + finally: + timer.cancel() + + @overload + async def receive_event( + self, name: str, *, timeout: timedelta | None = ..., + data_type: type[T], + ) -> T: + ... + + @overload + async def receive_event( + self, name: str, *, timeout: timedelta | None = ..., + data_type: None = ..., + ) -> Any: + ... + + async def receive_event( + self, name: str, *, timeout: timedelta | None = None, + data_type: type | None = None, + ) -> Any: + """Receive one event through a private, cancellation-owned subscription. + + A finite relative timeout starts when the coroutine executes and raises + ``TimeoutError`` if its durable timer wins. Buffered events win + immediately, even for zero/negative timeouts, without creating a timer. + Otherwise nonpositive timeouts fail immediately without a timer. + Pending races follow history order. On every exit, the subscription + and any outstanding timeout timer are withdrawn locally, so a later + receive cannot lose an event to this abandoned subscription. + + For a shared subscription use the original context's + ``wait_for_external_event()`` and borrowed ``Task.wait()`` instead. + """ + from durabletask.task import when_any + + if timeout is not None and not math.isfinite(timeout.total_seconds()): + raise ValueError("Event receive timeouts must be finite.") + event: task.CancellableTask[Any] = self._context.wait_for_external_event(name, data_type=data_type) + timer: task.TimerTask | None = None + try: + if event.is_complete: + return event.get_result() + if timeout is None: + return await event + if timeout <= timedelta(0): + raise TimeoutError(f"Timed out waiting for event {name!r}.") + timer = self._context.create_timer(timeout) + winner = await when_any([event, timer]) + if winner is timer: + raise TimeoutError(f"Timed out waiting for event {name!r}.") + return event.get_result() + finally: + event.cancel() + if timer is not None: + timer.cancel() diff --git a/durabletask/internal/asyncio_loop.py b/durabletask/internal/asyncio_loop.py new file mode 100644 index 00000000..490626eb --- /dev/null +++ b/durabletask/internal/asyncio_loop.py @@ -0,0 +1,422 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Dispatch-local asyncio scheduling on top of the shared durable history engine.""" + +from __future__ import annotations + +import asyncio +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") + + +class _CompletionObserver(task.CompositeTask[None]): + def __init__(self, durable: task.Task[Any], callback: Callable[[], None]) -> None: + self._callback = callback + super().__init__([durable]) + + def on_child_completed(self, task: task.Task[Any]) -> None: + if not self._is_complete: + self._is_complete = True + self._notify_parent() + self._callback() + + def detach(self) -> None: + self._notify_parent() + + +class _DurableAsyncTask(asyncio.Task[T]): + def __await__(self) -> Generator[Any, None, T]: + try: + return (yield from super().__await__()) + finally: + self._retire_if_observed() + + __iter__ = __await__ + + def result(self) -> T: + try: + return super().result() + finally: + self._retire_if_observed() + + def exception(self) -> BaseException | None: + try: + return super().exception() + finally: + self._retire_if_observed() + + def _retire_if_observed(self) -> None: + loop = self.get_loop() + assert isinstance(loop, DurableEventLoop) + loop.retire_task(self) + + def cancel(self, msg: Any = None) -> bool: + loop = self.get_loop() + assert isinstance(loop, DurableEventLoop) + previous = loop.cancelling_task + loop.cancelling_task = self + try: + return super().cancel(msg) + finally: + loop.cancelling_task = previous + + +class _DurableTimerHandle(asyncio.TimerHandle): + def __eq__(self, other: object) -> bool: + return self is other + + __hash__ = object.__hash__ + + def __init__( + self, when: float, callback: Callable[..., Any], args: tuple[Any, ...], + loop: DurableEventLoop, context: contextvars.Context | None, + ) -> None: + super().__init__(when, callback, args, loop, context=context) + self.durable_timer: task.TimerTask | None = None + + def cancel(self) -> None: + super().cancel() + if self.durable_timer is not None: + self.durable_timer.cancel() + + +class DurableEventLoop(asyncio.AbstractEventLoop): + """A FIFO microtask loop; it never polls a selector or advances wall time. + + Only history events can complete durable waits/timers. asyncio owns coroutine + stepping, gather and structured cancellation; this loop owns their ordering. + """ + + def __init__( + self, ctx: task.OrchestrationContext, *, is_complete: Callable[[], bool], + logger: logging.Logger | None = None, + ) -> None: + self.ctx = ctx + self._is_complete = is_complete + self._logger = logger or logging.getLogger(__name__) + self._ready: deque[asyncio.Handle] = deque() + self._cancel_ready: list[tuple[int, int, asyncio.Handle]] = [] + self._tasks: dict[asyncio.Task[Any], int] = {} + self._task_sequence = 0 + self.cancelling_task: asyncio.Task[Any] | None = None + self._waits: dict[int, tuple[task.Task[Any], list[asyncio.Future[Any]]]] = {} + 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: + # A fixed UTC epoch avoids local-time and platform timestamp differences. + return (self.ctx.current_utc_datetime - datetime(1970, 1, 1)).total_seconds() + + def get_debug(self) -> bool: + return False + + def is_running(self) -> bool: + return asyncio._get_running_loop() is self + + def is_closed(self) -> bool: + return self._closed + + def create_future(self) -> asyncio.Future[Any]: + return asyncio.Future(loop=self) + + def get_task_factory(self) -> None: + return None + + def create_task( + self, coro: Coroutine[Any, Any, T] | Generator[Any, None, T], *, name: str | None = None, + context: contextvars.Context | None = None, **kwargs: Any, + ) -> asyncio.Task[T]: + if self._closed: + coro.close() + raise RuntimeError("The durable orchestration dispatch is closed.") + if kwargs: + coro.close() + raise TypeError("Custom/eager task execution is not supported in durable orchestrators.") + future = _DurableAsyncTask(coro, loop=self, name=name, context=context) + self._tasks[future] = self._task_sequence + self._task_sequence += 1 + future.add_done_callback(self.retire_task) + return future + + def retire_task(self, future: asyncio.Task[Any]) -> None: + # A failed task may be observed after its done callbacks have run. + # Inspect without retrieving exceptions so unobserved failures still log. + if not self._closing and future.done() and not getattr(future, "_log_traceback", False): + self._tasks.pop(future, None) + + def call_soon( + self, callback: Callable[..., Any], *args: Any, + context: contextvars.Context | None = None, + ) -> asyncio.Handle: + handle = asyncio.Handle(callback, args, self, context=context) + if self.cancelling_task is not None: + # TaskGroup cancels a set of tasks. Do not let object hashes + # determine the order in which their finally blocks execute. + owner = getattr(callback, "__self__", None) + index = self._tasks[self.cancelling_task] + if isinstance(owner, asyncio.Task): + index = self._tasks.get(cast(asyncio.Task[Any], owner), index) + self._cancel_ready.append((index, len(self._cancel_ready), handle)) + else: + self._ready.append(handle) + return handle + + def call_later( + self, delay: float, callback: Callable[..., Any], *args: Any, + context: contextvars.Context | None = None, + ) -> asyncio.TimerHandle: + if not math.isfinite(delay): + raise ValueError("Durable delays must be finite.") + return self.call_at(self.time() + max(0, delay), callback, *args, context=context) + + def call_at( + self, when: float, callback: Callable[..., Any], *args: Any, + context: contextvars.Context | None = None, + ) -> asyncio.TimerHandle: + if not math.isfinite(when): + raise ValueError("Durable deadlines must be finite.") + handle = _DurableTimerHandle(when, callback, args, self, context) + if when <= self.time(): + self._ready.append(handle) + else: + if self._closing: + raise RuntimeError("Dispatch cleanup cannot schedule durable timers.") + deadline = datetime(1970, 1, 1) + timedelta(seconds=when) + if (deadline - datetime(1970, 1, 1)).total_seconds() < when: + deadline += timedelta(microseconds=1) + handle.durable_timer = self.ctx.create_timer(deadline) + self._timers[handle] = _CompletionObserver( + handle.durable_timer, + lambda: self._completions.append(lambda: self._complete_timer(handle))) + return handle + + def _timer_handle_cancelled(self, handle: asyncio.TimerHandle) -> None: + if isinstance(handle, _DurableTimerHandle): + observer = self._timers.pop(handle, None) + if observer is not None: + observer.detach() + + def _complete_timer(self, handle: _DurableTimerHandle) -> None: + self._timers.pop(handle, None) + if not handle.cancelled(): + 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() + entry = self._waits.setdefault(id(durable), (durable, [])) + entry[1].append(future) + if id(durable) not in self._observers: + self._observers[id(durable)] = _CompletionObserver( + durable, lambda: self._completions.append(lambda: self._complete_wait(durable))) + try: + return await future + finally: + # Cancelling one waiter never cancels a shared durable operation. + entry[1].remove(future) + if not entry[1]: + self._waits.pop(id(durable), None) + self._observers.pop(id(durable)).detach() + + def _complete_wait(self, durable: task.Task[Any]) -> None: + entry = self._waits.get(id(durable)) + if entry is not None: + for future in tuple(entry[1]): + if not future.done(): + try: + future.set_result(durable.get_result()) + except Exception as error: + future.set_exception(error) + + def _deliver_completions(self) -> None: + while self._completions: + self._completions.popleft()() + + def start(self, coro: Coroutine[Any, Any, Any]) -> None: + self._root = self.create_task(coro, name="durable-orchestrator") + self.resume() + + def resume(self) -> None: + if self._is_complete(): + return + self._deliver_completions() + self._drain() + assert self._root is not None + if self._root.done(): + if self._root.cancelled(): + raise task.TaskCancelledError( + "The orchestrator coroutine was cancelled. Catch local cancellation " + "inside the orchestrator; terminate the durable instance through the client.") + # 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(): + raise RuntimeError( + "The async orchestrator is blocked without durable work. " + "Await durable activities, events or timers; external I/O is unsupported.") + + 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: + 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(): + 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._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: + 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): + failure = RuntimeError("An orchestration callback raised StopIteration.") + failure.__cause__ = error + 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: + 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() + 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(): + 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(): + 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: + self._closed = True + self._ready.clear() + self._cancel_ready.clear() + for observer in (*self._observers.values(), *self._timers.values()): + observer.detach() + self._completions.clear() + self._observers.clear() + 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.") + + def call_soon_threadsafe(self, callback: Any, *args: Any, context: Any = None) -> asyncio.Handle: + raise RuntimeError("Cross-thread scheduling is unsupported in orchestrators; use an activity.") + + async def _unsupported_io(self, *args: Any, **kwargs: Any) -> Any: + raise RuntimeError("External I/O is unsupported in orchestrators; move it to an ordinary activity.") + + create_connection = _unsupported_io + create_server = _unsupported_io + create_datagram_endpoint = _unsupported_io + connect_read_pipe = _unsupported_io + connect_write_pipe = _unsupported_io + getaddrinfo = _unsupported_io + getnameinfo = _unsupported_io + sock_connect = _unsupported_io + sock_recv = _unsupported_io + sock_sendall = _unsupported_io + subprocess_exec = _unsupported_io + subprocess_shell = _unsupported_io diff --git a/durabletask/internal/exceptions.py b/durabletask/internal/exceptions.py index a04a304f..83193ec7 100644 --- a/durabletask/internal/exceptions.py +++ b/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__() diff --git a/durabletask/task.py b/durabletask/task.py index 17cc7f3d..3ff84ecd 100644 --- a/durabletask/task.py +++ b/durabletask/task.py @@ -6,12 +6,14 @@ import logging import math +import asyncio from abc import ABC, abstractmethod -from collections.abc import Callable, Generator, Sequence +from collections.abc import Callable, Coroutine, Generator, Sequence from dataclasses import dataclass from datetime import UTC, datetime, timedelta from typing import Any, Generic, TypeAlias, TypeVar, cast, overload +from durabletask.async_context import AsyncOrchestrationContext as AsyncOrchestrationContext from durabletask.entities import DurableEntity, EntityInstanceId, EntityLock, EntityContext import durabletask.internal.helpers as pbh import durabletask.internal.orchestrator_service_pb2 as pb @@ -23,6 +25,14 @@ class OrchestrationContext(ABC): + def as_async(self) -> AsyncOrchestrationContext: + """Return a cold, coroutine-native view of this context. + + Existing context methods remain eager and return reusable durable task + handles. The async view schedules work only when its coroutines start. + """ + return AsyncOrchestrationContext(self) + @property @abstractmethod def instance_id(self) -> str: @@ -612,6 +622,33 @@ def __init__(self) -> None: # Dict insertion order preserves subscription order for shared races. self._parents = {} + def __await__(self) -> Generator[Task[T], Any, T]: + """Wait for this task in a replay-driven ``async def`` orchestrator. + + A local asyncio waiter does not own or cancel the durable operation. + """ + from durabletask.internal.asyncio_loop import DurableEventLoop + + loop = asyncio._get_running_loop() + if isinstance(loop, DurableEventLoop): + return (yield from loop.wait_durable(self).__await__()) + if loop is not None: + raise RuntimeError("Durable tasks can only be awaited inside a durable orchestrator.") + # Allow mock-based orchestration drivers to send results without an event loop. + if not self.is_complete: + yield self + return self.get_result() + + async def wait(self) -> T: + """Borrow this operation's result through a fresh native coroutine. + + Suitable for ``asyncio.create_task`` and task groups. Calling this + method does not schedule another durable operation. Each call creates + an independent waiter; cancelling it does not cancel the underlying + operation, timer, or event subscription, or any other waiter. + """ + return await self + @property def is_complete(self) -> bool: """Returns True if the task has completed, False otherwise.""" @@ -942,8 +979,11 @@ def task_id(self) -> int: return self._task_id -# Orchestrators are generators that yield tasks, receive any type, and return TOutput -Orchestrator: TypeAlias = Callable[[OrchestrationContext, TInput], Generator[Task[Any], Any, TOutput] | TOutput] +# Orchestrators yield or await durable tasks, or return an immediate result. +Orchestrator: TypeAlias = Callable[ + [OrchestrationContext, TInput], + Generator[Task[Any], Any, TOutput] | Coroutine[Any, Any, TOutput] | TOutput, +] # Activities are simple functions that can be scheduled by orchestrators Activity: TypeAlias = Callable[[ActivityContext, TInput], TOutput] diff --git a/durabletask/worker.py b/durabletask/worker.py index e4099042..aa5fd3e1 100644 --- a/durabletask/worker.py +++ b/durabletask/worker.py @@ -8,12 +8,12 @@ import logging import os import time -from collections.abc import Callable, Generator, Sequence +from collections.abc import Callable, Coroutine, Generator, Sequence from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from datetime import UTC, datetime, timedelta from threading import Event, Lock, Thread -from types import GeneratorType +from types import CoroutineType, GeneratorType from enum import Enum from functools import partial from typing import Any, Self, TypeVar, cast, overload @@ -50,6 +50,7 @@ from durabletask.payload import helpers as payload_helpers from durabletask import task from durabletask.internal.grpc_interceptor import DefaultClientInterceptorImpl +from durabletask.internal.asyncio_loop import DurableEventLoop from durabletask.payload.store import PayloadStore from durabletask.serialization import DataConverter, JsonDataConverter @@ -62,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. @@ -1626,22 +1619,44 @@ def __init__(self, self._data_converter = data_converter self._exception_properties_provider = exception_properties_provider self._logger = logger - - def run(self, generator: Generator[task.Task[Any], Any, Any]) -> None: - self._generator = generator - # TODO: Do something with this task - task = next(generator) # this starts the generator - # TODO: Check if the task is null? - self._previous_task = task + 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): + self._async_loop = DurableEventLoop( + self, is_complete=lambda: self._is_complete, logger=self._logger) + self._async_loop.start(orchestrator) + return + self._generator = orchestrator + self._set_next_task(orchestrator.send(None)) + + 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() + + def _set_next_task(self, next_task: Any) -> None: + if not isinstance(next_task, task.Task): + raise TypeError( + "Generator orchestrators must yield durabletask.Task objects. " + "Use async def and await for asyncio composition." + ) + self._previous_task = next_task def resume(self) -> None: + if self._async_loop is not None: + self._async_loop.resume() + return if self._generator is None: # This is never expected unless maybe there's an issue with the history raise TypeError( - "The orchestrator generator is not initialized! Was the orchestration history corrupted?" + "The orchestrator is not initialized! Was the orchestration history corrupted?" ) - # We can resume the generator only if the previously yielded task + # We can resume the generator only if the previous task # has reached a completed state. The only time this won't be the # case is if the user yielded on a WhenAll task and there are still # outstanding child tasks that need to be completed. @@ -1656,9 +1671,7 @@ def resume(self) -> None: # This will either return a Task or raise StopIteration if it's done. next_task = self._generator.send(self._previous_task.get_result()) - if not isinstance(next_task, task.Task): - raise TypeError("The orchestrator generator yielded a non-Task object") - self._previous_task = next_task + self._set_next_task(next_task) def set_complete( self, @@ -1787,6 +1800,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 @@ -2431,7 +2446,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.", @@ -2465,7 +2480,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: @@ -2481,34 +2496,33 @@ def execute( finally: # Close after building the response, so the orchestrator's finally # blocks do not change the actions or custom status we return. - if ctx._generator is not None: # pyright: ignore[reportPrivateUsage] - try: - ctx._generator.close() # pyright: ignore[reportPrivateUsage] - except (Exception, asyncio.CancelledError): - cleanup_name = next( - (e.executionStarted.name for e in new_events if e.HasField("executionStarted")), - orchestration_name, - ) - self._logger.warning( - f"{instance_id}: Failed to close orchestrator {cleanup_name} generator.", - exc_info=True, - ) - except BaseException as ex: - self._logger.error( - f"{instance_id}: Orchestrator {orchestration_name} generator teardown exited unexpectedly.", - exc_info=True, - ) - if not dispatch_exited: - if _is_control_flow_exception(ex) or result is None: - raise - # Keep an existing failure rather than mask its cause. - if not any( - action.HasField("completeOrchestration") - and action.completeOrchestration.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED - for action in result.actions - ): - result = result.as_failed(ph.new_failure_details( - ex, self._exception_properties_provider, self._logger)) + try: + ctx.close() + except (Exception, asyncio.CancelledError): + cleanup_name = next( + (e.executionStarted.name for e in new_events if e.HasField("executionStarted")), + orchestration_name, + ) + self._logger.warning( + f"{instance_id}: Failed to close orchestrator {cleanup_name}.", + exc_info=True, + ) + except BaseException as ex: + self._logger.error( + f"{instance_id}: Orchestrator {orchestration_name} teardown exited unexpectedly.", + exc_info=True, + ) + if not dispatch_exited: + if pe.is_control_flow_exception(ex) or result is None: + raise + # Keep an existing failure rather than mask its cause. + if not any( + action.HasField("completeOrchestration") + and action.completeOrchestration.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + for action in result.actions + ): + result = result.as_failed(ph.new_failure_details( + ex, self._exception_properties_provider, self._logger)) assert result is not None return result @@ -2685,10 +2699,15 @@ def process_event( result = fn( ctx, input - ) # this does not execute the generator, only creates it - if isinstance(result, GeneratorType): - # Start the orchestrator's generator function - ctx.run(cast(Generator[task.Task[Any], Any, Any], result)) + ) # this creates the generator or coroutine without executing it + if isinstance(result, (GeneratorType, CoroutineType)): + # Both protocols suspend on durable tasks and finish with StopIteration. + ctx.run(cast(Generator[task.Task[Any], Any, Any] | Coroutine[Any, Any, Any], result)) + elif inspect.isasyncgen(result): + raise TypeError( + "Async generator orchestrators are not supported. " + "Use async def with await, not yield." + ) else: # This is an orchestrator that doesn't schedule any tasks ctx.set_complete(result, pb.ORCHESTRATION_STATUS_COMPLETED) diff --git a/tests/azure-functions-durable/e2e/apps/dtask_style/orchestrators.py b/tests/azure-functions-durable/e2e/apps/dtask_style/orchestrators.py index 34b44d98..8ffd9ba2 100644 --- a/tests/azure-functions-durable/e2e/apps/dtask_style/orchestrators.py +++ b/tests/azure-functions-durable/e2e/apps/dtask_style/orchestrators.py @@ -13,6 +13,7 @@ """ from datetime import timedelta +import asyncio import logging from typing import Any @@ -35,6 +36,28 @@ def activity_chain(ctx: task.OrchestrationContext, _: Any): return [first, second, third] +@bp.orchestration_trigger(context_name="context") +async def async_await_probe(ctx: task.OrchestrationContext, _: Any) -> dict[str, Any]: + greeting = await ctx.call_activity("say_hello", input="Tokyo", return_type=str) + values = await asyncio.gather(*(ctx.call_activity("square", input=i, return_type=int) for i in (2, 3))) + await asyncio.sleep(0) + await asyncio.sleep(0.01) + child: list[str] = await ctx.call_sub_orchestrator("activity_chain") + return {"greeting": greeting, "total": sum(values), "child": child} + + +@bp.orchestration_trigger(context_name="context") +async def async_cold_probe(ctx: task.OrchestrationContext, _: Any) -> dict[str, Any]: + ops = ctx.as_async() + greeting = await asyncio.create_task(ops.call_activity("say_hello", input="Tokyo", return_type=str)) + async with asyncio.TaskGroup() as group: + first = group.create_task(ops.call_activity("square", input=2, return_type=int)) + second = group.create_task(ops.call_activity("square", input=3, return_type=int)) + await ops.create_timer(timedelta(milliseconds=10)) + child: list[str] = await ops.call_sub_orchestrator("activity_chain") + return {"greeting": greeting, "total": first.result() + second.result(), "child": child} + + @bp.orchestration_trigger(context_name="context") def logging_filtered(ctx: task.OrchestrationContext, _: Any): _LOGGER.info("worker-info-filter-anchor %s", ctx.instance_id) diff --git a/tests/azure-functions-durable/e2e/apps/v1_style/orchestrators.py b/tests/azure-functions-durable/e2e/apps/v1_style/orchestrators.py index a96467c7..08ce5fda 100644 --- a/tests/azure-functions-durable/e2e/apps/v1_style/orchestrators.py +++ b/tests/azure-functions-durable/e2e/apps/v1_style/orchestrators.py @@ -13,6 +13,8 @@ """ from datetime import timedelta +import asyncio +from typing import Any import azure.durable_functions as df @@ -31,6 +33,28 @@ def activity_chain(context: df.DurableOrchestrationContext): return [first, second, third] +@bp.orchestration_trigger(context_name="context") +async def async_await_probe(context: df.DurableOrchestrationContext) -> dict[str, Any]: + greeting = await context.call_activity("say_hello", "Tokyo") + values = await asyncio.gather(*(context.call_activity("square", i) for i in (2, 3))) + await asyncio.sleep(0) + await asyncio.sleep(0.01) + child = await context.call_sub_orchestrator("activity_chain") + return {"greeting": greeting, "total": sum(values), "child": child} + + +@bp.orchestration_trigger(context_name="context") +async def async_cold_probe(context: df.DurableOrchestrationContext) -> dict[str, Any]: + ops = context.as_async() + greeting = await asyncio.create_task(ops.call_activity("say_hello", input="Tokyo", return_type=str)) + async with asyncio.TaskGroup() as group: + first = group.create_task(ops.call_activity("square", input=2, return_type=int)) + second = group.create_task(ops.call_activity("square", input=3, return_type=int)) + await ops.create_timer(timedelta(milliseconds=10)) + child: list[str] = await ops.call_sub_orchestrator("activity_chain") + return {"greeting": greeting, "total": first.result() + second.result(), "child": child} + + @bp.orchestration_trigger(context_name="context") def fan_out_fan_in(context: df.DurableOrchestrationContext): count = context.get_input() or 5 diff --git a/tests/azure-functions-durable/e2e/test_coroutine_e2e.py b/tests/azure-functions-durable/e2e/test_coroutine_e2e.py new file mode 100644 index 00000000..6e5f1629 --- /dev/null +++ b/tests/azure-functions-durable/e2e/test_coroutine_e2e.py @@ -0,0 +1,25 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Real-host smoke coverage for both async orchestrator authoring styles.""" + +import pytest + +from ._harness import FunctionApp + + +pytestmark = pytest.mark.functions_e2e + + +@pytest.mark.parametrize("app_fixture", ["dtask_app", "v1_app"]) +@pytest.mark.parametrize("orchestrator", ["async_await_probe", "async_cold_probe"]) +def test_async_await_orchestrator(request: pytest.FixtureRequest, app_fixture: str, orchestrator: str) -> None: + app: FunctionApp = request.getfixturevalue(app_fixture) + instance_id = app.start_orchestration(orchestrator) + status = app.wait_for_completion(instance_id) + assert status["runtimeStatus"].upper() == "COMPLETED", status + assert status["output"] == { + "greeting": "Hello Tokyo!", + "total": 13, + "child": ["Hello Tokyo!", "Hello Seattle!", "Hello London!"], + } diff --git a/tests/azure-functions-durable/test_functions_coroutine_orchestrators.py b/tests/azure-functions-durable/test_functions_coroutine_orchestrators.py new file mode 100644 index 00000000..bcceaf82 --- /dev/null +++ b/tests/azure-functions-durable/test_functions_coroutine_orchestrators.py @@ -0,0 +1,222 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Coroutine replay through native and compatibility Functions entry points.""" + +import base64 +import asyncio +import json +from collections.abc import Callable +from datetime import datetime, timedelta +from typing import Any + +import pytest + +import azure.durable_functions as df +from azure.durable_functions.orchestrator import Orchestrator +from azure.durable_functions.testing import orchestrator_generator_wrapper +from azure.durable_functions.worker import DurableFunctionsWorker +from durabletask import task +from durabletask.internal import helpers +from durabletask.internal import orchestrator_service_pb2 as pb + + +@pytest.mark.parametrize("native", [False, True], ids=["compat", "native"]) +@pytest.mark.parametrize("entry", ["worker-sync", "worker-async", "decorator", "v1-handle"]) +@pytest.mark.parametrize("cold", [False, True], ids=["handles", "cold"]) +async def test_coroutine_replay_through_functions_entry_points(native: bool, entry: str, cold: bool) -> None: + async def modern(ctx: task.OrchestrationContext, value: int) -> list[Any]: + if cold: + ops = ctx.as_async() + async with asyncio.TaskGroup() as group: + first = group.create_task(ops.call_activity("work", input=value, return_type=str)) + second = group.create_task(ops.call_activity("work", input=value + 1, return_type=str)) + return [first.result(), second.result()] + return list(await asyncio.gather( + ctx.call_activity("work", input=value, return_type=str), + ctx.call_activity("work", input=value + 1, return_type=str), + )) + + async def compatible(ctx: df.DurableOrchestrationContext) -> list[Any]: + value = ctx.get_input() + if cold: + ops = ctx.as_async() + async with asyncio.TaskGroup() as group: + first = group.create_task(ops.call_activity("work", input=value, return_type=str)) + second = group.create_task(ops.call_activity("work", input=value + 1, return_type=str)) + return [first.result(), second.result()] + return list(await asyncio.gather( + ctx.call_activity("work", value), + ctx.call_activity("work", value + 1), + )) + + fn: Callable[..., Any] = modern if native else compatible + worker = DurableFunctionsWorker() + handle = Orchestrator.create(fn) + app = df.DFApp() + builder = app.orchestration_trigger(context_name="ctx")(fn) + registered = builder.build().get_user_function() + instance_id = "async-functions" + + async def invoke(request: pb.OrchestratorRequest) -> pb.OrchestratorResponse: + encoded = base64.b64encode(request.SerializeToString()).decode() + if entry == "worker-sync": + output = worker.execute_orchestration_request(fn, encoded) + elif entry == "worker-async": + output = await worker.execute_orchestration_request_async(fn, encoded) + elif entry == "decorator": + output = await registered(ctx=encoded) + else: + output = handle(encoded) + response = pb.OrchestratorResponse() + response.ParseFromString(base64.b64decode(output)) + return response + + start = [ + helpers.new_orchestrator_started_event(), + helpers.new_execution_started_event(fn.__name__, instance_id, "3"), + ] + first = await invoke(pb.OrchestratorRequest(instanceId=instance_id, newEvents=start)) + assert [a.scheduleTask.name for a in first.actions] == ["work", "work"] + assert [json.loads(a.scheduleTask.input.value) for a in first.actions] == [3, 4] + assert [a.id for a in first.actions] == [1, 2] + + response = await invoke(pb.OrchestratorRequest( + instanceId=instance_id, + pastEvents=start + [helpers.new_task_scheduled_event(i, "work") for i in (1, 2)], + newEvents=[ + helpers.new_task_completed_event(2, '"second"'), + helpers.new_task_completed_event(1, '"first"'), + ], + )) + assert len(response.actions) == 1 + complete = response.actions[0].completeOrchestration + assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED + assert json.loads(complete.result.value) == ["first", "second"] + + +@pytest.mark.parametrize("native", [False, True], ids=["compat", "native"]) +def test_owned_event_timeout_replacement_through_functions_worker(native: bool) -> None: + async def receive(ops: task.AsyncOrchestrationContext) -> str: + try: + return await ops.receive_event("approval", timeout=timedelta(seconds=1), data_type=str) + except TimeoutError: + return await ops.receive_event("approval", data_type=str) + + async def modern(ctx: task.OrchestrationContext, _: Any) -> str: + return await receive(ctx.as_async()) + + async def compatible(ctx: df.DurableOrchestrationContext) -> str: + return await receive(ctx.as_async()) + + fn: Callable[..., Any] = modern if native else compatible + worker = DurableFunctionsWorker() + instance_id = "owned-event-functions" + now = datetime(2026, 1, 1) + + def invoke(old: list[pb.HistoryEvent], new: list[pb.HistoryEvent]) -> pb.OrchestratorResponse: + request = pb.OrchestratorRequest(instanceId=instance_id, pastEvents=old, newEvents=new) + encoded = base64.b64encode(request.SerializeToString()).decode() + response = pb.OrchestratorResponse() + response.ParseFromString(base64.b64decode(worker.execute_orchestration_request(fn, encoded))) + return response + + start = [ + helpers.new_orchestrator_started_event(now), + helpers.new_execution_started_event(fn.__name__, instance_id, "null"), + ] + initial = invoke([], start) + assert len(initial.actions) == 1 + assert initial.actions[0].createTimer.fireAt.ToDatetime() == now + timedelta(seconds=1) + old = start + [helpers.new_timer_created_event(1, now + timedelta(seconds=1))] + timeout = [ + helpers.new_orchestrator_started_event(now + timedelta(seconds=1)), + helpers.new_timer_fired_event(1, now + timedelta(seconds=1)), + ] + assert not invoke(old, timeout).actions + response = invoke(old + timeout, [helpers.new_event_raised_event("approval", '"approved"')]) + assert len(response.actions) == 1 + complete = response.actions[0].completeOrchestration + assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED, complete.failureDetails + assert json.loads(complete.result.value) == "approved" + + +@pytest.mark.parametrize("fail", [False, True]) +def test_testing_wrapper_drives_coroutine(fail: bool) -> None: + pending = task.CompletableTask[str]() + + async def orchestrator() -> str: + try: + return await pending + except task.TaskFailedError: + return "caught" + + values = orchestrator_generator_wrapper(orchestrator()) + assert next(values) is pending + if fail: + pending.fail("boom", ValueError("boom")) + else: + pending.complete("done") + assert list(values) == ["caught" if fail else "done"] + + +def test_testing_wrapper_drives_immediate_coroutine() -> None: + async def orchestrator() -> str: + return "done" + + assert list(orchestrator_generator_wrapper(orchestrator())) == ["done"] + + +@pytest.mark.parametrize("fail", [False, True]) +async def test_testing_wrapper_isolates_each_step_from_async_test_loop(fail: bool) -> None: + loop = asyncio.get_running_loop() + caller = asyncio.current_task() + first = task.CompletableTask[str]() + second = task.CompletableTask[str]() + + async def orchestrator() -> str: + assert asyncio._get_running_loop() is None + try: + await first + except task.TaskFailedError: + assert asyncio._get_running_loop() is None + assert asyncio._get_running_loop() is None + return await second + + values = orchestrator_generator_wrapper(orchestrator()) + assert next(values) is first + assert asyncio.get_running_loop() is loop + assert asyncio.current_task() is caller + if fail: + first.fail("boom", ValueError("boom")) + else: + first.complete("first") + assert next(values) is second + assert asyncio.get_running_loop() is loop + assert asyncio.current_task() is caller + second.complete("done") + assert list(values) == ["done"] + assert asyncio.get_running_loop() is loop + assert asyncio.current_task() is caller + await asyncio.sleep(0) + + +@pytest.mark.parametrize("error", [ValueError("test failure"), SystemExit("test exit")]) +async def test_testing_wrapper_restores_async_test_loop_on_exception(error: BaseException) -> None: + loop = asyncio.get_running_loop() + caller = asyncio.current_task() + pending = task.CompletableTask[str]() + + async def orchestrator() -> None: + await pending + raise error + + values = orchestrator_generator_wrapper(orchestrator()) + assert next(values) is pending + pending.complete("done") + with pytest.raises(type(error)) as caught: + next(values) + assert caught.value is error + assert asyncio.get_running_loop() is loop + assert asyncio.current_task() is caller + await asyncio.sleep(0) diff --git a/tests/azure-functions-durable/typing/import_package.py b/tests/azure-functions-durable/typing/import_package.py index 8d69f92d..cf7dbe3f 100644 --- a/tests/azure-functions-durable/typing/import_package.py +++ b/tests/azure-functions-durable/typing/import_package.py @@ -1,12 +1,14 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. +import asyncio from datetime import timedelta from typing import assert_type import azure.durable_functions as df from azure.durable_functions.client import DurableFunctionsClient, SyncDurableFunctionsClient from azure.durable_functions.worker import DurableFunctionsWorker +from durabletask import task def create_app() -> df.DFApp: @@ -39,3 +41,21 @@ def check_sync_context_managers( async def check_async_context_manager(client: DurableFunctionsClient) -> None: async with client as entered_client: assert_type(entered_client, DurableFunctionsClient) + + +async def check_cold_operations( + context: df.DurableOrchestrationContext, + native: task.OrchestrationContext, +) -> None: + ops = context.as_async() + assert_type(ops, task.AsyncOrchestrationContext) + assert_type(native.as_async(), task.AsyncOrchestrationContext) + pending = asyncio.create_task(ops.call_activity("activity", return_type=int)) + assert_type(pending, asyncio.Task[int]) + assert_type(await pending, int) + event = asyncio.create_task(ops.receive_event( + "approval", timeout=timedelta(seconds=1), data_type=str)) + assert_type(await event, str) + operation = native.call_activity("shared", return_type=int) + borrower = asyncio.create_task(operation.wait()) + assert_type(borrower, asyncio.Task[int]) diff --git a/tests/durabletask-azuremanaged/test_dts_coroutine_orchestrators.py b/tests/durabletask-azuremanaged/test_dts_coroutine_orchestrators.py new file mode 100644 index 00000000..b82112d9 --- /dev/null +++ b/tests/durabletask-azuremanaged/test_dts_coroutine_orchestrators.py @@ -0,0 +1,103 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Exercise both workers over gRPC and optionally against the DTS emulator.""" + +import json +import asyncio +import os +import socket +from collections.abc import Generator +from datetime import timedelta +from typing import Any + +import pytest + +from durabletask import client, task, worker +from durabletask.azuremanaged.client import DurableTaskSchedulerClient +from durabletask.azuremanaged.worker import DurableTaskSchedulerWorker +from durabletask.testing import create_test_backend + + +@pytest.mark.parametrize("transport", [ + "core", + "azuremanaged", + pytest.param("dts", marks=pytest.mark.dts), +]) +@pytest.mark.parametrize("coroutine_parent", [False, True], ids=["generator-parent", "coroutine-parent"]) +@pytest.mark.parametrize("cold", [False, True], ids=["handles", "cold"]) +def test_mixed_orchestrators_over_grpc(transport: str, coroutine_parent: bool, cold: bool) -> None: + def increment(ctx: task.ActivityContext, value: int) -> int: + return value + 1 + + async def coroutine_child(ctx: task.OrchestrationContext, value: int) -> int: + if cold: + ops = ctx.as_async() + await ops.create_timer(timedelta(milliseconds=10)) + return await ops.call_activity(increment, input=value, return_type=int) + await asyncio.sleep(0) + await asyncio.sleep(0.01) + return await ctx.call_activity(increment, input=value, return_type=int) + + def generator_child(ctx: task.OrchestrationContext, value: int) -> Generator[task.Task[Any], Any, int]: + return (yield ctx.call_activity(increment, input=value)) + + async def coroutine(ctx: task.OrchestrationContext, value: int) -> list[Any]: + if cold: + ops = ctx.as_async() + async with asyncio.TaskGroup() as group: + first = group.create_task(ops.call_sub_orchestrator(coroutine_child, input=value)) + second = group.create_task(ops.call_sub_orchestrator(generator_child, input=value + 1)) + return [first.result(), second.result(), await ops.receive_event("go")] + values = list(await asyncio.gather( + ctx.call_sub_orchestrator(coroutine_child, input=value), + ctx.call_sub_orchestrator(generator_child, input=value + 1), + )) + event = await ctx.wait_for_external_event("go") + return values + [event] + + def generator(ctx: task.OrchestrationContext, value: int) -> Generator[task.Task[Any], Any, list[Any]]: + values = yield task.when_all([ + ctx.call_sub_orchestrator(coroutine_child, input=value), + ctx.call_sub_orchestrator(generator_child, input=value + 1), + ]) + event = yield ctx.wait_for_external_event("go") + return values + [event] + + if transport == "dts": + backend = None + host = os.environ["ENDPOINT"] + taskhub = os.environ["TASKHUB"] + else: + with socket.socket() as listener: + listener.bind(("127.0.0.1", 0)) + port = listener.getsockname()[1] + backend = create_test_backend(port=port) + host = f"localhost:{port}" + taskhub = "test" + if transport != "core": + w = DurableTaskSchedulerWorker( + host_address=host, secure_channel=False, taskhub=taskhub, token_credential=None) + c = DurableTaskSchedulerClient( + host_address=host, secure_channel=False, taskhub=taskhub, token_credential=None) + else: + w = worker.TaskHubGrpcWorker(host_address=host) + c = client.TaskHubGrpcClient(host_address=host) + try: + with w, c: + parent = coroutine if coroutine_parent else generator + for fn in (parent, coroutine_child, generator_child): + w.add_orchestrator(fn) + w.add_activity(increment) + w.start() + instance_id = c.schedule_new_orchestration(parent, input=3) + c.raise_orchestration_event(instance_id, "go", data="approved") + state = c.wait_for_orchestration_completion(instance_id, timeout=15) + assert state is not None + assert state.runtime_status == client.OrchestrationStatus.COMPLETED, state.failure_details + assert state.serialized_output is not None + assert json.loads(state.serialized_output) == [4, 5, "approved"] + finally: + if backend is not None: + backend.stop() + backend.reset() diff --git a/tests/durabletask/test_async_context.py b/tests/durabletask/test_async_context.py new file mode 100644 index 00000000..bdf8c534 --- /dev/null +++ b/tests/durabletask/test_async_context.py @@ -0,0 +1,446 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Cold coroutine and ownership semantics through the shared replay executor.""" + +import asyncio +import inspect +import json +import logging +from collections.abc import Coroutine, Generator +from datetime import timedelta +from typing import Any + +import pytest + +from durabletask import entities, task +from durabletask.async_context import AsyncOrchestrationContext +from durabletask.internal import helpers +from durabletask.internal import orchestrator_service_pb2 as pb +from tests.durabletask.test_coroutine_orchestrators import ( + NOW, + _completion, # pyright: ignore[reportPrivateUsage] + _execute, # pyright: ignore[reportPrivateUsage] + _start, # pyright: ignore[reportPrivateUsage] +) + + +def test_additive_public_api_has_no_new_abstract_requirements() -> None: + assert "as_async" not in task.OrchestrationContext.__abstractmethods__ + assert "wait" not in task.Task.__abstractmethods__ + assert task.AsyncOrchestrationContext is AsyncOrchestrationContext + + +@pytest.mark.parametrize("cancel", [False, True]) +def test_all_operations_are_cold_and_prestart_cancellation_schedules_nothing(cancel: bool) -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any: + ops = ctx.as_async() + entity = entities.EntityInstanceId("counter", "one") + invalid_input = object() + operations: list[Coroutine[Any, Any, Any]] = [ + ops.call_activity("unused", input=invalid_input), + ops.call_sub_orchestrator("unused", input=invalid_input), + ops.call_entity(entity, "unused", invalid_input), + ops.lock_entities([entity]), + ops.create_timer(timedelta(days=1)), + ops.receive_event("unused", timeout=timedelta(days=1)), + ] + assert all(inspect.iscoroutine(operation) for operation in operations) + if cancel: + pending = [asyncio.create_task(operation) for operation in operations] + for waiter in pending: + waiter.cancel() + await asyncio.gather(*pending, return_exceptions=True) + else: + for operation in operations: + operation.close() + return await ctx.call_activity("actual", input=3) + + actions = _execute(orchestrator, [], _start()).actions + assert len(actions) == 1 + assert actions[0].id == 1 + assert actions[0].scheduleTask.name == "actual" + + +@pytest.mark.parametrize("cancel", [False, True]) +def test_cold_receive_does_not_consume_buffered_event(cancel: bool) -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> str: + await ctx.wait_for_external_event("go") + unused = ctx.as_async().receive_event("value") + if cancel: + pending = asyncio.create_task(unused) + pending.cancel() + await asyncio.gather(pending, return_exceptions=True) + else: + unused.close() + return await ctx.as_async().receive_event("value", data_type=str) + + result = _execute(orchestrator, _start(), [ + helpers.new_event_raised_event("value", '"kept"'), + helpers.new_event_raised_event("go", None), + ]) + assert json.loads(_completion(result).result.value) == "kept" + + +def test_create_task_and_task_group_accept_cold_operations() -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> list[int]: + ops = ctx.as_async() + first = asyncio.create_task(ops.call_activity("first", return_type=int)) + async with asyncio.TaskGroup() as group: + second = group.create_task(ops.call_activity("second", return_type=int)) + return [await first, second.result()] + + start = _start() + initial = _execute(orchestrator, [], start) + assert [(a.id, a.scheduleTask.name) for a in initial.actions] == [(1, "first"), (2, "second")] + old = start + [ + helpers.new_task_scheduled_event(1, "first"), + helpers.new_task_scheduled_event(2, "second"), + ] + result = _execute(orchestrator, old, [ + helpers.new_task_completed_event(2, "2"), + helpers.new_task_completed_event(1, "1"), + ]) + assert json.loads(_completion(result).result.value) == [1, 2] + + +def test_serialization_and_scheduling_follow_coroutine_execution_not_construction() -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any: + ops = ctx.as_async() + payload = {"value": 1} + first = ops.call_activity("first", input=payload) + second = ops.call_activity("second", input=payload) + payload["value"] = 2 + try: + await second + return await first + finally: + first.close() + + actions = _execute(orchestrator, [], _start()).actions + assert len(actions) == 1 + assert actions[0].id == 1 + assert actions[0].scheduleTask.name == "second" + assert json.loads(actions[0].scheduleTask.input.value) == {"value": 2} + + +def test_cold_relative_timer_deadline_is_evaluated_when_started() -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + timer = ctx.as_async().create_timer(timedelta(hours=1)) + try: + await ctx.wait_for_external_event("go") + await timer + finally: + timer.close() + + later = NOW + timedelta(hours=2) + result = _execute(orchestrator, _start(), [ + helpers.new_orchestrator_started_event(later), + helpers.new_event_raised_event("go", None), + ]) + assert len(result.actions) == 1 + assert result.actions[0].createTimer.fireAt.ToDatetime() == later + timedelta(hours=1) + + +def test_borrowed_waiters_are_independent_reusable_and_do_not_cancel_operations() -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> list[str]: + operation = ctx.call_activity("shared", return_type=str) + cancelled = asyncio.create_task(operation.wait()) + remaining = asyncio.create_task(operation.wait()) + await asyncio.sleep(0) + cancelled.cancel() + await asyncio.gather(cancelled, return_exceptions=True) + return [await remaining, await operation.wait(), await operation] + + old = _start() + [helpers.new_task_scheduled_event(1, "shared")] + result = _execute(orchestrator, old, [helpers.new_task_completed_event(1, '"ok"')]) + assert json.loads(_completion(result).result.value) == ["ok", "ok", "ok"] + + +def test_borrowed_event_cancellation_preserves_shared_subscription() -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> int: + event = ctx.wait_for_external_event("value", data_type=int) + first = asyncio.create_task(event.wait()) + second = asyncio.create_task(event.wait()) + await asyncio.sleep(0) + first.cancel() + await asyncio.gather(first, return_exceptions=True) + return await second + + result = _execute(orchestrator, _start(), [helpers.new_event_raised_event("value", "42")]) + assert json.loads(_completion(result).result.value) == 42 + + +def test_borrowed_cancelled_subscription_still_receives_before_owned_replacement() -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> list[int]: + shared = ctx.wait_for_external_event("value", data_type=int) + pending = asyncio.create_task(shared.wait()) + await asyncio.sleep(0) + pending.cancel() + await asyncio.gather(pending, return_exceptions=True) + next_value = await ctx.as_async().receive_event("value", data_type=int) + return [await shared.wait(), next_value] + + result = _execute(orchestrator, _start(), [ + helpers.new_event_raised_event("value", "1"), + helpers.new_event_raised_event("value", "2"), + ]) + assert json.loads(_completion(result).result.value) == [1, 2] + + +def test_cold_coroutines_are_single_use_not_shared_operation_handles() -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> int: + operation = ctx.as_async().call_activity("work", return_type=int) + value = await operation + with pytest.raises(RuntimeError, match="cannot reuse"): + await operation + return value + + old = _start() + [helpers.new_task_scheduled_event(1, "work")] + result = _execute(orchestrator, old, [helpers.new_task_completed_event(1, "42")]) + assert json.loads(_completion(result).result.value) == 42 + + +def test_owned_timeout_withdraws_subscription_before_replacement() -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> int: + ops = ctx.as_async() + try: + await ops.receive_event("value", timeout=timedelta(seconds=5)) + except TimeoutError: + return await ops.receive_event("value", data_type=int) + raise AssertionError("Expected timeout") + + deadline = NOW + timedelta(seconds=5) + initial = _execute(orchestrator, [], _start()) + assert initial.actions[0].id == 1 + old = _start() + [helpers.new_timer_created_event(1, deadline)] + expired = [helpers.new_orchestrator_started_event(deadline), helpers.new_timer_fired_event(1, deadline)] + assert _execute(orchestrator, old, expired).actions == [] + result = _execute(orchestrator, old + expired, [helpers.new_event_raised_event("value", "42")]) + assert json.loads(_completion(result).result.value) == 42 + + +@pytest.mark.parametrize("with_timeout", [False, True]) +def test_owned_cancellation_withdraws_subscription_and_uncommitted_timer(with_timeout: bool) -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> int: + ops = ctx.as_async() + pending = asyncio.create_task(ops.receive_event( + "value", timeout=timedelta(seconds=5) if with_timeout else None)) + await asyncio.sleep(0) + pending.cancel() + await asyncio.gather(pending, return_exceptions=True) + return await ops.receive_event("value", data_type=int) + + assert _execute(orchestrator, [], _start()).actions == [] + result = _execute(orchestrator, _start(), [helpers.new_event_raised_event("value", "42")]) + assert json.loads(_completion(result).result.value) == 42 + + +@pytest.mark.parametrize("buffered", [False, True]) +@pytest.mark.parametrize("seconds", [-1, 0, 5]) +def test_buffered_events_and_immediate_timeouts(buffered: bool, seconds: int) -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any: + await ctx.wait_for_external_event("go") + try: + return await ctx.as_async().receive_event( + "value", timeout=timedelta(seconds=seconds), data_type=int) + except TimeoutError: + return "timeout" + + new = ([helpers.new_event_raised_event("value", "42")] if buffered else []) + [ + helpers.new_event_raised_event("go", None)] + result = _execute(orchestrator, _start(), new) + if not buffered and seconds > 0: + assert len(result.actions) == 1 + assert result.actions[0].HasField("createTimer") + else: + assert len(result.actions) == 1 + assert json.loads(_completion(result).result.value) == (42 if buffered else "timeout") + + +@pytest.mark.parametrize("event_first", [False, True]) +def test_owned_receive_race_follows_history_order(event_first: bool) -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> str: + ops = ctx.as_async() + try: + value = await ops.receive_event("value", timeout=timedelta(seconds=5), data_type=int) + return f"event:{value}" + except TimeoutError: + return f"timeout:{await ops.receive_event('value', data_type=int)}" + + deadline = NOW + timedelta(seconds=5) + old = _start() + [helpers.new_timer_created_event(1, deadline)] + arrival = helpers.new_event_raised_event("value", "42") + fired = helpers.new_timer_fired_event(1, deadline) + result = _execute(orchestrator, old, [arrival, fired] if event_first else [fired, arrival]) + assert json.loads(_completion(result).result.value) == ("event:42" if event_first else "timeout:42") + + +@pytest.mark.parametrize("seconds", [float("inf"), float("-inf"), float("nan")]) +def test_nonfinite_timeout_rejected_before_consuming_event(seconds: float) -> None: + class InfiniteTimeout(timedelta): + def total_seconds(self) -> float: + return seconds + + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> int: + await ctx.wait_for_external_event("go") + ops = ctx.as_async() + with pytest.raises(ValueError, match="finite"): + await ops.receive_event("value", timeout=InfiniteTimeout()) + return await ops.receive_event("value", data_type=int) + + result = _execute(orchestrator, _start(), [ + helpers.new_event_raised_event("value", "42"), + helpers.new_event_raised_event("go", None), + ]) + assert json.loads(_completion(result).result.value) == 42 + + +def test_owned_timer_cancellation_does_not_cancel_borrowed_timer() -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + borrowed = ctx.create_timer(timedelta(seconds=10)) + owned = asyncio.create_task(ctx.as_async().create_timer(timedelta(seconds=5))) + waiter = asyncio.create_task(borrowed.wait()) + await asyncio.sleep(0) + owned.cancel() + waiter.cancel() + await asyncio.gather(owned, waiter, return_exceptions=True) + await borrowed + + result = _execute(orchestrator, [], _start()) + assert len(result.actions) == 1 + assert result.actions[0].id == 1 + assert result.actions[0].createTimer.fireAt.ToDatetime() == NOW + timedelta(seconds=10) + + +def test_owned_timer_completes_through_history() -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> str: + await asyncio.create_task(ctx.as_async().create_timer(timedelta(seconds=5))) + return "done" + + deadline = NOW + timedelta(seconds=5) + old = _start() + [helpers.new_timer_created_event(1, deadline)] + result = _execute(orchestrator, old, [helpers.new_timer_fired_event(1, deadline)]) + assert json.loads(_completion(result).result.value) == "done" + + +def test_owned_event_completion_cancels_only_its_uncommitted_timeout_timer() -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any: + await ctx.as_async().receive_event("value", timeout=timedelta(seconds=5)) + return await ctx.call_activity("next") + + result = _execute(orchestrator, _start(), [helpers.new_event_raised_event("value", "42")]) + assert len(result.actions) == 1 + assert result.actions[0].id == 2 + assert result.actions[0].scheduleTask.name == "next" + + +def test_dispatch_disposal_of_owned_resources_preserves_response_and_has_no_cleanup_errors( + caplog: pytest.LogCaptureFixture, +) -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + ops = ctx.as_async() + await asyncio.gather( + ops.create_timer(timedelta(seconds=10)), + ops.receive_event("value", timeout=timedelta(seconds=5)), + ops.receive_event("other"), + ) + + with caplog.at_level(logging.WARNING): + result = _execute(orchestrator, [], _start()) + assert [(a.id, a.createTimer.fireAt.ToDatetime()) for a in result.actions] == [ + (1, NOW + timedelta(seconds=10)), (2, NOW + timedelta(seconds=5))] + assert not caplog.records + + +def test_sequential_cold_operations_preserve_generator_actions() -> None: + async def coroutine(ctx: task.OrchestrationContext, _: Any) -> int: + ops = ctx.as_async() + first = await ops.call_activity("increment", input=3, return_type=int) + return await ops.call_activity("increment", input=first, return_type=int) + + def generator(ctx: task.OrchestrationContext, _: Any) -> Generator[task.Task[Any], Any, int]: + first = yield ctx.call_activity("increment", input=3, return_type=int) + return (yield ctx.call_activity("increment", input=first, return_type=int)) + + batches = [ + _start(), + [helpers.new_task_scheduled_event(1, "increment"), helpers.new_task_completed_event(1, "4")], + [helpers.new_task_scheduled_event(2, "increment"), helpers.new_task_completed_event(2, "5")], + ] + old: list[pb.HistoryEvent] = [] + result = None + for new in batches: + result = _execute(coroutine, old, new) + assert result.actions == _execute(generator, old, new).actions + old += new + assert result is not None + assert json.loads(_completion(result).result.value) == 5 + + +def test_entity_lock_and_typed_call_use_existing_replay_protocol() -> None: + entity = entities.EntityInstanceId("counter", "one") + + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> int: + ops = ctx.as_async() + with await ops.lock_entities([entity]): + return await ops.call_entity(entity, "get", return_type=int) + + old = _start() + lock = _execute(orchestrator, [], old).actions[0].sendEntityMessage.entityLockRequested + new = [ + helpers.new_event_sent_event(1, str(entity), json.dumps({"id": lock.criticalSectionId})), + helpers.new_event_raised_event(lock.criticalSectionId, None), + ] + actions = _execute(orchestrator, old, new).actions + call = actions[0].sendEntityMessage.entityOperationCalled + old += new + [helpers.new_event_sent_event(2, str(entity), json.dumps({"id": call.requestId}))] + complete = _completion(_execute( + orchestrator, old, [helpers.new_event_raised_event(call.requestId, json.dumps({"result": "42"}))])) + assert json.loads(complete.result.value) == 42 + + +@pytest.mark.parametrize("sub_orchestration", [False, True]) +def test_cold_retry_and_call_options_preserve_eager_actions(sub_orchestration: bool) -> None: + policy = task.RetryPolicy(first_retry_interval=timedelta(seconds=2), max_number_of_attempts=3) + + async def cold(ctx: task.OrchestrationContext, _: Any) -> int: + if sub_orchestration: + return await ctx.as_async().call_sub_orchestrator( + "child", input=3, instance_id="child-id", retry_policy=policy, + version="v1", return_type=int) + return await ctx.as_async().call_activity( + "work", input=3, retry_policy=policy, tags={"tag": "value"}, return_type=int) + + def eager(ctx: task.OrchestrationContext, _: Any) -> Generator[task.Task[Any], Any, int]: + if sub_orchestration: + return (yield ctx.call_sub_orchestrator( + "child", input=3, instance_id="child-id", retry_policy=policy, + version="v1", return_type=int)) + return (yield ctx.call_activity( + "work", input=3, retry_policy=policy, tags={"tag": "value"}, return_type=int)) + + start = _start() + assert _execute(cold, [], start).actions == _execute(eager, [], start).actions + if sub_orchestration: + scheduled = helpers.new_sub_orchestration_created_event(1, "child", "child-id") + failed = helpers.new_sub_orchestration_failed_event(1, ValueError("retry me")) + completed = helpers.new_sub_orchestration_completed_event(1, "42") + else: + scheduled = helpers.new_task_scheduled_event(1, "work") + failed = helpers.new_task_failed_event(1, ValueError("retry me")) + completed = helpers.new_task_completed_event(1, "42") + old = start + [scheduled] + retry = _execute(cold, old, [failed]) + assert retry.actions == _execute(eager, old, [failed]).actions + assert len(retry.actions) == 1 + assert retry.actions[0].id == 2 + deadline = retry.actions[0].createTimer.fireAt.ToDatetime() + old += [failed, helpers.new_timer_created_event(2, deadline)] + fired = [helpers.new_orchestrator_started_event(deadline), helpers.new_timer_fired_event(2, deadline)] + restarted = _execute(cold, old, fired) + assert restarted.actions == _execute(eager, old, fired).actions + assert restarted.actions[0].id == 1 + old += fired + [scheduled] + complete = _completion(_execute(cold, old, [completed])) + assert json.loads(complete.result.value) == 42 diff --git a/tests/durabletask/test_asyncio_orchestrators.py b/tests/durabletask/test_asyncio_orchestrators.py new file mode 100644 index 00000000..67518483 --- /dev/null +++ b/tests/durabletask/test_asyncio_orchestrators.py @@ -0,0 +1,886 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Deterministic asyncio behavior through the real history executor.""" + +import asyncio +import gc +import json +from contextlib import asynccontextmanager +from contextvars import ContextVar +from dataclasses import dataclass +from datetime import timedelta +from typing import Any +from weakref import ReferenceType, ref + +import pytest + +from durabletask import task +from durabletask.internal import helpers +from durabletask.internal import orchestrator_service_pb2 as pb +from tests.durabletask.test_coroutine_orchestrators import ( + NOW, + _completion, # pyright: ignore[reportPrivateUsage] + _execute, # pyright: ignore[reportPrivateUsage] + _start, # pyright: ignore[reportPrivateUsage] +) + + +@pytest.mark.parametrize("fail_branch", [False, True]) +def test_task_group_branches_inherit_context_and_replay_stable_actions(fail_branch: bool) -> None: + context: ContextVar[task.OrchestrationContext] = ContextVar("orchestration") + + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any: + token = context.set(ctx) + try: + results: dict[str, Any] = {} + + async def branch(name: str) -> None: + current = context.get() + assert current is ctx + results[name] = await current.call_activity(name, input=current.instance_id) + + async with asyncio.TaskGroup() as group: + group.create_task(branch("first")) + group.create_task(branch("second")) + return await context.get().call_activity("finish", input=[results["first"], results["second"]]) + finally: + context.reset(token) + + def replay(old: list[pb.HistoryEvent], new: list[pb.HistoryEvent]): + first = _execute(orchestrator, old, new) + assert _execute(orchestrator, old, new).actions == first.actions + with pytest.raises(LookupError): + context.get() + return first + + start = _start() + initial = replay([], start) + assert [action.scheduleTask.name for action in initial.actions] == ["first", "second"] + assert all(json.loads(action.scheduleTask.input.value) == "coroutine-test" for action in initial.actions) + old = start + [helpers.new_task_scheduled_event(1, "first"), + helpers.new_task_scheduled_event(2, "second")] + second_done = [helpers.new_task_completed_event(2, '"second result"')] + assert replay(old, second_done).actions == [] + first_done = [helpers.new_task_failed_event(1, ValueError("recorded branch failure"))] if fail_branch else [ + helpers.new_task_completed_event(1, '"first result"')] + result = replay(old + second_done, first_done) + if fail_branch: + complete = _completion(result) + assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + assert "recorded branch failure" in str(complete.failureDetails) + else: + assert len(result.actions) == 1 + assert result.actions[0].scheduleTask.name == "finish" + assert json.loads(result.actions[0].scheduleTask.input.value) == ["first result", "second result"] + complete = _completion(replay(old + second_done + first_done, [ + helpers.new_task_scheduled_event(3, "finish"), + helpers.new_task_completed_event(3, '"done"')])) + assert json.loads(complete.result.value) == "done" + + +def test_gather_resumes_each_branch_without_waiting_for_batch() -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> list[Any]: + async def branch(name: str) -> Any: + value: Any = await ctx.call_activity(name) + await asyncio.sleep(0) + return await ctx.call_activity(name + "-next", input=value) + + return list(await asyncio.gather(branch("a"), branch("b"))) + + start = _start() + actions = _execute(orchestrator, [], start).actions + assert [a.scheduleTask.name for a in actions] == ["a", "b"] + old = start + [helpers.new_task_scheduled_event(i, name) for i, name in [(1, "a"), (2, "b")]] + second = [helpers.new_task_completed_event(2, '"b"')] + actions = _execute(orchestrator, old, second).actions + assert [(a.id, a.scheduleTask.name) for a in actions] == [(3, "b-next")] + third = [helpers.new_task_scheduled_event(3, "b-next"), + helpers.new_task_completed_event(1, '"a"'), + helpers.new_task_scheduled_event(4, "a-next"), + helpers.new_task_completed_event(4, '"a-done"'), + helpers.new_task_completed_event(3, '"b-done"')] + assert json.loads(_completion(_execute(orchestrator, old + second, third)).result.value) == [ + "a-done", "b-done"] + + +@pytest.mark.parametrize("return_exceptions", [False, True]) +def test_gather_failure_does_not_cancel_sibling(return_exceptions: bool) -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any: + async def sibling() -> str: + await ctx.call_activity("slow") + return await ctx.call_activity("continued", return_type=str) + + pending = asyncio.create_task(sibling()) + try: + values = await asyncio.gather( + ctx.call_activity("fail", return_type=str), pending, return_exceptions=return_exceptions) + return [isinstance(values[0], task.TaskFailedError), values[1]] + except task.TaskFailedError: + assert not pending.cancelled() + return await pending + + old = _start() + [helpers.new_task_scheduled_event(1, "fail"), + helpers.new_task_scheduled_event(2, "slow")] + failed = [helpers.new_task_failed_event(1, ValueError("boom"))] + assert _execute(orchestrator, old, failed).actions == [] + completed = [helpers.new_task_completed_event(2, "null")] + assert _execute(orchestrator, old + failed, completed).actions[0].scheduleTask.name == "continued" + result = _completion(_execute(orchestrator, old + failed + completed, [ + helpers.new_task_scheduled_event(3, "continued"), helpers.new_task_completed_event(3, '"ok"')])) + assert json.loads(result.result.value) == ([True, "ok"] if return_exceptions else "ok") + + +def test_task_group_failure_cancels_waiters_not_durable_work() -> None: + cleaned: list[str] = [] + + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def branch(name: str) -> None: + try: + await ctx.call_activity(name) + finally: + await asyncio.sleep(0) + cleaned.append(name) + + async with asyncio.TaskGroup() as group: + group.create_task(branch("fail")) + group.create_task(branch("slow")) + + old = _start() + [helpers.new_task_scheduled_event(1, "fail"), + helpers.new_task_scheduled_event(2, "slow")] + result = _completion(_execute(orchestrator, old, [helpers.new_task_failed_event(1, ValueError("boom"))])) + assert result.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + assert cleaned == ["fail", "slow"] + + +@pytest.mark.parametrize("kind", ["sleep", "wait_for", "timeout"]) +def test_durable_timer_and_timeout(kind: str) -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> str: + try: + if kind == "sleep": + await asyncio.sleep(2) + elif kind == "wait_for": + await asyncio.wait_for(ctx.wait_for_external_event("never"), timeout=2) + else: + async with asyncio.timeout(2): + await ctx.wait_for_external_event("never") + except TimeoutError: + return "timed out" + return "slept" + + actions = _execute(orchestrator, [], _start()).actions + assert len(actions) == 1 + deadline = actions[0].createTimer.fireAt.ToDatetime() + assert deadline == NOW + timedelta(seconds=2) + old = _start() + [helpers.new_timer_created_event(1, deadline)] + new = [helpers.new_orchestrator_started_event(deadline), helpers.new_timer_fired_event(1, deadline)] + assert json.loads(_completion(_execute(orchestrator, old, new)).result.value) == ( + "slept" if kind == "sleep" else "timed out") + + +def test_sleep_zero_fifo_without_timer() -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> list[int]: + observed: list[int] = [] + + async def branch(index: int) -> None: + observed.append(index) + await asyncio.sleep(0) + observed.append(index + 2) + + await asyncio.gather(branch(0), branch(1)) + return observed + + result = _execute(orchestrator, [], _start()) + assert len(result.actions) == 1 + assert json.loads(_completion(result).result.value) == [0, 1, 2, 3] + + +def test_discarded_gather_batches_release_decoded_results_during_replay() -> None: + references: list[ReferenceType[Any]] = [] + + @dataclass + class Payload: + value: int + + def __post_init__(self) -> None: + references.append(ref(self)) + + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> str: + for batch in range(8): + await asyncio.gather(*(ctx.call_activity("work", return_type=Payload) for _ in range(4))) + await asyncio.sleep(0) + gc.collect() + assert len(references) == (batch + 1) * 4 + # The history handler still holds the most recent decoded result + # while resume() drains this batch's callbacks. + assert sum(reference() is not None for reference in references) <= 1 + return "done" + + history = _start() + for batch in range(8): + ids = range(batch * 4 + 1, batch * 4 + 5) + history += [helpers.new_task_scheduled_event(index, "work") for index in ids] + history += [helpers.new_task_completed_event(index, json.dumps({"value": index})) for index in ids] + complete = _completion(_execute(orchestrator, history[:-1], history[-1:])) + assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED, complete.failureDetails + assert json.loads(complete.result.value) == "done" + gc.collect() + assert all(reference() is None for reference in references) + + +@pytest.mark.parametrize("observe", [ + "await", "gather", "gather-exceptions", "result", "exception", "task-group", +]) +def test_observed_failures_release_traceback_payloads_during_replay(observe: str) -> None: + references: list[ReferenceType[Any]] = [] + + @dataclass + class Payload: + value: int + + def __post_init__(self) -> None: + references.append(ref(self)) + + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> str: + async def branch() -> None: + payload = await ctx.call_activity("work", return_type=Payload) + raise ValueError(f"Failed with payload {payload.value}") + + for batch in range(8): + if observe == "task-group": + group = asyncio.TaskGroup() + try: + async with group: + group.create_task(branch()) + except* ValueError: + pass + del group + else: + pending = asyncio.create_task(branch()) + await asyncio.sleep(0) + # Observe only after the failed task's initial retirement callback. + await ctx.call_activity("gate") + if observe == "gather-exceptions": + await asyncio.gather(pending, return_exceptions=True) + elif observe == "exception": + assert isinstance(pending.exception(), ValueError) + else: + try: + if observe == "await": + await pending + elif observe == "gather": + await asyncio.gather(pending) + else: + pending.result() + except ValueError: + pass + del pending + await asyncio.sleep(0) + gc.collect() + assert len(references) == batch + 1 + assert sum(reference() is not None for reference in references) <= 1 + return "done" + + history = _start() + sequence = 0 + for batch in range(8): + sequence += 1 + work_id = sequence + history.append(helpers.new_task_scheduled_event(work_id, "work")) + if observe != "task-group": + sequence += 1 + history.append(helpers.new_task_scheduled_event(sequence, "gate")) + history.append(helpers.new_task_completed_event(work_id, json.dumps({"value": batch}))) + if observe != "task-group": + history.append(helpers.new_task_completed_event(sequence, "null")) + complete = _completion(_execute(orchestrator, history[:-1], history[-1:])) + assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED, complete.failureDetails + assert json.loads(complete.result.value) == "done" + gc.collect() + assert all(reference() is None for reference in references) + + +def test_retired_tasks_do_not_reuse_cancellation_order() -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> str: + async def completed() -> None: + pass + + async def branch(index: int) -> None: + try: + await ctx.wait_for_external_event(str(index)) + finally: + await ctx.call_activity("compensate", input=index) + + asyncio.create_task(completed()) + first = asyncio.create_task(branch(0)) + await asyncio.sleep(0) + await asyncio.sleep(0) + second = asyncio.create_task(branch(1)) + await asyncio.sleep(0) + second.cancel() + first.cancel() + await asyncio.gather(first, second, return_exceptions=True) + return "done" + + initial = _execute(orchestrator, [], _start()) + assert [json.loads(action.scheduleTask.input.value) for action in initial.actions] == [0, 1] + assert _execute(orchestrator, [], _start()).actions == initial.actions + old = _start() + [helpers.new_task_scheduled_event(index, "compensate") for index in (1, 2)] + new = [helpers.new_task_completed_event(index, "null") for index in (1, 2)] + assert json.loads(_completion(_execute(orchestrator, old, new)).result.value) == "done" + + +@pytest.mark.parametrize("under_loop", [False, True]) +def test_dispatch_cleanup_unwinds_scopes_contexts_asyncgens(under_loop: bool) -> None: + cleaned: list[str] = [] + references: list[Any] = [] + + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + references.extend([ref(ctx), ref(asyncio.get_running_loop()), ref(asyncio.current_task())]) + + async def resource(): + try: + yield 1 + finally: + await asyncio.sleep(0) + cleaned.append("asyncgen") + + opened = resource() + await anext(opened) + try: + async with asyncio.TaskGroup(): + await ctx.call_activity("wait") + finally: + await asyncio.sleep(0) + cleaned.append("root") + + def dispatch() -> None: + result = _execute(orchestrator, [], _start()) + assert [a.scheduleTask.name for a in result.actions] == ["wait"] + assert cleaned == ["root", "asyncgen"] + + async def outer() -> None: + before = asyncio.current_task() + dispatch() + assert asyncio.current_task() is before + await asyncio.sleep(0) + + if under_loop: + asyncio.run(outer()) + else: + dispatch() + gc.collect() + assert all(reference() is None for reference in references) + + +def test_shared_operation_survives_local_waiter_cancellation() -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any: + durable = ctx.call_activity("shared", return_type=str) + + async def wait() -> Any: + return await durable + + first = asyncio.create_task(wait()) + second = asyncio.create_task(wait()) + await asyncio.sleep(0) + first.cancel() + await asyncio.gather(first, return_exceptions=True) + return await second + + old = _start() + [helpers.new_task_scheduled_event(1, "shared")] + assert json.loads(_completion(_execute(orchestrator, old, [ + helpers.new_task_completed_event(1, '"ok"')])).result.value) == "ok" + + +def test_local_durable_cancellation_wakes_waiter_without_history() -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> str: + timer = ctx.create_timer(timedelta(days=1)) + + async def waiter() -> None: + await timer + + pending = asyncio.create_task(waiter()) + await asyncio.sleep(0) + timer.cancel() + try: + await pending + except task.TaskCancelledError: + return "cancelled" + return "unexpected" + + result = _execute(orchestrator, [], _start()) + assert len(result.actions) == 1 + assert json.loads(_completion(result).result.value) == "cancelled" + + +@pytest.mark.parametrize("terminal", ["suspend", "terminate", "continue", "history-failure"]) +def test_async_teardown_preserves_response_and_cannot_schedule_work( + terminal: str, caplog: pytest.LogCaptureFixture) -> None: + cleaned: list[str] = [] + + @asynccontextmanager + async def resource(): + try: + yield + finally: + await asyncio.sleep(0) + cleaned.append("resource") + + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async with resource(): + try: + ctx.set_custom_status("waiting") + if terminal == "continue": + ctx.continue_as_new("next") + await ctx.call_activity("wait") + finally: + try: + await asyncio.sleep(0) + cleaned.append("inner") + ctx.call_activity("forbidden-cleanup") + finally: + await asyncio.sleep(0) + cleaned.append("outer") + + old = _start() + if terminal == "terminate": + new = [pb.HistoryEvent(executionTerminated=pb.ExecutionTerminatedEvent())] + elif terminal == "history-failure": + new = [helpers.new_task_scheduled_event(1, "wrong")] + else: + old, new = [], old + result = _execute(orchestrator, old, new) + assert cleaned == ["inner", "outer", "resource"] + assert result.encoded_custom_status == '"waiting"' + assert all(action.scheduleTask.name != "forbidden-cleanup" for action in result.actions) + if terminal == "terminate": + assert _completion(result).orchestrationStatus == pb.ORCHESTRATION_STATUS_TERMINATED + elif terminal == "history-failure": + assert _completion(result).failureDetails.errorType.endswith("NonDeterminismError") + elif terminal == "continue": + assert _completion(result).orchestrationStatus == pb.ORCHESTRATION_STATUS_CONTINUED_AS_NEW + else: + assert [action.scheduleTask.name for action in result.actions] == ["wait"] + assert "Dispatch cleanup cannot schedule durable work" in caplog.text + + +def test_unhandled_local_root_cancellation_fails_not_abandons() -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + raise asyncio.CancelledError + + complete = _completion(_execute(orchestrator, [], _start())) + assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + assert complete.failureDetails.errorType.endswith("TaskCancelledError") + + +@pytest.mark.parametrize("delay", [0.0000003, 0.000001, 0.0000013, 2.25]) +def test_timer_rounding_never_precedes_loop_deadline(delay: float) -> None: + observed: list[float] = [] + + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + loop = asyncio.get_running_loop() + observed.append(loop.time() + delay) + await asyncio.sleep(delay) + + actions = _execute(orchestrator, [], _start()).actions + assert len(actions) == 1 + deadline = actions[0].createTimer.fireAt.ToDatetime() + from datetime import datetime + assert (deadline - datetime(1970, 1, 1)).total_seconds() >= observed[0] + + +def test_async_sleep_uses_shared_long_timer_chunking() -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> str: + await asyncio.sleep(7 * 86400) + return "done" + + old = _start() + for timer_id, days in [(1, 3), (2, 6), (3, 7)]: + result = _execute(orchestrator, [], old) + assert len(result.actions) == 1 + action = result.actions[0] + assert action.id == timer_id + deadline = NOW + timedelta(days=days) + assert action.createTimer.fireAt.ToDatetime() == deadline + old += [helpers.new_timer_created_event(timer_id, deadline), + helpers.new_orchestrator_started_event(deadline), + helpers.new_timer_fired_event(timer_id, deadline)] + assert json.loads(_completion(_execute(orchestrator, old[:-1], old[-1:])).result.value) == "done" + + +@pytest.mark.parametrize("api", ["thread", "io", "foreign-loop"]) +def test_unsupported_work_fails_actionably(api: str) -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + if api == "thread": + await asyncio.to_thread(lambda: None) + elif api == "io": + await asyncio.open_connection("localhost", 1) + else: + foreign = asyncio.new_event_loop() + try: + await foreign.create_future() + finally: + foreign.close() + + complete = _completion(_execute(orchestrator, [], _start())) + assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + assert any(text in complete.failureDetails.errorMessage for text in ( + "use an activity", "ordinary activity", "different loop")) + + +@pytest.mark.parametrize("async_generator", [False, True]) +@pytest.mark.parametrize("unexpected", [False, True]) +def test_cleanup_failure_policy_includes_async_generator_tasks( + async_generator: bool, unexpected: bool, caplog: pytest.LogCaptureFixture) -> None: + class CleanupAbort(BaseException): + pass + + error = CleanupAbort("cleanup failed") if unexpected else ValueError("cleanup failed") + + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def resource(): + try: + yield 1 + finally: + await asyncio.sleep(0) + raise error + + opened = resource() + await anext(opened) + try: + await ctx.call_activity("wait") + finally: + if not async_generator: + await opened.aclose() + + result = _execute(orchestrator, [], _start()) + if unexpected: + complete = _completion(result) + assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + assert complete.failureDetails.errorMessage == "cleanup failed" + assert complete.failureDetails.errorType.endswith("CleanupAbort") + else: + assert [action.scheduleTask.name for action in result.actions] == ["wait"] + assert "cleanup failed" in caplog.text + + +def test_task_group_cancellation_durable_cleanup_has_stable_action_order() -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def branch(index: int) -> None: + try: + await ctx.call_activity("work", input=index) + finally: + # This is real task-group cancellation during replay, not + # dispatch teardown: durable compensation is still allowed. + await ctx.call_activity("compensate", input=index) + + async with asyncio.TaskGroup() as group: + for index in range(4): + group.create_task(branch(index)) + + old = _start() + [helpers.new_task_scheduled_event(index, "work") for index in range(1, 5)] + failure = [helpers.new_task_failed_event(1, ValueError("boom"))] + # The failing branch first enters its durable finally, so deliver that + # completion before task-group sibling cancellation starts. + first = _execute(orchestrator, old, failure) + assert [(a.id, a.scheduleTask.name, json.loads(a.scheduleTask.input.value)) for a in first.actions] == [ + (5, "compensate", 0)] + old += failure + [helpers.new_task_scheduled_event(5, "compensate"), + helpers.new_task_completed_event(5, "null")] + expected = [helpers.new_schedule_task_action(index + 6, "compensate", json.dumps(index + 1), None) + for index in range(3)] + for _ in range(10): + assert _execute(orchestrator, [], old).actions == expected + + +def test_cancelling_last_loop_timer_cannot_hide_foreign_wait() -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + handle = asyncio.get_running_loop().call_later(1, lambda: None) + handle.cancel() + await asyncio.Future() + + complete = _completion(_execute(orchestrator, [], _start())) + assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + assert "blocked without durable work" in complete.failureDetails.errorMessage + + +@pytest.mark.parametrize("observed", [False, True]) +@pytest.mark.parametrize("probe_pending", [False, True]) +def test_background_failure_is_reported_unless_observed( + observed: bool, probe_pending: bool, caplog: pytest.LogCaptureFixture) -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> str: + async def fail() -> None: + raise ValueError("background failure") + + pending = asyncio.create_task(fail()) + if probe_pending: + with pytest.raises(asyncio.InvalidStateError): + pending.result() + with pytest.raises(asyncio.InvalidStateError): + pending.exception() + await asyncio.sleep(0) + if observed: + await asyncio.gather(pending, return_exceptions=True) + del pending + await asyncio.sleep(0) + gc.collect() + return "done" + + assert json.loads(_completion(_execute(orchestrator, [], _start())).result.value) == "done" + assert ("background failure" in caplog.text) is not observed + + +def test_unexpected_cleanup_failure_does_not_hide_other_task_errors( + caplog: pytest.LogCaptureFixture) -> None: + class CleanupAbort(BaseException): + pass + + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def branch(error: BaseException) -> None: + try: + await ctx.wait_for_external_event("never") + finally: + await asyncio.sleep(0) + raise error + + asyncio.create_task(branch(CleanupAbort("first cleanup failure"))) + asyncio.create_task(branch(ValueError("second cleanup failure"))) + await ctx.call_activity("wait") + + complete = _completion(_execute(orchestrator, [], _start())) + assert complete.failureDetails.errorMessage == "first cleanup failure" + assert "second cleanup failure" in caplog.text + + +@pytest.mark.parametrize("control_flow", [False, True]) +def test_cleanup_callback_failure_still_drains_tasks_and_asyncgens( + control_flow: bool, caplog: pytest.LogCaptureFixture) -> None: + cleaned: list[str] = [] + retained: list[asyncio.Task[Any]] = [] + + def fail() -> None: + raise SystemExit("cleanup exit") if control_flow else ValueError("cleanup callback") + + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def resource(): + try: + yield 1 + finally: + await asyncio.sleep(0) + cleaned.append("asyncgen") + + async def sibling() -> None: + try: + await ctx.wait_for_external_event("never") + finally: + await asyncio.sleep(0) + cleaned.append("sibling") + + opened = resource() + await anext(opened) + retained.append(asyncio.create_task(sibling())) + current = asyncio.current_task() + assert current is not None + retained.append(current) + try: + await ctx.call_activity("wait") + finally: + if control_flow: + fail() + asyncio.get_running_loop().call_soon(fail) + await asyncio.sleep(0) + cleaned.append("root") + + if control_flow: + with pytest.raises(SystemExit, match="cleanup exit"): + _execute(orchestrator, [], _start()) + else: + result = _execute(orchestrator, [], _start()) + assert [action.scheduleTask.name for action in result.actions] == ["wait"] + assert "cleanup callback" in caplog.text + assert "root" in cleaned + assert "sibling" in cleaned + assert "asyncgen" in cleaned + assert all(pending.done() for pending in retained) + + +@pytest.mark.parametrize("count", [128, 4096]) +def test_fanout_waiter_inspections_scale_linearly( + count: int, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture) -> None: + inspections = 0 + + def is_complete(durable: task.Task[Any]) -> bool: + nonlocal inspections + inspections += 1 + return durable._is_complete # pyright: ignore[reportPrivateUsage] + + monkeypatch.setattr(task.Task, "is_complete", property(is_complete)) + + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> list[Any]: + async def branch(index: int) -> Any: + return await ctx.call_activity("work", input=index) + + return await asyncio.gather(*(branch(index) for index in range(count))) + + result = _execute(orchestrator, [], _start()) + assert [action.id for action in result.actions] == list(range(1, count + 1)) + assert inspections < count * 20 + assert not caplog.records + + +def test_nonquiescent_cleanup_does_not_skip_asyncgen_closure( + caplog: pytest.LogCaptureFixture) -> None: + cleaned: list[str] = [] + + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def resource(): + try: + yield 1 + finally: + await asyncio.sleep(0) + cleaned.append("asyncgen") + + opened = resource() + await anext(opened) + try: + await ctx.call_activity("wait") + finally: + while True: + await asyncio.sleep(0) + + result = _execute(orchestrator, [], _start()) + assert [action.scheduleTask.name for action in result.actions] == ["wait"] + assert cleaned == ["asyncgen"] + assert "cleanup did not quiesce" in caplog.text + + +@pytest.mark.parametrize("grouped", [False, True]) +@pytest.mark.parametrize("custom_error", [False, True]) +def test_forced_cleanup_preserves_process_control_after_other_errors(grouped: bool, custom_error: bool) -> None: + class CleanupAbort(BaseException): + pass + + cleaned: list[str] = [] + exit_error: BaseException = SystemExit("forced cleanup exit") + if grouped: + exit_error = BaseExceptionGroup("forced cleanup group", [exit_error, ValueError("nested")]) + + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def branch(name: str, error: BaseException | None) -> None: + try: + await ctx.wait_for_external_event(name) + finally: + try: + while True: + try: + await asyncio.Future() + except asyncio.CancelledError: + # Force teardown to close a cancellation-resistant coroutine. + pass + finally: + cleaned.append(name) + if error is not None: + raise error + + first_error = CleanupAbort("custom cleanup error") if custom_error else ValueError("ordinary cleanup error") + asyncio.create_task(branch("ordinary", first_error)) + asyncio.create_task(branch("exit", exit_error)) + asyncio.create_task(branch("last", None)) + await ctx.call_activity("wait") + + with pytest.raises(type(exit_error)) as caught: + _execute(orchestrator, [], _start()) + assert caught.value is exit_error + assert cleaned == ["ordinary", "exit", "last"] + + +@pytest.mark.parametrize("timer", [False, True]) +def test_callback_stop_iteration_fails_instead_of_completing_orchestrator(timer: bool) -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + def fail() -> None: + raise StopIteration("callback stopped") + + loop = asyncio.get_running_loop() + if timer: + loop.call_later(1, fail) + else: + loop.call_soon(fail) + await ctx.wait_for_external_event("never") + + start = _start() + if timer: + deadline = NOW + timedelta(seconds=1) + old = start + [helpers.new_timer_created_event(1, deadline)] + new = [helpers.new_orchestrator_started_event(deadline), helpers.new_timer_fired_event(1, deadline)] + else: + old, new = [], start + complete = _completion(_execute(orchestrator, old, new)) + assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + assert complete.failureDetails.errorType == "builtins.RuntimeError" + assert complete.failureDetails.innerFailure.errorType == "builtins.StopIteration" + assert complete.failureDetails.innerFailure.errorMessage == "callback stopped" + + +@pytest.mark.parametrize("cancelled", [None, 0, 1]) +def test_identical_timer_callbacks_complete_and_cancel_independently(cancelled: int | None) -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> int: + loop = asyncio.get_running_loop() + result = loop.create_future() + completed: list[None] = [] + expected = 2 if cancelled is None else 1 + + def callback() -> None: + completed.append(None) + if len(completed) == expected: + result.set_result(len(completed)) + + handles = [loop.call_at(loop.time() + 1, callback) for _ in range(2)] + assert handles[0] != handles[1] + assert len(set(handles)) == 2 + if cancelled is not None: + handles[cancelled].cancel() + return await result + + start = _start() + initial = _execute(orchestrator, [], start) + remaining = [index + 1 for index in range(2) if index != cancelled] + assert [action.id for action in initial.actions] == remaining + deadline = NOW + timedelta(seconds=1) + old = start + [helpers.new_timer_created_event(index, deadline) for index in remaining] + for offset, timer_id in enumerate(remaining): + new = [helpers.new_orchestrator_started_event(deadline), helpers.new_timer_fired_event(timer_id, deadline)] + result = _execute(orchestrator, old, new) + assert _execute(orchestrator, old, new).actions == result.actions + if offset < len(remaining) - 1: + assert result.actions == [] + else: + assert json.loads(_completion(result).result.value) == len(remaining) + old += new + + +def test_cooperative_cleanup_prioritizes_grouped_process_control() -> None: + class CleanupAbort(BaseException): + pass + + exit_error = BaseExceptionGroup("cleanup exit", [SystemExit("stop"), ValueError("nested")]) + cleaned: list[str] = [] + + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def branch(name: str, error: BaseException | None) -> None: + try: + await ctx.wait_for_external_event(name) + finally: + await asyncio.sleep(0) + cleaned.append(name) + if error is not None: + raise error + + asyncio.create_task(branch("custom", CleanupAbort("first failure"))) + asyncio.create_task(branch("exit", exit_error)) + asyncio.create_task(branch("last", None)) + await ctx.call_activity("wait") + + with pytest.raises(BaseExceptionGroup) as caught: + _execute(orchestrator, [], _start()) + assert caught.value is exit_error + assert cleaned == ["custom", "exit", "last"] diff --git a/tests/durabletask/test_coroutine_orchestrators.py b/tests/durabletask/test_coroutine_orchestrators.py new file mode 100644 index 00000000..9f32aa18 --- /dev/null +++ b/tests/durabletask/test_coroutine_orchestrators.py @@ -0,0 +1,338 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Replay-level coverage for coroutine orchestrators and generator history parity.""" + +import asyncio +import json +import logging +from collections.abc import AsyncGenerator, Generator +from datetime import datetime, timedelta +from typing import Any + +import pytest + +from durabletask import entities, task, worker +from durabletask.internal import helpers +from durabletask.internal import orchestrator_service_pb2 as pb +from durabletask.serialization import JsonDataConverter + + +INSTANCE_ID = "coroutine-test" +NOW = datetime(2026, 1, 1) + + +def _start() -> list[pb.HistoryEvent]: + return [ + helpers.new_orchestrator_started_event(NOW), + helpers.new_execution_started_event("test", INSTANCE_ID, json.dumps(3)), + ] + + +def _execute( + fn: task.Orchestrator[Any, Any], + old: list[pb.HistoryEvent], + new: list[pb.HistoryEvent], +) -> worker.ExecutionResults: + registry = worker._Registry() # pyright: ignore[reportPrivateUsage] + registry.add_named_orchestrator("test", fn) + executor = worker._OrchestrationExecutor( # pyright: ignore[reportPrivateUsage] + registry, logging.getLogger(__name__), JsonDataConverter()) + return executor.execute(INSTANCE_ID, old, new) + + +def _completion(result: worker.ExecutionResults) -> pb.CompleteOrchestrationAction: + completions = [a.completeOrchestration for a in result.actions if a.HasField("completeOrchestration")] + assert len(completions) == 1 + return completions[0] + + +def test_coroutine_immediate_return() -> None: + async def orchestrator(ctx: task.OrchestrationContext, value: int) -> int: + return value + + complete = _completion(_execute(orchestrator, [], _start())) + assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED + assert json.loads(complete.result.value) == 3 + + +def test_await_is_concrete_and_completed_tasks_do_not_suspend() -> None: + assert "__await__" not in task.Task.__abstractmethods__ + + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> list[int]: + completed = task.CompletableTask[int]() + completed.complete(7) + return [await completed, await completed] + + assert json.loads(_completion(_execute(orchestrator, [], _start())).result.value) == [7, 7] + + +@pytest.mark.parametrize("after_activity", [False, True]) +def test_generator_invalid_yield_recommends_async_authoring(after_activity: bool) -> None: + def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Generator[Any, Any, None]: + if after_activity: + yield ctx.call_activity("work") + yield None + + old = _start() + [helpers.new_task_scheduled_event(1, "work")] if after_activity else [] + new = [helpers.new_task_completed_event(1, "null")] if after_activity else _start() + complete = _completion(_execute(orchestrator, old, new)) + assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + assert complete.failureDetails.errorType == "builtins.TypeError" + assert "Generator orchestrators must yield durabletask.Task objects" in complete.failureDetails.errorMessage + assert "Use async def and await for asyncio composition" in complete.failureDetails.errorMessage + + +def test_precompleted_failure_and_cancellation_are_catchable() -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> list[str]: + failed = task.CompletableTask[int]() + failed.fail("failed", ValueError("boom")) + cancelled = task.TimerTask() + cancelled.cancel() + results: list[str] = [] + try: + await failed + except task.TaskFailedError: + results.append("failed") + try: + await cancelled + except task.TaskCancelledError: + results.append("cancelled") + return results + + assert json.loads(_completion(_execute(orchestrator, [], _start())).result.value) == [ + "failed", "cancelled"] + + +def test_nested_async_helpers_replay_with_generator_action_parity() -> None: + async def helper(ctx: task.OrchestrationContext, value: int) -> int: + return await ctx.call_activity("increment", input=value, return_type=int) + + async def coroutine(ctx: task.OrchestrationContext, value: int) -> int: + first = await helper(ctx, value) + return await helper(ctx, first) + + def generator(ctx: task.OrchestrationContext, value: int) -> Generator[task.Task[Any], Any, int]: + first = yield ctx.call_activity("increment", input=value, return_type=int) + return (yield ctx.call_activity("increment", input=first, return_type=int)) + + batches = [ + _start(), + [helpers.new_task_scheduled_event(1, "increment"), + helpers.new_task_completed_event(1, json.dumps(4))], + [helpers.new_task_scheduled_event(2, "increment"), + helpers.new_task_completed_event(2, json.dumps(5))], + ] + old: list[pb.HistoryEvent] = [] + result: worker.ExecutionResults | None = None + for new in batches: + result = _execute(coroutine, old, new) + assert result.actions == _execute(generator, old, new).actions + old += new + assert result is not None + assert json.loads(_completion(result).result.value) == 5 + + +def test_fan_out_waits_for_all_and_preserves_input_order() -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> list[str]: + pending: list[task.Task[str]] = [ + ctx.call_activity("work", input=i, return_type=str) for i in range(2)] + return await task.when_all(pending) + + old = _start() + assert len(_execute(orchestrator, [], old).actions) == 2 + old += [helpers.new_task_scheduled_event(i, "work") for i in (1, 2)] + second = [helpers.new_task_completed_event(2, '"second"')] + assert _execute(orchestrator, old, second).actions == [] + complete = _completion(_execute( + orchestrator, old + second, [helpers.new_task_completed_event(1, '"first"')])) + assert json.loads(complete.result.value) == ["first", "second"] + + +@pytest.mark.parametrize("caught", [False, True]) +def test_activity_failure_is_thrown_into_coroutine(caught: bool) -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> str: + try: + await ctx.call_activity("boom") + except task.TaskFailedError: + if not caught: + raise + return "caught" + return "unexpected" + + old = _start() + [helpers.new_task_scheduled_event(1, "boom")] + complete = _completion(_execute( + orchestrator, old, [helpers.new_task_failed_event(1, ValueError("boom"))])) + assert complete.orchestrationStatus == ( + pb.ORCHESTRATION_STATUS_COMPLETED if caught else pb.ORCHESTRATION_STATUS_FAILED) + if caught: + assert json.loads(complete.result.value) == "caught" + else: + assert complete.failureDetails.innerFailure.errorMessage == "boom" + + +@pytest.mark.parametrize("event_wins", [False, True]) +def test_when_any_timer_and_external_event(event_wins: bool) -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any: + event = ctx.wait_for_external_event("approval") + timer = ctx.create_timer(timedelta(seconds=1)) + winner = await task.when_any([event, timer]) + if winner is event: + timer.cancel() + return await event + return "timeout" + + deadline = NOW + timedelta(seconds=1) + old = _start() + [helpers.new_timer_created_event(1, deadline)] + new = [helpers.new_event_raised_event("approval", '"approved"')] if event_wins else [ + helpers.new_timer_fired_event(1, deadline)] + complete = _completion(_execute(orchestrator, old, new)) + assert json.loads(complete.result.value) == ("approved" if event_wins else "timeout") + + +def test_sub_orchestration_and_replay_flags() -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> list[Any]: + before = ctx.is_replaying + value = await ctx.call_sub_orchestrator("child", instance_id="child-test", return_type=str) + return [value, before, ctx.is_replaying] + + old = _start() + [helpers.new_sub_orchestration_created_event(1, "child", "child-test")] + complete = _completion(_execute( + orchestrator, old, [helpers.new_sub_orchestration_completed_event(1, '"done"')])) + assert json.loads(complete.result.value) == ["done", True, False] + + +def test_retry_timer_resumes_same_coroutine() -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any: + return await ctx.call_activity( + "retry", retry_policy=task.RetryPolicy( + first_retry_interval=timedelta(seconds=1), max_number_of_attempts=2)) + + old = _start() + [helpers.new_task_scheduled_event(1, "retry")] + failure = [helpers.new_task_failed_event(1, ValueError("retry me"))] + retry = _execute(orchestrator, old, failure).actions + assert len(retry) == 1 + assert retry[0].id == 2 + deadline = retry[0].createTimer.fireAt.ToDatetime() + old += failure + [helpers.new_timer_created_event(2, deadline)] + fired = [helpers.new_orchestrator_started_event(deadline), helpers.new_timer_fired_event(2, deadline)] + scheduled = _execute(orchestrator, old, fired).actions + assert len(scheduled) == 1 + assert scheduled[0].id == 1 + assert scheduled[0].scheduleTask.name == "retry" + old += fired + [helpers.new_task_scheduled_event(1, "retry")] + complete = _completion(_execute( + orchestrator, old, [helpers.new_task_completed_event(1, '"success"')])) + assert json.loads(complete.result.value) == "success" + + +@pytest.mark.parametrize("sub_orchestration", [False, True]) +@pytest.mark.parametrize("replayed", [False, True]) +@pytest.mark.parametrize("caught", [False, True]) +def test_coroutine_non_retriable_failure( + sub_orchestration: bool, replayed: bool, caught: bool) -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any: + retry = task.RetryPolicy(first_retry_interval=timedelta(seconds=1), max_number_of_attempts=3) + try: + if sub_orchestration: + return await ctx.call_sub_orchestrator("child", retry_policy=retry) + return await ctx.call_activity("child", retry_policy=retry) + except task.TaskFailedError as error: + if not caught: + raise + assert error.details.message == "permanent" + return "caught" + + old = _start() + if sub_orchestration: + old.append(helpers.new_sub_orchestration_created_event(1, "child", "child-instance")) + failed = helpers.new_sub_orchestration_failed_event(1, ValueError("permanent")) + details = failed.subOrchestrationInstanceFailed.failureDetails + else: + old.append(helpers.new_task_scheduled_event(1, "child")) + failed = helpers.new_task_failed_event(1, ValueError("permanent")) + details = failed.taskFailed.failureDetails + details.isNonRetriable = True + new = [helpers.new_orchestrator_started_event(NOW)] + if replayed: + old.append(failed) + else: + new.append(failed) + result = _execute(orchestrator, old, new) + assert len(result.actions) == 1 + complete = _completion(result) + if caught: + assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED + assert json.loads(complete.result.value) == "caught" + else: + assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + assert complete.failureDetails.innerFailure == details + + +def test_entity_lock_and_call_suspend_coroutine() -> None: + entity = entities.EntityInstanceId("counter", "one") + + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any: + with await ctx.lock_entities([entity]): + return await ctx.call_entity(entity, "get") + + old = _start() + lock = _execute(orchestrator, [], old).actions[0].sendEntityMessage.entityLockRequested + request = json.dumps({"id": lock.criticalSectionId}) + new = [ + helpers.new_event_sent_event(1, str(entity), request), + helpers.new_event_raised_event(lock.criticalSectionId, None), + ] + actions = _execute(orchestrator, old, new).actions + assert len(actions) == 1 + call = actions[0].sendEntityMessage.entityOperationCalled + assert call.targetInstanceId.value == str(entity) + old += new + [helpers.new_event_sent_event(2, str(entity), json.dumps({"id": call.requestId}))] + complete = _completion(_execute( + orchestrator, old, [helpers.new_event_raised_event(call.requestId, json.dumps({"result": "42"}))])) + assert json.loads(complete.result.value) == 42 + + +def test_continue_as_new_from_coroutine() -> None: + async def orchestrator(ctx: task.OrchestrationContext, value: int) -> None: + await ctx.wait_for_external_event("next") + ctx.continue_as_new(value + 1) + + complete = _completion(_execute( + orchestrator, _start(), [helpers.new_event_raised_event("next", None)])) + assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_CONTINUED_AS_NEW + assert json.loads(complete.result.value) == 4 + + +def test_nondeterminism_still_detected_for_coroutines() -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any: + return await ctx.call_activity("changed") + + complete = _completion(_execute( + orchestrator, _start(), [helpers.new_task_scheduled_event(1, "original")])) + assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + assert complete.failureDetails.errorType.endswith("NonDeterminismError") + + +@pytest.mark.parametrize("after_activity", [False, True]) +def test_non_durable_suspension_is_rejected(after_activity: bool) -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + if after_activity: + await ctx.call_activity("work") + await asyncio.Future() + + old = _start() + [helpers.new_task_scheduled_event(1, "work")] if after_activity else [] + new = [helpers.new_task_completed_event(1, "null")] if after_activity else _start() + complete = _completion(_execute(orchestrator, old, new)) + assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + assert complete.failureDetails.errorType == "builtins.RuntimeError" + assert "blocked without durable work" in complete.failureDetails.errorMessage + + +def test_async_generator_is_rejected() -> None: + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> AsyncGenerator[task.Task[Any], None]: + yield ctx.call_activity("work", return_type=int) + + complete = _completion(_execute(orchestrator, [], _start())) + assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + assert "Async generator orchestrators are not supported" in complete.failureDetails.errorMessage diff --git a/tests/durabletask/test_orchestration_generator_cleanup.py b/tests/durabletask/test_orchestration_generator_cleanup.py index 80adfd3e..4b348d38 100644 --- a/tests/durabletask/test_orchestration_generator_cleanup.py +++ b/tests/durabletask/test_orchestration_generator_cleanup.py @@ -347,7 +347,7 @@ def test_teardown_cancellation_preserves_response( assert len(caplog.records) == 1 record = caplog.records[0] assert record.levelno == logging.WARNING - assert record.message == f"{TEST_INSTANCE_ID}: Failed to close orchestrator orchestrator generator." + assert record.message == f"{TEST_INSTANCE_ID}: Failed to close orchestrator orchestrator." assert record.exc_info is not None assert record.exc_info[0] is asyncio.CancelledError diff --git a/tests/durabletask/typing/async_context.py b/tests/durabletask/typing/async_context.py new file mode 100644 index 00000000..42ee5767 --- /dev/null +++ b/tests/durabletask/typing/async_context.py @@ -0,0 +1,45 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""The async surface returns native coroutines with inferred result types.""" + +import asyncio +from collections.abc import Coroutine +from datetime import timedelta +from typing import Any, assert_type + +from durabletask import task +from durabletask.entities import EntityInstanceId, EntityLock + + +def increment(ctx: task.ActivityContext, value: int) -> int: + return value + 1 + + +async def child(ctx: task.OrchestrationContext, value: int) -> str: + return str(value) + + +async def check_cold_operations(ctx: task.OrchestrationContext, entity: EntityInstanceId) -> None: + ops = ctx.as_async() + assert_type(ops, task.AsyncOrchestrationContext) + activity: Coroutine[Any, Any, int] = ops.call_activity(increment, input=1) + assert_type(await activity, int) + assert_type(await ops.call_activity("named", return_type=str), str) + assert_type(await ops.call_activity("named"), Any) + assert_type(await ops.call_sub_orchestrator(child, input=1), str) + assert_type(await ops.call_sub_orchestrator("child", return_type=int), int) + assert_type(await ops.call_sub_orchestrator("child"), Any) + assert_type(await ops.call_entity(entity, "get", return_type=int), int) + assert_type(await ops.call_entity(entity, "get"), Any) + assert_type(await ops.create_timer(timedelta(seconds=1)), None) + assert_type(await ops.receive_event("event", data_type=str), str) + assert_type(await ops.receive_event("event"), Any) + assert_type(await ops.lock_entities([entity]), EntityLock) + assert_type(asyncio.create_task(ops.call_activity(increment, input=1)), asyncio.Task[int]) + async with asyncio.TaskGroup() as group: + assert_type(group.create_task(ops.receive_event("event", data_type=str)), asyncio.Task[str]) + operation = ctx.call_activity(increment, input=1) + assert_type(operation, task.CompletableTask[int]) + assert_type(await operation.wait(), int) + assert_type(asyncio.create_task(operation.wait()), asyncio.Task[int])