diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index f1194285..8a06c5c7 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -67,7 +67,7 @@ jobs: # Initializes the CodeQL tools for scanning. - name: Initialize CodeQL - uses: github/codeql-action/init@2892aa5e19bbd11bc0cff5427e3b750a04d9e3c2 # v4.38.2 + uses: github/codeql-action/init@6f2ac81f5cd81fe351a197ba0b00d91d03cd2aee # v4.39.0 with: languages: ${{ matrix.language }} build-mode: ${{ matrix.build-mode }} @@ -96,6 +96,6 @@ jobs: exit 1 - name: Perform CodeQL Analysis - uses: github/codeql-action/analyze@2892aa5e19bbd11bc0cff5427e3b750a04d9e3c2 # v4.38.2 + uses: github/codeql-action/analyze@6f2ac81f5cd81fe351a197ba0b00d91d03cd2aee # v4.39.0 with: category: "/language:${{matrix.language}}" diff --git a/.github/workflows/durabletask-azurefunctions.yml b/.github/workflows/durabletask-azurefunctions.yml index 4b04765a..9552fb0d 100644 --- a/.github/workflows/durabletask-azurefunctions.yml +++ b/.github/workflows/durabletask-azurefunctions.yml @@ -57,7 +57,7 @@ jobs: needs: run-tests runs-on: ubuntu-latest # The end-to-end suite launches a real Azure Functions host, which requires - # ``azure-functions>=2.3.0b2``. That build is published to PyPI and pulled + # ``azure-functions>=2.4.0b1``. That build is published to PyPI and pulled # in as a declared dependency of ``azure-functions-durable``. # Run the suite in both payload-typing modes: the default (loose) codec and # strict typing (``AZURE_FUNCTIONS_DURABLE_STRICT_TYPING=1``), which requires @@ -78,7 +78,7 @@ jobs: python-version: "3.13" - name: Set up Node.js (needed for Azurite and Functions Core Tools) - uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7.0.0 + uses: actions/setup-node@949feb2413d6458794dcd2491c4babbbce0c15c1 # v7.1.0 with: node-version: "22.x" @@ -106,7 +106,7 @@ jobs: # pytest summary does not surface, so upload it for diagnosis. - name: Upload Functions host logs if: failure() - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7.0.2 with: name: func-host-logs-strict-${{ matrix.strict-typing }} path: tests/azure-functions-durable/e2e/apps/*/_func_host.log diff --git a/.github/workflows/durabletask.yml b/.github/workflows/durabletask.yml index 4cef9ed4..2720e25a 100644 --- a/.github/workflows/durabletask.yml +++ b/.github/workflows/durabletask.yml @@ -14,6 +14,29 @@ permissions: contents: read jobs: + async-runtime-windows: + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "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@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 + with: + python-version: ${{ matrix.python-version }} + - name: Install runtime test dependencies + run: python -m pip install -r requirements.txt -e . + - name: Check asyncio, AnyIO, replay, and dispatch cleanup + run: >- + python -m pytest + tests\durabletask\test_asyncio_orchestrators.py + tests\durabletask\test_coroutine_orchestrators.py + tests\durabletask\test_pydantic_ai_orchestrator.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: @@ -46,7 +69,7 @@ jobs: python-version: ${{ matrix.python-version }} - name: Set up Node.js (needed for Azurite) - uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7.0.0 + uses: actions/setup-node@949feb2413d6458794dcd2491c4babbbce0c15c1 # v7.1.0 with: node-version: '22.x' diff --git a/CHANGELOG.md b/CHANGELOG.md index 16194e95..657bd191 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,42 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## Unreleased +ADDED + +- Added `OrchestrationQuery.task_hub_names` to filter orchestration instances +by one or more task hubs. +- 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. + +FIXED + +- Implemented `__hash__` and `__repr__` for `EntityInstanceId`, making identity attributes immutable to guarantee hash stability when used in sets or dict keys. +- Large-payload externalization now checks every eligible payload's maximum +size before uploading or replacing any protobuf values, including carryover +events and entity payloads, for both synchronous and asynchronous calls. +Oversized messages remain unchanged and cause no uploads; their size error +now takes precedence over hypothetical earlier storage failures. +- Orchestrations that exceed the configured external payload storage limit now +fail without retrying and release their entity locks, allowing other callers to +progress. Transient payload-store failures remain retryable. +- Previously unhandled non-cancellation `BaseException` exits, including result +serialization and generator teardown failures, now fail with a clear reason instead of leaving an accepted dispatch +unfinished. Cancellation and process-control signals retain their semantics. +The gRPC worker abandons work that exits before an outcome is sent and resets +the originating work-item stream after an uncertain completion or abandonment +RPC, allowing scheduler ownership recovery without retrying an uncertain outcome. +- Orchestrator generators are now closed at the end of each dispatch, including +while waiting for durable work, releasing resources held by async bridges. +`finally` blocks and context-manager exits can run before orchestration completion +and repeat across dispatches; see [cleanup semantics](docs/supported-patterns.md#cleanup-in-orchestrator-finally-blocks). +- Async cancellation raised by generator teardown no longer replaces dispatch +responses or version-rejection abandonment. + ## v1.12.0 ADDED diff --git a/README.md b/README.md index 2bd8e7fb..5d6f6a7b 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,46 @@ and [Azure Durable Functions](https://learn.microsoft.com/azure/azure-functions/ ## Optional Features +### Async Orchestrators + +All three packages support `async def` orchestrators and nested async helpers. +Existing generator orchestrators continue to work. + +```python +from durabletask import task + + +async def sequence(ctx: task.OrchestrationContext, value: int) -> list[int]: + first = await ctx.call_activity("plus_one", input=value, return_type=int) + pending: list[task.Task[int]] = [ + ctx.call_activity("plus_one", input=first, return_type=int), + ctx.call_activity("plus_one", input=value, return_type=int), + ] + return await task.when_all(pending) +``` + +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(...)`. + +> [!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. + +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, framework adapters, and coordinated release rules. + ### Large Payload Externalization Install the `azure-blob-payloads` extra to automatically offload oversized orchestration payloads to diff --git a/azure-functions-durable/CHANGELOG.md b/azure-functions-durable/CHANGELOG.md index b17a3654..45871657 100644 --- a/azure-functions-durable/CHANGELOG.md +++ b/azure-functions-durable/CHANGELOG.md @@ -7,6 +7,40 @@ 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; +asyncio scheduling requires replay/backend tests. +Requires a coordinated release after the core SDK is published; dependency +minimums are intentionally unchanged in this feature PR. +- Added `DFApp.configure_serialization(ensure_ascii=False)` to opt in to compact +Unicode JSON across durable clients, orchestration and entity workers, activity +bindings, and registered blueprints. The default remains ASCII-escaped JSON. +The opt-in uses the shared codec's `ensure_ascii` option available in +`azure-functions` v2.4.0b1; unsupported SDKs fail explicitly at configuration time. + +CHANGED + +- Updated the minimum `azure-functions` dependency to v2.4.0b1. + +FIXED + +- With the corresponding core `durabletask` SDK fix, oversized messages fail +large-payload validation before any uploads or payload replacements. Size +errors now take precedence over hypothetical earlier storage failures. +- With the corresponding core `durabletask` SDK fix, previously unhandled non-cancellation +`BaseException` exits from orchestration execution, result serialization, or teardown now produce a failed response +with the exception's type, message, and stack trace instead of losing the response. +- With the corresponding core `durabletask` SDK fix, orchestrator generators +are closed at the end of each dispatch, including while waiting for durable work. +See [cleanup semantics](../docs/supported-patterns.md#cleanup-in-orchestrator-finally-blocks) +for repeated finalizers, discarded teardown actions/status, and yielding restrictions. +- With the corresponding core `durabletask` SDK fix, async cancellation raised by +generator teardown no longer replaces dispatch responses or version-rejection abandonment. + ## v2.0.0rc3 ADDED diff --git a/azure-functions-durable/MIGRATION_GUIDE.md b/azure-functions-durable/MIGRATION_GUIDE.md index 5c0762c5..20d992da 100644 --- a/azure-functions-durable/MIGRATION_GUIDE.md +++ b/azure-functions-durable/MIGRATION_GUIDE.md @@ -66,7 +66,7 @@ later. Replace the 1.x package constraint in `requirements.txt`: ```text -azure-functions>=2.3.0b2 +azure-functions>=2.4.0b1 azure-functions-durable>=2.0.0b1,<3 ``` diff --git a/azure-functions-durable/README.md b/azure-functions-durable/README.md index e441e75d..125516b9 100644 --- a/azure-functions-durable/README.md +++ b/azure-functions-durable/README.md @@ -32,6 +32,70 @@ Key capabilities include durable orchestrations and sub-orchestrations, durable timers, external events, durable entities, retries, versioning, durable HTTP calls (`context.call_http(...)`), recurring scheduled tasks, and history export. +## 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. + +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 +ASCII-escaped JSON representation. With `azure-functions>=2.4.0b1`, opt in to +literal Unicode JSON on the root `DFApp` at startup: + +```python +import azure.durable_functions as df + +app = df.DFApp() +app.configure_serialization(ensure_ascii=False) +``` + +The shared `df_dumps(..., ensure_ascii=...)` API was released in +[`azure-functions` 2.4.0b1](https://github.com/Azure/azure-functions-python-library/releases/tag/2.4.0b1), +which is the minimum required shared SDK version. Dependency installation +selects a compatible release automatically. If an older SDK is manually +installed despite that requirement, configuring the opt-in raises `RuntimeError` +rather than silently ignoring it. Leaving the option unset, or explicitly +choosing `ensure_ascii=True`, preserves the existing JSON encoding. + +There is one encoding policy per Python worker process. It applies to both +synchronous and asynchronous durable clients, orchestration and entity workers, +activity binding outputs, and all registered blueprints, including handlers +created before configuration. Configure every scaled-out worker consistently +before invocations. Repeating the same setting is allowed; a conflicting setting +or changing the policy after the first durable JSON serialization raises +`ValueError`. + +Literal Unicode reduces escaping overhead for non-Latin text. For example, +a JSON object containing 100,000 Cyrillic characters occupies approximately +200,012 UTF-8 bytes rather than 600,012 ASCII-escaped bytes. Actual reduction +depends on the payload. The shared codec preserves surrogate-code-unit and +ASCII DEL escapes, custom-object envelopes, and strict/loose typing semantics. +This option does not normalize Unicode or change the decoder; existing escaped +payloads remain readable. + +> [!WARNING] +> Enabling or disabling this option changes serialized bytes for non-ASCII +> text. Review applications deriving hashes, signatures, IDs, or replay +> decisions from serialized JSON before rolling out or rolling back a change. +> Do not change the setting inside orchestrator code or during invocations. + +The policy applies before large-payload externalization, so uploads use the +selected JSON encoding and size thresholds use its serialized UTF-8 bytes. +Reduced escaping does not eliminate aggregate request limits or replace +externalization. Protocol envelopes, reference tokens, management responses, and +JSON serialized explicitly by application code are not reformatted by this +setting. + ## Large payloads Configure a `durabletask.payload.PayloadStore` once at app startup to store large diff --git a/azure-functions-durable/azure/durable_functions/decorators/durable_app.py b/azure-functions-durable/azure/durable_functions/decorators/durable_app.py index 70ce34c5..2e82c4f3 100644 --- a/azure-functions-durable/azure/durable_functions/decorators/durable_app.py +++ b/azure-functions-durable/azure/durable_functions/decorators/durable_app.py @@ -493,6 +493,23 @@ class DFApp(Blueprint, FunctionRegister): Exports the decorators required to declare and index DF Function-types. """ + def configure_serialization(self, *, ensure_ascii: bool = True) -> None: + """Configure JSON encoding for this app and its registered blueprints. + + The default preserves ASCII-escaped JSON. Set ``ensure_ascii=False`` + to emit ordinary Unicode literally with ``azure-functions>=2.4.0b1``. + An older SDK raises RuntimeError when this option is requested. + + Call at startup in every worker process, before durable payloads are + serialized. Clients, orchestrations, entities, and activity bindings + share one policy per process, including handlers created earlier. + Repeating the same setting is allowed. A conflicting setting, or + changing the policy after serialization starts, raises ValueError. + """ + from ..internal.serialization import configure_serialization + + configure_serialization(ensure_ascii=ensure_ascii) + def configure_large_payloads(self, *, payload_store: PayloadStore) -> None: """Enable payload externalization for this app and its blueprints. diff --git a/azure-functions-durable/azure/durable_functions/internal/converters.py b/azure-functions-durable/azure/durable_functions/internal/converters.py index 34eadd38..faf2319f 100644 --- a/azure-functions-durable/azure/durable_functions/internal/converters.py +++ b/azure-functions-durable/azure/durable_functions/internal/converters.py @@ -32,7 +32,6 @@ from azure.functions._durable_functions import ( EntityContext, OrchestrationContext, - df_dumps, df_loads, ) @@ -43,6 +42,7 @@ ORCHESTRATION_TRIGGER, ) from .payloads import ActivityPayload, get_transport_payload_store +from .serialization import serialize_json _TriggerMetadata = Optional[Mapping[str, meta.Datum]] @@ -162,7 +162,7 @@ def encode(cls, obj: Any, *, if isinstance(obj, ActivityPayload): return meta.Datum(type='json', value=obj.value) try: - result = df_dumps(obj) + result = serialize_json(obj) except TypeError as e: raise ValueError( f'activity trigger output must be json serializable ({obj})') from e diff --git a/azure-functions-durable/azure/durable_functions/internal/serialization.py b/azure-functions-durable/azure/durable_functions/internal/serialization.py index b8f296c0..3f2257a8 100644 --- a/azure-functions-durable/azure/durable_functions/internal/serialization.py +++ b/azure-functions-durable/azure/durable_functions/internal/serialization.py @@ -10,13 +10,72 @@ """ from __future__ import annotations +import inspect +from threading import Lock from azure.functions._durable_functions import df_dumps, df_loads -from typing import Any, override +from typing import Any, Callable, Protocol, TypeGuard, override from durabletask.serialization import JsonDataConverter +class _ConfigurableJsonEncoder(Protocol): + def __call__(self, value: Any, *, ensure_ascii: bool = True) -> str: + ... + + +_serialization_lock = Lock() +_ensure_ascii: bool | None = None +_unicode_encoder: _ConfigurableJsonEncoder | None = None + + +def _supports_ensure_ascii( + encoder: Callable[[Any], str]) -> TypeGuard[_ConfigurableJsonEncoder]: + parameter = inspect.signature(encoder).parameters.get("ensure_ascii") + return parameter is not None and parameter.kind in ( + inspect.Parameter.KEYWORD_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD, + ) + + +def configure_serialization(*, ensure_ascii: object = True) -> None: + """Select one JSON encoding policy before durable payload serialization.""" + global _ensure_ascii, _unicode_encoder + if not isinstance(ensure_ascii, bool): + raise TypeError("ensure_ascii must be a bool") + with _serialization_lock: + if _ensure_ascii is not None and _ensure_ascii != ensure_ascii: + raise ValueError( + "JSON serialization is already configured with " + f"ensure_ascii={_ensure_ascii} in this process; configure " + "serialization once before serializing durable payloads") + if not ensure_ascii: + encoder = df_dumps + if not _supports_ensure_ascii(encoder): + raise RuntimeError( + "ensure_ascii=False requires azure-functions>=2.4.0b1, whose " + "df_dumps supports the ensure_ascii keyword argument. " + "Install a compatible azure-functions SDK before opting in.") + _unicode_encoder = encoder + _ensure_ascii = ensure_ascii + + +def serialize_json(value: Any) -> str: + """Apply the process-wide policy using the shared Functions JSON codec.""" + global _ensure_ascii + with _serialization_lock: + if _ensure_ascii is None: + _ensure_ascii = True + ensure_ascii = _ensure_ascii + encoder = _unicode_encoder + if ensure_ascii: + # Older shared SDKs do not accept the optional keyword argument. + return df_dumps(value) + if encoder is None: + raise RuntimeError("Unicode JSON serialization has no configured encoder") + return encoder(value, ensure_ascii=False) + + class FunctionsDataConverter(JsonDataConverter): """:class:`DataConverter` that serializes via azure-functions' codec. @@ -38,7 +97,7 @@ class FunctionsDataConverter(JsonDataConverter): def serialize(self, value: Any) -> str | None: if value is None: return None - return df_dumps(value) + return serialize_json(value) @override def deserialize(self, data: str | None, target_type: type | None = None) -> Any: 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..61408762 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,17 @@ # 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 +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 +21,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 only. 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. @@ -34,13 +39,12 @@ def orchestrator_generator_wrapper( orchestrator's final return value. """ try: - previous = next(generator) + previous = generator.send(None) 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 diff --git a/azure-functions-durable/pyproject.toml b/azure-functions-durable/pyproject.toml index 6e19bf95..5cdecffc 100644 --- a/azure-functions-durable/pyproject.toml +++ b/azure-functions-durable/pyproject.toml @@ -32,7 +32,7 @@ readme = "README.md" dependencies = [ "durabletask[opentelemetry]>=1.12.0", "azure-identity>=1.19.0", - "azure-functions>=2.3.0b2" + "azure-functions>=2.4.0b1" ] [project.urls] diff --git a/azure-functions-durable/samples/durable-entities/requirements.txt b/azure-functions-durable/samples/durable-entities/requirements.txt index 7ef4b6a5..1abc822d 100644 --- a/azure-functions-durable/samples/durable-entities/requirements.txt +++ b/azure-functions-durable/samples/durable-entities/requirements.txt @@ -1,2 +1,2 @@ -azure-functions>=2.3.0b2 +azure-functions>=2.4.0b1 azure-functions-durable>=2.0.0b1,<3 diff --git a/azure-functions-durable/samples/fan-out-fan-in/requirements.txt b/azure-functions-durable/samples/fan-out-fan-in/requirements.txt index 7ef4b6a5..1abc822d 100644 --- a/azure-functions-durable/samples/fan-out-fan-in/requirements.txt +++ b/azure-functions-durable/samples/fan-out-fan-in/requirements.txt @@ -1,2 +1,2 @@ -azure-functions>=2.3.0b2 +azure-functions>=2.4.0b1 azure-functions-durable>=2.0.0b1,<3 diff --git a/azure-functions-durable/samples/function-chaining/requirements.txt b/azure-functions-durable/samples/function-chaining/requirements.txt index 7ef4b6a5..1abc822d 100644 --- a/azure-functions-durable/samples/function-chaining/requirements.txt +++ b/azure-functions-durable/samples/function-chaining/requirements.txt @@ -1,2 +1,2 @@ -azure-functions>=2.3.0b2 +azure-functions>=2.4.0b1 azure-functions-durable>=2.0.0b1,<3 diff --git a/azure-functions-durable/samples/human-interaction/requirements.txt b/azure-functions-durable/samples/human-interaction/requirements.txt index 7ef4b6a5..1abc822d 100644 --- a/azure-functions-durable/samples/human-interaction/requirements.txt +++ b/azure-functions-durable/samples/human-interaction/requirements.txt @@ -1,2 +1,2 @@ -azure-functions>=2.3.0b2 +azure-functions>=2.4.0b1 azure-functions-durable>=2.0.0b1,<3 diff --git a/docs/async-orchestrators.md b/docs/async-orchestrators.md new file mode 100644 index 00000000..db692fd8 --- /dev/null +++ b/docs/async-orchestrators.md @@ -0,0 +1,174 @@ +# 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 durabletask import task + + +async def workflow(ctx: task.OrchestrationContext, values: list[int]) -> list[int]: + async def branch(value: int) -> int: + first = await ctx.call_activity("increment", input=value, return_type=int) + await asyncio.sleep(1) + return await ctx.call_activity("increment", input=first, return_type=int) + + return await asyncio.gather(*(branch(value) for value in values)) +``` + +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. + +## 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`, and `asyncio.sleep`. On Python 3.11+, +it supports `asyncio.TaskGroup` and `asyncio.timeout` / `timeout_at`. +`asyncio.wait_for`, AnyIO's asyncio task groups, cancel scopes, shielded cleanup, +and `fail_after` use the same task/future scheduling and durable timers. +AnyIO support is bounded to these paths, not all backends or APIs. + +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. + +Task groups cancel and drain local sibling tasks on failure and use their +library's normal exception grouping rules. AnyIO uses level cancellation: +awaited cleanup in a cancelled scope may require +`with anyio.CancelScope(shield=True)`. Python-version-specific exception +grouping and cancel-count behavior remain those of asyncio/AnyIO. + +Cancelling a local task or timing out its wait detaches that waiter; it does +**not** cancel the durable activity, sub-orchestration, entity operation, or +external-event subscription it was waiting on. Another waiter can still await +the same durable task. Keep the durable task if you intend to await it again. +Use explicit SDK cancellation for cancellable durable timers/event waits. +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. 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. + +Pydantic AI's real `Agent.run` / parallel-tool AnyIO path is exercised with a +deterministic `TestModel` in +[`test_pydantic_ai_orchestrator.py`](../tests/durabletask/test_pydantic_ai_orchestrator.py). +The orchestrator directly awaits the agent; no adapter event loop, trampoline, +or per-batch durable join is needed. The proof adapter still registers ordinary +model/tool activities, encodes/decodes framework payloads, supplies replay-stable +IDs and a deadline clock, and explicitly canonicalizes volatile message +timestamps. These are **adapter responsibilities**, not SDK-wide interception. +This proof is not a complete production Pydantic adapter or compatibility +promise for arbitrary agents, model providers, tools, or framework versions. +No paid model calls are needed. + +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. + +The Functions `orchestrator_generator_wrapper` helper supports direct durable +awaits for 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. + +## Coordinated delivery + +This feature is one draft PR/release unit across all packages. Do not ship a +limited intermediate contract. Release the core package first; then use a +dedicated coordinated release PR to update provider dependency minimums, +versions, and release notes together. Feature development does not change +package versions or dependency minimums. diff --git a/docs/supported-patterns.md b/docs/supported-patterns.md index 660044c4..d8d0f3e2 100644 --- a/docs/supported-patterns.md +++ b/docs/supported-patterns.md @@ -2,6 +2,10 @@ 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/AnyIO composition, durable clocks/timeouts, and dispatch-local cleanup. + ## Function chaining An orchestration can chain a sequence of function calls using the following syntax: @@ -22,6 +26,78 @@ def sequence(ctx: task.OrchestrationContext, _): See the full [function chaining example](../examples/activity_sequence.py). +## Cleanup in orchestrator finally blocks + +The SDK closes an orchestrator generator at the end of every dispatch, including +when it is waiting for durable work. Closing injects `GeneratorExit`, running +active `finally` blocks and context-manager exits. Replay creates a fresh +generator, so this cleanup may run many times before the orchestration completes. + +Use these hooks to release dispatch-local resources, not as orchestration-completion +or exactly-once compensation hooks. Do not release resources that an activity +still needs: that activity may still be running when the dispatch ends. + +The response is built before this forced teardown. Durable actions, timer +cancellation, custom-status changes, and `continue_as_new()` issued only during +teardown do not change that response. Keep workflow-level cleanup in normal +orchestration control flow, as a yielded activity: + +```python +from io import StringIO +from typing import Any, Generator +from durabletask import task + + +def orchestrator( + ctx: task.OrchestrationContext, value: str, +) -> Generator[task.Task[Any], Any, Any]: + scratch = StringIO(value) + try: + result = yield ctx.call_activity("process", input=scratch.getvalue()) + finally: + scratch.close() # Local buffer cleanup runs at each dispatch boundary. + + yield ctx.call_activity("cleanup_workflow", input=value) + return result +``` + +`ctx.is_replaying` is normally false after new events are processed. Neither +`not ctx.is_replaying` nor replay-safe logging makes teardown completion-only. +Yielding during forced close raises `RuntimeError: generator ignored GeneratorExit`; +the SDK logs a warning, but the generator remains open and may retain resources. +A `finally` reached through normal execution, rather than forced close, can +still yield durable work. + +Garbage collection could already trigger finalization; completion-only timing +was never guaranteed. Explicit dispatch teardown does not guarantee cleanup +across process failure. + +Unexpected non-cancellation `BaseException` exits from orchestration execution, +result serialization, or forced generator teardown fail the orchestration with the exception's type, +message, and stack trace. An existing orchestration failure is retained if +teardown also fails. Outstanding acquired entity locks are released with the +failed response; durable work and status changes made only by teardown are not +committed. + +The gRPC worker does not turn `asyncio.CancelledError`, `GeneratorExit`, +`SystemExit`, or `KeyboardInterrupt` into orchestration failures. This also applies +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 +work-item stream so the scheduler can recover ownership. It does not retry an +uncertain RPC or send a conflicting completion or abandonment. Other work owned +by that stream can also be redelivered. This is not an exactly-once execution +guarantee or a bounded recovery-time guarantee. Cancelling an asyncio await of a +still-running executor thread does not abandon that thread's dispatch. + ## Fan-out/fan-in An orchestration can fan-out a dynamic number of function calls in parallel and then fan-in the diff --git a/durabletask-azuremanaged/CHANGELOG.md b/durabletask-azuremanaged/CHANGELOG.md index 338b183f..caabcd59 100644 --- a/durabletask-azuremanaged/CHANGELOG.md +++ b/durabletask-azuremanaged/CHANGELOG.md @@ -7,6 +7,35 @@ 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). +Requires a coordinated release after the core SDK is published; dependency +minimums are intentionally unchanged in this feature PR. + +FIXED + +- With the corresponding core `durabletask` SDK fix, oversized messages fail +large-payload validation before any uploads or payload replacements. Size +errors now take precedence over hypothetical earlier storage failures. +- With the corresponding core `durabletask` SDK fix, orchestrations that exceed +the configured external payload storage limit fail without retrying and release +their entity locks. Transient payload-store failures remain retryable. +- With the corresponding core `durabletask` SDK fix, previously unhandled +non-cancellation `BaseException` exits, including result serialization failures, +now fail with a clear reason. Cancellation and process-control +exits recover dispatch ownership, and uncertain completion or abandonment RPCs +reset the originating work-item stream rather than retaining unfinished work. +- With the corresponding core `durabletask` SDK fix, orchestrator generators +are closed at the end of each dispatch, including while waiting for durable work. +See [cleanup semantics](../docs/supported-patterns.md#cleanup-in-orchestrator-finally-blocks) +for repeated finalizers, discarded teardown actions/status, and yielding restrictions. +- With the corresponding core `durabletask` SDK fix, async cancellation raised by +generator teardown no longer replaces dispatch responses or version-rejection abandonment. + ## v1.12.0 ADDED diff --git a/durabletask/client.py b/durabletask/client.py index 4bec371a..b4dae185 100644 --- a/durabletask/client.py +++ b/durabletask/client.py @@ -197,6 +197,7 @@ class OrchestrationQuery: max_instance_count: int | None = (1 << 31) - 1 fetch_inputs_and_outputs: bool = False instance_id_prefix: str | None = None + task_hub_names: list[str] | None = None @dataclass diff --git a/durabletask/entities/entity_instance_id.py b/durabletask/entities/entity_instance_id.py index 9c984877..e195736c 100644 --- a/durabletask/entities/entity_instance_id.py +++ b/durabletask/entities/entity_instance_id.py @@ -1,12 +1,42 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. + class EntityInstanceId: + __slots__ = ("_entity", "_key") + def __init__(self, entity: str, key: str): EntityInstanceId.validate_entity_name(entity) EntityInstanceId.validate_key(key) - self.entity = entity.lower() - self.key = key + object.__setattr__(self, "_entity", entity.lower()) + object.__setattr__(self, "_key", key) + + @property + def entity(self) -> str: + return self._entity + + @property + def key(self) -> str: + return self._key + + def __setattr__(self, name: str, value: object) -> None: + raise AttributeError( + f"'{type(self).__name__}' object does not support attribute assignment" + ) + + def __delattr__(self, name: str) -> None: + raise AttributeError( + f"'{type(self).__name__}' object does not support attribute deletion" + ) + + def __copy__(self) -> "EntityInstanceId": + return self + + def __deepcopy__(self, memo: object) -> "EntityInstanceId": + return self + + def __reduce__(self) -> tuple[type["EntityInstanceId"], tuple[str, str]]: + return (type(self), (self._entity, self._key)) def __str__(self) -> str: return f"@{self.entity}@{self.key}" @@ -21,6 +51,12 @@ def __lt__(self, other: object) -> bool: return NotImplemented return str(self) < str(other) + def __hash__(self) -> int: + return hash((self.entity, self.key)) + + def __repr__(self) -> str: + return f"EntityInstanceId(entity={self.entity!r}, key={self.key!r})" + @staticmethod def parse(entity_id: str) -> "EntityInstanceId": """Parse a string representation of an entity ID into an EntityInstanceId object. diff --git a/durabletask/internal/asyncio_loop.py b/durabletask/internal/asyncio_loop.py new file mode 100644 index 00000000..ca9a0cf8 --- /dev/null +++ b/durabletask/internal/asyncio_loop.py @@ -0,0 +1,332 @@ +# 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, TYPE_CHECKING, cast + +from durabletask import task + +if TYPE_CHECKING: + from durabletask.worker import _RuntimeOrchestrationContext # pyright: ignore[reportPrivateUsage] + +T = TypeVar("T") + + +class _DurableAsyncTask(asyncio.Task[T]): + 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 __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: _RuntimeOrchestrationContext) -> None: + self.ctx = ctx + self._ready: deque[asyncio.Handle] = deque() + self._cancel_ready: list[tuple[int, int, asyncio.Handle]] = [] + self._tasks: dict[asyncio.Task[Any], int] = {} + self.cancelling_task: asyncio.Task[Any] | None = None + self._waits: dict[int, tuple[task.Task[Any], list[asyncio.Future[Any]]]] = {} + self._timers: list[_DurableTimerHandle] = [] + self._asyncgens: dict[int, Any] = {} + self._closed = False + self._closing = False + self._fatal: 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.") + if sys.version_info >= (3, 11): + future = _DurableAsyncTask(coro, loop=self, name=name, context=context) + else: + if context is not None: + coro.close() + raise TypeError("Task contexts require Python 3.11 or later.") + future = _DurableAsyncTask(coro, loop=self, name=name) + self._tasks[future] = len(self._tasks) + return future + + 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 and AnyIO cancel sets 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.append(handle) + return handle + + def _timer_handle_cancelled(self, handle: asyncio.TimerHandle) -> None: + # No selector heap to maintain; _DurableTimerHandle owns cancellation. + pass + + 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) + 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) + + def _deliver_completions(self) -> None: + for durable, futures in tuple(self._waits.values()): + if durable.is_complete: + for future in tuple(futures): + if not future.done(): + try: + future.set_result(durable.get_result()) + except Exception as error: + future.set_exception(error) + for handle in tuple(self._timers): + if handle.cancelled(): + self._timers.remove(handle) + elif handle.durable_timer is not None and handle.durable_timer.is_complete: + self._timers.remove(handle) + self._ready.append(handle) + + 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.ctx._is_complete: # pyright: ignore[reportPrivateUsage] + 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 any( + not handle.cancelled() for handle in 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 + while self._ready or self._cancel_ready: + self._ready.extend(handle for _, _, handle in sorted(self._cancel_ready)) + self._cancel_ready.clear() + handle = self._ready.popleft() + if not handle.cancelled(): + handle._run() + if self._fatal is not None: + error, self._fatal = self._fatal, None + raise error + if not self._closing and self.ctx._is_complete: # pyright: ignore[reportPrivateUsage] + 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 >= 10000: + raise RuntimeError("Async cleanup did not quiesce; cleanup must cooperate with cancellation.") + 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") + self._fatal = error if isinstance(error, BaseException) else RuntimeError(context["message"]) + + def close(self) -> None: + self._closing = True + self.ctx._async_cleanup = True # pyright: ignore[reportPrivateUsage] + 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 = None + for future in tuple(self._tasks): + if not future.done(): + try: + future.get_coro().close() + except BaseException as error: + 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): + if not isinstance(error, Exception): + failure = failure or error + logger = self.ctx._logger or logging.getLogger(__name__) # pyright: ignore[reportPrivateUsage] + logger.warning( + "%s: Async %s task failed.", self.ctx.instance_id, + "background" if future in completed_at_close else "cleanup", exc_info=error) + if failure is not None: + raise failure + finally: + self._closed = True + self._ready.clear() + self._cancel_ready.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/client_helpers.py b/durabletask/internal/client_helpers.py index cd277429..06913c18 100644 --- a/durabletask/internal/client_helpers.py +++ b/durabletask/internal/client_helpers.py @@ -103,6 +103,7 @@ def build_query_instances_req( runtimeStatus=[status.value for status in orchestration_query.runtime_status] if orchestration_query.runtime_status else None, createdTimeFrom=helpers.new_timestamp(orchestration_query.created_time_from) if orchestration_query.created_time_from else None, createdTimeTo=helpers.new_timestamp(orchestration_query.created_time_to) if orchestration_query.created_time_to else None, + taskHubNames=[wrappers_pb2.StringValue(value=name) for name in orchestration_query.task_hub_names] if orchestration_query.task_hub_names else None, instanceIdPrefix=helpers.get_string_value(orchestration_query.instance_id_prefix), maxInstanceCount=orchestration_query.max_instance_count, fetchInputsAndOutputs=orchestration_query.fetch_inputs_and_outputs, diff --git a/durabletask/internal/exceptions.py b/durabletask/internal/exceptions.py index 276cbe60..a04a304f 100644 --- a/durabletask/internal/exceptions.py +++ b/durabletask/internal/exceptions.py @@ -12,3 +12,7 @@ def __init__(self, error_details: pb.TaskFailureDetails) -> None: class AbandonOrchestrationError(Exception): pass + + +class PayloadSizeExceededError(ValueError): + """A raw payload exceeds the configured external storage limit.""" diff --git a/durabletask/internal/helpers.py b/durabletask/internal/helpers.py index 0e170371..6e0b1a0a 100644 --- a/durabletask/internal/helpers.py +++ b/durabletask/internal/helpers.py @@ -215,7 +215,7 @@ def failure_details_from_protobuf(details: pb.TaskFailureDetails) -> "FailureDet def new_failure_details( - ex: Exception, + ex: BaseException, exception_properties_provider: ExceptionPropertiesProvider | None = None, logger: logging.Logger | None = None, _visited: set[int] | None = None) -> pb.TaskFailureDetails: @@ -226,7 +226,7 @@ def new_failure_details( if len(_visited) > 10 or (inner and id(inner) in _visited) or not isinstance(inner, Exception): inner = None properties: dict[str, struct_pb2.Value] | None = None - if exception_properties_provider is not None: + if exception_properties_provider is not None and isinstance(ex, Exception): try: provider_properties = exception_properties_provider.get_exception_properties(ex) except Exception: @@ -257,7 +257,9 @@ def new_failure_details( failure_details = pb.TaskFailureDetails( errorType=get_qualified_name(type(ex)), errorMessage=str(ex), - stackTrace=wrappers_pb2.StringValue(value=''.join(traceback.format_tb(ex.__traceback__))), + stackTrace=wrappers_pb2.StringValue(value=''.join( + traceback.format_tb(ex.__traceback__) if isinstance(ex, Exception) + else traceback.format_exception(type(ex), ex, ex.__traceback__))), innerFailure=( new_failure_details( inner, exception_properties_provider, logger, _visited) diff --git a/durabletask/internal/orchestration_entity_context.py b/durabletask/internal/orchestration_entity_context.py index e47d2ab5..3366590d 100644 --- a/durabletask/internal/orchestration_entity_context.py +++ b/durabletask/internal/orchestration_entity_context.py @@ -19,6 +19,8 @@ def __init__(self, instance_id: str): self.critical_section_id: str | None = None self.critical_section_locks: list[EntityInstanceId] = [] self.available_locks: list[EntityInstanceId] = [] + self._lock_release_obligations: list[EntityInstanceId] = [] + self._pending_lock_releases: list[pb.SendEntityMessageAction] = [] @property def is_inside_critical_section(self) -> bool: @@ -61,19 +63,40 @@ def recover_lock_after_call(self, target_instance_id: EntityInstanceId): if self.is_inside_critical_section: self.available_locks.append(target_instance_id) + def get_lock_release_messages(self) -> Generator[pb.SendEntityMessageAction, None, None]: + yield from self._pending_lock_releases + if self.is_inside_critical_section: + # Preserve the historical shared-list order, then include held + # locks removed from that list by calls which are still pending. + release_order = self.critical_section_locks + [ + entity_id for entity_id in self._lock_release_obligations + if entity_id not in self.critical_section_locks + ] + for entity_id in release_order: + yield self._new_lock_release_message(entity_id) + + def _new_lock_release_message(self, entity_id: EntityInstanceId) -> pb.SendEntityMessageAction: + return pb.SendEntityMessageAction(entityUnlockSent=pb.EntityUnlockSentEvent( + criticalSectionId=self.critical_section_id, + targetInstanceId=get_string_value(str(entity_id)), + parentInstanceId=get_string_value(self.instance_id) + )) + def emit_lock_release_messages(self) -> Generator[pb.SendEntityMessageAction, None, None]: if self.is_inside_critical_section: + # Normal emission must retain historical action IDs. Missing held + # locks are kept only for a later terminal-failure snapshot. + self._pending_lock_releases.extend( + self._new_lock_release_message(entity_id) + for entity_id in self._lock_release_obligations + if entity_id not in self.critical_section_locks + ) for entity_id in self.critical_section_locks: - unlock_event = pb.SendEntityMessageAction(entityUnlockSent=pb.EntityUnlockSentEvent( - criticalSectionId=self.critical_section_id, - targetInstanceId=get_string_value(str(entity_id)), - parentInstanceId=get_string_value(self.instance_id) - )) - yield unlock_event - - self.critical_section_locks = [] - self.available_locks = [] - self.critical_section_id = None + yield self._new_lock_release_message(entity_id) + self.critical_section_locks = [] + self.available_locks = [] + self._lock_release_obligations = [] + self.critical_section_id = None def emit_request_message(self, target: Any, operation_name: str, one_way: bool, operation_id: str, scheduled_time_utc: datetime, input: str | None, @@ -106,6 +129,7 @@ def emit_acquire_message( self.critical_section_id = critical_section_id self.critical_section_locks = entity_ids_dedup + self._lock_release_obligations = list(entity_ids_dedup) self.lock_acquisition_pending = True return request, target diff --git a/durabletask/payload/helpers.py b/durabletask/payload/helpers.py index b42ed119..4fe52ac1 100644 --- a/durabletask/payload/helpers.py +++ b/durabletask/payload/helpers.py @@ -12,11 +12,14 @@ from __future__ import annotations import logging +from collections.abc import Iterator from typing import TYPE_CHECKING, Any from google.protobuf import message as proto_message from google.protobuf import wrappers_pb2 +from durabletask.internal.exceptions import PayloadSizeExceededError + if TYPE_CHECKING: from durabletask.payload.store import PayloadStore @@ -38,11 +41,23 @@ def externalize_payloads( Any ``StringValue`` whose UTF-8 byte length exceeds ``store.options.threshold_bytes`` is uploaded and its value replaced - with the token returned by the store. + with the token returned by the store. Empty values and known tokens + are skipped. + + All eligible payloads are snapshotted and checked against + ``store.options.max_stored_payload_bytes`` before any upload or mutation. + A size violation therefore raises ``ValueError`` before any hypothetical + earlier upload failure. Upload failures can still leave earlier uploads + and token replacements in place; this is not a storage transaction. """ - threshold = store.options.threshold_bytes - max_bytes = store.options.max_stored_payload_bytes - _walk_and_externalize(msg, store, threshold, max_bytes, instance_id) + payloads = _collect_payloads_to_externalize(msg, store) + for field_name, sv, payload_bytes in payloads: + token = store.upload(payload_bytes, instance_id=instance_id) + sv.value = token + logger.debug( + "Externalized %d-byte payload in field '%s' -> %s", + len(payload_bytes), field_name, token, + ) def deexternalize_payloads( @@ -66,10 +81,19 @@ async def externalize_payloads_async( *, instance_id: str | None = None, ) -> None: - """Async version of :func:`externalize_payloads`.""" - threshold = store.options.threshold_bytes - max_bytes = store.options.max_stored_payload_bytes - await _walk_and_externalize_async(msg, store, threshold, max_bytes, instance_id) + """Async version of :func:`externalize_payloads`, with the same preflight. + + Uploads remain sequential; upload errors and cancellation propagate without + rolling back earlier uploads or token replacements. + """ + payloads = _collect_payloads_to_externalize(msg, store) + for field_name, sv, payload_bytes in payloads: + token = await store.upload_async(payload_bytes, instance_id=instance_id) + sv.value = token + logger.debug( + "Externalized %d-byte payload in field '%s' -> %s", + len(payload_bytes), field_name, token, + ) async def deexternalize_payloads_async( @@ -81,7 +105,7 @@ async def deexternalize_payloads_async( # ------------------------------------------------------------------ -# Internal recursive walkers – sync +# Internal preflight and recursive walkers – sync # ------------------------------------------------------------------ def _is_map_field(fd: Any) -> bool: @@ -90,88 +114,50 @@ def _is_map_field(fd: Any) -> bool: return mt is not None and fd.is_repeated and mt.GetOptions().map_entry -def _walk_and_externalize( +def _iter_payload_fields( msg: proto_message.Message, - store: PayloadStore, - threshold: int, - max_bytes: int, - instance_id: str | None, -) -> None: +) -> Iterator[tuple[str, wrappers_pb2.StringValue]]: for fd in msg.DESCRIPTOR.fields: if fd.message_type is None: continue if _is_map_field(fd): - # Map fields: iterate values. ScalarMap values are not - # messages and will be skipped by the isinstance check. - for map_value in getattr(msg, fd.name).values(): - if isinstance(map_value, proto_message.Message): - if isinstance(map_value, wrappers_pb2.StringValue): - _try_externalize_field( - fd.name, map_value, store, - threshold, max_bytes, instance_id, - ) - else: - _walk_and_externalize( - map_value, store, threshold, max_bytes, instance_id - ) + values = getattr(msg, fd.name).values() elif fd.is_repeated: - value = getattr(msg, fd.name) - for item in value: - if isinstance(item, proto_message.Message): - if isinstance(item, wrappers_pb2.StringValue): - _try_externalize_field( - fd.name, item, store, - threshold, max_bytes, instance_id, - ) - else: - _walk_and_externalize( - item, store, threshold, max_bytes, instance_id - ) + values = getattr(msg, fd.name) else: - # Singular message field — only recurse if actually set if not msg.HasField(fd.name): continue - value = getattr(msg, fd.name) + values = (getattr(msg, fd.name),) + + for value in values: if isinstance(value, wrappers_pb2.StringValue): - _try_externalize_field( - fd.name, value, store, - threshold, max_bytes, instance_id, - ) - else: - _walk_and_externalize( - value, store, threshold, max_bytes, instance_id - ) + yield fd.name, value + elif isinstance(value, proto_message.Message): + yield from _iter_payload_fields(value) -def _try_externalize_field( - field_name: str, - sv: wrappers_pb2.StringValue, +def _collect_payloads_to_externalize( + msg: proto_message.Message, store: PayloadStore, - threshold: int, - max_bytes: int, - instance_id: str | None, -) -> None: - val = sv.value - if not val: - return - # Already a token – skip - if store.is_known_token(val): - return - payload_bytes = val.encode("utf-8") - if len(payload_bytes) <= threshold: - return - if len(payload_bytes) > max_bytes: - raise ValueError( - f"Payload size {len(payload_bytes)} bytes exceeds the maximum " - f"allowed size of {max_bytes} bytes." - ) - token = store.upload(payload_bytes, instance_id=instance_id) - sv.value = token - logger.debug( - "Externalized %d-byte payload in field '%s' -> %s", - len(payload_bytes), field_name, token, - ) +) -> list[tuple[str, wrappers_pb2.StringValue, bytes]]: + threshold = store.options.threshold_bytes + max_bytes = store.options.max_stored_payload_bytes + payloads: list[tuple[str, wrappers_pb2.StringValue, bytes]] = [] + for field_name, sv in _iter_payload_fields(msg): + val = sv.value + if not val or store.is_known_token(val): + continue + payload_bytes = val.encode("utf-8") + if len(payload_bytes) <= threshold: + continue + if len(payload_bytes) > max_bytes: + raise PayloadSizeExceededError( + f"Payload size {len(payload_bytes)} bytes exceeds the maximum " + f"allowed size of {max_bytes} bytes." + ) + payloads.append((field_name, sv, payload_bytes)) + return payloads def _walk_and_deexternalize( @@ -223,87 +209,6 @@ def _try_deexternalize_field( # Internal recursive walkers – async # ------------------------------------------------------------------ -async def _walk_and_externalize_async( - msg: proto_message.Message, - store: PayloadStore, - threshold: int, - max_bytes: int, - instance_id: str | None, -) -> None: - for fd in msg.DESCRIPTOR.fields: - if fd.message_type is None: - continue - - if _is_map_field(fd): - for map_value in getattr(msg, fd.name).values(): - if isinstance(map_value, proto_message.Message): - if isinstance(map_value, wrappers_pb2.StringValue): - await _try_externalize_field_async( - fd.name, map_value, store, - threshold, max_bytes, instance_id, - ) - else: - await _walk_and_externalize_async( - map_value, store, threshold, max_bytes, instance_id, - ) - elif fd.is_repeated: - value = getattr(msg, fd.name) - for item in value: - if isinstance(item, proto_message.Message): - if isinstance(item, wrappers_pb2.StringValue): - await _try_externalize_field_async( - fd.name, item, store, - threshold, max_bytes, instance_id, - ) - else: - await _walk_and_externalize_async( - item, store, threshold, max_bytes, instance_id, - ) - else: - if not msg.HasField(fd.name): - continue - value = getattr(msg, fd.name) - if isinstance(value, wrappers_pb2.StringValue): - await _try_externalize_field_async( - fd.name, value, store, - threshold, max_bytes, instance_id, - ) - else: - await _walk_and_externalize_async( - value, store, threshold, max_bytes, instance_id, - ) - - -async def _try_externalize_field_async( - field_name: str, - sv: wrappers_pb2.StringValue, - store: PayloadStore, - threshold: int, - max_bytes: int, - instance_id: str | None, -) -> None: - val = sv.value - if not val: - return - # Already a token – skip - if store.is_known_token(val): - return - payload_bytes = val.encode("utf-8") - if len(payload_bytes) <= threshold: - return - if len(payload_bytes) > max_bytes: - raise ValueError( - f"Payload size {len(payload_bytes)} bytes exceeds the maximum " - f"allowed size of {max_bytes} bytes." - ) - token = await store.upload_async(payload_bytes, instance_id=instance_id) - sv.value = token - logger.debug( - "Externalized %d-byte payload in field '%s' -> %s", - len(payload_bytes), field_name, token, - ) - - async def _walk_and_deexternalize_async( msg: proto_message.Message, store: PayloadStore, diff --git a/durabletask/serialization.py b/durabletask/serialization.py index 52c3dfed..c2d916f0 100644 --- a/durabletask/serialization.py +++ b/durabletask/serialization.py @@ -42,7 +42,8 @@ import types import typing from abc import ABC, abstractmethod -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Set as AbstractSet from types import SimpleNamespace from typing import Any, cast @@ -68,7 +69,6 @@ class DataConverter(ABC): @abstractmethod def serialize(self, value: Any) -> str | None: """Serialize ``value`` to a string, or ``None`` when ``value`` is ``None``.""" - pass @abstractmethod def deserialize(self, data: str | None, target_type: type | None = None) -> Any: @@ -85,7 +85,6 @@ def deserialize(self, data: str | None, target_type: type | None = None) -> Any: :class:`JsonDataConverter` is best-effort and falls back; a validating converter may instead raise. """ - pass @abstractmethod def coerce(self, value: Any, target_type: type | None = None) -> Any: @@ -99,7 +98,6 @@ def coerce(self, value: Any, target_type: type | None = None) -> Any: (strict vs. best-effort) that an implementation applies in :meth:`deserialize` should apply here. """ - pass def can_reconstruct(self, target_type: Any) -> bool: """Return True if this converter can rebuild ``target_type`` from a payload. @@ -136,7 +134,9 @@ class JsonDataConverter(DataConverter): ``SimpleNamespace`` handling, and nested values are encoded recursively, so a dataclass (or any object) with a ``to_json`` hook -- including one nested inside another dataclass, ``list``, or ``dict`` -- round-trips through its - own hooks. + own hooks. Sets and frozensets serialize to JSON arrays ordered + deterministically by their canonical JSON representation (using lexical + string order, so numeric sets like ``{2, 10}`` serialize as ``[10, 2]``). Deserialization (and value-level :meth:`coerce`) is **best-effort**: when a ``target_type`` is supplied and the value cannot be coerced to it, the raw @@ -167,6 +167,18 @@ class JsonDataConverter(DataConverter): > payload; the destination type is always supplied by the caller. """ + def __init__( + self, + *, + enable_generic_sets: bool = False, + enable_set_coercion: bool = False, + ) -> None: + self.enable_generic_sets: bool = enable_generic_sets or enable_set_coercion + + @property + def enable_set_coercion(self) -> bool: + return self.enable_generic_sets + def serialize(self, value: Any) -> str | None: if value is None: return None @@ -221,15 +233,40 @@ def _log_coercion_fallback(target_type: type, error: Exception) -> None: # --------------------------------------------------------------------------- +def _is_supported_set_element(converter: DataConverter, elem_type: Any) -> bool: + if elem_type is None or elem_type is typing.Any: + return True + elem_origin = typing.get_origin(elem_type) + if elem_origin is not None: + if elem_origin in (list, dict, Sequence, Mapping, set, AbstractSet): + # Mutable collections are unhashable + return False + if elem_origin is typing.Union or elem_origin is types.UnionType: + non_none = [a for a in typing.get_args(elem_type) if a is not type(None)] + return bool(non_none) and all(_is_supported_set_element(converter, a) for a in non_none) + if elem_origin is frozenset: + args = typing.get_args(elem_type) + return not args or _is_supported_set_element(converter, args[0]) + return False + if not isinstance(elem_type, type): + return False + if issubclass(elem_type, (list, dict, set)): + return False + return True + + def _can_reconstruct(converter: DataConverter, target_type: Any) -> bool: """:class:`JsonDataConverter`'s reconstruction policy. Recognizes dataclasses and ``from_json()``-capable types, plus ``Optional`` / ``list`` / ``Sequence`` hints wrapping them; builtins and unknown - annotations are excluded. Recurses through ``converter.can_reconstruct`` - (not itself) so a :class:`JsonDataConverter` subclass that overrides - ``can_reconstruct`` still participates in the element-type checks of - ``Optional`` / ``list`` hints. + annotations are excluded. When set support is enabled on the converter + (via ``enable_generic_sets=True``), ``set`` / ``frozenset`` / ``Set`` + containers with supported element types are also recognized. Recurses + through ``converter.can_reconstruct`` (not itself) so a + :class:`JsonDataConverter` subclass that overrides ``can_reconstruct`` + still participates in the element-type checks of ``Optional`` / ``list`` + hints. """ origin = typing.get_origin(target_type) if origin is not None: @@ -240,11 +277,19 @@ def _can_reconstruct(converter: DataConverter, target_type: Any) -> bool: ) if origin in (list, Sequence): return any(converter.can_reconstruct(a) for a in args) + if origin in (set, frozenset, AbstractSet, typing.AbstractSet): + if not getattr(converter, "enable_generic_sets", False): + return False + if not args: + return True + return _is_supported_set_element(converter, args[0]) return False if not isinstance(target_type, type): return False if dataclasses.is_dataclass(target_type): return True + if target_type in (set, frozenset, AbstractSet): + return getattr(converter, "enable_generic_sets", False) return callable(getattr(cast(Any, target_type), "from_json", None)) @@ -253,9 +298,12 @@ def _to_json(obj: Any) -> str: Builtins serialize to plain JSON. Objects exposing a ``to_json()`` method, dataclasses, and ``SimpleNamespace`` instances are serialized to plain JSON - as well (without any type marker), recursively. Custom objects can be - reconstructed on the receiving side by passing ``expected_type`` to - :func:`_from_json`. + as well (without any type marker), recursively. Sets and frozensets are + serialized as JSON arrays with elements deterministically ordered by their + canonical JSON string representation (note that numeric set elements follow + lexical JSON ordering, e.g. ``{2, 10}`` serializes as ``[10, 2]``). Custom + objects can be reconstructed on the receiving side by passing ``expected_type`` + to :func:`_from_json`. """ try: return json.dumps(obj, default=_encode_custom_object) @@ -319,6 +367,31 @@ def _encode_custom_object(o: Any) -> Any: to_json_hook = getattr(cast(Any, type(o)), "to_json", None) if callable(to_json_hook): return to_json_hook(o) + if isinstance(o, (set, frozenset, AbstractSet)): + items: AbstractSet[Any] = cast("AbstractSet[Any]", o) + # Sets and frozensets are unordered collections. To ensure deterministic + # serialization across processes, worker replays, and PYTHONHASHSEED + # variations, elements are individually encoded to canonical JSON + # (with sort_keys=True), sorted lexically by their encoded JSON string, + # and emitted as normalized parsed values. + # + # Emitting the normalized parsed values rather than original objects + # guarantees identical order and eliminates leaks from dict insertion + # order variations, NaN/non-finite comparison anomalies, or partial + # ordering in frozensets. A consequence of sorting canonical JSON strings + # is lexical ordering for numeric sets (e.g., {2, 10} serializes as + # [10, 2]). + try: + encoded_items = sorted( + json.dumps(item, default=_encode_custom_object, sort_keys=True) + for item in items + ) + return [json.loads(encoded) for encoded in encoded_items] + except Exception as e: + raise TypeError( + f"Cannot deterministically serialize set of type '{type(items).__name__}' " + f"because its elements cannot be compared or deterministically encoded: {e}" + ) from e if isinstance(o, enum.Enum): # Emit the member's underlying value (a primitive). Reconstruction is # type-directed: passing the enum type to ``deserialize`` rebuilds the @@ -477,6 +550,19 @@ def _coerce_generic(value: Any, expected_type: Any, origin: Any, if origin in (list, Sequence) and isinstance(value, list): elem_type = args[0] if args else None return [_coerce_to_type(item, elem_type, converter) for item in cast(list[Any], value)] + if origin in (set, frozenset, AbstractSet, typing.AbstractSet) and isinstance( + value, (list, set, frozenset, AbstractSet) + ): + if not getattr(converter, "enable_generic_sets", False): + return cast(Any, value) + elem_type = args[0] if args else None + coerced: list[Any] = [ + _coerce_to_type(item, elem_type, converter) + for item in cast("Iterable[Any]", value) + ] + if origin is frozenset: + return frozenset(coerced) + return set(coerced) if origin in (dict, Mapping) and isinstance(value, dict): # JSON object keys are always strings, so only the values can carry a # reconstructable type. Keys are passed through unchanged. @@ -595,12 +681,12 @@ def _resolve_forward_refs(tp: Any, globalns: dict[str, Any]) -> Any: """ if isinstance(tp, str): try: - tp = eval(tp, globalns) # noqa: S307 - resolves caller-authored annotations + tp = eval(tp, globalns) except Exception: return tp elif isinstance(tp, typing.ForwardRef): try: - tp = eval(tp.__forward_arg__, globalns) # noqa: S307 + tp = eval(tp.__forward_arg__, globalns) except Exception: return tp diff --git a/durabletask/task.py b/durabletask/task.py index b2ca7b32..2c123b94 100644 --- a/durabletask/task.py +++ b/durabletask/task.py @@ -6,8 +6,9 @@ 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 datetime, timedelta, timezone from typing import TYPE_CHECKING, Any, Generic, TypeAlias, TypeVar, cast, overload @@ -631,6 +632,22 @@ 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.") + if not self.is_complete: + yield self + return self.get_result() + @property def is_complete(self) -> bool: """Returns True if the task has completed, False otherwise.""" @@ -961,8 +978,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 25010a17..613bea6b 100644 --- a/durabletask/worker.py +++ b/durabletask/worker.py @@ -7,14 +7,16 @@ import json import logging import os +import sys 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 datetime, timedelta, timezone 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, TypeVar, cast, overload import uuid from packaging.version import InvalidVersion, parse @@ -49,6 +51,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 @@ -61,6 +64,14 @@ _WorkItem = tuple[Callable[..., Any], Callable[..., Any], tuple[Any, ...], dict[str, Any]] +def _is_control_flow_exception(error: BaseException) -> bool: + if isinstance(error, (asyncio.CancelledError, GeneratorExit, KeyboardInterrupt, SystemExit)): + return True + if sys.version_info >= (3, 11) and isinstance(error, BaseExceptionGroup): + return any(_is_control_flow_exception(inner) for inner in error.exceptions) + return False + + class ConcurrencyOptions: """Configuration options for controlling concurrency of different work item types and the thread pool size. @@ -935,7 +946,8 @@ def should_invalidate_connection(rpc_error: grpc.RpcError) -> bool: get_work_items_request.workItemFilters.CopyFrom( self._work_item_filters._to_grpc() # pyright: ignore[reportPrivateUsage] ) - self._response_stream = stub.GetWorkItems(get_work_items_request) + response_stream = stub.GetWorkItems(get_work_items_request) + self._response_stream = response_stream self._logger.info( f"Successfully connected to {self._host_address}. Waiting for work items..." ) @@ -946,13 +958,8 @@ def should_invalidate_connection(rpc_error: grpc.RpcError) -> bool: work_item_queue: queue.Queue[Any] = queue.Queue() saw_message = False - def stream_reader() -> None: + def stream_reader(response_stream: Any = response_stream) -> None: try: - response_stream = self._response_stream - if response_stream is None: - work_item_queue.put(_STREAM_CLOSED_SENTINEL) - return - for work_item in response_stream: work_item_queue.put(work_item) work_item_queue.put(_STREAM_CLOSED_SENTINEL) @@ -1004,8 +1011,14 @@ def stream_reader() -> None: if work_item.HasField("orchestratorRequest"): submit_work_item( self._async_worker_manager.submit_orchestration, - self._execute_orchestrator, - self._cancel_orchestrator, + partial( + self._execute_orchestrator, + _on_outcome_uncertain=response_stream.cancel, + ), + partial( + self._cancel_orchestrator, + _on_outcome_uncertain=response_stream.cancel, + ), work_item.orchestratorRequest, stub, work_item.completionToken, @@ -1148,7 +1161,66 @@ def _execute_orchestrator( req: pb.OrchestratorRequest, stub: Any, completionToken: Any, + *, + _on_outcome_uncertain: Callable[[], Any] | None = None, ) -> None: + outcome_attempted = False + try: + try: + res = self._get_orchestrator_response(req, completionToken) + except pe.AbandonOrchestrationError: + outcome_attempted = True + self._cancel_orchestrator(req, stub, completionToken) + return + + outcome_attempted = True + stub.CompleteOrchestratorTask(res) + except BaseException as error: + recovery_error: BaseException | None = None + outcome_uncertain = outcome_attempted + self._logger.error( + f"Orchestrator dispatch for '{req.instanceId}' exited without a confirmed outcome.", + exc_info=True, + ) + if not outcome_attempted: + try: + self._cancel_orchestrator(req, stub, completionToken) + except BaseException as abandonment_error: + recovery_error = abandonment_error + outcome_uncertain = True + # A failed RPC may already have been accepted. Do not retry it or + # send a conflicting outcome; release ownership through its stream. + if outcome_uncertain and _on_outcome_uncertain is not None: + try: + self._reset_orchestrator_stream(req.instanceId, _on_outcome_uncertain) + except BaseException as reset_error: + if recovery_error is None or ( + isinstance(recovery_error, Exception) and not isinstance(reset_error, Exception) + ): + recovery_error = reset_error + if not isinstance(error, Exception): + raise + if recovery_error is not None and not isinstance(recovery_error, Exception): + raise recovery_error + + def _reset_orchestrator_stream(self, instance_id: str, reset: Callable[[], Any]) -> None: + self._logger.error( + f"Resetting the owning work-item stream for '{instance_id}' after an uncertain outcome." + ) + try: + reset() + except BaseException: + self._logger.error( + f"Failed to reset the owning work-item stream for '{instance_id}'.", + exc_info=True, + ) + raise + + def _get_orchestrator_response( + self, + req: pb.OrchestratorRequest, + completionToken: Any, + ) -> pb.OrchestratorResponse: instance_id = req.instanceId # De-externalize any large-payload tokens in the incoming request @@ -1179,6 +1251,7 @@ def _execute_orchestrator( if (req.HasField("orchestrationTraceContext") and req.orchestrationTraceContext.HasField("spanID") and req.orchestrationTraceContext.spanID.value): persisted_orch_span_id = req.orchestrationTraceContext.spanID.value + result: ExecutionResults | None = None try: executor = _OrchestrationExecutor( self._registry, self._logger, @@ -1192,30 +1265,6 @@ def _execute_orchestrator( result = executor.execute( instance_id, req.pastEvents, req.newEvents) - # Determine completion status for span - is_complete = False - is_failed = False - failure_details = None - for action in result.actions: - if action.HasField("completeOrchestration"): - is_complete = True - orch_status = action.completeOrchestration.orchestrationStatus - if orch_status == pb.ORCHESTRATION_STATUS_FAILED: - is_failed = True - failure_details = action.completeOrchestration.failureDetails - - if is_complete and self._emit_trace_spans: - # Orchestration finished — emit a single span covering its lifetime - tracing.emit_orchestration_span( - orchestration_name, - instance_id, - start_time_ns, - is_failed, - failure_details=failure_details, - parent_trace_context=parent_trace_ctx, - orchestration_trace_context=result._orchestration_trace_context, # pyright: ignore[reportPrivateUsage] - ) - # Include the span ID in the orchestration trace context # so it persists across dispatches. orch_span_id = None @@ -1239,16 +1288,7 @@ def _execute_orchestrator( ), ) except pe.AbandonOrchestrationError: - # Abandoned — no span needed - self._logger.info( - f"Abandoning orchestration. InstanceId = '{instance_id}'. Completion token = '{completionToken}'" - ) - stub.AbandonTaskOrchestratorWorkItem( - pb.AbandonOrchestrationTaskRequest( - completionToken=completionToken - ) - ) - return + raise except Exception as ex: # Unhandled error — emit a failed span if self._emit_trace_spans: @@ -1276,29 +1316,69 @@ def _execute_orchestrator( completionToken=completionToken, ) + # Externalize before attempting the completion RPC. try: - # Externalize any large payloads in the response if self._payload_store is not None: payload_helpers.externalize_payloads( res, self._payload_store, instance_id=instance_id, ) - stub.CompleteOrchestratorTask(res) - except Exception as ex: - self._logger.exception( - f"Failed to deliver orchestrator response for '{req.instanceId}' to sidecar: {ex}" + except pe.PayloadSizeExceededError as ex: + self._logger.warning( + f"Orchestration '{instance_id}' exceeded the payload storage limit: {ex}" ) + # Use the original actions, not a partially externalized response. + original = result if result is not None else ExecutionResults(list(res.actions), None) + failed = original.as_failed(pb.TaskFailureDetails( + errorType=ph.get_qualified_name(type(ex)), + errorMessage=str(ex)[:1024], + isNonRetriable=True, + )) + res.ClearField("actions") + res.actions.extend(failed.actions) + res.ClearField("customStatus") + # Keep this fallback inline: it must not need the failing store. + + if result is not None and self._emit_trace_spans: + for action in res.actions: + if action.HasField("completeOrchestration"): + completion = action.completeOrchestration + tracing.emit_orchestration_span( + orchestration_name, + instance_id, + start_time_ns, + completion.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED, + failure_details=completion.failureDetails if completion.HasField("failureDetails") else None, + parent_trace_context=parent_trace_ctx, + orchestration_trace_context=result._orchestration_trace_context, # pyright: ignore[reportPrivateUsage] + ) + return res def _cancel_orchestrator( self, req: pb.OrchestratorRequest, stub: Any, completionToken: Any, + *, + _on_outcome_uncertain: Callable[[], Any] | None = None, ) -> None: - stub.AbandonTaskOrchestratorWorkItem( - pb.AbandonOrchestrationTaskRequest( - completionToken=completionToken + try: + stub.AbandonTaskOrchestratorWorkItem( + pb.AbandonOrchestrationTaskRequest( + completionToken=completionToken + ) ) - ) + except BaseException as error: + self._logger.error( + f"Failed to abandon orchestrator dispatch for '{req.instanceId}'.", + exc_info=True, + ) + if _on_outcome_uncertain is not None: + try: + self._reset_orchestrator_stream(req.instanceId, _on_outcome_uncertain) + except BaseException as reset_error: + if isinstance(error, Exception) and not isinstance(reset_error, Exception): + raise + raise self._logger.info(f"Cancelled orchestration task for invocation ID: {req.instanceId}") def _on_activity_execution_started(self, req: pb.ActivityRequest) -> None: @@ -1503,7 +1583,7 @@ def _cancel_entity_batch( class _RuntimeOrchestrationContext(task.OrchestrationContext): - _generator: Generator[task.Task[Any], Any, Any] | None + _generator: Generator[task.Task[Any], Any, Any] | Coroutine[Any, Any, Any] | None _previous_task: task.Task[Any] | None def __init__(self, @@ -1548,22 +1628,36 @@ def __init__(self, self._data_converter = data_converter self._exception_properties_provider = exception_properties_provider self._logger = logger + self._async_loop: DurableEventLoop | None = None + self._async_cleanup = False - def run(self, generator: Generator[task.Task[Any], Any, Any]) -> None: + def run(self, generator: Generator[task.Task[Any], Any, Any] | Coroutine[Any, Any, Any]) -> None: + if isinstance(generator, CoroutineType): + self._async_loop = DurableEventLoop(self) + self._async_loop.start(generator) + return 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._set_next_task(generator.send(None)) + + def _set_next_task(self, next_task: Any) -> None: + if not isinstance(next_task, task.Task): + raise TypeError( + "Orchestrators must only yield or await durabletask.Task objects. " + "Use durable timers and task.when_all/when_any instead of asyncio APIs." + ) + 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 or coroutine 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. @@ -1578,9 +1672,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, @@ -1595,28 +1687,38 @@ def set_complete( if self._entity_context.is_inside_critical_section: self._exit_critical_section() - self._is_complete = True - self._completion_status = status # This is probably a bug - an orchestrator may complete with some actions remaining that the user still # wants to execute - for example, signaling an entity. So we shouldn't clear the pending actions here. # self._pending_actions.clear() # Cancel any pending actions + # Block reentrant terminal transitions from serialization callbacks, + # but leave failure conversion available if no completion was built. + self._is_complete = True + try: + result_json: str | None = None + if result is not None: + try: + result_json = result if is_result_encoded else self._data_converter.serialize(result) + except StopIteration as ex: + # This is a serialization failure, not generator completion. + self._is_complete = False + self.set_failed(ex) + return + except (ValueError, TypeError): + self._is_complete = False + self.set_failed(JsonEncodeOutputException(result)) + return + action = ph.new_complete_orchestration_action( + self.next_sequence_number(), status, result_json + ) + self._pending_actions[action.id] = action + except BaseException: + self._is_complete = False + raise self._result = result - result_json: str | None = None - if result is not None: - try: - result_json = result if is_result_encoded else self._data_converter.serialize(result) - except (ValueError, TypeError): - self._is_complete = False - self._result = None - self.set_failed(JsonEncodeOutputException(result)) - return - action = ph.new_complete_orchestration_action( - self.next_sequence_number(), status, result_json - ) - self._pending_actions[action.id] = action + self._completion_status = status - def set_failed(self, ex: Exception | pb.TaskFailureDetails): + def set_failed(self, ex: BaseException | pb.TaskFailureDetails): if self._is_complete: return @@ -1635,7 +1737,7 @@ def set_failed(self, ex: Exception | pb.TaskFailureDetails): ex, self._exception_properties_provider, self._logger) failure_details.innerFailure.CopyFrom( ex._failure_details) # pyright: ignore[reportPrivateUsage] - elif isinstance(ex, Exception): + elif isinstance(ex, BaseException): failure_details = ph.new_failure_details( ex, self._exception_properties_provider, self._logger) else: @@ -1699,6 +1801,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 @@ -2194,14 +2298,31 @@ class ExecutionResults: actions: list[pb.OrchestratorAction] encoded_custom_status: str | None _orchestration_trace_context: pb.TraceContext | None + entity_lock_releases: list[pb.SendEntityMessageAction] def __init__( self, actions: list[pb.OrchestratorAction], encoded_custom_status: str | None, orchestration_trace_context: pb.TraceContext | None = None, + entity_lock_releases: list[pb.SendEntityMessageAction] | None = None, ): self.actions = actions self.encoded_custom_status = encoded_custom_status self._orchestration_trace_context = orchestration_trace_context + self.entity_lock_releases = entity_lock_releases if entity_lock_releases is not None else [] + + def as_failed(self, failure_details: pb.TaskFailureDetails) -> "ExecutionResults": + actions = [ + action for action in self.actions + if action.HasField("sendEntityMessage") and action.sendEntityMessage.HasField("entityUnlockSent") + ] + next_id = max((action.id for action in self.actions), default=-1) + 1 + for message in self.entity_lock_releases: + actions.append(pb.OrchestratorAction(id=next_id, sendEntityMessage=message)) + next_id += 1 + actions.append(ph.new_complete_orchestration_action( + next_id, pb.ORCHESTRATION_STATUS_FAILED, None, failure_details)) + return ExecutionResults( + actions, None, orchestration_trace_context=self._orchestration_trace_context) class _OrchestrationExecutor: @@ -2288,65 +2409,125 @@ def execute( exception_properties_provider=self._exception_properties_provider, logger=self._logger, ) + result: ExecutionResults | None = None + dispatch_exited = False + unexpected_failure = False try: - # Rebuild local state by replaying old history into the orchestrator function - self._logger.debug( - f"{instance_id}: Rebuilding local state with {len(old_events)} history event..." - ) - ctx._is_replaying = True # pyright: ignore[reportPrivateUsage] - for old_event in old_events: - self.process_event(ctx, old_event) - - # Get new actions by executing newly received events into the orchestrator function - if self._logger.level <= logging.DEBUG: - summary = _get_new_event_summary(new_events) + try: + # Rebuild local state by replaying old history into the orchestrator function self._logger.debug( - f"{instance_id}: Processing {len(new_events)} new event(s): {summary}" + f"{instance_id}: Rebuilding local state with {len(old_events)} history event..." ) - ctx._is_replaying = False # pyright: ignore[reportPrivateUsage] - for new_event in new_events: - self.process_event(ctx, new_event) - - except pe.VersionFailureException as ex: - if self._registry.versioning and self._registry.versioning.failure_strategy == VersionFailureStrategy.FAIL: - if ex.error_details: - ctx.set_failed(ex.error_details) - else: - ctx.set_failed(ex) - elif self._registry.versioning and self._registry.versioning.failure_strategy == VersionFailureStrategy.REJECT: - raise pe.AbandonOrchestrationError + ctx._is_replaying = True # pyright: ignore[reportPrivateUsage] + for old_event in old_events: + self.process_event(ctx, old_event) + + # Get new actions by executing newly received events into the orchestrator function + if self._logger.level <= logging.DEBUG: + summary = _get_new_event_summary(new_events) + self._logger.debug( + f"{instance_id}: Processing {len(new_events)} new event(s): {summary}" + ) + ctx._is_replaying = False # pyright: ignore[reportPrivateUsage] + for new_event in new_events: + self.process_event(ctx, new_event) + + except pe.VersionFailureException as ex: + if self._registry.versioning and self._registry.versioning.failure_strategy == VersionFailureStrategy.FAIL: + if ex.error_details: + ctx.set_failed(ex.error_details) + else: + ctx.set_failed(ex) + elif self._registry.versioning and self._registry.versioning.failure_strategy == VersionFailureStrategy.REJECT: + raise pe.AbandonOrchestrationError - except Exception as ex: - # Unhandled exceptions fail the orchestration - self._logger.debug(f"{instance_id}: Orchestration {orchestration_name} failed") - ctx.set_failed(ex) - - if not ctx._is_complete: # pyright: ignore[reportPrivateUsage] - task_count = len(ctx._pending_tasks) # pyright: ignore[reportPrivateUsage] - event_count = len(ctx._pending_events) # pyright: ignore[reportPrivateUsage] - self._logger.info( - f"{instance_id}: Orchestrator {orchestration_name} yielded with {task_count} task(s) " - f"and {event_count} event(s) outstanding." - ) - elif ( - ctx._completion_status and ctx._completion_status is not pb.ORCHESTRATION_STATUS_CONTINUED_AS_NEW # pyright: ignore[reportPrivateUsage] - ): - completion_status_str = ph.get_orchestration_status_str( - ctx._completion_status # pyright: ignore[reportPrivateUsage] - ) - self._logger.info( - f"{instance_id}: Orchestration {orchestration_name} completed with status: {completion_status_str}" - ) + except Exception as ex: + # Unhandled exceptions fail the orchestration + self._logger.debug(f"{instance_id}: Orchestration {orchestration_name} failed") + ctx.set_failed(ex) + + except BaseException as ex: + if _is_control_flow_exception(ex): + raise + self._logger.error( + f"{instance_id}: Orchestrator {orchestration_name} exited unexpectedly.", + exc_info=True, + ) + unexpected_failure = not ctx._is_complete # pyright: ignore[reportPrivateUsage] + ctx.set_failed(ex) + + if not ctx._is_complete: # pyright: ignore[reportPrivateUsage] + task_count = len(ctx._pending_tasks) # pyright: ignore[reportPrivateUsage] + event_count = len(ctx._pending_events) # pyright: ignore[reportPrivateUsage] + self._logger.info( + f"{instance_id}: Orchestrator {orchestration_name} yielded with {task_count} task(s) " + f"and {event_count} event(s) outstanding." + ) + elif ( + ctx._completion_status and ctx._completion_status is not pb.ORCHESTRATION_STATUS_CONTINUED_AS_NEW # pyright: ignore[reportPrivateUsage] + ): + completion_status_str = ph.get_orchestration_status_str( + ctx._completion_status # pyright: ignore[reportPrivateUsage] + ) + self._logger.info( + f"{instance_id}: Orchestration {orchestration_name} completed with status: {completion_status_str}" + ) - actions = ctx.get_actions() - if self._logger.level <= logging.DEBUG: - self._logger.debug( - f"{instance_id}: Returning {len(actions)} action(s): {_get_action_summary(actions)}" + actions = ctx.get_actions() + if self._logger.level <= logging.DEBUG: + self._logger.debug( + f"{instance_id}: Returning {len(actions)} action(s): {_get_action_summary(actions)}" + ) + 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. + entity_lock_releases=list(ctx._entity_context.get_lock_release_messages()), # pyright: ignore[reportPrivateUsage] ) - return ExecutionResults( - actions=actions, encoded_custom_status=ctx._encoded_custom_status, # pyright: ignore[reportPrivateUsage] - orchestration_trace_context=ctx._orchestration_trace_context, # pyright: ignore[reportPrivateUsage] - ) + if unexpected_failure and result.entity_lock_releases: + failure_details = next( + action.completeOrchestration.failureDetails for action in result.actions + if action.HasField("completeOrchestration") + and action.completeOrchestration.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + ) + result = result.as_failed(failure_details) + except BaseException: + dispatch_exited = True + raise + finally: + # Close after building the response, so the orchestrator's finally + # blocks do not change the actions or custom status we return. + closeable = ctx._async_loop or ctx._generator # pyright: ignore[reportPrivateUsage] + if closeable is not None: + try: + closeable.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} 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)) + assert result is not None + return result def _build_rewind_result( self, @@ -2521,10 +2702,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/noxfile.py b/noxfile.py index 1e09cf7c..252b8d51 100644 --- a/noxfile.py +++ b/noxfile.py @@ -18,7 +18,7 @@ nox -s functions_unit-3.13 nox -s functions_e2e -``azure-functions>=2.3.0b2`` is published to PyPI and installed as a declared +``azure-functions>=2.4.0b1`` is published to PyPI and installed as a declared dependency of ``azure-functions-durable``, so no local build is required. """ diff --git a/requirements.txt b/requirements.txt index ea61d217..2d865693 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,6 +3,8 @@ grpcio>=1.60.0 # 1.60.0 is the version introducing protobuf 1.25.X support, newe protobuf pytest pytest-asyncio +anyio +pydantic-ai-slim==2.60.0; python_version >= "3.11" # Framework proof, no model providers pytest-cov azure-identity azure-functions 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..ccb003c3 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,16 @@ 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") 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..28d70500 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,16 @@ 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") 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..06215119 --- /dev/null +++ b/tests/azure-functions-durable/e2e/test_coroutine_e2e.py @@ -0,0 +1,24 @@ +# 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"]) +def test_async_await_orchestrator(request: pytest.FixtureRequest, app_fixture: str) -> None: + app: FunctionApp = request.getfixturevalue(app_fixture) + instance_id = app.start_orchestration("async_await_probe") + 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_client_compat.py b/tests/azure-functions-durable/test_client_compat.py index 8cf96308..6be86e97 100644 --- a/tests/azure-functions-durable/test_client_compat.py +++ b/tests/azure-functions-durable/test_client_compat.py @@ -2,6 +2,7 @@ # Licensed under the MIT License. import json +import pickle from datetime import datetime, timedelta, timezone from types import SimpleNamespace from unittest.mock import AsyncMock, Mock, patch @@ -927,6 +928,33 @@ def test_entity_id_maps_to_entity_instance_id(): assert str(entity_id) == "@counter@one" +def test_entity_id_pickle_round_trip(): + with pytest.warns(DeprecationWarning): + entity_id = df.EntityId("Counter", "one") + + restored = pickle.loads(pickle.dumps(entity_id)) + assert type(restored) is df.EntityId + assert isinstance(restored, df.EntityId) + assert restored.name == "counter" + assert restored.entity == "counter" + assert restored.key == "one" + assert str(restored) == "@counter@one" + assert restored == entity_id + assert hash(restored) == hash(entity_id) + + # Immutability guarantees remain intact after unpickling + with pytest.raises(AttributeError): + restored.entity = "other" + with pytest.raises(AttributeError): + restored.key = "other" + with pytest.raises(AttributeError): + restored.name = "other" + with pytest.raises(AttributeError): + del restored.entity + with pytest.raises(AttributeError): + del restored.key + + def test_entity_id_url_path(): with pytest.warns(DeprecationWarning): entity_id = df.EntityId("Counter", "one") 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..cdb157e0 --- /dev/null +++ b/tests/azure-functions-durable/test_functions_coroutine_orchestrators.py @@ -0,0 +1,107 @@ +# 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 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"]) +async def test_coroutine_replay_through_functions_entry_points(native: bool, entry: str) -> None: + async def modern(ctx: task.OrchestrationContext, value: int) -> list[Any]: + return 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() + return 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("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"] diff --git a/tests/azure-functions-durable/test_serialization_configuration.py b/tests/azure-functions-durable/test_serialization_configuration.py new file mode 100644 index 00000000..0bd120b4 --- /dev/null +++ b/tests/azure-functions-durable/test_serialization_configuration.py @@ -0,0 +1,384 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""App-level encoding policy and shared-codec compatibility coverage.""" + +import base64 +import json +from unittest.mock import AsyncMock, Mock + +import pytest +from azure.functions import meta +from azure.functions._durable_functions import df_dumps as shared_dumps + +import azure.durable_functions as df +from azure.durable_functions.internal import payloads, serialization +from azure.durable_functions.internal.compat.activity import wrap_activity_payloads +from azure.durable_functions.internal.converters import ActivityTriggerConverter +from azure.durable_functions.worker import DurableFunctionsWorker +from durabletask.entities import EntityInstanceId +from durabletask.internal import helpers +from durabletask.internal import orchestrator_service_pb2 as pb + + +VALUE = {"\u043a\u043b\u044e\u0447": "\u041f\u0440\u0438\u0432\u0435\u0442 \U0001f680"} +CLIENT_CONFIG = json.dumps({ + "taskHubName": "unicode-tests", + "rpcBaseUrl": "http://localhost:8080/", +}) + + +@pytest.fixture(autouse=True) +def reset_serialization(monkeypatch): + monkeypatch.setattr(serialization, "_ensure_ascii", None) + monkeypatch.setattr(serialization, "_unicode_encoder", None) + monkeypatch.setattr(payloads, "_payload_store", None) + + +@pytest.fixture +def configurable_codec(monkeypatch): + calls = [] + + def dumps(value, *, ensure_ascii=True): + calls.append(ensure_ascii) + return shared_dumps(value, ensure_ascii=ensure_ascii) + + monkeypatch.setattr(serialization, "df_dumps", dumps) + return calls + + +def test_legacy_sdk_keeps_default_and_explicit_true(monkeypatch): + calls = [] + + def legacy_dumps(value): + calls.append(value) + return shared_dumps(value) + + monkeypatch.setattr(serialization, "df_dumps", legacy_dumps) + converter = serialization.FunctionsDataConverter() + assert converter.serialize(VALUE) == shared_dumps(VALUE) + df.DFApp().configure_serialization(ensure_ascii=True) + assert ActivityTriggerConverter.encode(VALUE, expected_type=None).value == shared_dumps(VALUE) + assert calls == [VALUE, VALUE] + + +def test_legacy_sdk_rejects_optin_without_changing_policy(monkeypatch): + def legacy_dumps(value): + return shared_dumps(value) + + monkeypatch.setattr(serialization, "df_dumps", legacy_dumps) + app = df.DFApp() + with pytest.raises(RuntimeError, match=r"requires azure-functions>=2\.4\.0b1"): + app.configure_serialization(ensure_ascii=False) + assert serialization._ensure_ascii is None + app.configure_serialization() + assert serialization.serialize_json(VALUE) == shared_dumps(VALUE) + + +@pytest.mark.parametrize("invalid", [None, 0, 1, "false", [], {}]) +def test_configuration_requires_bool(invalid): + with pytest.raises(TypeError, match="ensure_ascii must be a bool"): + df.DFApp().configure_serialization(ensure_ascii=invalid) + assert serialization._ensure_ascii is None + + +def test_configuration_option_is_keyword_only(): + with pytest.raises(TypeError): + df.DFApp().configure_serialization(False) + + +@pytest.mark.parametrize("ensure_ascii", [True, False]) +def test_configuration_is_idempotent_but_rejects_conflicts(configurable_codec, ensure_ascii): + first = df.DFApp() + second = df.DFApp() + first.configure_serialization(ensure_ascii=ensure_ascii) + second.configure_serialization(ensure_ascii=ensure_ascii) + with pytest.raises(ValueError, match="already configured"): + second.configure_serialization(ensure_ascii=not ensure_ascii) + serialized = serialization.serialize_json(VALUE) + assert json.loads(serialized) == VALUE + assert configurable_codec == [ensure_ascii] + + +def test_first_serialization_freezes_legacy_default(configurable_codec): + assert serialization.serialize_json(VALUE) == shared_dumps(VALUE) + with pytest.raises(ValueError, match="before serializing durable payloads"): + df.DFApp().configure_serialization(ensure_ascii=False) + df.DFApp().configure_serialization(ensure_ascii=True) + + +@pytest.mark.parametrize("strict", [False, True]) +def test_shared_default_and_optin_round_trip(configurable_codec, monkeypatch, strict): + monkeypatch.setenv("AZURE_FUNCTIONS_DURABLE_STRICT_TYPING", "1" if strict else "0") + converter = serialization.FunctionsDataConverter() + df.DFApp().configure_serialization(ensure_ascii=False) + for value in [ + VALUE, + {"text": "\u0645\u0631\u062d\u0628\u0627 \u4f60\u597d e\u0301 \u2028\u2029"}, + {"text": '\x00\n\t"\\ literal \\u0416'}, + {"text": "ASCII only"}, + ]: + encoded = converter.serialize(value) + assert encoded == json.dumps(value, ensure_ascii=False) + assert converter.deserialize(encoded, dict) == value + assert configurable_codec == [False] * 4 + + +class Label: + def __init__(self, text): + self.text = text + + def to_json(self): + return {"text": self.text} + + @classmethod + def from_json(cls, data): + return cls(data["text"]) + + +@pytest.mark.parametrize("strict", [False, True]) +def test_optin_preserves_custom_object_rules(configurable_codec, monkeypatch, strict): + monkeypatch.setenv("AZURE_FUNCTIONS_DURABLE_STRICT_TYPING", "1" if strict else "0") + df.DFApp().configure_serialization(ensure_ascii=False) + converter = serialization.FunctionsDataConverter() + value = Label("\u0416\U0001f680") + encoded = converter.serialize(value) + restored = converter.deserialize(encoded, Label) + assert restored.text == value.text + nested = {"label": value} + if strict: + with pytest.raises(TypeError): + converter.serialize(nested) + else: + assert converter.deserialize(converter.serialize(nested), dict)["label"].text == value.text + + +def test_blueprints_and_existing_workers_use_app_policy(configurable_codec): + blueprint = df.Blueprint() + + @blueprint.orchestration_trigger(context_name="context") + def unicode_orchestrator(context): + return context.get_input() + + app = df.DFApp() + app.register_functions(blueprint) + worker = DurableFunctionsWorker() + app.configure_serialization(ensure_ascii=False) + request = pb.OrchestratorRequest(instanceId="unicode-orchestrator") + request.newEvents.extend([ + helpers.new_orchestrator_started_event(), + helpers.new_execution_started_event( + "unicode_orchestrator", request.instanceId, shared_dumps(VALUE)), + ]) + encoded = base64.b64encode(request.SerializeToString()).decode("ascii") + result = worker.execute_orchestration_request( + unicode_orchestrator._function._func.orchestrator_function, encoded) + response = pb.OrchestratorResponse.FromString(base64.b64decode(result)) + completion = response.actions[0].completeOrchestration + assert completion.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED + assert completion.result.value == json.dumps(VALUE, ensure_ascii=False) + + +def test_activity_converter_uses_optin(configurable_codec): + df.DFApp().configure_serialization(ensure_ascii=False) + encoded = ActivityTriggerConverter.encode(VALUE, expected_type=None) + assert encoded.type == "json" + assert encoded.value == json.dumps(VALUE, ensure_ascii=False) + for payload in (shared_dumps(VALUE), encoded.value): + assert ActivityTriggerConverter.decode( + meta.Datum(type="json", value=payload), trigger_metadata=None) == VALUE + assert configurable_codec == [False] + + +@pytest.mark.parametrize("initial_ascii", [False, True]) +@pytest.mark.parametrize("escaped_result", [False, True]) +def test_replay_across_encoding_rollout_and_rollback( + configurable_codec, monkeypatch, initial_ascii, escaped_result): + df.DFApp().configure_serialization(ensure_ascii=initial_ascii) + + def unicode_replay(context): + value = context.get_input() + result = yield context.call_activity("echo", value) + result = yield context.call_sub_orchestrator( + "child", result, instance_id="unicode-child") + context.set_custom_status(result) + return result + + def execute(request): + encoded = base64.b64encode(request.SerializeToString()).decode("ascii") + result = DurableFunctionsWorker().execute_orchestration_request(unicode_replay, encoded) + return pb.OrchestratorResponse.FromString(base64.b64decode(result)) + + request = pb.OrchestratorRequest(instanceId="unicode-replay") + request.newEvents.extend([ + helpers.new_orchestrator_started_event(), + helpers.new_execution_started_event( + "unicode_replay", request.instanceId, shared_dumps(VALUE)), + ]) + initial = execute(request) + scheduled = initial.actions[0] + assert scheduled.scheduleTask.input.value == json.dumps(VALUE, ensure_ascii=initial_ascii) + request.pastEvents.extend(request.newEvents) + request.ClearField("newEvents") + request.pastEvents.append(helpers.new_task_scheduled_event( + scheduled.id, "echo", scheduled.scheduleTask.input.value)) + result_payload = json.dumps(VALUE, ensure_ascii=escaped_result) + request.newEvents.extend([ + helpers.new_orchestrator_started_event(), + helpers.new_task_completed_event(scheduled.id, result_payload), + ]) + + # Simulate a fresh worker process deploying the other encoding policy. + monkeypatch.setattr(serialization, "_ensure_ascii", None) + monkeypatch.setattr(serialization, "_unicode_encoder", None) + df.DFApp().configure_serialization(ensure_ascii=not initial_ascii) + resumed = execute(request) + assert len(resumed.actions) == 1 + child = resumed.actions[0] + assert child.createSubOrchestration.input.value == json.dumps( + VALUE, ensure_ascii=not initial_ascii) + request.pastEvents.extend(request.newEvents) + request.ClearField("newEvents") + request.pastEvents.append(helpers.new_sub_orchestration_created_event( + child.id, "child", "unicode-child", child.createSubOrchestration.input.value)) + request.newEvents.extend([ + helpers.new_orchestrator_started_event(), + helpers.new_sub_orchestration_completed_event(child.id, result_payload), + ]) + completed = execute(request) + assert len(completed.actions) == 1 + completion = completed.actions[0].completeOrchestration + assert completion.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED + expected = json.dumps(VALUE, ensure_ascii=not initial_ascii) + assert completion.result.value == expected + assert completed.customStatus.value == expected + + +@pytest.mark.parametrize("use_async", [False, True]) +async def test_clients_created_before_configuration_use_optin( + configurable_codec, monkeypatch, use_async): + client = (df.DurableFunctionsClient(CLIENT_CONFIG) if use_async + else df.SyncDurableFunctionsClient(CLIENT_CONFIG)) + stub = Mock() + stub.StartInstance = (AsyncMock if use_async else Mock)( + return_value=pb.CreateInstanceResponse(instanceId="unicode-client")) + stub.SignalEntity = (AsyncMock if use_async else Mock)( + return_value=pb.SignalEntityResponse()) + if use_async: + monkeypatch.setattr(client, "_get_stub", lambda: stub) + else: + monkeypatch.setattr(client, "_stub", stub) + df.DFApp().configure_serialization(ensure_ascii=False) + try: + if use_async: + await client.schedule_new_orchestration("echo", input=VALUE) + await client.signal_entity(EntityInstanceId("counter", "key"), "set", VALUE) + else: + client.schedule_new_orchestration("echo", input=VALUE) + client.signal_entity(EntityInstanceId("counter", "key"), "set", VALUE) + assert stub.StartInstance.call_args.args[0].input.value == json.dumps( + VALUE, ensure_ascii=False) + assert stub.SignalEntity.call_args.args[0].input.value == json.dumps( + VALUE, ensure_ascii=False) + assert configurable_codec == [False, False] + finally: + if use_async: + await client.close() + else: + client.close() + + +@pytest.mark.parametrize("native", [False, True]) +def test_entity_state_result_and_actions_use_optin(configurable_codec, native): + worker = DurableFunctionsWorker() + df.DFApp().configure_serialization(ensure_ascii=False) + + def compat_entity(context): + assert context.get_input() == VALUE + assert context.get_state() == VALUE + context.set_state(VALUE) + context.set_result(VALUE) + + def native_entity(context, input): + assert input == VALUE + assert context.get_state() == VALUE + context.set_state(VALUE) + context.signal_entity(EntityInstanceId("other", "key"), "set", VALUE) + return VALUE + + unicode_entity = native_entity if native else compat_entity + + request = pb.EntityBatchRequest(instanceId=f"@{unicode_entity.__name__}@key") + request.entityState.value = shared_dumps(VALUE) + request.operations.add( + requestId="unicode-operation", operation="set", + input=helpers.get_string_value(shared_dumps(VALUE))) + encoded = base64.b64encode(request.SerializeToString()).decode("ascii") + response = pb.EntityBatchResult.FromString(base64.b64decode( + worker.execute_entity_batch_request(unicode_entity, encoded))) + expected = json.dumps(VALUE, ensure_ascii=False) + assert response.entityState.value == expected + assert response.results[0].success.result.value == expected + if native: + assert response.actions[0].sendSignal.input.value == expected + + +@pytest.mark.parametrize("use_async", [False, True]) +async def test_externalized_activity_bytes_and_threshold_use_optin( + configurable_codec, monkeypatch, payload_store_factory, use_async): + value = {"text": "\u0416" * 100} + store = payload_store_factory(threshold_bytes=300) + monkeypatch.setattr(payloads, "_payload_store", store) + df.DFApp().configure_serialization(ensure_ascii=False) + + def activity(payload): + return value + + async def async_activity(payload): + return value + + wrapped = wrap_activity_payloads(async_activity if use_async else activity, "payload") + input_payload = ActivityTriggerConverter.decode( + meta.Datum(type="json", value="null"), trigger_metadata=None) + result = await wrapped(input_payload) if use_async else wrapped(input_payload) + encoded = ActivityTriggerConverter.encode(result, expected_type=None) + assert encoded.value == json.dumps(value, ensure_ascii=False) + assert store._blobs == {} + larger = {"text": "\u0416" * 200} + + def large_activity(payload): + return larger + + async def large_async_activity(payload): + return larger + + wrapped = wrap_activity_payloads( + large_async_activity if use_async else large_activity, "payload") + result = await wrapped(input_payload) if use_async else wrapped(input_payload) + encoded = ActivityTriggerConverter.encode(result, expected_type=None) + token = json.loads(encoded.value) + assert store.download(token) == json.dumps(larger, ensure_ascii=False).encode("utf-8") + + +@pytest.mark.parametrize("strict", [False, True]) +def test_optin_with_actual_shared_codec(monkeypatch, strict): + monkeypatch.setenv("AZURE_FUNCTIONS_DURABLE_STRICT_TYPING", "1" if strict else "0") + df.DFApp().configure_serialization(ensure_ascii=False) + for value in [VALUE, {"text": "\ud800\udc00\udfff\x7f"}]: + encoded = serialization.serialize_json(value) + assert encoded == shared_dumps(value, ensure_ascii=False) + wire = pb.ActivityResponse(result=helpers.get_string_value(encoded)).SerializeToString() + assert pb.ActivityResponse.FromString(wire).result.value == encoded + + +@pytest.mark.parametrize("strict", [False, True]) +def test_released_codec_reduces_serialized_utf8_size(monkeypatch, strict): + monkeypatch.setenv("AZURE_FUNCTIONS_DURABLE_STRICT_TYPING", "1" if strict else "0") + value = {"text": "\u0416" * 100_000} + legacy = shared_dumps(value) + assert legacy == shared_dumps(value, ensure_ascii=True) + assert len(legacy.encode("utf-8")) == 600_012 + df.DFApp().configure_serialization(ensure_ascii=False) + encoded = serialization.serialize_json(value) + assert len(encoded.encode("utf-8")) == 200_012 + assert serialization.FunctionsDataConverter().deserialize(encoded, dict) == value diff --git a/tests/azure-functions-durable/test_worker_compat.py b/tests/azure-functions-durable/test_worker_compat.py index 5e20f3a2..1529b91c 100644 --- a/tests/azure-functions-durable/test_worker_compat.py +++ b/tests/azure-functions-durable/test_worker_compat.py @@ -56,6 +56,34 @@ def test_worker_uses_propagate_only_tracing(): assert worker.emit_trace_spans is False +@pytest.mark.parametrize("use_async", [False, True]) +async def test_worker_returns_failed_response_for_unexpected_teardown(monkeypatch, use_async): + monkeypatch.setattr(payloads, "_payload_store", None) + + class TeardownExit(BaseException): + pass + + def orchestrator(context): + try: + yield context.call_activity("work") + finally: + raise TeardownExit("deterministic teardown failure") + + worker = DurableFunctionsWorker() + request = _encode_orchestrator_request("teardown") + if use_async: + encoded = await worker.execute_orchestration_request_async(orchestrator, request) + else: + encoded = worker.execute_orchestration_request(orchestrator, request) + response = _decode_orchestrator_response(encoded) + completion = _get_completion_action(response) + + assert completion.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + assert completion.failureDetails.errorType.endswith(".TeardownExit") + assert completion.failureDetails.errorMessage == "deterministic teardown failure" + assert len(response.actions) == 1 + + def test_worker_created_before_configuration_hydrates_and_externalizes(monkeypatch, payload_store_factory): monkeypatch.setattr(payloads, "_payload_store", None) worker = DurableFunctionsWorker() diff --git a/tests/azure-functions-durable/typing/import_package.py b/tests/azure-functions-durable/typing/import_package.py index fe91818c..e9395042 100644 --- a/tests/azure-functions-durable/typing/import_package.py +++ b/tests/azure-functions-durable/typing/import_package.py @@ -7,7 +7,9 @@ def create_app() -> df.DFApp: - return df.DFApp() + app = df.DFApp() + app.configure_serialization(ensure_ascii=True) + return app def create_retry_policy() -> df.RetryPolicy: diff --git a/tests/durabletask-azuremanaged/test_azuremanaged_grpc_resiliency.py b/tests/durabletask-azuremanaged/test_azuremanaged_grpc_resiliency.py index 4a4d0f28..0a96eaf9 100644 --- a/tests/durabletask-azuremanaged/test_azuremanaged_grpc_resiliency.py +++ b/tests/durabletask-azuremanaged/test_azuremanaged_grpc_resiliency.py @@ -1,10 +1,13 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -from unittest.mock import patch +from unittest.mock import Mock, patch from collections.abc import Mapping from typing import Any +from durabletask import task +from durabletask.internal import helpers +from durabletask.internal import orchestrator_service_pb2 as pb from durabletask.azuremanaged.client import ( AsyncDurableTaskSchedulerClient, DurableTaskSchedulerClient, @@ -40,6 +43,34 @@ def test_dts_worker_passes_resiliency_options_to_base_worker(): assert mock_init.call_args.kwargs["resiliency_options"] is resiliency +def test_dts_worker_completes_unexpected_teardown_failure(): + class TeardownExit(BaseException): + pass + + def orchestrator(ctx: task.OrchestrationContext, _: Any): + try: + yield ctx.call_activity("work") + finally: + raise TeardownExit("deterministic teardown failure") + + with DurableTaskSchedulerWorker( + host_address="localhost:4001", taskhub="hub", token_credential=None) as worker: + name = worker.add_orchestrator(orchestrator) + request = pb.OrchestratorRequest(instanceId="instance", newEvents=[ + helpers.new_orchestrator_started_event(), + helpers.new_execution_started_event(name, "instance"), + ]) + stub = Mock() + worker._execute_orchestrator(request, stub, "original-token") + + stub.CompleteOrchestratorTask.assert_called_once() + stub.AbandonTaskOrchestratorWorkItem.assert_not_called() + response = stub.CompleteOrchestratorTask.call_args.args[0] + assert response.completionToken == "original-token" + assert response.actions[0].completeOrchestration.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + assert response.actions[0].completeOrchestration.failureDetails.errorMessage == "deterministic teardown failure" + + def test_dts_worker_passes_exception_properties_provider_to_base_worker(): class PropertiesProvider: def get_exception_properties(self, exception: Exception) -> Mapping[str, Any] | None: 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..2ac8ab01 --- /dev/null +++ b/tests/durabletask-azuremanaged/test_dts_coroutine_orchestrators.py @@ -0,0 +1,91 @@ +# 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 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"]) +def test_mixed_orchestrators_over_grpc(transport: str, coroutine_parent: bool) -> None: + def increment(ctx: task.ActivityContext, value: int) -> int: + return value + 1 + + async def coroutine_child(ctx: task.OrchestrationContext, value: int) -> 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]: + values = 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/entities/test_entity_id_parsing.py b/tests/durabletask/entities/test_entity_id_parsing.py index 77bb339e..1a9fb942 100644 --- a/tests/durabletask/entities/test_entity_id_parsing.py +++ b/tests/durabletask/entities/test_entity_id_parsing.py @@ -1,6 +1,9 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. +import copy +import pickle + import pytest from durabletask.entities import EntityInstanceId @@ -51,3 +54,129 @@ def test_entity_id_parsing_failures(): # Test multiple @ symbols with pytest.raises(ValueError): EntityInstanceId.parse("invalid@symbol@placement") + + +def test_entity_instance_id_hash_and_equality(): + id1 = EntityInstanceId("MyEntity", "instance1") + id2 = EntityInstanceId("myentity", "instance1") + id3 = EntityInstanceId("myentity", "instance2") + id4 = EntityInstanceId("other", "instance1") + + assert hash(id1) == hash(id2) + assert hash(id1) != hash(id3) + assert hash(id1) != hash(id4) + + +def test_entity_instance_id_in_set(): + id1 = EntityInstanceId("MyEntity", "key1") + id2 = EntityInstanceId("myentity", "key1") + id3 = EntityInstanceId("myentity", "key2") + + s = {id1, id2, id3} + assert len(s) == 2 + assert EntityInstanceId("myentity", "key1") in s + assert EntityInstanceId("myentity", "key2") in s + + +def test_entity_instance_id_as_dict_key(): + id1 = EntityInstanceId("Counter", "key1") + id2 = EntityInstanceId("counter", "key1") + + mapping = {id1: "value1"} + assert mapping[id2] == "value1" + mapping[id2] = "value2" + assert len(mapping) == 1 + assert mapping[id1] == "value2" + + +def test_entity_instance_id_repr(): + eid = EntityInstanceId("MyEntity", "InstanceKey") + expected = "EntityInstanceId(entity='myentity', key='InstanceKey')" + assert repr(eid) == expected + + +def test_entity_instance_id_immutable(): + """Verify that entity and key cannot be reassigned or deleted after construction, + ensuring hash stability when instances are stored in sets or dicts.""" + eid = EntityInstanceId("Counter", "key1") + + with pytest.raises(AttributeError): + eid.entity = "other" + + with pytest.raises(AttributeError): + eid.key = "other_key" + + with pytest.raises(AttributeError): + eid.custom_attr = "value" + + with pytest.raises(AttributeError): + del eid.entity + + with pytest.raises(AttributeError): + del eid.key + + # Verify the original values are unchanged + assert eid.entity == "counter" + assert eid.key == "key1" + + +def test_entity_instance_id_immutable_dict_stability(): + """Verify that immutability preserves dict/set lookup integrity.""" + eid = EntityInstanceId("Counter", "key1") + mapping = {eid: "value"} + s = {eid} + + # Attempt mutation - should be rejected + with pytest.raises(AttributeError): + eid.entity = "other" + + # Lookups with an equal instance must still succeed + lookup = EntityInstanceId("counter", "key1") + assert mapping[lookup] == "value" + assert lookup in s + + +def test_entity_instance_id_copy_and_pickle(): + """Verify that immutable EntityInstanceId instances can be copied and pickled.""" + eid = EntityInstanceId("Counter", "key1") + + # Copying should preserve identity / immutability + assert copy.copy(eid) is eid + assert copy.deepcopy(eid) is eid + + # Pickling should preserve equality and hash + pickled = pickle.loads(pickle.dumps(eid)) + assert pickled == eid + assert hash(pickled) == hash(eid) + assert pickled.entity == "counter" + assert pickled.key == "key1" + + +class _SubclassEntityId(EntityInstanceId): + @property + def custom_name(self) -> str: + return self.entity.upper() + + +def test_entity_instance_id_subclass_pickle(): + """Verify that unpickling a subclass of EntityInstanceId preserves the concrete type.""" + custom = _SubclassEntityId("Counter", "key1") + assert isinstance(custom, _SubclassEntityId) + assert custom.custom_name == "COUNTER" + + restored = pickle.loads(pickle.dumps(custom)) + assert type(restored) is _SubclassEntityId + assert isinstance(restored, _SubclassEntityId) + assert restored.custom_name == "COUNTER" + assert restored.entity == "counter" + assert restored.key == "key1" + assert restored == custom + assert hash(restored) == hash(custom) + + # Immutability guarantees remain intact after unpickling + with pytest.raises(AttributeError): + restored.entity = "other" + with pytest.raises(AttributeError): + restored.key = "other" + with pytest.raises(AttributeError): + del restored.entity diff --git a/tests/durabletask/test_asyncio_orchestrators.py b/tests/durabletask/test_asyncio_orchestrators.py new file mode 100644 index 00000000..552a618d --- /dev/null +++ b/tests/durabletask/test_asyncio_orchestrators.py @@ -0,0 +1,465 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Deterministic asyncio and AnyIO behavior through the real history executor.""" + +import asyncio +import gc +import json +import sys +from contextlib import asynccontextmanager +from datetime import timedelta +from typing import Any +from weakref import ref + +import anyio +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 _completion, _execute, _start, NOW + + +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 = await ctx.call_activity(name) + await asyncio.sleep(0) + return await ctx.call_activity(name + "-next", input=value) + + return 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"), 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") + + +@pytest.mark.parametrize("native", [False, True]) +def test_task_group_failure_cancels_waiters_not_durable_work(native: bool) -> None: + if native and sys.version_info < (3, 11): + pytest.skip("asyncio.TaskGroup requires Python 3.11") + cleaned: list[str] = [] + + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def branch(name: str) -> None: + try: + await ctx.call_activity(name) + finally: + with anyio.CancelScope(shield=True): + await asyncio.sleep(0) + cleaned.append(name) + + if native: + async with asyncio.TaskGroup() as group: + group.create_task(branch("fail")) + group.create_task(branch("slow")) + else: + async with anyio.create_task_group() as group: + group.start_soon(branch, "fail") + group.start_soon(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", "anyio", "timeout"]) +def test_durable_timer_and_timeout(kind: str) -> None: + if kind == "timeout" and sys.version_info < (3, 11): + pytest.skip("asyncio.timeout requires Python 3.11") + + 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) + elif kind == "timeout": + async with asyncio.timeout(2): + await ctx.wait_for_external_event("never") + else: + with anyio.fail_after(2): + await ctx.wait_for_external_event("never") + except (TimeoutError, asyncio.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] + + +@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 anyio.create_task_group(): + 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") + + 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. + with anyio.CancelScope(shield=True): + await ctx.call_activity("compensate", input=index) + + async with anyio.create_task_group() as group: + for index in range(4): + group.start_soon(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]) +def test_background_failure_is_reported_unless_observed( + observed: 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()) + await asyncio.sleep(0) + if observed: + await asyncio.gather(pending, return_exceptions=True) + 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 diff --git a/tests/durabletask/test_batch_actions.py b/tests/durabletask/test_batch_actions.py index 24989746..8b90e822 100644 --- a/tests/durabletask/test_batch_actions.py +++ b/tests/durabletask/test_batch_actions.py @@ -197,12 +197,36 @@ def test_orchestration_query_serializes_instance_id_prefix(): assert request.query.instanceIdPrefix.value == "prefix-" +def test_orchestration_query_serializes_task_hub_names(): + request = build_query_instances_req( + client.OrchestrationQuery(task_hub_names=["hub1", "hub2"]), + continuation_token=None, + ) + + assert [name.value for name in request.query.taskHubNames] == ["hub1", "hub2"] + + +def test_orchestration_query_serializes_task_hub_names_and_instance_id_prefix(): + request = build_query_instances_req( + client.OrchestrationQuery( + instance_id_prefix="prefix-", + task_hub_names=["hub1"], + ), + continuation_token=None, + ) + + assert request.query.HasField("instanceIdPrefix") + assert request.query.instanceIdPrefix.value == "prefix-" + assert [name.value for name in request.query.taskHubNames] == ["hub1"] + + def test_orchestration_query_preserves_positional_argument_order(): query = client.OrchestrationQuery(None, None, None, None, True) assert query.max_instance_count is None assert query.fetch_inputs_and_outputs is True assert query.instance_id_prefix is None + assert query.task_hub_names is None def test_get_orchestration_state_pagination_succeeds(backend): diff --git a/tests/durabletask/test_coroutine_orchestrators.py b/tests/durabletask/test_coroutine_orchestrators.py new file mode 100644 index 00000000..be9af370 --- /dev/null +++ b/tests/durabletask/test_coroutine_orchestrators.py @@ -0,0 +1,279 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Replay-level coverage for the async/await orchestration prototype.""" + +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] + + +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" + + +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_data_converter.py b/tests/durabletask/test_data_converter.py index 19909138..8186280f 100644 --- a/tests/durabletask/test_data_converter.py +++ b/tests/durabletask/test_data_converter.py @@ -12,13 +12,14 @@ from unittest.mock import patch import pytest - from durabletask.internal.entity_state_shim import StateShim from durabletask.serialization import ( DEFAULT_DATA_CONVERTER, DataConverter, JsonDataConverter, ) +from durabletask.internal import type_discovery +from durabletask.worker import _ActivityExecutor, _Registry @dataclass @@ -219,7 +220,7 @@ def __init__(self, /, **kwargs): self.self = kwargs.get("self", 0) converter = JsonDataConverter() - encoded = converter.serialize(PositionalReceiver(**{"self": 42})) + encoded = converter.serialize(PositionalReceiver(self=42)) restored = converter.deserialize(encoded, PositionalReceiver) assert isinstance(restored, PositionalReceiver) assert restored.self == 42 @@ -375,3 +376,102 @@ def coerce(self, value: Any, target_type: type | None = None) -> Any: assert conv.serialize("hello") == '"HELLO"' assert conv.deserialize('"HELLO"') == "HELLO" assert conv.coerce("HELLO") == "HELLO" + + +# ----- set and frozenset container support ----- + + +@dataclass(frozen=True) +class TaggedOrder: + item: str + tags: set[str] + flags: frozenset[int] + + +def test_roundtrip_dataclass_with_sets(): + converter = JsonDataConverter(enable_generic_sets=True) + original = TaggedOrder(item="laptop", tags={"electronics", "sale"}, flags=frozenset({1, 2})) + serialized = converter.serialize(original) + assert serialized is not None + restored = converter.deserialize(serialized, TaggedOrder) + assert isinstance(restored, TaggedOrder) + assert restored.item == "laptop" + assert isinstance(restored.tags, set) + assert restored.tags == {"electronics", "sale"} + assert isinstance(restored.flags, frozenset) + assert restored.flags == frozenset({1, 2}) + + +def test_can_reconstruct_set_and_frozenset(): + # By default (enable_generic_sets=False), generic sets are not reconstructed to + # preserve backward compatibility and existing-instance replay behavior. + default_converter = JsonDataConverter() + assert default_converter.can_reconstruct(set[PricedOrder]) is False + assert default_converter.can_reconstruct(frozenset[PricedOrder]) is False + assert default_converter.can_reconstruct(set[int]) is False + assert default_converter.can_reconstruct(set[str]) is False + + # When opted in, supported sets (including primitive elements) return True, + # while unsupported unhashable element types return False. + opt_in_converter = JsonDataConverter(enable_generic_sets=True) + assert opt_in_converter.can_reconstruct(set[PricedOrder]) is True + assert opt_in_converter.can_reconstruct(frozenset[PricedOrder]) is True + assert opt_in_converter.can_reconstruct(set[int]) is True + assert opt_in_converter.can_reconstruct(set[str]) is True + assert opt_in_converter.can_reconstruct(frozenset[int]) is True + assert opt_in_converter.can_reconstruct(set[list[str]]) is False + assert opt_in_converter.can_reconstruct(set[dict[str, int]]) is False + + +def test_existing_instance_replay_compatibility_preserves_list(): + """Preserve existing-instance replay behavior for recorded arrays under set annotations. + + When generic set coercion is not explicitly enabled, historical JSON arrays + containing duplicates and order (e.g. {"tags": ["b", "a", "b"]}) remain raw + lists, avoiding NonDeterminismError during orchestration replay. + """ + default_converter = JsonDataConverter() + recorded_input = '{"item": "laptop", "tags": ["b", "a", "b"], "flags": [1, 2, 1]}' + restored = default_converter.deserialize(recorded_input, TaggedOrder) + assert isinstance(restored, TaggedOrder) + assert isinstance(restored.tags, list) + assert restored.tags == ["b", "a", "b"] + assert len(restored.tags) == 3 + assert isinstance(restored.flags, list) + assert restored.flags == [1, 2, 1] + assert len(restored.flags) == 3 + + # With opt-in enabled, set coercion deduplicates and converts to sets + opt_in_converter = JsonDataConverter(enable_generic_sets=True) + opt_in_restored = opt_in_converter.deserialize(recorded_input, TaggedOrder) + assert isinstance(opt_in_restored, TaggedOrder) + assert isinstance(opt_in_restored.tags, set) + assert opt_in_restored.tags == {"a", "b"} + assert len(opt_in_restored.tags) == 2 + assert isinstance(opt_in_restored.flags, frozenset) + assert opt_in_restored.flags == frozenset({1, 2}) + assert len(opt_in_restored.flags) == 2 + + +def test_type_discovery_and_activity_execution_with_sets(): + """Verify that type discovery recognizes set annotations when opted in, + enabling activities to receive sets rather than raw lists.""" + def tag_activity(ctx: Any, tags: set[str]) -> int: + assert isinstance(tags, set) + tags.add("processed") + return len(tags) + + reg = _Registry() + reg.add_activity(tag_activity) + logger = logging.getLogger("test_type_discovery_sets") + + # With opt-in, activity input is discovered as set[str] and executed + conv = JsonDataConverter(enable_generic_sets=True) + assert type_discovery.activity_input_type(tag_activity, conv) == set[str] + executor = _ActivityExecutor(reg, logger, conv) + output = executor.execute("orch-1", "tag_activity", 1, '["item1", "item2"]') + assert output == "3" + + # Default converter preserves input as None (raw JSON), passing list + default_conv = JsonDataConverter() + assert type_discovery.activity_input_type(tag_activity, default_conv) is None diff --git a/tests/durabletask/test_data_converter_roundtrip.py b/tests/durabletask/test_data_converter_roundtrip.py index 465aeb3a..804f81bf 100644 --- a/tests/durabletask/test_data_converter_roundtrip.py +++ b/tests/durabletask/test_data_converter_roundtrip.py @@ -545,8 +545,6 @@ class HasUnion: @pytest.mark.parametrize("value", [ datetime(2020, 1, 1, 12, 0, 0), Decimal("1.5"), - {1, 2, 3}, - frozenset({1, 2}), b"bytes", ]) def test_unsupported_types_raise_clear_typeerror(conv, value): diff --git a/tests/durabletask/test_large_payload.py b/tests/durabletask/test_large_payload.py index dc0f32d5..ace3af61 100644 --- a/tests/durabletask/test_large_payload.py +++ b/tests/durabletask/test_large_payload.py @@ -12,6 +12,7 @@ from google.protobuf import wrappers_pb2 import durabletask.internal.orchestrator_service_pb2 as pb +from durabletask.internal.exceptions import PayloadSizeExceededError from durabletask.payload.helpers import ( deexternalize_payloads, deexternalize_payloads_async, @@ -82,6 +83,16 @@ def sv(text: str) -> wrappers_pb2.StringValue: return wrappers_pb2.StringValue(value=text) +@pytest.fixture(params=["sync", "async"]) +def externalize(request): + def run(msg, store, *, instance_id=None): + if request.param == "async": + asyncio.run(externalize_payloads_async(msg, store, instance_id=instance_id)) + else: + externalize_payloads(msg, store, instance_id=instance_id) + return run + + # ------------------------------------------------------------------ # Tests: externalize_payloads # ------------------------------------------------------------------ @@ -142,8 +153,9 @@ def test_exceeds_max_raises_error(self): name="MyOrch", input=sv(large_data), ) - with pytest.raises(ValueError, match="exceeds the maximum"): + with pytest.raises(PayloadSizeExceededError, match="exceeds the maximum") as failure: externalize_payloads(req, store, instance_id="test-1") + assert isinstance(failure.value, ValueError) def test_empty_value_not_externalized(self): """Empty StringValue should not be externalized.""" @@ -237,6 +249,178 @@ def test_activity_response_externalized(self): assert res.result.value.startswith(FakePayloadStore.TOKEN_PREFIX) +class TestExternalizePreflight: + @pytest.mark.parametrize("message_kind", ["actions", "carryover", "entity-results", "entity-state"]) + @pytest.mark.parametrize("upload_fails", [False, True]) + def test_later_oversized_payload_has_no_side_effects(self, externalize, message_kind, upload_fails): + store = FakePayloadStore(threshold_bytes=8, max_stored_payload_bytes=32) + if message_kind == "actions": + response = pb.OrchestratorResponse( + instanceId="size-preflight", + actions=[ + pb.OrchestratorAction(id=1, scheduleTask=pb.ScheduleTaskAction(input=sv("a" * 20))), + pb.OrchestratorAction(id=2, scheduleTask=pb.ScheduleTaskAction(input=sv("b" * 33))), + ], + ) + elif message_kind == "carryover": + response = pb.OrchestratorResponse( + instanceId="size-preflight", + actions=[pb.OrchestratorAction( + id=1, + completeOrchestration=pb.CompleteOrchestrationAction( + result=sv("a" * 20), + carryoverEvents=[pb.HistoryEvent( + eventId=2, + eventRaised=pb.EventRaisedEvent(name="event", input=sv("b" * 33)), + )], + ), + )], + ) + elif message_kind == "entity-results": + response = pb.EntityBatchResult(results=[ + pb.OperationResult(success=pb.OperationResultSuccess(result=sv("a" * 20))), + pb.OperationResult(success=pb.OperationResultSuccess(result=sv("b" * 33))), + ]) + else: + response = pb.EntityBatchResult( + results=[pb.OperationResult(success=pb.OperationResultSuccess(result=sv("a" * 20)))], + actions=[pb.OperationAction(sendSignal=pb.SendSignalAction(input=sv("a" * 20)))], + entityState=sv("b" * 33), + ) + before = response.SerializeToString() + error = OSError("must not upload") if upload_fails else None + store.upload = MagicMock(wraps=store.upload, side_effect=error) + store.upload_async = AsyncMock(wraps=store.upload_async, side_effect=error) + + with pytest.raises(ValueError, match="Payload size 33 bytes exceeds the maximum allowed size of 32 bytes"): + externalize(response, store, instance_id="size-preflight") + + store.upload.assert_not_called() + store.upload_async.assert_not_called() + assert not store._blobs + assert response.SerializeToString() == before + + @pytest.mark.parametrize("payload", ["a" * 33, "\u00e9" * 16 + "a", "\U0001f600" * 8 + "a"]) + def test_utf8_maximum_uses_bytes(self, externalize, payload): + store = FakePayloadStore(threshold_bytes=8, max_stored_payload_bytes=32) + response = pb.OrchestrationState(input=sv("a" * 20), output=sv(payload)) + before = response.SerializeToString() + + with pytest.raises(ValueError, match="Payload size 33 bytes exceeds the maximum"): + externalize(response, store) + + assert not store._blobs + assert response.SerializeToString() == before + + @pytest.mark.parametrize("payload", [ + "a" * 7, "a" * 8, "a" * 9, "a" * 32, + "\u00e9" * 4, "\u00e9" * 4 + "a", "\u00e9" * 16, "\U0001f600" * 8, + ]) + def test_valid_utf8_threshold_and_maximum_boundaries(self, externalize, payload): + store = FakePayloadStore(threshold_bytes=8, max_stored_payload_bytes=32) + response = pb.ActivityResponse(result=sv(payload)) + + externalize(response, store) + + if len(payload.encode("utf-8")) > 8: + assert store._blobs[response.result.value] == payload.encode("utf-8") + assert len(store._blobs) == 1 + else: + assert response.result.value == payload + assert not store._blobs + + def test_all_valid_payloads_preserve_order_tokens_and_presence(self, externalize): + store = FakePayloadStore(threshold_bytes=8, max_stored_payload_bytes=32) + token = f"{FakePayloadStore.TOKEN_PREFIX}existing-blob" + response = pb.OrchestratorResponse(actions=[ + pb.OrchestratorAction(id=1, scheduleTask=pb.ScheduleTaskAction(input=sv("a" * 20))), + pb.OrchestratorAction(id=2, scheduleTask=pb.ScheduleTaskAction(input=sv("\u00e9" * 16))), + pb.OrchestratorAction(id=3, scheduleTask=pb.ScheduleTaskAction(input=sv(token))), + pb.OrchestratorAction(id=4, scheduleTask=pb.ScheduleTaskAction(input=sv(""))), + pb.OrchestratorAction(id=5, scheduleTask=pb.ScheduleTaskAction()), + ]) + upload = store.upload = MagicMock(wraps=store.upload) + + externalize(response, store, instance_id="size-preflight") + + assert list(store._blobs.values()) == [b"a" * 20, ("\u00e9" * 16).encode("utf-8")] + assert [call.kwargs for call in upload.call_args_list] == [ + {"instance_id": "size-preflight"}, {"instance_id": "size-preflight"}, + ] + assert response.actions[2].scheduleTask.input.value == token + assert response.actions[3].scheduleTask.HasField("input") + assert response.actions[3].scheduleTask.input.value == "" + assert not response.actions[4].scheduleTask.HasField("input") + before = response.SerializeToString() + externalize(response, store) + assert upload.call_count == 2 + assert response.SerializeToString() == before + + def test_inline_payload_above_maximum_is_still_ignored(self, externalize): + store = FakePayloadStore(threshold_bytes=32, max_stored_payload_bytes=8) + response = pb.ActivityResponse(result=sv("a" * 20)) + before = response.SerializeToString() + + externalize(response, store) + + assert not store._blobs + assert response.SerializeToString() == before + + def test_uploads_use_snapshots(self, externalize): + store = FakePayloadStore(threshold_bytes=8, max_stored_payload_bytes=32) + response = pb.OrchestrationState(input=sv("a" * 20), output=sv("b" * 20)) + upload = store.upload + + def mutate_later_field(data, *, instance_id=None): + response.output.value = "changed during upload" + return upload(data, instance_id=instance_id) + + store.upload = MagicMock(side_effect=mutate_later_field) + + externalize(response, store) + + assert list(store._blobs.values()) == [b"a" * 20, b"b" * 20] + assert store._blobs[response.output.value] == b"b" * 20 + + @pytest.mark.parametrize("error", [OSError("storage failure"), ValueError("custom store failure")]) + def test_upload_errors_propagate_without_rollback(self, externalize, error): + store = FakePayloadStore(threshold_bytes=8, max_stored_payload_bytes=32) + response = pb.OrchestrationState(input=sv("a" * 20), output=sv("b" * 20)) + upload = store.upload + + def upload_or_fail(data, *, instance_id=None): + if store._blobs: + raise error + return upload(data, instance_id=instance_id) + + store.upload = MagicMock(side_effect=upload_or_fail) + + with pytest.raises(type(error)) as raised: + externalize(response, store) + + assert raised.value is error + assert len(store._blobs) == 1 + assert store.is_known_token(response.input.value) + assert response.output.value == "b" * 20 + assert store.upload.call_count == 2 + + @pytest.mark.asyncio + async def test_async_cancellation_propagates_without_later_uploads(self): + store = FakePayloadStore(threshold_bytes=8, max_stored_payload_bytes=32) + response = pb.OrchestrationState(input=sv("a" * 20), output=sv("b" * 20)) + before = response.SerializeToString() + error = asyncio.CancelledError("cancelled upload") + store.upload_async = AsyncMock(side_effect=error) + + with pytest.raises(asyncio.CancelledError) as raised: + await externalize_payloads_async(response, store) + + assert raised.value is error + assert store.upload_async.await_count == 1 + assert not store._blobs + assert response.SerializeToString() == before + + # ------------------------------------------------------------------ # Tests: deexternalize_payloads # ------------------------------------------------------------------ @@ -358,6 +542,40 @@ async def test_async_externalize_and_deexternalize(self): await deexternalize_payloads_async(req, store) assert req.input.value == original + @pytest.mark.asyncio + async def test_async_cap_failure_is_a_value_error(self): + store = FakePayloadStore(threshold_bytes=1, max_stored_payload_bytes=10) + req = pb.CreateInstanceRequest(input=sv("x" * 11)) + with pytest.raises(PayloadSizeExceededError) as failure: + await externalize_payloads_async(req, store) + assert isinstance(failure.value, ValueError) + assert not store._blobs + + @pytest.mark.asyncio + @pytest.mark.parametrize("use_async", [False, True]) + async def test_token_larger_than_cap_is_not_reuploaded(self, use_async): + store = FakePayloadStore(threshold_bytes=1, max_stored_payload_bytes=10) + token = f"{FakePayloadStore.TOKEN_PREFIX}existing-blob" + assert len(token.encode("utf-8")) > store.options.max_stored_payload_bytes + req = pb.CreateInstanceRequest(input=sv(token)) + if use_async: + await externalize_payloads_async(req, store) + else: + externalize_payloads(req, store) + assert req.input.value == token + assert not store._blobs + + @pytest.mark.asyncio + @pytest.mark.parametrize("error", [OSError("unavailable"), ValueError("store error"), asyncio.CancelledError()]) + async def test_async_store_error_propagates_unchanged(self, error): + store = FakePayloadStore(threshold_bytes=1) + store.upload_async = AsyncMock(side_effect=error) + req = pb.CreateInstanceRequest(input=sv("large enough")) + with pytest.raises(type(error)) as failure: + await externalize_payloads_async(req, store) + assert failure.value is error + assert req.input.value == "large enough" + # ------------------------------------------------------------------ # Tests: worker capability flag diff --git a/tests/durabletask/test_orchestration_dispatch_recovery.py b/tests/durabletask/test_orchestration_dispatch_recovery.py new file mode 100644 index 00000000..4d4f7ee6 --- /dev/null +++ b/tests/durabletask/test_orchestration_dispatch_recovery.py @@ -0,0 +1,691 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +# pyright: reportPrivateUsage=false + +import asyncio +import json +import logging +import sys +from threading import Event +from typing import Any, Generator +from unittest.mock import Mock + +import pytest + +from durabletask import client, task, worker +from durabletask.entities import EntityInstanceId +from durabletask.internal import exceptions, helpers +from durabletask.internal import orchestrator_service_pb2 as pb +from durabletask.serialization import JsonDataConverter +from durabletask.testing import create_test_backend +from tests.durabletask._port_utils import find_free_port + +INSTANCE_ID = "dispatch-recovery" +LOGGER = logging.getLogger("tests.dispatch_recovery") + + +class TeardownExit(BaseException): + pass + + +class _FailingOutput: + def __init__(self, error: BaseException): + self.error = error + + def to_json(self) -> str: + raise self.error + + +def _request() -> pb.OrchestratorRequest: + return pb.OrchestratorRequest(instanceId=INSTANCE_ID, newEvents=[ + helpers.new_orchestrator_started_event(), + helpers.new_execution_started_event("orchestrator", INSTANCE_ID), + ]) + + +def _executor(orchestrator: task.Orchestrator[Any, Any]) -> worker._OrchestrationExecutor: + registry = worker._Registry() + registry.add_named_orchestrator("orchestrator", orchestrator) + return worker._OrchestrationExecutor(registry, LOGGER, JsonDataConverter()) + + +@pytest.mark.parametrize("during_teardown", [False, True]) +def test_unexpected_base_exception_fails_orchestration_without_reprocessing( + during_teardown: bool, caplog: pytest.LogCaptureFixture): + error = TeardownExit("deterministic dispatch failure") + + def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Generator[task.Task[Any], Any, None]: + if not during_teardown: + raise error + try: + ctx.set_custom_status("waiting") + yield ctx.call_activity("work") + finally: + raise error + + with worker.TaskHubGrpcWorker(logger=LOGGER) as host: + host._registry = _executor(orchestrator)._registry + stub = Mock() + reset = Mock() + with caplog.at_level(logging.ERROR, logger=LOGGER.name): + host._execute_orchestrator( + _request(), stub, "original-token", _on_outcome_uncertain=reset) + + stub.CompleteOrchestratorTask.assert_called_once() + stub.AbandonTaskOrchestratorWorkItem.assert_not_called() + reset.assert_not_called() + response: pb.OrchestratorResponse = stub.CompleteOrchestratorTask.call_args.args[0] + assert response.instanceId == INSTANCE_ID + assert response.completionToken == "original-token" + completion = next(a.completeOrchestration for a in response.actions if a.HasField("completeOrchestration")) + assert completion.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + assert completion.failureDetails.errorType.endswith(".TeardownExit") + assert completion.failureDetails.errorMessage == str(error) + assert "raise error" in completion.failureDetails.stackTrace.value + if during_teardown: + assert len(response.actions) == 1 + assert not response.HasField("customStatus") + assert any(record.exc_info and record.exc_info[1] is error for record in caplog.records) + + +@pytest.mark.parametrize("as_generator", [False, True]) +@pytest.mark.parametrize("error", [ + TeardownExit("output serialization exited"), + RuntimeError("output serialization failed"), + StopIteration("output serialization stopped"), + ValueError("invalid output"), + TypeError("invalid output type"), +]) +def test_output_serialization_exit_submits_failed_completion(as_generator: bool, error: BaseException): + output = _FailingOutput(error) + + def episode(ctx: task.OrchestrationContext) -> Generator[task.Task[Any], Any, _FailingOutput]: + yield ctx.wait_for_external_event("ready") + return output + + def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any: + return episode(ctx) if as_generator else output + + request = _request() + request.newEvents.append(helpers.new_event_raised_event("ready", None)) + with worker.TaskHubGrpcWorker(logger=LOGGER) as host: + host._registry = _executor(orchestrator)._registry + stub = Mock() + reset = Mock() + host._execute_orchestrator(request, stub, "original-token", _on_outcome_uncertain=reset) + + stub.CompleteOrchestratorTask.assert_called_once() + stub.AbandonTaskOrchestratorWorkItem.assert_not_called() + reset.assert_not_called() + response: pb.OrchestratorResponse = stub.CompleteOrchestratorTask.call_args.args[0] + assert response.instanceId == INSTANCE_ID + assert response.completionToken == "original-token" + assert len(response.actions) == 1 + completion = response.actions[0].completeOrchestration + assert completion.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + if isinstance(error, (TypeError, ValueError)): + assert completion.failureDetails.errorType.endswith(".JsonEncodeOutputException") + else: + assert completion.failureDetails.errorType.endswith(f".{type(error).__name__}") + assert completion.failureDetails.errorMessage == str(error) + assert "raise self.error" in completion.failureDetails.stackTrace.value + + +@pytest.mark.parametrize("error", [ + None, + ValueError("invalid output"), + TypeError("invalid output type"), + TeardownExit("output serialization exited"), +]) +def test_output_serialization_cannot_reenter_terminal_transition(error: BaseException | None): + def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any: + class Output: + def to_json(self) -> str: + ctx.continue_as_new("replacement") + if error is not None: + raise error + return "original" + + return Output() + + result = _executor(orchestrator).execute(INSTANCE_ID, [], _request().newEvents) + + assert len(result.actions) == 1 + completion = result.actions[0].completeOrchestration + if error is None: + assert completion.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED + assert completion.result.value == '"original"' + else: + assert completion.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + if isinstance(error, (ValueError, TypeError)): + assert completion.failureDetails.errorType.endswith(".JsonEncodeOutputException") + else: + assert completion.failureDetails.errorMessage == str(error) + + +def test_output_serialization_exit_preserves_entity_lock_release(): + entity_id = EntityInstanceId("Counter", "key") + error = TeardownExit("output serialization exited with a lock") + + def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Generator[task.Task[Any], Any, _FailingOutput]: + yield ctx.lock_entities([entity_id]) + return _FailingOutput(error) + + executor = _executor(orchestrator) + request = _request() + initial = executor.execute(INSTANCE_ID, [], request.newEvents) + lock_id = initial.actions[0].sendEntityMessage.entityLockRequested.criticalSectionId + result = executor.execute(INSTANCE_ID, request.newEvents, [ + helpers.new_event_sent_event(1, str(entity_id), json.dumps({"id": lock_id})), + helpers.new_event_raised_event(lock_id, None), + ]) + + assert len(result.actions) == 2 + release = result.actions[0].sendEntityMessage.entityUnlockSent + assert release.criticalSectionId == lock_id + assert release.targetInstanceId.value == str(entity_id) + completion = result.actions[1].completeOrchestration + assert completion.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + assert completion.failureDetails.errorMessage == str(error) + + +@pytest.mark.parametrize("fail_during_teardown", [False, True]) +def test_terminal_failure_after_pending_call_timeout_preserves_missing_release(fail_during_teardown: bool): + entity_a = EntityInstanceId("Counter", "a") + entity_b = EntityInstanceId("Counter", "b") + error = TeardownExit("terminal failure after pending call timeout") + + def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Generator[task.Task[Any], Any, Any]: + after_timeout = False + try: + with (yield ctx.lock_entities([entity_a, entity_b])): + call = ctx.call_entity(entity_a, "get") + yield task.when_any([call, ctx.wait_for_external_event("timeout")]) + after_timeout = True + yield ctx.call_activity("after_timeout") + if fail_during_teardown: + yield ctx.wait_for_external_event("never") + return _FailingOutput(error) + finally: + if fail_during_teardown and after_timeout: + raise error + + executor = _executor(orchestrator) + started = list(_request().newEvents) + initial = executor.execute(INSTANCE_ID, [], started) + lock_id = initial.actions[0].sendEntityMessage.entityLockRequested.criticalSectionId + granted = [ + helpers.new_event_sent_event(1, str(entity_a), json.dumps({"id": lock_id})), + helpers.new_event_raised_event(lock_id, None), + ] + calling = executor.execute(INSTANCE_ID, started, granted) + call_id = calling.actions[0].sendEntityMessage.entityOperationCalled.requestId + history = started + granted + [ + helpers.new_event_sent_event(2, str(entity_a), json.dumps({"id": call_id})), + helpers.new_event_raised_event("timeout", None), + helpers.new_event_sent_event(3, str(entity_b), None), + helpers.new_task_scheduled_event(4, "after_timeout"), + ] + result = executor.execute(INSTANCE_ID, history, [helpers.new_task_completed_event(4, '"done"')]) + + assert len(result.actions) == 2 + release = result.actions[0].sendEntityMessage.entityUnlockSent + assert release.targetInstanceId.value == str(entity_a) + assert release.criticalSectionId == lock_id + completion = result.actions[1].completeOrchestration + assert completion.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + assert completion.failureDetails.errorMessage == str(error) + + +@pytest.mark.parametrize("during_serialization", [False, True]) +def test_unexpected_exit_commits_failure_and_allows_next_dispatch(during_serialization: bool): + def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Generator[task.Task[Any], Any, None]: + try: + yield ctx.call_activity("unused") + finally: + raise TeardownExit("deterministic teardown failure") + + def unserializable(ctx: task.OrchestrationContext, _: Any) -> _FailingOutput: + return _FailingOutput(TeardownExit("output serialization exited")) + + def healthy(ctx: task.OrchestrationContext, _: Any) -> str: + return "healthy" + + port = find_free_port() + address = f"localhost:{port}" + backend = create_test_backend(port=port) + try: + with worker.TaskHubGrpcWorker( + host_address=address, + concurrency_options=worker.ConcurrencyOptions( + maximum_concurrent_orchestration_work_items=1, + maximum_thread_pool_workers=1, + )) as host: + host.add_orchestrator(orchestrator) + host.add_orchestrator(unserializable) + host.add_orchestrator(healthy) + host.start() + with client.TaskHubGrpcClient(host_address=address) as connection: + failed_id = connection.schedule_new_orchestration( + unserializable if during_serialization else orchestrator) + failed = connection.wait_for_orchestration_completion(failed_id, timeout=10) + assert failed is not None + assert failed.runtime_status == client.OrchestrationStatus.FAILED + assert failed.failure_details is not None + assert failed.failure_details.error_type.endswith(".TeardownExit") + assert failed.failure_details.message == ( + "output serialization exited" if during_serialization else "deterministic teardown failure") + next_id = connection.schedule_new_orchestration(healthy) + completed = connection.wait_for_orchestration_completion(next_id, timeout=10) + assert completed is not None + assert completed.runtime_status == client.OrchestrationStatus.COMPLETED + assert completed.serialized_output == '"healthy"' + finally: + backend.stop() + + +def test_unexpected_cleanup_does_not_replace_existing_failure(): + def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Generator[task.Task[Any], Any, None]: + try: + yield ctx.call_activity("expected") + finally: + raise TeardownExit("secondary cleanup failure") + + request = _request() + request.newEvents.append(helpers.new_task_scheduled_event(1, "wrong")) + result = _executor(orchestrator).execute(INSTANCE_ID, [], request.newEvents) + + completion = next(a.completeOrchestration for a in result.actions if a.HasField("completeOrchestration")) + assert completion.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + assert "NonDeterminismError" in completion.failureDetails.errorType + assert "secondary cleanup failure" not in completion.failureDetails.errorMessage + + +@pytest.mark.parametrize("wait", ["locked-event", "released-event", "entity-call"]) +def test_unexpected_teardown_failure_preserves_entity_lock_releases(wait: str): + entity_id = EntityInstanceId("Counter", "key") + + def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Generator[task.Task[Any], Any, None]: + acquired = False + try: + with (yield ctx.lock_entities([entity_id])): + acquired = True + if wait == "entity-call": + yield ctx.call_entity(entity_id, "get") + elif wait == "locked-event": + yield ctx.wait_for_external_event("ready") + yield ctx.wait_for_external_event("ready") + finally: + if acquired: + ctx.set_custom_status("discarded") + raise TeardownExit("cleanup failed with acquired lock") + + executor = _executor(orchestrator) + request = _request() + initial = executor.execute(INSTANCE_ID, [], request.newEvents) + lock_id = initial.actions[0].sendEntityMessage.entityLockRequested.criticalSectionId + result = executor.execute(INSTANCE_ID, request.newEvents, [ + helpers.new_event_sent_event(1, str(entity_id), json.dumps({"id": lock_id})), + helpers.new_event_raised_event(lock_id, None), + ]) + + assert len(result.actions) == 2 + release = result.actions[0].sendEntityMessage.entityUnlockSent + assert release.criticalSectionId == lock_id + assert release.targetInstanceId.value == str(entity_id) + assert release.parentInstanceId.value == INSTANCE_ID + assert result.actions[1].completeOrchestration.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + assert result.actions[0].id < result.actions[1].id + assert result.encoded_custom_status is None + + +@pytest.mark.parametrize("call_completed", [False, True]) +def test_teardown_failure_preserves_compatible_lock_release_order(call_completed: bool): + entity_a = EntityInstanceId("Counter", "a") + entity_b = EntityInstanceId("Counter", "b") + + def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Generator[task.Task[Any], Any, None]: + fail_teardown = False + try: + with (yield ctx.lock_entities([entity_a, entity_b])): + fail_teardown = not call_completed + yield ctx.call_entity(entity_a, "get") + fail_teardown = True + yield ctx.wait_for_external_event("ready") + finally: + if fail_teardown: + raise TeardownExit("cleanup with entity locks") + + executor = _executor(orchestrator) + request = _request() + initial = executor.execute(INSTANCE_ID, [], request.newEvents) + lock_id = initial.actions[0].sendEntityMessage.entityLockRequested.criticalSectionId + lock_events = [ + helpers.new_event_sent_event(1, str(entity_a), json.dumps({"id": lock_id})), + helpers.new_event_raised_event(lock_id, None), + ] + calling = executor.execute(INSTANCE_ID, request.newEvents, lock_events) + if call_completed: + call_id = calling.actions[0].sendEntityMessage.entityOperationCalled.requestId + history = list(request.newEvents) + lock_events + [ + helpers.new_event_sent_event(2, str(entity_a), json.dumps({"id": call_id})), + ] + result = executor.execute(INSTANCE_ID, history, [ + helpers.new_event_raised_event(call_id, json.dumps({"result": "42"})), + ]) + else: + result = calling + + assert [action.sendEntityMessage.entityUnlockSent.targetInstanceId.value + for action in result.actions[:-1]] == [str(entity_b), str(entity_a)] + assert result.actions[-1].completeOrchestration.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + + +def test_cleanup_does_not_replace_pending_control_flow_error(monkeypatch: pytest.MonkeyPatch): + primary = asyncio.CancelledError("original cancellation") + + def orchestrator(ctx: task.OrchestrationContext, _: Any): + try: + yield ctx.wait_for_external_event("ready") + finally: + raise TeardownExit("secondary cleanup failure") + + executor = _executor(orchestrator) + process_event = executor.process_event + + def process_then_cancel(ctx: worker._RuntimeOrchestrationContext, event: pb.HistoryEvent): + process_event(ctx, event) + if event.HasField("executionStarted"): + raise primary + + monkeypatch.setattr(executor, "process_event", process_then_cancel) + with pytest.raises(asyncio.CancelledError) as caught: + executor.execute(INSTANCE_ID, [], _request().newEvents) + assert caught.value is primary + + +@pytest.mark.parametrize("error", [TeardownExit("unexpected"), SystemExit("stop"), asyncio.CancelledError("cancel")]) +def test_caller_exception_state_does_not_suppress_teardown_outcome(error: BaseException): + def orchestrator(ctx: task.OrchestrationContext, _: Any): + try: + yield ctx.wait_for_external_event("ready") + finally: + raise error + + try: + raise ValueError("unrelated caller exception") + except ValueError: + executor = _executor(orchestrator) + if isinstance(error, SystemExit): + with pytest.raises(SystemExit) as caught: + executor.execute(INSTANCE_ID, [], _request().newEvents) + assert caught.value is error + else: + result = executor.execute(INSTANCE_ID, [], _request().newEvents) + if isinstance(error, TeardownExit): + assert result.actions[0].completeOrchestration.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + else: + assert result.actions == [] + + +def _control_flow_errors() -> list[BaseException]: + errors: list[BaseException] = [ + asyncio.CancelledError("cancel"), + SystemExit("exit"), + KeyboardInterrupt("interrupt"), + GeneratorExit("close"), + ] + if sys.version_info >= (3, 11): + errors.extend([ + BaseExceptionGroup("grouped cancellation", [asyncio.CancelledError("cancel")]), + BaseExceptionGroup("mixed", [ValueError("ordinary"), asyncio.CancelledError("cancel")]), + BaseExceptionGroup("nested", [ + TeardownExit("unexpected"), + BaseExceptionGroup("inner", [SystemExit("exit")]), + ]), + ]) + return errors + + +@pytest.mark.parametrize("error", _control_flow_errors(), ids=str) +def test_output_serialization_control_flow_propagates_after_abandonment(error: BaseException): + def orchestrator(ctx: task.OrchestrationContext, _: Any) -> _FailingOutput: + return _FailingOutput(error) + + with worker.TaskHubGrpcWorker(logger=LOGGER) as host: + host._registry = _executor(orchestrator)._registry + stub = Mock() + reset = Mock() + with pytest.raises(type(error)) as caught: + host._execute_orchestrator( + _request(), stub, "original-token", _on_outcome_uncertain=reset) + + assert caught.value is error + stub.CompleteOrchestratorTask.assert_not_called() + stub.AbandonTaskOrchestratorWorkItem.assert_called_once_with( + pb.AbandonOrchestrationTaskRequest(completionToken="original-token")) + reset.assert_not_called() + + +@pytest.mark.parametrize("error", _control_flow_errors(), ids=str) +@pytest.mark.parametrize("during_teardown", [False, True]) +def test_control_flow_propagates_after_same_token_abandonment( + error: BaseException, during_teardown: bool): + if during_teardown and isinstance(error, (asyncio.CancelledError, GeneratorExit)): + pytest.skip("Direct teardown cancellation and GeneratorExit retain their existing cleanup semantics.") + + def orchestrator(ctx: task.OrchestrationContext, _: Any): + if not during_teardown: + raise error + try: + yield ctx.wait_for_external_event("ready") + finally: + raise error + + with worker.TaskHubGrpcWorker(logger=LOGGER) as host: + host._registry = _executor(orchestrator)._registry + stub = Mock() + reset = Mock() + with pytest.raises(type(error)) as caught: + host._execute_orchestrator( + _request(), stub, "original-token", _on_outcome_uncertain=reset) + + assert caught.value is error + stub.CompleteOrchestratorTask.assert_not_called() + stub.AbandonTaskOrchestratorWorkItem.assert_called_once_with( + pb.AbandonOrchestrationTaskRequest(completionToken="original-token")) + reset.assert_not_called() + + +@pytest.mark.parametrize("stage", ["prepare", "complete", "abandon"]) +@pytest.mark.parametrize("error", [RuntimeError("rpc failed"), asyncio.CancelledError("rpc cancelled")]) +def test_rpc_uncertainty_resets_stream_without_duplicate_outcomes( + stage: str, error: BaseException, monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture): + with worker.TaskHubGrpcWorker(logger=LOGGER) as host: + prepare = Mock(return_value=pb.OrchestratorResponse( + instanceId=INSTANCE_ID, completionToken="original-token")) + monkeypatch.setattr(host, "_get_orchestrator_response", prepare) + stub = Mock() + reset = Mock() + if stage == "prepare": + prepare.side_effect = error + stub.AbandonTaskOrchestratorWorkItem.side_effect = RuntimeError("abandon failed") + elif stage == "complete": + stub.CompleteOrchestratorTask.side_effect = error + else: + prepare.side_effect = exceptions.AbandonOrchestrationError() + stub.AbandonTaskOrchestratorWorkItem.side_effect = error + + with caplog.at_level(logging.ERROR, logger=LOGGER.name): + if isinstance(error, Exception): + host._execute_orchestrator( + _request(), stub, "original-token", _on_outcome_uncertain=reset) + else: + with pytest.raises(type(error)) as caught: + host._execute_orchestrator( + _request(), stub, "original-token", _on_outcome_uncertain=reset) + assert caught.value is error + + assert stub.CompleteOrchestratorTask.call_count == (stage == "complete") + assert stub.AbandonTaskOrchestratorWorkItem.call_count == (stage != "complete") + if stage != "complete": + assert stub.AbandonTaskOrchestratorWorkItem.call_args.args[0].completionToken == "original-token" + reset.assert_called_once_with() + assert any("uncertain outcome" in record.message for record in caplog.records) + assert any(record.exc_info and record.exc_info[1] is error for record in caplog.records) + + +@pytest.mark.parametrize("secondary", [RuntimeError("secondary"), asyncio.CancelledError("secondary")]) +def test_failed_recovery_preserves_original_control_flow( + secondary: BaseException, monkeypatch: pytest.MonkeyPatch): + original = asyncio.CancelledError("primary") + with worker.TaskHubGrpcWorker(logger=LOGGER) as host: + monkeypatch.setattr(host, "_get_orchestrator_response", Mock(side_effect=original)) + stub = Mock() + stub.AbandonTaskOrchestratorWorkItem.side_effect = secondary + reset = Mock(side_effect=TeardownExit("reset failure")) + with pytest.raises(asyncio.CancelledError) as caught: + host._execute_orchestrator( + _request(), stub, "original-token", _on_outcome_uncertain=reset) + assert caught.value is original + reset.assert_called_once() + + +def test_recovery_process_control_is_not_swallowed(monkeypatch: pytest.MonkeyPatch): + stop = SystemExit("stop during recovery") + with worker.TaskHubGrpcWorker(logger=LOGGER) as host: + monkeypatch.setattr(host, "_get_orchestrator_response", Mock(side_effect=ValueError("preparation"))) + stub = Mock() + stub.AbandonTaskOrchestratorWorkItem.side_effect = stop + reset = Mock() + with pytest.raises(SystemExit) as caught: + host._execute_orchestrator( + _request(), stub, "original-token", _on_outcome_uncertain=reset) + assert caught.value is stop + reset.assert_called_once() + + +@pytest.mark.parametrize("reset_error", _control_flow_errors(), ids=str) +def test_reset_control_flow_is_not_hidden_by_ordinary_abandonment_error( + reset_error: BaseException, monkeypatch: pytest.MonkeyPatch): + with worker.TaskHubGrpcWorker(logger=LOGGER) as host: + monkeypatch.setattr(host, "_get_orchestrator_response", Mock(side_effect=ValueError("preparation"))) + stub = Mock() + stub.AbandonTaskOrchestratorWorkItem.side_effect = RuntimeError("abandonment failed") + reset = Mock(side_effect=reset_error) + with pytest.raises(type(reset_error)) as caught: + host._execute_orchestrator( + _request(), stub, "original-token", _on_outcome_uncertain=reset) + assert caught.value is reset_error + reset.assert_called_once() + + +def test_failure_details_do_not_pass_base_exception_to_public_properties_provider(): + provider = Mock() + error = TeardownExit("unexpected") + details = helpers.new_failure_details(error, provider, LOGGER) + assert details.errorType.endswith(".TeardownExit") + assert details.errorMessage == "unexpected" + provider.get_exception_properties.assert_not_called() + + +@pytest.mark.skipif(sys.version_info < (3, 11), reason="Exception groups require Python 3.11+.") +def test_grouped_unexpected_failure_keeps_nested_reasons(): + if sys.version_info >= (3, 11): + def orchestrator(ctx: task.OrchestrationContext, _: Any): + try: + yield ctx.wait_for_external_event("ready") + finally: + raise BaseExceptionGroup("teardown", [TeardownExit("root cause"), ValueError("second cause")]) + + result = _executor(orchestrator).execute(INSTANCE_ID, [], _request().newEvents) + failure = result.actions[0].completeOrchestration.failureDetails + assert failure.errorType.endswith(".BaseExceptionGroup") + assert "TeardownExit: root cause" in failure.stackTrace.value + assert "ValueError: second cause" in failure.stackTrace.value + + +@pytest.mark.asyncio +async def test_cancelling_executor_await_does_not_abandon_running_dispatch(): + started = Event() + finish = Event() + + def orchestrator(ctx: task.OrchestrationContext, _: Any): + started.set() + assert finish.wait(timeout=5) + return "finished" + + with worker.TaskHubGrpcWorker(logger=LOGGER) as host: + host._registry = _executor(orchestrator)._registry + manager = host._async_worker_manager + queue: asyncio.Queue[Any] = asyncio.Queue() + semaphore = asyncio.Semaphore(0) + queue.put_nowait(None) + stub = Mock() + reset = Mock() + + def dispatch(): + host._execute_orchestrator( + _request(), stub, "original-token", _on_outcome_uncertain=reset) + + running = asyncio.create_task(manager._process_work_item( + semaphore, queue, dispatch, host._cancel_orchestrator, (), {})) + try: + assert await asyncio.to_thread(started.wait, 2) + running.cancel() + with pytest.raises(asyncio.CancelledError): + await running + stub.CompleteOrchestratorTask.assert_not_called() + stub.AbandonTaskOrchestratorWorkItem.assert_not_called() + reset.assert_not_called() + await asyncio.wait_for(queue.join(), timeout=2) + await asyncio.wait_for(semaphore.acquire(), timeout=2) + finally: + finish.set() + manager.thread_pool.shutdown(wait=True) + + stub.CompleteOrchestratorTask.assert_called_once() + stub.AbandonTaskOrchestratorWorkItem.assert_not_called() + reset.assert_not_called() + + +@pytest.mark.asyncio +async def test_base_exception_recovery_keeps_queue_and_semaphore_live(): + with worker.TaskHubGrpcWorker(logger=LOGGER) as host: + manager = host._async_worker_manager + queue: asyncio.Queue[Any] = asyncio.Queue() + semaphore = asyncio.Semaphore(1) + stub = Mock() + error = asyncio.CancelledError("dispatch cancelled") + + def orchestrator(ctx: task.OrchestrationContext, _: Any): + raise error + + host._registry = _executor(orchestrator)._registry + queue.put_nowait(None) + await semaphore.acquire() + with pytest.raises(asyncio.CancelledError): + await manager._process_work_item( + semaphore, queue, host._execute_orchestrator, host._cancel_orchestrator, + (_request(), stub, "original-token"), {}) + await asyncio.wait_for(queue.join(), timeout=2) + await asyncio.wait_for(semaphore.acquire(), timeout=2) + + def healthy(ctx: task.OrchestrationContext, _: Any): + return "healthy" + + host._registry = _executor(healthy)._registry + queue.put_nowait(None) + await manager._process_work_item( + semaphore, queue, host._execute_orchestrator, host._cancel_orchestrator, + (_request(), stub, "next-token"), {}) + await asyncio.wait_for(queue.join(), timeout=2) + await asyncio.wait_for(semaphore.acquire(), timeout=2) + semaphore.release() + manager.thread_pool.shutdown(wait=True) + + stub.AbandonTaskOrchestratorWorkItem.assert_called_once() + stub.CompleteOrchestratorTask.assert_called_once() + assert stub.CompleteOrchestratorTask.call_args.args[0].completionToken == "next-token" diff --git a/tests/durabletask/test_orchestration_generator_cleanup.py b/tests/durabletask/test_orchestration_generator_cleanup.py new file mode 100644 index 00000000..5c1e17e8 --- /dev/null +++ b/tests/durabletask/test_orchestration_generator_cleanup.py @@ -0,0 +1,480 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +import asyncio +import gc +import logging +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timedelta +from typing import Any, Generator +from unittest.mock import Mock +from weakref import ReferenceType, ref + +import anyio +import pytest + +from durabletask import task, worker +from durabletask.internal import exceptions, helpers +from durabletask.internal import orchestrator_service_pb2 as pb +from durabletask.serialization import JsonDataConverter + +TEST_INSTANCE_ID = "generator-cleanup" +TEST_LOGGER = logging.getLogger("tests.generator_cleanup") +START_TIME = datetime(2026, 1, 1) + + +def _executor( + orchestrator: task.Orchestrator[Any, Any], +) -> worker._OrchestrationExecutor: # pyright: ignore[reportPrivateUsage] + registry = worker._Registry() # pyright: ignore[reportPrivateUsage] + registry.add_named_orchestrator("orchestrator", orchestrator) + return worker._OrchestrationExecutor( # pyright: ignore[reportPrivateUsage] + registry, TEST_LOGGER, JsonDataConverter()) + + +def _start_events() -> list[pb.HistoryEvent]: + return [ + helpers.new_orchestrator_started_event(START_TIME), + helpers.new_execution_started_event("orchestrator", TEST_INSTANCE_ID), + ] + + +def test_suspended_generator_closes_after_response_snapshot(): + closed: list[str] = [] + generators: list[Generator[task.Task[Any], Any, Any]] = [] + + def episode(ctx: task.OrchestrationContext): + timer = ctx.create_timer(timedelta(seconds=5)) + ctx.set_custom_status("waiting") + try: + yield timer + finally: + closed.append("closed") + timer.cancel() + ctx.call_activity("cleanup") + ctx.set_custom_status("cleanup") + ctx.continue_as_new("cleanup") + + def orchestrator(ctx: task.OrchestrationContext, _: Any): + generator = episode(ctx) + generators.append(generator) + return generator + + result = _executor(orchestrator).execute(TEST_INSTANCE_ID, [], _start_events()) + + assert closed == ["closed"] + assert len(result.actions) == 1 + assert result.actions[0].HasField("createTimer") + assert result.actions[0].createTimer.fireAt.ToDatetime() == START_TIME + timedelta(seconds=5) + assert result.encoded_custom_status == '"waiting"' + generators[0].close() + generators[0].close() + assert closed == ["closed"] + + +@pytest.mark.parametrize("fail", [False, True]) +@pytest.mark.parametrize("suspend_first", [False, True]) +def test_exhausted_generator_cleanup_is_not_repeated(fail: bool, suspend_first: bool): + closed: list[str] = [] + generators: list[Generator[task.Task[Any], Any, Any]] = [] + + def episode(ctx: task.OrchestrationContext): + try: + if suspend_first: + yield ctx.wait_for_external_event("ready") + if fail: + raise ValueError("orchestrator failed") + return "done" + finally: + closed.append("closed") + + def orchestrator(ctx: task.OrchestrationContext, _: Any): + generator = episode(ctx) + generators.append(generator) + return generator + + result = _executor(orchestrator).execute( + TEST_INSTANCE_ID, [], _start_events() + [helpers.new_event_raised_event("ready", None)]) + + assert closed == ["closed"] + generators[0].close() + assert closed == ["closed"] + assert len(result.actions) == 1 + completion = result.actions[0].completeOrchestration + if fail: + assert completion.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + assert completion.failureDetails.errorMessage == "orchestrator failed" + else: + assert completion.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED + assert completion.result.value == '"done"' + + +def test_generator_stays_open_across_old_and_new_history_events(): + observed: list[str] = [] + + def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Generator[task.Task[Any], str, None]: + try: + observed.append((yield ctx.call_activity("first", return_type=str))) + observed.append((yield ctx.call_activity("second", return_type=str))) + assert observed == ["first result", "second result"] + yield ctx.call_activity("third", return_type=str) + finally: + observed.append("closed") + + old_events = _start_events() + [ + helpers.new_task_scheduled_event(1, "first"), + helpers.new_task_completed_event(1, '"first result"'), + helpers.new_task_scheduled_event(2, "second"), + ] + new_events = [helpers.new_task_completed_event(2, '"second result"')] + result = _executor(orchestrator).execute(TEST_INSTANCE_ID, old_events, new_events) + + assert observed == ["first result", "second result", "closed"] + assert len(result.actions) == 1 + assert result.actions[0].id == 3 + assert result.actions[0].scheduleTask.name == "third" + + +def test_continue_as_new_closes_suspended_generator(): + closed: list[str] = [] + + def orchestrator(ctx: task.OrchestrationContext, _: Any): + try: + ctx.continue_as_new("next") + yield ctx.wait_for_external_event("unused") + finally: + closed.append("closed") + ctx.call_activity("cleanup") + + result = _executor(orchestrator).execute(TEST_INSTANCE_ID, [], _start_events()) + + assert closed == ["closed"] + assert len(result.actions) == 1 + completion = result.actions[0].completeOrchestration + assert completion.orchestrationStatus == pb.ORCHESTRATION_STATUS_CONTINUED_AS_NEW + assert completion.result.value == '"next"' + + +def test_termination_closes_suspended_generator(): + closed: list[str] = [] + + def orchestrator(ctx: task.OrchestrationContext, _: Any): + try: + yield ctx.wait_for_external_event("ready") + finally: + closed.append("closed") + + result = _executor(orchestrator).execute( + TEST_INSTANCE_ID, _start_events(), + [pb.HistoryEvent(executionTerminated=pb.ExecutionTerminatedEvent())]) + + assert closed == ["closed"] + assert len(result.actions) == 1 + assert result.actions[0].completeOrchestration.orchestrationStatus == pb.ORCHESTRATION_STATUS_TERMINATED + + +def test_response_construction_failure_still_closes_generator(monkeypatch: pytest.MonkeyPatch): + closed: list[str] = [] + + def orchestrator(ctx: task.OrchestrationContext, _: Any): + try: + yield ctx.wait_for_external_event("ready") + finally: + closed.append("closed") + + def fail_get_actions(ctx: worker._RuntimeOrchestrationContext): # pyright: ignore[reportPrivateUsage] + raise ValueError("response failed") + + monkeypatch.setattr( + worker._RuntimeOrchestrationContext, "get_actions", fail_get_actions) # pyright: ignore[reportPrivateUsage] + with pytest.raises(ValueError, match="response failed"): + _executor(orchestrator).execute(TEST_INSTANCE_ID, [], _start_events()) + assert closed == ["closed"] + + +@pytest.mark.parametrize("cleanup_fails", [False, True]) +def test_history_failure_closes_suspended_generator( + cleanup_fails: bool, caplog: pytest.LogCaptureFixture): + closed: list[str] = [] + + def orchestrator(ctx: task.OrchestrationContext, _: Any): + try: + yield ctx.call_activity("expected", return_type=str) + finally: + closed.append("closed") + if cleanup_fails: + raise ValueError("cleanup failed") + + old_events = _start_events() + with caplog.at_level(logging.WARNING, logger=TEST_LOGGER.name): + result = _executor(orchestrator).execute( + TEST_INSTANCE_ID, old_events, [helpers.new_task_scheduled_event(1, "wrong")]) + + assert closed == ["closed"] + assert len(result.actions) == 1 + completion = result.actions[0].completeOrchestration + assert completion.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + assert "NonDeterminismError" in completion.failureDetails.errorType + assert "cleanup failed" not in completion.failureDetails.errorMessage + assert bool(caplog.records) == cleanup_fails + + +@pytest.mark.parametrize("strategy", list(worker.VersionFailureStrategy)) +@pytest.mark.parametrize("cleanup_fails", [False, True]) +def test_version_failure_closes_generator_without_changing_dispatch_outcome( + strategy: worker.VersionFailureStrategy, cleanup_fails: bool, + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture): + closed: list[str] = [] + + def orchestrator(ctx: task.OrchestrationContext, _: Any): + try: + yield ctx.wait_for_external_event("ready") + finally: + closed.append("closed") + if cleanup_fails: + raise ValueError("cleanup failed") + + executor = _executor(orchestrator) + executor._registry.versioning = worker.VersioningOptions( # pyright: ignore[reportPrivateUsage] + failure_strategy=strategy, match_strategy=worker.VersionMatchStrategy.NONE) + process_event = executor.process_event + failure = pb.TaskFailureDetails(errorType="VersionMismatch", errorMessage="version failed") + + def process_with_version_failure( + ctx: worker._RuntimeOrchestrationContext, # pyright: ignore[reportPrivateUsage] + event: pb.HistoryEvent): + process_event(ctx, event) + if event.HasField("eventRaised"): + raise exceptions.VersionFailureException(failure) + + monkeypatch.setattr(executor, "process_event", process_with_version_failure) + events = _start_events() + [helpers.new_event_raised_event("unrelated", None)] + with caplog.at_level(logging.WARNING, logger=TEST_LOGGER.name): + if strategy == worker.VersionFailureStrategy.REJECT: + with pytest.raises(exceptions.AbandonOrchestrationError): + executor.execute(TEST_INSTANCE_ID, [], events) + else: + result = executor.execute(TEST_INSTANCE_ID, [], events) + assert len(result.actions) == 1 + assert result.actions[0].completeOrchestration.failureDetails == failure + assert closed == ["closed"] + assert bool(caplog.records) == cleanup_fails + + +@pytest.mark.parametrize("yield_during_close", [False, True]) +def test_cleanup_failure_is_logged_without_changing_returned_actions( + yield_during_close: bool, caplog: pytest.LogCaptureFixture): + generators: list[Generator[task.Task[Any], Any, Any]] = [] + + def episode(ctx: task.OrchestrationContext): + try: + yield ctx.create_timer(timedelta(seconds=5)) + finally: + if yield_during_close: + yield ctx.wait_for_external_event("invalid-cleanup") + else: + raise ValueError("cleanup failed") + + def orchestrator(ctx: task.OrchestrationContext, _: Any): + generator = episode(ctx) + generators.append(generator) + return generator + + try: + with caplog.at_level(logging.WARNING, logger=TEST_LOGGER.name): + result = _executor(orchestrator).execute(TEST_INSTANCE_ID, [], _start_events()) + assert len(result.actions) == 1 + assert result.actions[0].HasField("createTimer") + assert len(caplog.records) == 1 + record = caplog.records[0] + assert record.levelno == logging.WARNING + assert TEST_INSTANCE_ID in record.message + assert "orchestrator" in record.message + assert "" not in record.message + assert record.exc_info is not None + expected = "generator ignored GeneratorExit" if yield_during_close else "cleanup failed" + assert str(record.exc_info[1]) == expected + finally: + for generator in generators: + generator.close() + + +def _cancelling_orchestrator( + ctx: task.OrchestrationContext, _: Any, +) -> Generator[task.Task[Any], Any, None]: + loop = asyncio.new_event_loop() + + async def body(): + await asyncio.Future[None]() + + root = loop.create_task(body()) + loop.run_until_complete(asyncio.sleep(0)) + ctx.set_custom_status("waiting") + try: + yield ctx.create_timer(timedelta(seconds=5)) + finally: + root.cancel() + try: + loop.run_until_complete(root) + finally: + loop.close() + + +@pytest.mark.parametrize("via_worker", [False, True]) +def test_teardown_cancellation_preserves_response( + via_worker: bool, caplog: pytest.LogCaptureFixture): + executor = _executor(_cancelling_orchestrator) + with caplog.at_level(logging.WARNING, logger=TEST_LOGGER.name): + if via_worker: + with worker.TaskHubGrpcWorker(logger=TEST_LOGGER) as host: + host._registry = executor._registry # pyright: ignore[reportPrivateUsage] + stub = Mock() + request = pb.OrchestratorRequest(instanceId=TEST_INSTANCE_ID, newEvents=_start_events()) + host._execute_orchestrator(request, stub, "token") # pyright: ignore[reportPrivateUsage] + stub.CompleteOrchestratorTask.assert_called_once() + stub.AbandonTaskOrchestratorWorkItem.assert_not_called() + response: pb.OrchestratorResponse = stub.CompleteOrchestratorTask.call_args.args[0] + assert response.instanceId == TEST_INSTANCE_ID + assert response.completionToken == "token" + actions = list(response.actions) + custom_status = response.customStatus.value + else: + result = executor.execute(TEST_INSTANCE_ID, [], _start_events()) + actions = result.actions + custom_status = result.encoded_custom_status + + assert actions == [helpers.new_create_timer_action(1, START_TIME + timedelta(seconds=5))] + assert custom_status == '"waiting"' + 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.exc_info is not None + assert record.exc_info[0] is asyncio.CancelledError + + +@pytest.mark.parametrize("via_worker", [False, True]) +def test_teardown_cancellation_preserves_version_rejection( + via_worker: bool, monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture): + executor = _executor(_cancelling_orchestrator) + executor._registry.versioning = worker.VersioningOptions( # pyright: ignore[reportPrivateUsage] + failure_strategy=worker.VersionFailureStrategy.REJECT, + match_strategy=worker.VersionMatchStrategy.NONE) + process_event = worker._OrchestrationExecutor.process_event # pyright: ignore[reportPrivateUsage] + + def process_with_version_failure( + self: worker._OrchestrationExecutor, # pyright: ignore[reportPrivateUsage] + ctx: worker._RuntimeOrchestrationContext, # pyright: ignore[reportPrivateUsage] + event: pb.HistoryEvent): + process_event(self, ctx, event) + if event.HasField("eventRaised"): + raise exceptions.VersionFailureException(pb.TaskFailureDetails(errorType="VersionMismatch")) + + monkeypatch.setattr( + worker._OrchestrationExecutor, "process_event", process_with_version_failure) # pyright: ignore[reportPrivateUsage] + events = _start_events() + [helpers.new_event_raised_event("unrelated", None)] + with caplog.at_level(logging.WARNING, logger=TEST_LOGGER.name): + if via_worker: + with worker.TaskHubGrpcWorker(logger=TEST_LOGGER) as host: + host._registry = executor._registry # pyright: ignore[reportPrivateUsage] + stub = Mock() + request = pb.OrchestratorRequest(instanceId=TEST_INSTANCE_ID, newEvents=events) + host._execute_orchestrator(request, stub, "token") # pyright: ignore[reportPrivateUsage] + stub.CompleteOrchestratorTask.assert_not_called() + stub.AbandonTaskOrchestratorWorkItem.assert_called_once_with( + pb.AbandonOrchestrationTaskRequest(completionToken="token")) + else: + with pytest.raises(exceptions.AbandonOrchestrationError): + executor.execute(TEST_INSTANCE_ID, [], events) + + assert len(caplog.records) == 1 + record = caplog.records[0] + assert record.levelno == logging.WARNING + assert record.exc_info is not None + assert record.exc_info[0] is asyncio.CancelledError + + +@pytest.mark.parametrize("error_type", [SystemExit, KeyboardInterrupt]) +def test_process_control_exception_from_cleanup_still_propagates(error_type: type[BaseException]): + def orchestrator(ctx: task.OrchestrationContext, _: Any): + try: + yield ctx.wait_for_external_event("ready") + finally: + raise error_type("stop") + + with pytest.raises(error_type, match="stop"): + _executor(orchestrator).execute(TEST_INSTANCE_ID, [], _start_events()) + + +@pytest.mark.parametrize("thread_pool", [False, True]) +def test_async_bridge_resources_are_released(thread_pool: bool): + references: dict[str, ReferenceType[Any]] = {} + cleaned: list[str] = [] + + def episode(ctx: task.OrchestrationContext): + loop = asyncio.new_event_loop() + + async def resource(): + try: + yield + finally: + await asyncio.sleep(0) + cleaned.append("asyncgen closed") + + async def body(): + assert ctx.instance_id == TEST_INSTANCE_ID + opened_resource = resource() + await anext(opened_resource) + references["asyncgen"] = ref(opened_resource) + try: + async with anyio.create_task_group(): + await asyncio.Future() + finally: + await asyncio.sleep(0) + cleaned.append("closed") + + root = loop.create_task(body()) + references["loop"] = ref(loop) + references["root"] = ref(root) + loop.run_until_complete(asyncio.sleep(0)) + try: + yield ctx.create_timer(timedelta(seconds=5)) + finally: + root.cancel() + loop.run_until_complete(asyncio.gather(root, return_exceptions=True)) + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.close() + + def orchestrator(ctx: task.OrchestrationContext, _: Any): + generator = episode(ctx) + references["context"] = ref(ctx) + references["generator"] = ref(generator) + return generator + + def dispatch(): + result = _executor(orchestrator).execute(TEST_INSTANCE_ID, [], _start_events()) + assert len(result.actions) == 1 + assert result.actions[0].HasField("createTimer") + assert sorted(cleaned) == ["asyncgen closed", "closed"] + loop = references["loop"]() + if loop is not None: + assert loop.is_closed() + root = references["root"]() + if root is not None: + assert root.done() + + try: + if thread_pool: + with ThreadPoolExecutor(max_workers=1) as pool: + pool.submit(dispatch).result() + else: + dispatch() + gc.collect() + assert set(references) == {"context", "generator", "loop", "root", "asyncgen"} + assert all(weak_reference() is None for weak_reference in references.values()) + finally: + weak_generator = references.get("generator") + generator = weak_generator() if weak_generator else None + if generator is not None: + generator.close() diff --git a/tests/durabletask/test_orchestrator_payload_failure.py b/tests/durabletask/test_orchestrator_payload_failure.py new file mode 100644 index 00000000..1412e4c3 --- /dev/null +++ b/tests/durabletask/test_orchestrator_payload_failure.py @@ -0,0 +1,474 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Permanent payload failures must complete without stranding entity locks.""" + +import asyncio +import json +import logging +from threading import Event +from unittest.mock import MagicMock + +import pytest + +from durabletask import client, entities, task, worker +from durabletask.internal import helpers as ph +from durabletask.internal import orchestrator_service_pb2 as pb +from durabletask.internal.orchestration_entity_context import OrchestrationEntityContext +from durabletask.serialization import JsonDataConverter +from durabletask.testing import create_test_backend +from tests.durabletask._port_utils import find_free_port +from tests.durabletask.test_large_payload import FakePayloadStore + + +ENTITY_IDS = [ + entities.EntityInstanceId("counter", "a"), + entities.EntityInstanceId("counter", "b"), +] +LARGE_PAYLOAD = "x" * 1024 + + +def _fail_payload(ctx: task.OrchestrationContext, mode: str): + if mode == "status": + ctx.set_custom_status(LARGE_PAYLOAD) + elif mode == "call": + yield ctx.call_entity(ENTITY_IDS[0], "add", LARGE_PAYLOAD) + elif mode == "continue": + ctx.continue_as_new(LARGE_PAYLOAD) + return LARGE_PAYLOAD if mode == "result" else "done" + + +def _request(name: str) -> pb.OrchestratorRequest: + return pb.OrchestratorRequest( + instanceId="owner", + newEvents=[ + ph.new_orchestrator_started_event(), + ph.new_execution_started_event(name, "owner"), + ], + ) + + +@pytest.mark.parametrize("call_failed", [False, True]) +def test_legacy_multi_lock_release_history_replays(call_failed): + def orchestrator(ctx, _): + with (yield ctx.lock_entities(ENTITY_IDS)): + try: + yield ctx.call_entity(ENTITY_IDS[0], "get") + except task.TaskFailedError: + # Handle the expected failure so replay can verify lock releases. + pass + yield ctx.wait_for_external_event("next") + return "done" + + registry = worker._Registry() + name = registry.add_orchestrator(orchestrator) + executor = worker._OrchestrationExecutor(registry, logging.getLogger("test"), JsonDataConverter()) + started = list(_request(name).newEvents) + first = executor.execute("owner", [], started) + section = first.actions[0].sendEntityMessage.entityLockRequested.criticalSectionId + granted = [ + ph.new_event_sent_event(1, str(ENTITY_IDS[0]), json.dumps({"id": section})), + ph.new_event_raised_event(section, None), + ] + second = executor.execute("owner", started, granted) + request_id = second.actions[0].sendEntityMessage.entityOperationCalled.requestId + response = {"result": "42"} + if call_failed: + response["exceptionType"] = "ValueError" + called = [ + ph.new_event_sent_event(2, str(ENTITY_IDS[0]), json.dumps({"id": request_id})), + ph.new_event_raised_event(request_id, json.dumps(response)), + ] + third = executor.execute("owner", started + granted, called) + # Before acquired/available lock tracking was separated, completing a call + # moved that entity to the end of the shared list used for releases. + legacy_releases = [(3, str(ENTITY_IDS[1])), (4, str(ENTITY_IDS[0]))] + assert [ + (action.id, action.sendEntityMessage.entityUnlockSent.targetInstanceId.value) + for action in third.actions + ] == legacy_releases + released = [ph.new_event_sent_event(action_id, target, None) for action_id, target in legacy_releases] + history = started + granted + called + released + replay = executor.execute("owner", history, [ph.new_orchestrator_started_event()]) + assert not replay.actions + assert not replay.entity_lock_releases + completed = executor.execute("owner", history, [ph.new_event_raised_event("next", None)]) + assert len(completed.actions) == 1 + assert completed.actions[0].id == 5 + assert completed.actions[0].completeOrchestration.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED + assert completed.actions[0].completeOrchestration.result.value == '"done"' + + +def _timeout_history(executor, name): + started = list(_request(name).newEvents) + first = executor.execute("owner", [], started) + section = first.actions[0].sendEntityMessage.entityLockRequested.criticalSectionId + granted = [ + ph.new_event_sent_event(1, str(ENTITY_IDS[0]), json.dumps({"id": section})), + ph.new_event_raised_event(section, None), + ] + second = executor.execute("owner", started, granted) + call = second.actions[0].sendEntityMessage.entityOperationCalled + pending = ph.new_event_sent_event(2, str(ENTITY_IDS[0]), json.dumps({"id": call.requestId})) + return started + granted + [pending], section + + +def test_legacy_pending_call_timeout_history_replays_through_activity(): + def orchestrator(ctx, _): + with (yield ctx.lock_entities(ENTITY_IDS)): + call = ctx.call_entity(ENTITY_IDS[0], "get") + timeout = ctx.wait_for_external_event("timeout") + yield task.when_any([call, timeout]) + return (yield ctx.call_activity("after_timeout")) + + registry = worker._Registry() + name = registry.add_orchestrator(orchestrator) + executor = worker._OrchestrationExecutor(registry, logging.getLogger("test"), JsonDataConverter()) + history, _ = _timeout_history(executor, name) + timed_out = ph.new_event_raised_event("timeout", None) + result = executor.execute("owner", history, [timed_out]) + assert [action.id for action in result.actions] == [3, 4] + assert result.actions[0].sendEntityMessage.entityUnlockSent.targetInstanceId.value == str(ENTITY_IDS[1]) + assert result.actions[1].scheduleTask.name == "after_timeout" + assert [release.entityUnlockSent.targetInstanceId.value for release in result.entity_lock_releases] == [ + str(ENTITY_IDS[0]), + ] + + # Literal history from the old emitter: only b was released, so the next + # activity kept ID 4 even though the call to a was still outstanding. + history += [ + timed_out, + ph.new_event_sent_event(3, str(ENTITY_IDS[1]), None), + ph.new_task_scheduled_event(4, "after_timeout"), + ] + assert not executor.execute("owner", history, [ph.new_orchestrator_started_event()]).actions + completed = executor.execute("owner", history, [ph.new_task_completed_event(4, '"done"')]) + assert len(completed.actions) == 1 + assert completed.actions[0].id == 5 + assert completed.actions[0].completeOrchestration.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED + assert completed.actions[0].completeOrchestration.result.value == '"done"' + + +@pytest.mark.parametrize("mode", ["result", "status", "activity"]) +@pytest.mark.parametrize("replay_exit", [False, True]) +def test_payload_failure_after_normal_pending_call_exit_releases_all_locks(mode, replay_exit): + def orchestrator(ctx, _): + with (yield ctx.lock_entities(ENTITY_IDS)): + call = ctx.call_entity(ENTITY_IDS[0], "get") + yield task.when_any([call, ctx.wait_for_external_event("timeout")]) + if replay_exit: + yield ctx.wait_for_external_event("fail") + if mode == "result": + return LARGE_PAYLOAD + if mode == "status": + ctx.set_custom_status(LARGE_PAYLOAD) + yield ctx.call_activity("after_timeout", input=LARGE_PAYLOAD if mode == "activity" else None) + return None + + store = FakePayloadStore(threshold_bytes=100, max_stored_payload_bytes=200) + stub = MagicMock() + with worker.TaskHubGrpcWorker(payload_store=store) as w: + w.add_orchestrator(orchestrator) + executor = worker._OrchestrationExecutor(w._registry, logging.getLogger("test"), JsonDataConverter()) + history, section = _timeout_history(executor, "orchestrator") + events = [ph.new_event_raised_event("timeout", None)] + if replay_exit: + exited = executor.execute("owner", history, events) + assert len(exited.actions) == 1 + assert exited.actions[0].id == 3 + history += events + [ph.new_event_sent_event(3, str(ENTITY_IDS[1]), None)] + events = [ph.new_event_raised_event("fail", None)] + w._execute_orchestrator(pb.OrchestratorRequest( + instanceId="owner", pastEvents=history, + newEvents=events, + ), stub, "failure-token") + + response = stub.CompleteOrchestratorTask.call_args.args[0] + assert response.completionToken == "failure-token" + assert len(response.actions) == (2 if replay_exit else 3) + if not replay_exit: + assert response.actions[0].id == 3 + releases = [action.sendEntityMessage.entityUnlockSent for action in response.actions[:-1]] + expected_targets = [str(ENTITY_IDS[0])] if replay_exit else [str(ENTITY_IDS[1]), str(ENTITY_IDS[0])] + assert [release.targetInstanceId.value for release in releases] == expected_targets + assert all(release.parentInstanceId.value == "owner" for release in releases) + assert all(release.criticalSectionId == section for release in releases) + completion = response.actions[-1].completeOrchestration + assert completion.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + assert completion.failureDetails.isNonRetriable + assert not response.HasField("customStatus") + + +def test_release_snapshot_retains_pending_locks_after_completed_calls(): + entity_ctx = OrchestrationEntityContext("owner") + entity_ctx.emit_acquire_message("section", ENTITY_IDS) + entity_ctx.complete_acquire("section") + assert entity_ctx.validate_operation_transition(ENTITY_IDS[0], False) == (True, "") + entity_ctx.recover_lock_after_call(ENTITY_IDS[0]) + assert entity_ctx.validate_operation_transition(ENTITY_IDS[1], False) == (True, "") + releases = list(entity_ctx.get_lock_release_messages()) + assert [release.entityUnlockSent.targetInstanceId.value for release in releases] == [ + str(ENTITY_IDS[0]), str(ENTITY_IDS[1]), + ] + assert all(release.entityUnlockSent.parentInstanceId.value == "owner" for release in releases) + assert all(release.entityUnlockSent.criticalSectionId == "section" for release in releases) + assert entity_ctx.available_locks == [ENTITY_IDS[0]] + assert entity_ctx.critical_section_locks == [ENTITY_IDS[0]] + assert entity_ctx.is_inside_critical_section + + +@pytest.mark.parametrize("legacy_protocol", [False, True]) +@pytest.mark.parametrize("overlapping_call", [False, True]) +def test_pre_lock_call_can_complete_while_acquisition_is_pending(legacy_protocol, overlapping_call): + outside_entity = ENTITY_IDS[0] if overlapping_call else entities.EntityInstanceId("counter", "outside") + + def orchestrator(ctx, _): + pending = ctx.call_entity(outside_entity, "get") + with (yield ctx.lock_entities(ENTITY_IDS)): + result = yield pending + yield ctx.wait_for_external_event("next") + return result + + registry = worker._Registry() + name = registry.add_orchestrator(orchestrator) + executor = worker._OrchestrationExecutor(registry, logging.getLogger("test"), JsonDataConverter()) + started = list(_request(name).newEvents) + first = executor.execute("owner", [], started) + call = first.actions[0].sendEntityMessage.entityOperationCalled + lock = first.actions[1].sendEntityMessage.entityLockRequested + if legacy_protocol: + called = [ + ph.new_event_sent_event(1, str(outside_entity), json.dumps({"id": call.requestId})), + ph.new_event_sent_event(2, str(ENTITY_IDS[0]), json.dumps({"id": lock.criticalSectionId})), + ph.new_event_raised_event(call.requestId, json.dumps({"result": "42"})), + ] + granted = ph.new_event_raised_event(lock.criticalSectionId, None) + else: + called = [ + pb.HistoryEvent(eventId=1, entityOperationCalled=call), + pb.HistoryEvent(eventId=2, entityLockRequested=lock), + pb.HistoryEvent(entityOperationCompleted=pb.EntityOperationCompletedEvent( + requestId=call.requestId, output=ph.get_string_value("42"))), + ] + granted = pb.HistoryEvent(entityLockGranted=pb.EntityLockGrantedEvent( + criticalSectionId=lock.criticalSectionId)) + waiting = executor.execute("owner", started, called) + assert not waiting.actions + unlocked = executor.execute("owner", started + called, [granted]) + assert [action.id for action in unlocked.actions] == [3, 4] + assert [ + action.sendEntityMessage.entityUnlockSent.targetInstanceId.value for action in unlocked.actions + ] == [str(eid) for eid in ENTITY_IDS] + if legacy_protocol: + released = [ph.new_event_sent_event(action.id, str(eid), None) + for action, eid in zip(unlocked.actions, ENTITY_IDS)] + else: + released = [pb.HistoryEvent(eventId=action.id, entityUnlockSent=action.sendEntityMessage.entityUnlockSent) + for action in unlocked.actions] + history = started + called + [granted] + released + assert not executor.execute("owner", history, [ph.new_orchestrator_started_event()]).actions + completed = executor.execute("owner", history, [ph.new_event_raised_event("next", None)]) + assert len(completed.actions) == 1 + assert completed.actions[-1].id == 5 + assert completed.actions[-1].completeOrchestration.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED + assert completed.actions[-1].completeOrchestration.result.value == "42" + + +@pytest.mark.parametrize("mode", ["result", "status", "call", "continue"]) +@pytest.mark.parametrize("completed_call", [False, True]) +def test_permanent_payload_failure_releases_acquired_locks(mode, completed_call): + def orchestrator(ctx, _): + with (yield ctx.lock_entities(list(reversed(ENTITY_IDS)))): + if completed_call: + yield ctx.call_entity(ENTITY_IDS[0], "get") + return (yield from _fail_payload(ctx, mode)) + + store = FakePayloadStore(threshold_bytes=100, max_stored_payload_bytes=200) + provider = MagicMock() + stub = MagicMock() + with worker.TaskHubGrpcWorker(payload_store=store, exception_properties_provider=provider) as w: + w.add_orchestrator(orchestrator) + first = _request("orchestrator") + w._execute_orchestrator(first, stub, "first-token") + first_response = stub.CompleteOrchestratorTask.call_args.args[0] + lock_action = first_response.actions[0] + lock_request = lock_action.sendEntityMessage.entityLockRequested + stub.reset_mock() + + req = pb.OrchestratorRequest( + instanceId="owner", + pastEvents=[ + *first.newEvents, + pb.HistoryEvent(eventId=lock_action.id, entityLockRequested=lock_request), + ], + newEvents=[ + ph.new_orchestrator_started_event(), + pb.HistoryEvent(entityLockGranted=pb.EntityLockGrantedEvent( + criticalSectionId=lock_request.criticalSectionId)), + ], + ) + if completed_call: + w._execute_orchestrator(req, stub, "call-token") + call_action = stub.CompleteOrchestratorTask.call_args.args[0].actions[0] + call_message = call_action.sendEntityMessage.entityOperationCalled + req = pb.OrchestratorRequest( + instanceId="owner", + pastEvents=[ + *req.pastEvents, + *req.newEvents, + pb.HistoryEvent(eventId=call_action.id, entityOperationCalled=call_message), + ], + newEvents=[ + ph.new_orchestrator_started_event(), + pb.HistoryEvent(entityOperationCompleted=pb.EntityOperationCompletedEvent( + requestId=call_message.requestId, output=ph.get_string_value("42"))), + ], + ) + stub.reset_mock() + w._execute_orchestrator(req, stub, "failure-token") + + stub.CompleteOrchestratorTask.assert_called_once() + stub.AbandonTaskOrchestratorWorkItem.assert_not_called() + response = stub.CompleteOrchestratorTask.call_args.args[0] + assert response.instanceId == "owner" + assert response.completionToken == "failure-token" + assert not response.HasField("customStatus") + assert response.ByteSize() < 1024 + assert len(response.actions) == 3 + releases = [action.sendEntityMessage.entityUnlockSent for action in response.actions[:-1]] + release_order = list(reversed(ENTITY_IDS)) if completed_call or mode == "call" else ENTITY_IDS + assert [release.targetInstanceId.value for release in releases] == [str(eid) for eid in release_order] + assert all(release.parentInstanceId.value == "owner" for release in releases) + assert all(release.criticalSectionId == lock_request.criticalSectionId for release in releases) + assert [action.id for action in response.actions] == sorted({action.id for action in response.actions}) + completion = response.actions[-1].completeOrchestration + assert completion.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + assert not completion.HasField("result") + failure = completion.failureDetails + assert failure.isNonRetriable + assert failure.errorType == "durabletask.internal.exceptions.PayloadSizeExceededError" + assert "exceeds the maximum" in failure.errorMessage + assert not failure.HasField("stackTrace") + assert not failure.HasField("innerFailure") + assert not failure.properties + provider.get_exception_properties.assert_not_called() + + +@pytest.mark.parametrize("error", [OSError("unavailable"), ValueError("store error"), asyncio.CancelledError()]) +def test_storage_failure_and_cancellation_do_not_become_permanent(error): + def orchestrator(ctx, _): + return LARGE_PAYLOAD + + store = FakePayloadStore(threshold_bytes=100) + original_upload = store.upload + store.upload = MagicMock(side_effect=error) + stub = MagicMock() + with worker.TaskHubGrpcWorker(payload_store=store) as w: + w.add_orchestrator(orchestrator) + req = _request("orchestrator") + if isinstance(error, asyncio.CancelledError): + with pytest.raises(asyncio.CancelledError): + w._execute_orchestrator(req, stub, "token") + else: + w._execute_orchestrator(req, stub, "token") + stub.CompleteOrchestratorTask.assert_not_called() + + store.upload = original_upload + w._execute_orchestrator(req, stub, "retry-token") + + completion = stub.CompleteOrchestratorTask.call_args.args[0].actions[-1].completeOrchestration + assert completion.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED + assert not completion.HasField("failureDetails") + + +def test_failure_preserves_emitted_releases_before_outstanding_releases(monkeypatch): + entity_ctx = OrchestrationEntityContext("owner") + entity_ctx.emit_acquire_message("old-section", ENTITY_IDS) + emitted = [ + pb.OrchestratorAction(id=index, sendEntityMessage=release) + for index, release in enumerate(entity_ctx.emit_lock_release_messages(), start=2) + ] + entity_ctx.emit_acquire_message("new-section", ENTITY_IDS) + entity_ctx.complete_acquire("new-section") + result = worker.ExecutionResults( + actions=[ + *emitted, + ph.new_call_entity_action(4, "owner", ENTITY_IDS[0], "add", LARGE_PAYLOAD, "request"), + ], + encoded_custom_status="status to discard", + entity_lock_releases=list(entity_ctx.get_lock_release_messages()), + ) + monkeypatch.setattr(worker._OrchestrationExecutor, "execute", lambda *args: result) + store = FakePayloadStore(threshold_bytes=1, max_stored_payload_bytes=200) + stub = MagicMock() + with worker.TaskHubGrpcWorker(payload_store=store) as w: + w._execute_orchestrator(_request("unused"), stub, "token") + + response = stub.CompleteOrchestratorTask.call_args.args[0] + assert response.actions[:2] == emitted + releases = [action.sendEntityMessage.entityUnlockSent for action in response.actions[:-1]] + assert [(release.criticalSectionId, release.targetInstanceId.value) for release in releases] == [ + (section, str(eid)) for section in ("old-section", "new-section") for eid in ENTITY_IDS + ] + assert all(release.parentInstanceId.value == "owner" for release in releases) + assert [action.id for action in response.actions] == [2, 3, 5, 6, 7] + assert response.actions[-1].completeOrchestration.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + assert not response.HasField("customStatus") + + +@pytest.mark.parametrize("mode", ["result", "status", "call"]) +def test_next_lock_caller_progresses_after_permanent_payload_failure(monkeypatch, mode): + locked = Event() + second_queued = Event() + + def failing(ctx, _): + with (yield ctx.lock_entities(ENTITY_IDS)): + locked.set() + yield ctx.wait_for_external_event("fail") + return (yield from _fail_payload(ctx, mode)) + + def next_caller(ctx, _): + with (yield ctx.lock_entities(ENTITY_IDS)): + return (yield ctx.call_entity(ENTITY_IDS[0], "get")) + + def counter(ctx: entities.EntityContext, _): + return 42 + + port = find_free_port() + host = f"localhost:{port}" + backend = create_test_backend(port=port) + original_grant = backend._try_grant_lock + + def observe_grant(pending): + granted = original_grant(pending) + if pending.parent_instance_id == "second" and not granted: + second_queued.set() + return granted + + monkeypatch.setattr(backend, "_try_grant_lock", observe_grant) + store = FakePayloadStore(threshold_bytes=100, max_stored_payload_bytes=200) + try: + with worker.TaskHubGrpcWorker(host_address=host, payload_store=store) as w: + w.add_orchestrator(failing) + w.add_orchestrator(next_caller) + w.add_entity(counter) + w.start() + with client.TaskHubGrpcClient(host_address=host) as c: + c.schedule_new_orchestration(failing, instance_id="first") + assert locked.wait(10), "First orchestration did not acquire both locks" + c.schedule_new_orchestration(next_caller, instance_id="second") + assert second_queued.wait(10), "Second caller was not blocked by the first" + c.raise_orchestration_event("first", "fail") + failed = c.wait_for_orchestration_completion("first", timeout=10) + completed = c.wait_for_orchestration_completion("second", timeout=10) + + assert failed is not None + assert failed.runtime_status == client.OrchestrationStatus.FAILED + assert failed.failure_details is not None + assert failed.failure_details.error_type == "durabletask.internal.exceptions.PayloadSizeExceededError" + assert completed is not None + assert completed.runtime_status == client.OrchestrationStatus.COMPLETED + assert completed.serialized_output == "42" + finally: + backend.stop() + backend.reset() diff --git a/tests/durabletask/test_pydantic_ai_orchestrator.py b/tests/durabletask/test_pydantic_ai_orchestrator.py new file mode 100644 index 00000000..fa856410 --- /dev/null +++ b/tests/durabletask/test_pydantic_ai_orchestrator.py @@ -0,0 +1,131 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Real Pydantic AI parallel tools with ordinary activities and no trampoline. + +This small proof adapter retains model/tool registration and JSON codecs. It is +not a replacement for the framework adapter's full serialization/error policy. +""" + +import asyncio +import json +from contextvars import ContextVar +from datetime import timezone +from typing import Any + +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 _completion, _execute, _start + + +@pytest.mark.parametrize("fail_tool", [False, True]) +def test_pydantic_agent_parallel_tools_replay_without_bridge(fail_tool: bool) -> None: + pytest.importorskip("pydantic_ai", reason="The framework requires Python 3.11+") + from pydantic import TypeAdapter + from pydantic_ai import Agent + from pydantic_ai.capabilities import AbstractCapability + from pydantic_ai.messages import ModelMessage, ModelResponse + from pydantic_ai.models import ModelRequestParameters + from pydantic_ai.models.test import TestModel + from pydantic_ai.settings import ModelSettings + + context: ContextVar[task.OrchestrationContext] = ContextVar("context") + request_codec = TypeAdapter(tuple[list[ModelMessage], ModelSettings | None, ModelRequestParameters]) + response_codec = TypeAdapter(ModelResponse) + invoked: list[str] = [] + + class DurableModel(TestModel): + async def request(self, messages: Any, model_settings: Any, model_request_parameters: Any) -> ModelResponse: + payload = request_codec.dump_python((messages, model_settings, model_request_parameters), mode="json") + # This proof adapter explicitly canonicalizes framework message + # metadata; the SDK must not intercept wall-clock calls globally. + stamp = context.get().current_utc_datetime.replace(tzinfo=timezone.utc).isoformat() + for message in payload[0]: + message["timestamp"] = stamp + for part in message["parts"]: + if "timestamp" in part: + part["timestamp"] = stamp + result = await context.get().call_activity("model", input=payload) + return response_codec.validate_python(result) + + async def model(payload: Any) -> Any: + messages, settings, parameters = request_codec.validate_python(payload) + result = await TestModel(custom_output_text="done").request(messages, settings, parameters) + return response_codec.dump_python(result, mode="json") + + async def first_activity(value: int) -> str: + if fail_tool: + raise RuntimeError("recorded tool failure") + return f"first:{value}" + + async def second_activity(value: int) -> str: + return f"second:{value}" + + activities = {"model": model, "first": first_activity, "second": second_activity} + + class ReplayClock(AbstractCapability[None]): + def _run_clock(self) -> Any: + return lambda: context.get().current_utc_datetime.replace(tzinfo=timezone.utc) + + agent = Agent(DurableModel(), name="proof", capabilities=[ReplayClock()]) + + @agent.tool_plain + async def first(value: int) -> str: + return await context.get().call_activity("first", input=value, return_type=str) + + @agent.tool_plain + async def second(value: int) -> str: + return await context.get().call_activity("second", input=value, return_type=str) + + async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> str: + token = context.set(ctx) + try: + result = await agent.run( + "Call both tools.", run_id=str(ctx.new_uuid()), conversation_id=ctx.instance_id) + return result.output + finally: + context.reset(token) + + history = _start() + result = _execute(orchestrator, [], history) + batches: list[list[str]] = [] + while not any(action.HasField("completeOrchestration") for action in result.actions): + assert result.actions + scheduled = list(result.actions) + batches.append([action.scheduleTask.name for action in scheduled]) + events = [] + for action in scheduled: + assert action.HasField("scheduleTask") + name = action.scheduleTask.name + history.append(helpers.new_task_scheduled_event(action.id, name)) + payload = json.loads(action.scheduleTask.input.value) + invoked.append(name) + try: + output = asyncio.run(activities[name](payload)) + except RuntimeError as error: + events.append(helpers.new_task_failed_event(action.id, error)) + else: + events.append(helpers.new_task_completed_event(action.id, json.dumps(output))) + # Replay twice with identical recorded results: no model/tool body reruns. + before = list(invoked) + first_result = _execute(orchestrator, history, events) + second_result = _execute(orchestrator, history, events) + assert first_result.actions == second_result.actions + assert invoked == before + history += events + result = first_result + assert len(batches) < 10 + complete = _completion(result) + if fail_tool: + assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED + assert "recorded tool failure" in complete.failureDetails.errorMessage + assert batches == [["model"], ["first", "second"]] + assert invoked == ["model", "first", "second"] + return + assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED, complete.failureDetails + assert json.loads(complete.result.value) == "done" + assert batches == [["model"], ["first", "second"], ["model"]] + assert invoked == ["model", "first", "second", "model"] diff --git a/tests/durabletask/test_serialization.py b/tests/durabletask/test_serialization.py index 02861dcd..5bd7851a 100644 --- a/tests/durabletask/test_serialization.py +++ b/tests/durabletask/test_serialization.py @@ -9,20 +9,23 @@ """ import json +import os +import subprocess +import sys +from pathlib import Path from collections import namedtuple from dataclasses import dataclass from types import SimpleNamespace from typing import Any, List, Optional, Union, get_args import pytest - from durabletask.serialization import _AUTO_SERIALIZED as AUTO_SERIALIZED +from durabletask.serialization import JsonDataConverter from durabletask.serialization import _coerce_to_type as coerce_to_type from durabletask.serialization import _from_json as from_json from durabletask.serialization import _resolve_forward_refs as resolve_forward_refs from durabletask.serialization import _to_json as to_json - # ----- Test fixtures ----- @@ -548,7 +551,7 @@ def from_json(cls, value, strict=False): def test_resolve_forward_refs_resolves_string_element_type(): # ``list["Address"]`` evaluates to a generic whose arg is the raw string # "Address" -- exactly what 3.10 leaves behind. - resolved = resolve_forward_refs(list["Address"], {"Address": Address}) + resolved = resolve_forward_refs(List["Address"], {"Address": Address}) assert get_args(resolved) == (Address,) @@ -571,3 +574,222 @@ def test_resolve_forward_refs_leaves_unresolvable_name_untouched(): # the parsed JSON rather than raising. resolved = resolve_forward_refs(list["DoesNotExist"], {}) assert get_args(resolved) == ("DoesNotExist",) + + +# ----- set and frozenset container serialization and coercion ----- + + +@dataclass(frozen=True) +class FrozenAddress: + street: str + city: str + + +def test_to_json_serializes_set_and_frozenset(): + encoded_set = to_json({"b", "a"}) + assert json.loads(encoded_set) == ["a", "b"] + + encoded_frozenset = to_json(frozenset([2, 1])) + assert json.loads(encoded_frozenset) == [1, 2] + + # Numeric sets follow lexical JSON ordering (e.g. "10" < "2") + encoded_numeric_lexical = to_json({2, 10}) + assert json.loads(encoded_numeric_lexical) == [10, 2] + + # Verify disjoint and nested frozenset in-process serialization + encoded_disjoint = to_json({ + frozenset({"alpha"}), + frozenset({"beta"}), + frozenset({"gamma"}), + }) + assert json.loads(encoded_disjoint) == [["alpha"], ["beta"], ["gamma"]] + + encoded_nested = to_json({ + frozenset({frozenset({"b", "a"}), frozenset({"c"})}), + frozenset({frozenset({"d"})}), + }) + assert json.loads(encoded_nested) == [[["a", "b"], ["c"]], [["d"]]] + + +def test_set_serialization_non_finite_numbers(): + """Verify that sets with non-finite numbers (NaN, Infinity, -Infinity) + serialize deterministically with lexical JSON string sorting.""" + import math + + s = {float("nan"), float("inf"), float("-inf"), 2.0, 10.0} + encoded = to_json(s) + parsed = json.loads(encoded) + # Lexical order of canonical JSON strings: + # "-Infinity", "10.0", "2.0", "Infinity", "NaN" + assert parsed[0] == float("-inf") + assert parsed[1] == 10.0 + assert parsed[2] == 2.0 + assert parsed[3] == float("inf") + assert math.isnan(parsed[4]) + + +def test_set_serialization_equal_canonical_keys_differently_ordered_dicts(): + """Verify that custom objects with identical canonical JSON but differently ordered + dictionaries emit normalized values without order leaks.""" + class CustomDictObj: + def __init__(self, data: dict[str, Any], tag: str): + self.data = data + self.tag = tag + + def to_json(self) -> dict[str, Any]: + return self.data + + def __hash__(self) -> int: + return hash(self.tag) + + def __eq__(self, other: object) -> bool: + return isinstance(other, CustomDictObj) and self.tag == other.tag + + obj1 = CustomDictObj({"z": 100, "a": 1, "m": 50}, "obj1") + obj2 = CustomDictObj({"a": 1, "m": 50, "z": 100}, "obj2") + s = {obj1, obj2} + encoded = to_json(s) + parsed = json.loads(encoded) + # Both elements must be normalized to canonical sorted key order: {"a": 1, "m": 50, "z": 100} + assert len(parsed) == 2 + assert parsed[0] == {"a": 1, "m": 50, "z": 100} + assert parsed[1] == {"a": 1, "m": 50, "z": 100} + + +def test_set_serialization_deterministic_across_hash_seeds(): + """Verify that sets of disjoint frozensets, nested frozensets, unorderable custom objects, + non-finite numbers, and equal canonical keys with differently ordered dictionaries + are deterministically serialized across separate processes with differing PYTHONHASHSEED.""" + repo_root = str(Path(__file__).resolve().parents[2]) + code = """ +import json +from dataclasses import dataclass +from typing import Any +from durabletask.serialization import JsonDataConverter + +conv = JsonDataConverter() + +# Case A: Disjoint frozensets (exposes frozenset partial-ordering defects) +case_a = conv.serialize({ + frozenset({"alpha"}), + frozenset({"beta"}), + frozenset({"gamma"}), +}) + +# Case B: Nested frozensets (multi-element, multi-level nesting) +case_b = conv.serialize({ + frozenset({frozenset({"b", "a"}), frozenset({"c"})}), + frozenset({frozenset({"d"})}), +}) + +# Case C: Non-orderable custom objects (frozen dataclass without order=True) +@dataclass(frozen=True) +class Entry: + id: str + val: int + +case_c = conv.serialize({ + Entry("delta", 4), + Entry("alpha", 1), + Entry("charlie", 3), + Entry("bravo", 2), +}) + +# Case D: Numeric sets with lexical JSON ordering and non-finite floats +case_d = conv.serialize({float("nan"), float("inf"), float("-inf"), 2.0, 10.0}) + +# Case E: Equal canonical keys with differently ordered dictionaries +class DictItem: + def __init__(self, data: dict[str, Any], tag: str): + self.data = data + self.tag = tag + def to_json(self) -> dict[str, Any]: + return self.data + def __hash__(self) -> int: + return hash(self.tag) + def __eq__(self, other: object) -> bool: + return isinstance(other, DictItem) and self.tag == other.tag + +case_e = conv.serialize({ + DictItem({"z": 10, "a": 1}, "item1"), + DictItem({"a": 1, "z": 10}, "item2"), +}) + +print(json.dumps({"a": case_a, "b": case_b, "c": case_c, "d": case_d, "e": case_e})) +""" + seeds = ["0", "1", "2", "42", "999"] + outputs: dict[str, dict[str, str]] = {} + + for seed in seeds: + env = os.environ.copy() + env["PYTHONHASHSEED"] = seed + env["PYTHONPATH"] = ( + f"{repo_root}{os.pathsep}{env['PYTHONPATH']}" + if "PYTHONPATH" in env + else repo_root + ) + proc = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + env=env, + ) + assert proc.returncode == 0, ( + f"Subprocess failed for seed {seed}:\nstdout: {proc.stdout}\nstderr: {proc.stderr}" + ) + outputs[seed] = json.loads(proc.stdout.strip()) + + base = outputs[seeds[0]] + for seed in seeds[1:]: + assert outputs[seed] == base, ( + f"Deterministic serialization failed between seed {seeds[0]} and {seed}:\n" + f"Seed {seeds[0]}: {base}\nSeed {seed}: {outputs[seed]}" + ) + + expected_a = '[["alpha"], ["beta"], ["gamma"]]' + expected_b = '[[["a", "b"], ["c"]], [["d"]]]' + expected_c = ( + '[{"id": "alpha", "val": 1}, {"id": "bravo", "val": 2}, ' + '{"id": "charlie", "val": 3}, {"id": "delta", "val": 4}]' + ) + expected_d = '[-Infinity, 10.0, 2.0, Infinity, NaN]' + expected_e = '[{"a": 1, "z": 10}, {"a": 1, "z": 10}]' + + assert base["a"] == expected_a + assert base["b"] == expected_b + assert base["c"] == expected_c + assert base["d"] == expected_d + assert base["e"] == expected_e + + +def test_deserialize_coerces_set_when_opted_in(): + converter = JsonDataConverter(enable_generic_sets=True) + encoded = converter.serialize(["a", "b", "c"]) + assert encoded is not None + result = converter.deserialize(encoded, set[str]) + assert isinstance(result, set) + assert result == {"a", "b", "c"} + + encoded_int = converter.serialize([1, 2, 3]) + assert encoded_int is not None + result_frozen = converter.deserialize(encoded_int, frozenset[int]) + assert isinstance(result_frozen, frozenset) + assert result_frozen == frozenset({1, 2, 3}) + + encoded_dc = json.dumps([ + {"street": "1 Main St", "city": "Redmond"}, + {"street": "2 1st Ave", "city": "Seattle"}, + ]) + result_dc = converter.deserialize(encoded_dc, set[FrozenAddress]) + assert isinstance(result_dc, set) + assert all(isinstance(item, FrozenAddress) for item in result_dc) + assert FrozenAddress("1 Main St", "Redmond") in result_dc + + +def test_deserialize_preserves_raw_array_when_opted_out(): + default_converter = JsonDataConverter() + payload = '["b", "a", "b"]' + result = default_converter.deserialize(payload, set[str]) + # Preserves list and duplicates for replay compatibility + assert isinstance(result, list) + assert result == ["b", "a", "b"] diff --git a/tests/durabletask/test_worker_resiliency.py b/tests/durabletask/test_worker_resiliency.py index 28faa215..479816e4 100644 --- a/tests/durabletask/test_worker_resiliency.py +++ b/tests/durabletask/test_worker_resiliency.py @@ -746,6 +746,93 @@ def create_stub(channel): provided_channel.close.assert_not_called() +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["complete", "abandon"]) +async def test_orchestrator_uncertain_outcome_resets_only_originating_stream(monkeypatch, operation): + channel = MagicMock() + worker = TaskHubGrpcWorker(channel=channel) + manager = DummyWorkerManager() + worker._async_worker_manager = manager + monkeypatch.setattr(worker._shutdown, "wait", lambda timeout: False) + + first_stream = FakeResponseStream(items=[pb.WorkItem( + orchestratorRequest=pb.OrchestratorRequest(instanceId="instance-id"), + completionToken="original-token", + )]) + next_stream = FakeResponseStream() + first_stub = MagicMock() + first_stub.GetWorkItems.return_value = first_stream + first_stub.CompleteOrchestratorTask.side_effect = RuntimeError("completion outcome unknown") + first_stub.AbandonTaskOrchestratorWorkItem.side_effect = RuntimeError("abandonment outcome unknown") + next_stub = MagicMock() + next_stub.GetWorkItems.return_value = next_stream + stop_stub = MagicMock() + stop_stub.GetWorkItems.side_effect = _cancel_after_shutdown(worker) + stubs = [first_stub, next_stub, stop_stub] + monkeypatch.setattr("durabletask.worker.stubs.TaskHubSidecarServiceStub", lambda _: stubs.pop(0)) + await worker._async_run_loop() + + assert len(manager.submissions) == 1 + first_stream.cancelled = False + next_stream.cancelled = False + worker._response_stream = next_stream + _, submission = manager.submissions[0] + handler, cancel, request, original_stub, token = submission + if operation == "complete": + handler(request, original_stub, token) + first_stub.CompleteOrchestratorTask.assert_called_once() + assert first_stub.CompleteOrchestratorTask.call_args.args[0].completionToken == "original-token" + first_stub.AbandonTaskOrchestratorWorkItem.assert_not_called() + else: + with pytest.raises(RuntimeError, match="abandonment outcome unknown"): + cancel(request, original_stub, token) + first_stub.CompleteOrchestratorTask.assert_not_called() + first_stub.AbandonTaskOrchestratorWorkItem.assert_called_once() + assert first_stub.AbandonTaskOrchestratorWorkItem.call_args.args[0].completionToken == "original-token" + assert first_stream.cancelled + assert not next_stream.cancelled + channel.close.assert_not_called() + + +@pytest.mark.asyncio +async def test_orchestrator_delivery_failure_reconnects_live_work_item_stream(monkeypatch): + worker = TaskHubGrpcWorker( + channel=MagicMock(), + concurrency_options=ConcurrencyOptions( + maximum_concurrent_orchestration_work_items=1, + maximum_thread_pool_workers=1, + ), + ) + monkeypatch.setattr(worker._shutdown, "wait", lambda timeout: False) + cancelled = Event() + + class LiveStream: + def __iter__(self): + yield pb.WorkItem( + orchestratorRequest=pb.OrchestratorRequest(instanceId="instance-id"), + completionToken="original-token", + ) + assert cancelled.wait(timeout=5), "orphaned dispatch left its stream open" + + def cancel(self): + cancelled.set() + + first_stub = MagicMock() + first_stub.GetWorkItems.return_value = LiveStream() + first_stub.CompleteOrchestratorTask.side_effect = RuntimeError("completion outcome unknown") + next_stub = MagicMock() + next_stub.GetWorkItems.side_effect = _cancel_after_shutdown(worker) + stubs = [first_stub, next_stub] + monkeypatch.setattr("durabletask.worker.stubs.TaskHubSidecarServiceStub", lambda _: stubs.pop(0)) + + await asyncio.wait_for(worker._async_run_loop(), timeout=5) + + assert cancelled.is_set() + first_stub.CompleteOrchestratorTask.assert_called_once() + first_stub.AbandonTaskOrchestratorWorkItem.assert_not_called() + next_stub.GetWorkItems.assert_called_once() + + @pytest.mark.asyncio async def test_worker_uses_reconnect_backoff_helper_after_connection_failure(monkeypatch): worker = TaskHubGrpcWorker(