diff --git a/.github/scripts/install_otel_test_wheels.py b/.github/scripts/install_otel_test_wheels.py index 938996aca..b1332e5c2 100644 --- a/.github/scripts/install_otel_test_wheels.py +++ b/.github/scripts/install_otel_test_wheels.py @@ -16,6 +16,7 @@ ROOT = Path(__file__).resolve().parents[2] CORE = "aws-durable-execution-sdk-python" OTEL = CORE + "-otel" +TESTING = CORE + "-testing" def built_wheel(package: str) -> Path: @@ -72,7 +73,9 @@ def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--legacy-plugin", action="store_true") args = parser.parse_args() - packages = [CORE] if args.legacy_plugin else [CORE, OTEL] + # This PR also repairs the local service simulator. Validate its actual + # built artifact with the SDK/plugin, retaining every test assertion. + packages = [CORE] if args.legacy_plugin else [CORE, OTEL, TESTING] wheels = [built_wheel(package) for package in packages] subprocess.run( [ diff --git a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py index 034be671e..c27d0948d 100644 --- a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py +++ b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py @@ -1,8 +1,12 @@ +import re from pathlib import Path import yaml +SHARED_WORKFLOW_REF = "bdb4f1cd0f9252c1aaa978bb8b341b71f2b9d9dc" +CONFORMANCE_TEST_REF = "05960beb6c81d7b24b53114f2581b3b6f9fef6d4" + WORKFLOW_PATH = ( Path(__file__).parents[2] / "workflows" / "opentelemetry-conformance-tests.yml" ) @@ -18,7 +22,10 @@ def test_opentelemetry_conformance_caller_uses_current_workflow_contract() -> No "uses: aws/aws-durable-execution-conformance-tests/.github/workflows/" "opentelemetry-orchestrator.yml@" ) - assert orchestrator in workflow + pinned_ref = re.search(re.escape(orchestrator) + r"([0-9a-f]{40})", workflow) + assert pinned_ref is not None + assert pinned_ref.group(1) == SHARED_WORKFLOW_REF + assert f"default: {CONFORMANCE_TEST_REF}" in workflow assert "python-opentelemetry.yml@" not in workflow assert "\n otlp_endpoint:" not in workflow @@ -28,7 +35,7 @@ def test_opentelemetry_conformance_caller_uses_current_workflow_contract() -> No "resource_prefix: p", "sdk_repository: aws/aws-durable-execution-sdk-python", "sdk_ref: ${{ github.event.pull_request.head.sha || github.sha }}", - "conformance_test_ref: ${{ inputs.conformance_test_ref || 'main' }}", + f"conformance_test_ref: ${{{{ inputs.conformance_test_ref || '{CONFORMANCE_TEST_REF}' }}}}", "checkout_sdk: true", f"examples_dir: {EXAMPLES_DIR}", "adot_release_repository: aws-observability/aws-otel-python-instrumentation", diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 73b0413d9..c0e7bf84e 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -10,7 +10,7 @@ on: branches: [ main ] pull_request: - branches: [ main ] + branches: ["main"] jobs: lint-commits: diff --git a/.github/workflows/opentelemetry-conformance-tests.yml b/.github/workflows/opentelemetry-conformance-tests.yml index 6ce0099f5..c45d2e583 100644 --- a/.github/workflows/opentelemetry-conformance-tests.yml +++ b/.github/workflows/opentelemetry-conformance-tests.yml @@ -45,7 +45,7 @@ on: conformance_test_ref: description: Conformance test commit SHA or branch name required: true - default: main + default: 05960beb6c81d7b24b53114f2581b3b6f9fef6d4 type: string # Backend stacks are shared across PRs. Queue whole runs so reusable @@ -66,14 +66,14 @@ jobs: actions: write contents: read id-token: write - uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@e48ac54502d3c14e2e3fa216c5e200dbb4c14ac2 + uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@bdb4f1cd0f9252c1aaa978bb8b341b71f2b9d9dc with: language: python runs_on: codebuild-github-actions-runner-${{ github.run_id }}-${{ github.run_attempt }} resource_prefix: p sdk_repository: aws/aws-durable-execution-sdk-python sdk_ref: ${{ github.event.pull_request.head.sha || github.sha }} - conformance_test_ref: ${{ inputs.conformance_test_ref || 'main' }} + conformance_test_ref: ${{ inputs.conformance_test_ref || '05960beb6c81d7b24b53114f2581b3b6f9fef6d4' }} # Check the SDK out so the handlers and templates below are on disk. The handlers # themselves are installed from sdk_ref by src/requirements.txt during the SAM build. checkout_sdk: true diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index a3b6bc2f7..b583d7461 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -72,11 +72,11 @@ hatch run dev-examples:test # run examples tests only ### Installed package compatibility testing -Build the core and OTel distributions with `hatch build` in each package, then run +Build the core, OTel and testing-library distributions with `hatch build` in each package, then run these commands from the repository root: ```bash -hatch run test-wheel-otel:test # full OTel suite on the two built wheels +hatch run test-wheel-otel:test # full OTel suite on the three built wheels hatch run test-wheel-otel-legacy:test # released OTel 1.0.0 with the built core hatch run test-pypi-examples:test # examples against the published core ``` diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md index 5ff978c88..f8d59871d 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md @@ -33,7 +33,7 @@ template-long-running.yaml # otel-long-running suite tests/ # contract tests for the templates and handlers ``` -The 20 invocation and 20 execution requirements reuse the same scenario +The invocation and execution requirements reuse the same scenario handlers; the view is selected per function through the `OTEL_PLUGIN_MODE` environment variable, which `common.otel_plugin()` reads to pick `InvocationOtelPlugin` or `ExecutionOtelPlugin`. `template.yaml` deploys only the @@ -63,6 +63,12 @@ view named by its `OtelSuite` parameter. | `otel-invocation-18` | `otel_18_chained_invoke_failure.handler` | Verifies failed chained-invoke telemetry. | | `otel-invocation-19` | `otel_19_execution_failure.handler` | Verifies telemetry for a direct handler failure. | | `otel-invocation-20` | `otel_20_virtual_context.handler` | Verifies a virtual child-context span without context checkpoints. | +| `otel-invocation-21` | `otel_21_completed_step_replay.handler` | Replays a completed step after a successful wait/resume; each step body runs once. | +| `otel-invocation-22` | `otel_22_user_function_context.handler` | Creates ordinary user spans under the active handler, attempt, child, branch, and iteration contexts. | +| `otel-invocation-23` | `otel_23_callback_function_context.handler` | Verifies retry/check attempts, callback submitter, wrapped retry helper, and virtual-child callback parents. | +| `otel-invocation-24` | `otel_24_invocation_retry_status.handler` | Raises a retryable invocation error after a completed step, then resumes with its saved result. | +| `otel-invocation-25` | `otel_25_caught_callback_failure_replay.handler` | Catches the failed callback child, records `caught`, replays through two typed barriers, and returns `caught/one/two`. | +| `otel-invocation-26` | `otel_26_external_callback_completion_replay.handler` | Completes a root callback after suspension, saves its result in a step, and replays it through two callback barriers. | | `otel-execution-1` | `otel_1_success.handler` | Verifies the execution-view workflow, step, and attempt hierarchy. | | `otel-execution-2` | `otel_2_wait_resume.handler` | Verifies the execution view across a resumed invocation. | | `otel-execution-3` | `otel_3_retry.handler` | Verifies the execution view across retry attempts. | @@ -83,6 +89,12 @@ view named by its `OtelSuite` parameter. | `otel-execution-18` | `otel_18_chained_invoke_failure.handler` | Verifies source and target failed workflow roots. | | `otel-execution-19` | `otel_19_execution_failure.handler` | Verifies a failed invocation without a completed workflow. | | `otel-execution-20` | `otel_20_virtual_context.handler` | Verifies a virtual child-context span under the workflow root. | +| `otel-execution-21` | `otel_21_completed_step_replay.handler` | Verifies completed-operation spans are exported once across normal successful replay. | +| `otel-execution-22` | `otel_22_user_function_context.handler` | Observes active execution-view callback contexts without supplying or repairing parents. | +| `otel-execution-23` | `otel_23_callback_function_context.handler` | Verifies the same SDK-owned callback lifecycle parents in the execution view. | +| `otel-execution-24` | `otel_24_invocation_retry_status.handler` | Verifies invocation retry status independently of step retry and recovery re-exports. | +| `otel-execution-25` | `otel_25_caught_callback_failure_replay.handler` | Requires one `FAILED`/`ERROR` callback leaf without an invented exception event while the workflow succeeds after two replays. | +| `otel-execution-26` | `otel_26_external_callback_completion_replay.handler` | Requires one terminal root-callback export at first completion and no duplicate exports on two later replays. | | `otel-long-running-1` | `otel_long_running_1_wait.handler` | Verifies wait and resume telemetry across a long durable suspension. | | `otel-long-running-2` | `otel_long_running_2_retry.handler` | Verifies retry telemetry across a long durable backoff. | | `otel-long-running-3` | `otel_long_running_3_callback.handler` | Verifies callback telemetry when completion arrives after a long delay. | @@ -91,6 +103,65 @@ view named by its `OtelSuite` parameter. The runner discovers each mapping from `TestingMetadata.TestDescription` on the functions in the templates. +## External completion and status coverage + +The local testing library includes callback success, failure and timeout in the +next invocation's `UpdatedOperationIds`. The core delivers each terminal update +notification once per invocation, including when a resumed operation and a later +checkpoint response carry the same completion. It tracks actual notifications, +preserves first delivery when the update-ID metadata is absent, and clears that +tracking at invocation boundaries. Public runner regressions cover memory and +file stores, stored step results, failure payloads and two subsequent replays. +An exactly empty serialized callback error is represented as absent in both +stores, matching the SDK's service parser. Present fields remain intact, +including empty messages, types, data and stack lists; the enclosing failed +callback future still raises the same error as the file-store baseline. +Detailed callback-failure history retains the service's empty error `Payload` +object with `Truncated: false`, while SDK-facing state still has no error details. +Metadata-only history and nonempty errors retain their existing representation. + +Case 24 covers invocation `RETRY` becoming `RETRYING`/`UNSET`. Case 25 uses +`caught-callback-failure-replay`: it catches only `CallbackExternalError`, which +`wait_for_callback` restores from the failed child on replay, then checkpoints `caught` and awaits +two root barriers. The final result is `caught/one/two`, with three PENDING +invocations followed by SUCCEEDED. The failed callback remains `FAILED`/`ERROR` +without inventing an exception event. The marker body and failed submitter run +once; replay reconstructs the failed child without revisiting its callback body. +Marker and barrier failures escape the narrow catch. Case17 retains its +unhandled supplied-error behavior. + +This fixture must be paired with the Shared case25 catalog that expects a +successful execution and the three phase-gated callback actions. The previous +case25 catalog expects a failed execution and is not compatible with this new +routing. Cloud coverage requires the typed empty, untruncated callback error +payload and actual trace evidence in both views. The initial-active-invocation, +completion-during-flush and pruned-child-history boundaries are separate claims. +Other rare terminal statuses without details retain local coverage; no new +cloud path is claimed for them. + +Case 26 revisits a root-level public callback on every replay. The driver waits +for `InvocationCompleted` before completing the target and each barrier. The +target's terminal span must precede the observed step and must not be exported +again during the two later resumes. Invocation view also retains its initial +pending callback segment. Cloud validation must count the raw S3 export records +without deduplicating equal span IDs; local runner checks alone do not establish +cloud coverage. + +## Callback coverage boundary + +Cases 22 and 23 create normal user spans from the active context, without an +explicit parent or a copied execution ARN. They cover handler context, step and +condition attempts, child and branch bodies, callback submission, wrapped +`with_retry` body/strategy callbacks, and a virtual child. Case 23 places the +retry helper and virtual child after callback completion so ordinary successful +replay does not repeat their probes. + +This does not promise an operation/attempt parent for every arbitrary callback. +General retry/wait policies, serializers, summary generators, and item naming +have phase-specific caller scopes outside the wrapped user-function lifecycle; +these cases do not impose a new ownership policy on them. Instrumentation +extensions and user-created threads are outside this business-callback contract. + ## How a handler maps to a requirement ```yaml diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/Makefile b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/Makefile index 320577e7f..8382a49cb 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/Makefile +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/Makefile @@ -5,6 +5,8 @@ build-Otel14MapFailure build-Otel15WaitInterrupted build-Otel16WaitForConditionFailure \ build-Otel17WaitForCallbackFailure build-Otel18ChainedInvokeFailure \ build-Otel18InvokeTarget build-Otel19ExecutionFailure build-Otel20VirtualContext \ + build-Otel21CompletedStepReplay build-Otel22UserFunctionContext build-Otel23CallbackFunctionContext build-Otel24InvocationRetryStatus \ + build-Otel25CallbackFailureWithoutError build-Otel26ExternalCallbackReplay \ build-OtelExecution1Success \ build-OtelExecution2WaitResume build-OtelExecution3Retry \ build-OtelExecution4TerminalFailure build-OtelExecution5ChildContext \ @@ -17,6 +19,9 @@ build-OtelExecution17WaitForCallbackFailure build-OtelExecution18ChainedInvokeFailure \ build-OtelExecution18InvokeTarget build-OtelExecution19ExecutionFailure \ build-OtelExecution20VirtualContext \ + build-OtelExecution21CompletedStepReplay build-OtelExecution22UserFunctionContext build-OtelExecution23CallbackFunctionContext \ + build-OtelExecution24InvocationRetryStatus \ + build-OtelExecution25CallbackFailureWithoutError build-OtelExecution26ExternalCallbackReplay \ build-OtelLongRunning1Wait \ build-OtelLongRunning2Retry build-OtelLongRunning3Callback \ build-OtelLongRunning4ChainedInvoke build-OtelLongRunning4InvokeTarget @@ -28,6 +33,8 @@ build-Otel11InvokeTarget build-Otel12ChildContextFailure build-Otel13ParallelFai build-Otel14MapFailure build-Otel15WaitInterrupted build-Otel16WaitForConditionFailure \ build-Otel17WaitForCallbackFailure build-Otel18ChainedInvokeFailure \ build-Otel18InvokeTarget build-Otel19ExecutionFailure build-Otel20VirtualContext \ +build-Otel21CompletedStepReplay build-Otel22UserFunctionContext build-Otel23CallbackFunctionContext build-Otel24InvocationRetryStatus \ +build-Otel25CallbackFailureWithoutError build-Otel26ExternalCallbackReplay \ build-OtelExecution1Success \ build-OtelExecution2WaitResume build-OtelExecution3Retry \ build-OtelExecution4TerminalFailure build-OtelExecution5ChildContext \ @@ -40,6 +47,9 @@ build-OtelExecution16WaitForConditionFailure \ build-OtelExecution17WaitForCallbackFailure build-OtelExecution18ChainedInvokeFailure \ build-OtelExecution18InvokeTarget build-OtelExecution19ExecutionFailure \ build-OtelExecution20VirtualContext \ +build-OtelExecution21CompletedStepReplay build-OtelExecution22UserFunctionContext build-OtelExecution23CallbackFunctionContext \ +build-OtelExecution24InvocationRetryStatus \ +build-OtelExecution25CallbackFailureWithoutError build-OtelExecution26ExternalCallbackReplay \ build-OtelLongRunning1Wait \ build-OtelLongRunning2Retry build-OtelLongRunning3Callback \ build-OtelLongRunning4ChainedInvoke build-OtelLongRunning4InvokeTarget: diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_21_completed_step_replay.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_21_completed_step_replay.py new file mode 100644 index 000000000..2ea9ef4a0 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_21_completed_step_replay.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Normal successful replay of a completed step for OTel case 21.""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python import ( + DurableContext, + StepContext, + durable_execution, + durable_step, +) +from aws_durable_execution_sdk_python.config import Duration +from common import otel_plugin, require_scenario + + +@durable_step +def before_wait(_step_context: StepContext) -> str: + return "before" + + +@durable_step +def after_wait(_step_context: StepContext) -> str: + return "after" + + +@durable_execution(plugins=[otel_plugin()]) +def handler(event: dict[str, Any], context: DurableContext) -> str: + require_scenario(event, "completed-step-replay") + before = context.step(before_wait(), name="otel-before-wait") + context.wait(Duration.from_seconds(1), name="otel-replay-wait") + after = context.step(after_wait(), name="otel-after-wait") + return f"{before}-{after}" diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_22_user_function_context.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_22_user_function_context.py new file mode 100644 index 000000000..4f103ba4e --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_22_user_function_context.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Observe the real active SDK context inside public user callbacks.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +from aws_durable_execution_sdk_python import ( + DurableContext, + StepContext, + durable_execution, +) +from aws_durable_execution_sdk_python.config import ( + Duration, + MapConfig, + ParallelBranch, + ParallelConfig, +) +from common import otel_plugin, require_scenario +from opentelemetry import trace + + +def probe(label: str) -> None: + if not trace.get_current_span().get_span_context().is_valid: + raise RuntimeError(f"No active span for conformance.{label}") + span = trace.get_tracer("aws-durable-execution-conformance").start_span( + f"conformance.{label}", attributes={"conformance.callback": label} + ) + span.end() + + +def step_body(_context: StepContext) -> str: + probe("step") + return "step" + + +def child_step(_context: StepContext) -> str: + probe("child-step") + return "child-step" + + +def child_body(context: DurableContext) -> str: + probe("child") + context.step(child_step, name="otel-context-child-step") + probe("child-restored") + return "child" + + +def parallel_step_a(_context: StepContext) -> str: + probe("parallel-step-a") + return "a" + + +def parallel_step_b(_context: StepContext) -> str: + probe("parallel-step-b") + return "b" + + +def parallel_a(context: DurableContext) -> str: + probe("parallel-a") + return context.step(parallel_step_a, name="otel-context-branch-step-a") + + +def parallel_b(context: DurableContext) -> str: + probe("parallel-b") + return context.step(parallel_step_b, name="otel-context-branch-step-b") + + +def iteration_name(_item: int, index: int) -> str: + return ("otel-context-iteration-0", "otel-context-iteration-1")[index] + + +def mapper( + context: DurableContext, item: int, index: int, _items: Sequence[int] +) -> int: + probe(("map-0", "map-1")[index]) + + def map_step(_step_context: StepContext) -> int: + probe(("map-step-0", "map-step-1")[index]) + return item + + return context.step( + map_step, name=("otel-context-map-step-0", "otel-context-map-step-1")[index] + ) + + +@durable_execution(plugins=[otel_plugin()]) +def handler(event: dict[str, Any], context: DurableContext) -> str: + require_scenario(event, "user-function-context") + probe("handler") + context.step(step_body, name="otel-context-step") + context.run_in_child_context(child_body, name="otel-context-child") + context.parallel( + [ + ParallelBranch(parallel_a, name="otel-context-branch-a"), + ParallelBranch(parallel_b, name="otel-context-branch-b"), + ], + name="otel-context-parallel", + config=ParallelConfig(max_concurrency=2), + ) + context.map( + [0, 1], + mapper, + name="otel-context-map", + config=MapConfig(max_concurrency=2, item_namer=iteration_name), + ) + probe("handler-restored") + context.wait(Duration.from_seconds(1), name="otel-context-resume") + probe("handler-after-resume") + return "context-complete" diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_23_callback_function_context.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_23_callback_function_context.py new file mode 100644 index 000000000..c7fae30ec --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_23_callback_function_context.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Observe callback contexts at SDK-owned user-function lifecycle boundaries.""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python import ( + DurableContext, + StepContext, + durable_execution, +) +from aws_durable_execution_sdk_python.config import ( + ChildConfig, + Duration, + JitterStrategy, + StepConfig, +) +from aws_durable_execution_sdk_python.exceptions import ChildContextError +from aws_durable_execution_sdk_python.retries import ( + RetryDecision, + RetryStrategyConfig, + WithRetryConfig, + create_retry_strategy, + with_retry, +) +from aws_durable_execution_sdk_python.types import ( + WaitForCallbackContext, + WaitForConditionCheckContext, +) +from aws_durable_execution_sdk_python.waits import ( + WaitForConditionConfig, + WaitForConditionDecision, +) +from common import otel_plugin, require_scenario +from opentelemetry import trace + + +HELPER_FAILURE = "intentional-helper-failure" + + +def probe(label: str) -> None: + if not trace.get_current_span().get_span_context().is_valid: + raise RuntimeError(f"No active span for conformance.{label}") + span = trace.get_tracer("aws-durable-execution-conformance").start_span( + f"conformance.{label}", attributes={"conformance.callback": label} + ) + span.end() + + +def retry_step(context: StepContext) -> str: + probe(f"retry-attempt-{context.attempt}") + if context.attempt == 1: + raise RuntimeError("intentional-step-retry") + return "retried" + + +def check_condition(state: int, _context: WaitForConditionCheckContext) -> int: + next_state = state + 1 + probe(f"condition-check-{next_state}") + return next_state + + +def wait_strategy(state: int, _attempt: int) -> WaitForConditionDecision: + if state >= 2: + return WaitForConditionDecision.stop_polling() + return WaitForConditionDecision.continue_waiting(Duration.from_seconds(1)) + + +def submit_callback(_callback_id: str, _context: WaitForCallbackContext) -> None: + probe("callback-submitter") + + +def helper_body(_context: DurableContext, _attempt: int) -> str: + probe("with-retry-body") + raise RuntimeError(HELPER_FAILURE) + + +def helper_retry_strategy(_error: Exception, _attempt: int) -> RetryDecision: + probe("with-retry-strategy") + return RetryDecision.no_retry() + + +def virtual_child(_context: DurableContext) -> str: + probe("virtual-child") + return "virtual" + + +@durable_execution(plugins=[otel_plugin()]) +def handler(event: dict[str, Any], context: DurableContext) -> str: + require_scenario(event, "callback-function-context") + context.step( + retry_step, + name="otel-context-retry-step", + config=StepConfig( + retry_strategy=create_retry_strategy( + RetryStrategyConfig( + max_attempts=2, + initial_delay=Duration.from_seconds(1), + max_delay=Duration.from_seconds(1), + backoff_rate=1.0, + jitter_strategy=JitterStrategy.NONE, + retryable_error_types=[RuntimeError], + ) + ) + ), + ) + context.wait_for_condition( + check_condition, + name="otel-context-condition", + config=WaitForConditionConfig(initial_state=0, wait_strategy=wait_strategy), + ) + context.wait_for_callback(submit_callback, name="otel-context-callback") + + # These callbacks run after the last asynchronous wait, so normal replay + # does not repeat the helper body or the checkpointless virtual child. + try: + with_retry( + context, + helper_body, + WithRetryConfig(retry_strategy=helper_retry_strategy), + name="otel-context-with-retry", + ) + except ChildContextError as error: + if error.message != HELPER_FAILURE: + raise + else: + raise RuntimeError("Expected the intentional helper failure") + context.run_in_child_context( + virtual_child, + name="otel-context-virtual", + config=ChildConfig(is_virtual=True), + ) + return "callback-context-complete" diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_24_invocation_retry_status.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_24_invocation_retry_status.py new file mode 100644 index 000000000..568622517 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_24_invocation_retry_status.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""A real invocation retry for OTel status-mapping case 24.""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python import ( + DurableContext, + StepContext, + durable_execution, + durable_step, +) +from aws_durable_execution_sdk_python.exceptions import InvocationError +from common import otel_plugin, require_scenario + + +@durable_step +def before_invocation_retry(_step_context: StepContext) -> str: + return "saved" + + +@durable_execution(plugins=[otel_plugin()]) +def handler(event: dict[str, Any], context: DurableContext) -> str: + require_scenario(event, "invocation-retry-status") + # Capture this at entry: consuming the completed step can end replay mode. + entered_replay = context.is_replaying() + context.step(before_invocation_retry(), name="otel-before-invocation-retry") + if not entered_replay: + raise InvocationError("Conformance invocation retry after a completed step") + return "retry-complete" diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_25_caught_callback_failure_replay.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_25_caught_callback_failure_replay.py new file mode 100644 index 000000000..d9d1a2f45 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_25_caught_callback_failure_replay.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Catch a real callback failure, then replay it before a successful return.""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python import ( + DurableContext, + StepContext, + durable_execution, + durable_step, +) +from aws_durable_execution_sdk_python.config import WaitForCallbackConfig +from aws_durable_execution_sdk_python.exceptions import CallbackExternalError +from aws_durable_execution_sdk_python.serdes import JsonSerDes +from aws_durable_execution_sdk_python.types import LambdaContext, WaitForCallbackContext +from common import otel_plugin, require_scenario + + +def submit_callback(_callback_id: str, _context: WaitForCallbackContext) -> None: + return None + + +@durable_step +def observe_failure(_context: StepContext) -> str: + return "caught" + + +def _run(event: dict[str, Any], context: DurableContext) -> str: + require_scenario(event, "caught-callback-failure-replay") + config = WaitForCallbackConfig(serdes=JsonSerDes()) + try: + context.wait_for_callback( + submit_callback, name="otel-failed-callback", config=config + ) + except CallbackExternalError: + # wait_for_callback restores the callback-typed cause on both first + # consumption and replay of its persisted failed child context. + pass + else: + raise AssertionError("Expected the callback to fail") + + observed = context.step(observe_failure(), name="otel-failed-callback-observed") + one = context.wait_for_callback( + submit_callback, name="otel-failed-callback-barrier-one", config=config + ) + two = context.wait_for_callback( + submit_callback, name="otel-failed-callback-barrier-two", config=config + ) + return "/".join((observed, one, two)) + + +_durable_handler = durable_execution(_run, plugins=[otel_plugin()]) + + +def handler(event: Any, context: LambdaContext) -> Any: + return _durable_handler(event, context) diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_26_external_callback_completion_replay.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_26_external_callback_completion_replay.py new file mode 100644 index 000000000..4e562e2f5 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_26_external_callback_completion_replay.py @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Public callback completion followed by two controlled replay barriers.""" + +from __future__ import annotations + +from typing import Any, cast + +from aws_durable_execution_sdk_python import ( + DurableContext, + StepContext, + durable_execution, + durable_step, +) +from aws_durable_execution_sdk_python.config import ( + CallbackConfig, + WaitForCallbackConfig, +) +from aws_durable_execution_sdk_python.serdes import JsonSerDes +from aws_durable_execution_sdk_python.types import WaitForCallbackContext +from common import otel_plugin, require_scenario + + +def submit_callback(_callback_id: str, _context: WaitForCallbackContext) -> None: + return None + + +@durable_step +def observe_target(_context: StepContext, result: str) -> str: + return result + + +@durable_execution(plugins=[otel_plugin()]) +def handler(event: dict[str, Any], context: DurableContext) -> str: + require_scenario(event, "external-callback-completion-replay") + config = WaitForCallbackConfig(serdes=JsonSerDes()) + target = cast( + str, + context.create_callback( + name="otel-external-target", config=CallbackConfig(serdes=JsonSerDes()) + ).result(), + ) + observed = context.step( + observe_target(target), name="otel-external-target-observed" + ) + one = context.wait_for_callback( + submit_callback, name="otel-external-barrier-one", config=config + ) + two = context.wait_for_callback( + submit_callback, name="otel-external-barrier-two", config=config + ) + return "/".join((observed, one, two)) diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/template.yaml b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/template.yaml index be70e122d..819b253b7 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/template.yaml +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/template.yaml @@ -409,6 +409,90 @@ Resources: FunctionName: !Sub "${AWS::StackName}-otel-invocation-20" Role: !Ref LambdaExecutionRoleArn + Otel21CompletedStepReplay: + Type: AWS::Serverless::Function + Condition: DeployInvocationView + Metadata: + BuildMethod: makefile + TestingMetadata: + TestDescription: + - otel-invocation-21 + Properties: + CodeUri: src/ + Handler: otel_21_completed_step_replay.handler + FunctionName: !Sub "${AWS::StackName}-otel-invocation-21" + Role: !Ref LambdaExecutionRoleArn + + Otel22UserFunctionContext: + Type: AWS::Serverless::Function + Condition: DeployInvocationView + Metadata: + BuildMethod: makefile + TestingMetadata: + TestDescription: + - otel-invocation-22 + Properties: + CodeUri: src/ + Handler: otel_22_user_function_context.handler + FunctionName: !Sub "${AWS::StackName}-otel-invocation-22" + Role: !Ref LambdaExecutionRoleArn + + Otel23CallbackFunctionContext: + Type: AWS::Serverless::Function + Condition: DeployInvocationView + Metadata: + BuildMethod: makefile + TestingMetadata: + TestDescription: + - otel-invocation-23 + Properties: + CodeUri: src/ + Handler: otel_23_callback_function_context.handler + FunctionName: !Sub "${AWS::StackName}-otel-invocation-23" + Role: !Ref LambdaExecutionRoleArn + + Otel24InvocationRetryStatus: + Type: AWS::Serverless::Function + Condition: DeployInvocationView + Metadata: + BuildMethod: makefile + TestingMetadata: + TestDescription: + - otel-invocation-24 + Properties: + CodeUri: src/ + Handler: otel_24_invocation_retry_status.handler + FunctionName: !Sub "${AWS::StackName}-otel-invocation-24" + Role: !Ref LambdaExecutionRoleArn + + Otel25CallbackFailureWithoutError: + Type: AWS::Serverless::Function + Condition: DeployInvocationView + Metadata: + BuildMethod: makefile + TestingMetadata: + TestDescription: + - otel-invocation-25 + Properties: + CodeUri: src/ + Handler: otel_25_caught_callback_failure_replay.handler + FunctionName: !Sub "${AWS::StackName}-otel-invocation-25" + Role: !Ref LambdaExecutionRoleArn + + Otel26ExternalCallbackReplay: + Type: AWS::Serverless::Function + Condition: DeployInvocationView + Metadata: + BuildMethod: makefile + TestingMetadata: + TestDescription: + - otel-invocation-26 + Properties: + CodeUri: src/ + Handler: otel_26_external_callback_completion_replay.handler + FunctionName: !Sub "${AWS::StackName}-otel-invocation-26" + Role: !Ref LambdaExecutionRoleArn + OtelExecution1Success: Type: AWS::Serverless::Function Condition: DeployExecutionView @@ -781,3 +865,105 @@ Resources: Environment: Variables: OTEL_PLUGIN_MODE: execution + + OtelExecution21CompletedStepReplay: + Type: AWS::Serverless::Function + Condition: DeployExecutionView + Metadata: + BuildMethod: makefile + TestingMetadata: + TestDescription: + - otel-execution-21 + Properties: + CodeUri: src/ + Handler: otel_21_completed_step_replay.handler + FunctionName: !Sub "${AWS::StackName}-otel-execution-21" + Role: !Ref LambdaExecutionRoleArn + Environment: + Variables: + OTEL_PLUGIN_MODE: execution + + OtelExecution22UserFunctionContext: + Type: AWS::Serverless::Function + Condition: DeployExecutionView + Metadata: + BuildMethod: makefile + TestingMetadata: + TestDescription: + - otel-execution-22 + Properties: + CodeUri: src/ + Handler: otel_22_user_function_context.handler + FunctionName: !Sub "${AWS::StackName}-otel-execution-22" + Role: !Ref LambdaExecutionRoleArn + Environment: + Variables: + OTEL_PLUGIN_MODE: execution + + OtelExecution23CallbackFunctionContext: + Type: AWS::Serverless::Function + Condition: DeployExecutionView + Metadata: + BuildMethod: makefile + TestingMetadata: + TestDescription: + - otel-execution-23 + Properties: + CodeUri: src/ + Handler: otel_23_callback_function_context.handler + FunctionName: !Sub "${AWS::StackName}-otel-execution-23" + Role: !Ref LambdaExecutionRoleArn + Environment: + Variables: + OTEL_PLUGIN_MODE: execution + + OtelExecution24InvocationRetryStatus: + Type: AWS::Serverless::Function + Condition: DeployExecutionView + Metadata: + BuildMethod: makefile + TestingMetadata: + TestDescription: + - otel-execution-24 + Properties: + CodeUri: src/ + Handler: otel_24_invocation_retry_status.handler + FunctionName: !Sub "${AWS::StackName}-otel-execution-24" + Role: !Ref LambdaExecutionRoleArn + Environment: + Variables: + OTEL_PLUGIN_MODE: execution + + OtelExecution25CallbackFailureWithoutError: + Type: AWS::Serverless::Function + Condition: DeployExecutionView + Metadata: + BuildMethod: makefile + TestingMetadata: + TestDescription: + - otel-execution-25 + Properties: + CodeUri: src/ + Handler: otel_25_caught_callback_failure_replay.handler + FunctionName: !Sub "${AWS::StackName}-otel-execution-25" + Role: !Ref LambdaExecutionRoleArn + Environment: + Variables: + OTEL_PLUGIN_MODE: execution + + OtelExecution26ExternalCallbackReplay: + Type: AWS::Serverless::Function + Condition: DeployExecutionView + Metadata: + BuildMethod: makefile + TestingMetadata: + TestDescription: + - otel-execution-26 + Properties: + CodeUri: src/ + Handler: otel_26_external_callback_completion_replay.handler + FunctionName: !Sub "${AWS::StackName}-otel-execution-26" + Role: !Ref LambdaExecutionRoleArn + Environment: + Variables: + OTEL_PLUGIN_MODE: execution diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_callback_failure_without_error.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_callback_failure_without_error.py new file mode 100644 index 000000000..2c7f41b80 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_callback_failure_without_error.py @@ -0,0 +1,268 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Public callback failures preserve error-detail semantics across stores.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +import time +from pathlib import Path + +import pytest +from aws_durable_execution_sdk_python.lambda_service import ErrorObject +from aws_durable_execution_sdk_python.plugin import OperationEndInfo, OperationType +from aws_durable_execution_sdk_python_testing.model import ( + GetDurableExecutionHistoryResponse, + events_to_operations, +) +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from aws_durable_execution_sdk_python_testing.stores.filesystem import ( + FileSystemExecutionStore, +) +from opentelemetry import context +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +SRC_DIR = Path(__file__).resolve().parents[1] / "src" + + +def _run_public_callback_failure( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + plugin_type: type[ExecutionOtelPlugin] | type[InvocationOtelPlugin], + filesystem: bool, + error_kind: str, +) -> dict[str, object]: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = plugin_type( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + callback_errors: list[ErrorObject | None] = [] + original_end = plugin.on_operation_end + + def observe_error(info: OperationEndInfo) -> None: + if info.operation_type is OperationType.CALLBACK: + callback_errors.append(info.error) + original_end(info) + + monkeypatch.setattr(plugin, "on_operation_end", observe_error) + original_context = context.get_current() + common_spec = importlib.util.spec_from_file_location( + "common", SRC_DIR / "common.py" + ) + assert common_spec is not None and common_spec.loader is not None + common = importlib.util.module_from_spec(common_spec) + monkeypatch.setitem(sys.modules, "common", common) + common_spec.loader.exec_module(common) + monkeypatch.setattr(common, "otel_plugin", lambda: plugin) + spec = importlib.util.spec_from_file_location( + "otel_17_wait_for_callback_failure", + SRC_DIR / "otel_17_wait_for_callback_failure.py", + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + error = { + "omitted": None, + "empty": ErrorObject(message=None, type=None, data=None, stack_trace=None), + "rich": ErrorObject( + message="explicit failure", + type="CallbackFailure", + data=None, + stack_trace=None, + ), + "empty-message": ErrorObject( + message="", type=None, data=None, stack_trace=None + ), + "empty-stack": ErrorObject(message=None, type=None, data=None, stack_trace=[]), + "empty-type": ErrorObject(message=None, type="", data=None, stack_trace=None), + "empty-data": ErrorObject(message=None, type=None, data="", stack_trace=None), + }[error_kind] + try: + store = FileSystemExecutionStore.create(tmp_path) if filesystem else None + with DurableFunctionTestRunner( + handler=module.handler, + store=store, + poll_interval=0.01, + execution_timeout=20, + skip_time=False, + ) as runner: + arn = runner.run_async( + input=json.dumps({"scenario": "wait-for-callback-failure"}) + ) + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + events = runner.get_execution_history( + arn, include_execution_data=True + ).events + starts = [ + event for event in events if event.event_type == "CallbackStarted" + ] + if starts and any( + event.event_type == "InvocationCompleted" + and event.event_id > starts[0].event_id + for event in events + ): + details = starts[0].callback_started_details + assert details is not None and details.callback_id is not None + callback_id = details.callback_id + break + time.sleep(0.01) + else: + raise AssertionError("Callback did not suspend before failure") + if error_kind == "omitted": + runner.send_callback_failure(callback_id) + else: + runner.send_callback_failure(callback_id, error=error) + result = runner.wait_for_result(arn, timeout=10) + with_data = runner.get_execution_history(arn, include_execution_data=True) + without_data = runner.get_execution_history( + arn, include_execution_data=False + ) + + assert result.status.value == "FAILED" + spans = exporter.get_finished_spans() + leaves = [ + span + for span in spans + if (span.attributes or {}).get("durable.operation.type") == "CALLBACK" + and (span.attributes or {}).get("durable.operation.status") == "FAILED" + ] + assert len(leaves) == 1 + has_details = error_kind not in {"omitted", "empty"} + assert [ + item.to_dict() if item is not None else None for item in callback_errors + ] == [error.to_dict() if has_details and error is not None else None] + if has_details and not filesystem: + assert callback_errors[0] is error + assert leaves[0].status.status_code.name == "ERROR" + assert [event.name for event in leaves[0].events] == ( + ["exception"] if has_details else [] + ) + # The public failed future still raises, so its enclosing context fails. + parents = [ + span + for span in spans + if span.name == "otel-failed-callback" + and (span.attributes or {}).get("durable.operation.status") == "FAILED" + ] + assert len(parents) == 1 + assert parents[0].status.status_code.name == "ERROR" + invocations = sorted( + [span for span in spans if span.name == "Invocation"], + key=lambda span: span.start_time or 0, + ) + statuses = [ + (span.attributes or {}).get("durable.invocation.status") + for span in invocations + ] + assert statuses == ["PENDING", "FAILED"] + assert result.error is not None + assert context.get_current() == original_context + for history in (with_data, without_data): + wire = history.to_dict() + decoded = GetDurableExecutionHistoryResponse.from_dict(wire) + original_error = next( + event["CallbackFailedDetails"]["Error"] + for event in wire["Events"] + if event["EventType"] == "CallbackFailed" + ) + for candidate in (history, decoded): + assert ( + next( + event.to_dict()["CallbackFailedDetails"]["Error"] + for event in candidate.events + if event.event_type == "CallbackFailed" + ) + == original_error + ) + callback = next( + operation + for operation in events_to_operations(candidate.events) + if operation.callback_details is not None + ) + details = callback.callback_details + assert details is not None + assert ( + details.error.to_dict() if details.error is not None else None + ) == (error.to_dict() if has_details and error is not None else None) + return { + "status": result.status.value, + "caller_error": dict(result.error.to_dict()), + "invocation_statuses": statuses, + "submitted_error": dict(error.to_dict()) if error is not None else {}, + "history_error": next( + event.to_dict()["CallbackFailedDetails"]["Error"] + for event in with_data.events + if event.event_type == "CallbackFailed" + ), + "metadata_error": next( + event.to_dict()["CallbackFailedDetails"]["Error"] + for event in without_data.events + if event.event_type == "CallbackFailed" + ), + } + finally: + provider.shutdown() + + +@pytest.mark.parametrize("plugin_type", [ExecutionOtelPlugin, InvocationOtelPlugin]) +@pytest.mark.parametrize( + "error_kind", + [ + "omitted", + "empty", + "rich", + "empty-message", + "empty-stack", + "empty-type", + "empty-data", + ], +) +def test_public_callback_failure_preserves_error_details( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + plugin_type: type[ExecutionOtelPlugin] | type[InvocationOtelPlugin], + error_kind: str, +) -> None: + memory = _run_public_callback_failure( + monkeypatch, tmp_path / "memory", plugin_type, False, error_kind + ) + filesystem = _run_public_callback_failure( + monkeypatch, tmp_path / "filesystem", plugin_type, True, error_kind + ) + assert memory == filesystem + expected_payload = memory["submitted_error"] + for outcome in [memory, filesystem]: + # Actual AWS history retains an empty Payload object for this failure, + # independently of the SDK-facing absent callback error. + # Retain the existing flags for nonempty errors; only the observed + # no-details projection is changed here. + assert outcome["history_error"] == { + "Payload": expected_payload, + "Truncated": bool(expected_payload), + } + # Preserve this API's existing metadata-only projection. + assert outcome["metadata_error"] == ( + {"Payload": expected_payload, "Truncated": True} + if expected_payload + else {"Truncated": True} + ) diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_callback_function_context.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_callback_function_context.py new file mode 100644 index 000000000..2aecd8cc6 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_callback_function_context.py @@ -0,0 +1,122 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Verify actual SDK callback parents across retries and external completion.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +import time +from collections import Counter +from pathlib import Path + +import pytest +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry import context, trace +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +SRC_DIR = Path(__file__).resolve().parents[1] / "src" +EXPECTED_PARENTS = { + "retry-attempt-1": "otel-context-retry-step attempt 1", + "retry-attempt-2": "otel-context-retry-step attempt 2", + "condition-check-1": "otel-context-condition attempt 1", + "condition-check-2": "otel-context-condition attempt 2", + "callback-submitter": "otel-context-callback submitter attempt 1", + "with-retry-body": "otel-context-with-retry", + "with-retry-strategy": "otel-context-with-retry", + "virtual-child": "otel-context-virtual", +} + + +@pytest.mark.parametrize("plugin_type", [ExecutionOtelPlugin, InvocationOtelPlugin]) +def test_callback_probes_keep_their_actual_sdk_parent( + monkeypatch: pytest.MonkeyPatch, + plugin_type: type[ExecutionOtelPlugin] | type[InvocationOtelPlugin], +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + monkeypatch.setattr(trace, "get_tracer_provider", lambda: provider) + plugin = plugin_type( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + before_context = context.get_current() + common_spec = importlib.util.spec_from_file_location( + "common", SRC_DIR / "common.py" + ) + assert common_spec is not None and common_spec.loader is not None + common = importlib.util.module_from_spec(common_spec) + monkeypatch.setitem(sys.modules, "common", common) + common_spec.loader.exec_module(common) + monkeypatch.setattr(common, "otel_plugin", lambda: plugin) + spec = importlib.util.spec_from_file_location( + "otel_23_callback_function_context", + SRC_DIR / "otel_23_callback_function_context.py", + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + try: + with DurableFunctionTestRunner(handler=module.handler) as runner: + arn = runner.run_async( + input=json.dumps({"scenario": "callback-function-context"}), timeout=30 + ) + callback_id = runner.wait_for_callback( + arn, name="otel-context-callback create callback id", timeout=10 + ) + # External delivery delay, outside the durable handler. + time.sleep(1) + runner.send_callback_success(callback_id, result=b"callback-complete") + result = runner.wait_for_result(arn, timeout=15) + assert result.status.value == "SUCCEEDED" + assert result.result is not None + assert json.loads(result.result) == "callback-context-complete" + spans = exporter.get_finished_spans() + probes = [span for span in spans if span.name.startswith("conformance.")] + assert Counter(span.name for span in probes) == Counter( + f"conformance.{label}" for label in EXPECTED_PARENTS + ) + spans_by_id = { + (span.context.trace_id, span.context.span_id): span + for span in spans + if span.context is not None + } + sdk_traces = { + span.context.trace_id + for span in spans + if span.context is not None + and span.attributes is not None + and span.attributes.get("durable.execution.arn") == arn + } + for probe_span in probes: + assert probe_span.context is not None + assert probe_span.parent is not None + assert probe_span.attributes is not None + label = str(probe_span.attributes["conformance.callback"]) + assert "durable.execution.arn" not in probe_span.attributes + assert probe_span.context.trace_id in sdk_traces + parent = spans_by_id[ + (probe_span.context.trace_id, probe_span.parent.span_id) + ] + assert parent.name == EXPECTED_PARENTS[label] + assert parent.attributes is not None + assert parent.attributes["durable.execution.arn"] == arn + assert context.get_current() == before_context + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_caught_callback_failure_replay.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_caught_callback_failure_replay.py new file mode 100644 index 000000000..676227def --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_caught_callback_failure_replay.py @@ -0,0 +1,358 @@ +"""The real case25 handler catches only its intended failure across replay.""" + +import importlib.util +import json +import os +from pathlib import Path +import sys +import time + +import pytest +from aws_durable_execution_sdk_python import DurableContext +from aws_durable_execution_sdk_python.exceptions import ( + CallbackExternalError, + ChildContextError, +) +from aws_durable_execution_sdk_python.lambda_service import ErrorObject +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from aws_durable_execution_sdk_python_testing.stores.filesystem import ( + FileSystemExecutionStore, +) +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +SRC = Path(__file__).resolve().parents[1] / "src" +FAILED = "otel-failed-callback" +OBSERVED = FAILED + "-observed" +BARRIERS = [FAILED + "-barrier-one", FAILED + "-barrier-two"] +EXPECTED_HISTORY = [ + "ExecutionStarted", + "ContextStarted", + "CallbackStarted", + "StepStarted", + "StepSucceeded", + "InvocationCompleted", + "CallbackFailed", + "ContextFailed", + "StepStarted", + "StepSucceeded", + "ContextStarted", + "CallbackStarted", + "StepStarted", + "StepSucceeded", + "InvocationCompleted", + "CallbackSucceeded", + "ContextSucceeded", + "ContextStarted", + "CallbackStarted", + "StepStarted", + "StepSucceeded", + "InvocationCompleted", + "CallbackSucceeded", + "ContextSucceeded", + "InvocationCompleted", + "ExecutionSucceeded", +] + + +def _callback_after_suspend(runner, arn, name): + deadline = time.monotonic() + 8 + while time.monotonic() < deadline: + events = runner.get_execution_history(arn, include_execution_data=True).events + starts = [ + e + for e in events + if e.event_type == "CallbackStarted" + and e.name == name + " create callback id" + ] + ends = [ + e + for e in events + if starts + and e.event_type == "InvocationCompleted" + and e.event_id > starts[0].event_id + ] + if starts and ends: + return starts[0].callback_started_details.callback_id, ( + starts[0].event_id, + ends[0].event_id, + ) + time.sleep(0.01) + raise AssertionError(f"No creating invocation completion for {name}") + + +def _load(monkeypatch, plugin): + spec = importlib.util.spec_from_file_location("common", SRC / "common.py") + common = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, "common", common) + spec.loader.exec_module(common) + monkeypatch.setattr(common, "otel_plugin", lambda: plugin) + spec = importlib.util.spec_from_file_location( + "otel_25_caught_callback_failure_replay", + SRC / "otel_25_caught_callback_failure_replay.py", + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.mark.parametrize( + "view,mode", + [ + (ExecutionOtelPlugin, "normal"), + (InvocationOtelPlugin, "normal"), + (ExecutionOtelPlugin, "unexpected-success"), + (ExecutionOtelPlugin, "barrier-failure"), + ], +) +def test_case25_public_caught_failure_and_replay(monkeypatch, tmp_path, view, mode): + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + starts = [] + invocation_ends = [] + exported = [] + operations = [] + user_starts = [] + failures = [] + marker_calls = [] + + class Exporter(InMemorySpanExporter): + def export(self, spans): + exported.extend((s, len(starts)) for s in spans) + return super().export(spans) + + exporter = Exporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + + def record(method, collection): + original = getattr(plugin, method) + + def wrapped(info): + collection.append((len(starts), info)) + return original(info) + + monkeypatch.setattr(plugin, method, wrapped) + + original_start = plugin.on_invocation_start + + def on_start(info): + starts.append(info) + return original_start(info) + + monkeypatch.setattr(plugin, "on_invocation_start", on_start) + record("on_invocation_end", invocation_ends) + record("on_operation_end", operations) + record("on_user_function_start", user_starts) + original_wait = DurableContext.wait_for_callback + + def wait(context, *args, **kwargs): + try: + return original_wait(context, *args, **kwargs) + except (CallbackExternalError, ChildContextError) as error: + if kwargs.get("name") == FAILED: + failures.append( + { + "invocation": len(starts), + "class": type(error).__name__, + "error_type": error.error_type, + "message": error.message, + } + ) + raise + + monkeypatch.setattr(DurableContext, "wait_for_callback", wait) + module = _load(monkeypatch, plugin) + original_marker = module.observe_failure + + def counted_marker(): + step = original_marker() + + def run(context): + marker_calls.append(len(starts)) + return step(context) + + return run + + monkeypatch.setattr(module, "observe_failure", counted_marker) + gates = [] + try: + with DurableFunctionTestRunner( + handler=module.handler, + store=FileSystemExecutionStore.create(tmp_path / "store"), + skip_time=False, + poll_interval=0.01, + execution_timeout=30, + ) as runner: + arn = runner.run_async( + input=json.dumps({"scenario": "caught-callback-failure-replay"}) + ) + callback, gate = _callback_after_suspend(runner, arn, FAILED) + gates.append(gate) + if mode == "unexpected-success": + runner.send_callback_success(callback, result=b'"unexpected"') + else: + runner.send_callback_failure(callback) + for name, result in zip(BARRIERS, ("one", "two")): + callback, gate = _callback_after_suspend(runner, arn, name) + gates.append(gate) + if mode == "barrier-failure": + runner.send_callback_failure( + callback, error=ErrorObject.from_message("barrier failure") + ) + break + runner.send_callback_success( + callback, result=json.dumps(result).encode() + ) + result = runner.wait_for_result(arn, timeout=10) + history = runner.get_execution_history( + arn, include_execution_data=True + ).to_dict() + span_rows = [ + { + "name": s.name, + "invocation": n, + "trace_id": f"{s.context.trace_id:032x}", + "span_id": f"{s.context.span_id:016x}", + "parent_id": f"{s.parent.span_id:016x}" if s.parent else None, + "status": s.status.status_code.name, + "start_time": s.start_time, + "end_time": s.end_time, + "attributes": dict(s.attributes or {}), + "events": [ + { + "name": e.name, + "timestamp": e.timestamp, + "attributes": dict(e.attributes or {}), + } + for e in s.events + ], + "links": [ + { + "trace_id": f"{link.context.trace_id:032x}", + "span_id": f"{link.context.span_id:016x}", + "attributes": dict(link.attributes or {}), + } + for link in s.links + ], + } + for s, n in exported + ] + summary = { + "view": view.__name__, + "mode": mode, + "execution_status": result.status.value, + "result": result.result, + "error": result.error.to_dict() if result.error else None, + "invocation_statuses": [info.status.value for _, info in invocation_ends], + "gates": gates, + "caught_failures": failures, + "marker_body_invocations": marker_calls, + "operation_ends": [ + { + "invocation": n, + "id": i.operation_id, + "name": i.name, + "parent_id": i.parent_id, + "type": i.operation_type.value, + "status": i.status.value, + "error": i.error.to_dict() if i.error else None, + } + for n, i in operations + ], + "user_function_starts": [ + { + "invocation": n, + "id": i.operation_id, + "name": i.name, + "type": i.operation_type.value, + } + for n, i in user_starts + ], + } + proof_dir = Path(os.environ.get("PYTHON_CASE25_PROOF_DIR", str(tmp_path))) + proof_dir.mkdir(parents=True, exist_ok=True) + prefix = f"python-case25-{view.__name__}-{mode}" + for suffix, data in ( + ("history", history), + ("spans", span_rows), + ("summary", summary), + ): + (proof_dir / f"{prefix}-{suffix}.json").write_text( + json.dumps(data, indent=2, default=str) + ) + if mode != "normal": + assert result.status.value == "FAILED" + assert marker_calls == ([] if mode == "unexpected-success" else [2]) + assert not any(e.get("Name") == BARRIERS[1] for e in history["Events"]) + return + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == "caught/one/two" + assert summary["invocation_statuses"] == [ + "PENDING", + "PENDING", + "PENDING", + "SUCCEEDED", + ] + assert [e["EventType"] for e in history["Events"]] == EXPECTED_HISTORY + assert [e["EventId"] for e in history["Events"]] == list(range(1, 27)) + assert gates == [(3, 6), (12, 15), (19, 22)] + assert marker_calls == [2] + assert [f["invocation"] for f in failures] == [2, 3, 4] + assert {f["class"] for f in failures} == {"CallbackExternalError"} + assert len({f["error_type"] for f in failures}) == 1 + failed = history["Events"][6] + assert failed["CallbackFailedDetails"]["Error"] == { + "Payload": {}, + "Truncated": False, + } + target_id = failed["Id"] + leaves = [ + s + for s in span_rows + if s["attributes"].get("durable.operation.id") == target_id + and s["attributes"].get("durable.operation.status") == "FAILED" + ] + assert len(leaves) == 1 + leaf = leaves[0] + assert ( + leaf["invocation"] == 2 + and leaf["status"] == "ERROR" + and leaf["events"] == [] + ) + marker_spans = [ + s + for s in span_rows + if s["attributes"].get("durable.operation.name") == OBSERVED + ] + assert marker_spans and all(s["invocation"] == 2 for s in marker_spans) + assert leaf["end_time"] <= min(s["start_time"] for s in marker_spans) + workflow = [s for s in span_rows if s["name"] == "Workflow"] + assert len(workflow) == 1 and workflow[0]["status"] == "OK" + parent = next(s for s in span_rows if s["span_id"] == leaf["parent_id"]) + assert ( + parent["attributes"]["durable.operation.id"] == history["Events"][1]["Id"] + ) + if view is InvocationOtelPlugin: + invocation = next( + s for s in span_rows if s["span_id"] == parent["parent_id"] + ) + assert invocation["name"] == "Invocation" and invocation["invocation"] == 2 + submitter_id = history["Events"][3]["Id"] + assert sum(i.operation_id == submitter_id for _, i in user_starts) == 1 + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_completed_step_replay.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_completed_step_replay.py new file mode 100644 index 000000000..d331decca --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_completed_step_replay.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Exercise case 21 through the public decorator and local durable runner.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +from pathlib import Path +from typing import Any + +import pytest +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry import context +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +SRC_DIR = Path(__file__).resolve().parents[1] / "src" + + +@pytest.mark.parametrize("plugin_type", [ExecutionOtelPlugin, InvocationOtelPlugin]) +def test_completed_step_replay_uses_saved_result( + monkeypatch: pytest.MonkeyPatch, + plugin_type: type[ExecutionOtelPlugin] | type[InvocationOtelPlugin], +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = plugin_type( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + before_context = context.get_current() + common_spec = importlib.util.spec_from_file_location( + "common", SRC_DIR / "common.py" + ) + assert common_spec is not None and common_spec.loader is not None + common = importlib.util.module_from_spec(common_spec) + monkeypatch.setitem(sys.modules, "common", common) + common_spec.loader.exec_module(common) + monkeypatch.setattr(common, "otel_plugin", lambda: plugin) + spec = importlib.util.spec_from_file_location( + "otel_21_completed_step_replay", SRC_DIR / "otel_21_completed_step_replay.py" + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + calls: list[str] = [] + + def track(factory: Any, label: str) -> Any: + def make_step() -> Any: + step = factory() + + def counted(step_context: Any) -> str: + calls.append(label) + return step(step_context) + + return counted + + return make_step + + monkeypatch.setattr(module, "before_wait", track(module.before_wait, "before")) + monkeypatch.setattr(module, "after_wait", track(module.after_wait, "after")) + try: + with DurableFunctionTestRunner(handler=module.handler) as runner: + result = runner.run( + input=json.dumps({"scenario": "completed-step-replay"}), timeout=15 + ) + assert result.status.value == "SUCCEEDED" + assert result.result is not None + assert json.loads(result.result) == "before-after" + assert calls == ["before", "after"] + spans = exporter.get_finished_spans() + invocations = [span for span in spans if span.name == "Invocation"] + assert [ + span.attributes["durable.invocation.status"] + for span in invocations + if span.attributes is not None + ] == ["PENDING", "SUCCEEDED"] + if plugin_type is ExecutionOtelPlugin: + # Count raw exports, including duplicates that reuse a span ID. + for name in ("otel-before-wait", "otel-replay-wait", "otel-after-wait"): + assert len([span for span in spans if span.name == name]) == 1 + assert context.get_current() == before_context + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_external_callback_completion_replay.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_external_callback_completion_replay.py new file mode 100644 index 000000000..2ec2f34fa --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_external_callback_completion_replay.py @@ -0,0 +1,195 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""A real callback completion is exported before two later public replays.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +import time +from pathlib import Path +from typing import Any + +import pytest +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from aws_durable_execution_sdk_python_testing.stores.filesystem import ( + FileSystemExecutionStore, +) +from opentelemetry import context +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +SRC_DIR = Path(__file__).resolve().parents[1] / "src" + + +def _suspended_callback( + runner: DurableFunctionTestRunner, arn: str, name: str +) -> tuple[str, tuple[int, int]]: + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + events = runner.get_execution_history(arn, include_execution_data=True).events + starts = [ + event + for event in events + if event.event_type == "CallbackStarted" and event.name == name + ] + completions = [ + event + for event in events + if event.event_type == "InvocationCompleted" + and starts + and event.event_id > starts[0].event_id + ] + if starts and completions: + details = starts[0].callback_started_details + assert details is not None and details.callback_id is not None + return details.callback_id, (starts[0].event_id, completions[0].event_id) + time.sleep(0.01) + raise AssertionError(f"Callback {name} did not suspend") + + +@pytest.mark.parametrize("plugin_type", [ExecutionOtelPlugin, InvocationOtelPlugin]) +def test_external_callback_completion_precedes_later_replays( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + plugin_type: type[ExecutionOtelPlugin] | type[InvocationOtelPlugin], +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = plugin_type( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + before_context = context.get_current() + common_spec = importlib.util.spec_from_file_location( + "common", SRC_DIR / "common.py" + ) + assert common_spec is not None and common_spec.loader is not None + common = importlib.util.module_from_spec(common_spec) + monkeypatch.setitem(sys.modules, "common", common) + common_spec.loader.exec_module(common) + monkeypatch.setattr(common, "otel_plugin", lambda: plugin) + spec = importlib.util.spec_from_file_location( + "otel_26_external_callback_completion_replay", + SRC_DIR / "otel_26_external_callback_completion_replay.py", + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + calls: list[str] = [] + original_observe = module.observe_target + + def counted_observe(value: str) -> Any: + step = original_observe(value) + + def counted(step_context: Any) -> str: + calls.append(value) + return step(step_context) + + return counted + + monkeypatch.setattr(module, "observe_target", counted_observe) + gates: list[tuple[int, int]] = [] + try: + with DurableFunctionTestRunner( + handler=module.handler, + store=FileSystemExecutionStore.create(tmp_path), + poll_interval=0.01, + execution_timeout=25, + skip_time=False, + ) as runner: + arn = runner.run_async( + input=json.dumps({"scenario": "external-callback-completion-replay"}) + ) + for name, payload in ( + ("otel-external-target", "target"), + ("otel-external-barrier-one create callback id", "one"), + ("otel-external-barrier-two create callback id", "two"), + ): + callback_id, gate = _suspended_callback(runner, arn, name) + gates.append(gate) + runner.send_callback_success( + callback_id, result=json.dumps(payload).encode("utf-8") + ) + result = runner.wait_for_result(arn, timeout=10) + history = runner.get_execution_history(arn, include_execution_data=True) + + assert result.status.value == "SUCCEEDED" + assert result.result is not None + assert json.loads(result.result) == "target/one/two" + assert calls == ["target"] + assert gates == [(2, 3), (8, 11), (15, 18)] + assert [ + event.event_id + for event in history.events + if event.event_type == "InvocationCompleted" + ] == [3, 11, 18, 21] + assert history.events[-1].event_type == "ExecutionSucceeded" + assert history.events[-1].event_id == 22 + + spans = exporter.get_finished_spans() + assert all(span.context is not None for span in spans) + assert len({span.context.trace_id for span in spans if span.context}) == 1 + invocations = sorted( + [span for span in spans if span.name == "Invocation"], + key=lambda span: span.start_time or 0, + ) + assert len(invocations) == 4 + assert [ + (span.attributes or {}).get("durable.invocation.status") + for span in invocations + ] == ["PENDING", "PENDING", "PENDING", "SUCCEEDED"] + targets = [ + span + for span in spans + if span.name == "otel-external-target" + and (span.attributes or {}).get("durable.operation.status") == "SUCCEEDED" + ] + # Count raw exports, including duplicate records with identical span IDs. + assert len(targets) == 1 + target = targets[0] + assert target.status.status_code.name == "OK" + observed = [ + span + for span in spans + if span.name == "otel-external-target-observed attempt 1" + ] + assert len(observed) == 1 + assert target.end_time is not None and observed[0].start_time is not None + assert target.end_time <= observed[0].start_time + assert invocations[2].start_time is not None + assert invocations[3].start_time is not None + assert target.end_time < invocations[2].start_time < invocations[3].start_time + parent = ( + invocations[1] + if plugin_type is InvocationOtelPlugin + else next(span for span in spans if span.name == "Workflow") + ) + assert target.parent is not None and parent.context is not None + assert target.parent.span_id == parent.context.span_id + # Invocation view also exports the legitimate first-invocation segment. + pending_target = [ + span + for span in spans + if span.name == "otel-external-target" + and (span.attributes or {}).get("durable.operation.status") == "STARTED" + ] + assert len(pending_target) == (1 if plugin_type is InvocationOtelPlugin else 0) + assert context.get_current() == before_context + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_invocation_retry_status.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_invocation_retry_status.py new file mode 100644 index 000000000..d2851a709 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_invocation_retry_status.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Exercise a real invocation retry through the public decorator and local runner.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +from pathlib import Path +from typing import Any + +import pytest +from aws_durable_execution_sdk_python.plugin import InvocationEndInfo, InvocationStatus +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry import context +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +SRC_DIR = Path(__file__).resolve().parents[1] / "src" + + +@pytest.mark.parametrize("plugin_type", [ExecutionOtelPlugin, InvocationOtelPlugin]) +def test_invocation_retry_preserves_completed_step( + monkeypatch: pytest.MonkeyPatch, + plugin_type: type[ExecutionOtelPlugin] | type[InvocationOtelPlugin], +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = plugin_type( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + statuses: list[InvocationStatus] = [] + original_end = plugin.on_invocation_end + + def observe_end(info: InvocationEndInfo) -> None: + statuses.append(info.status) + original_end(info) + + monkeypatch.setattr(plugin, "on_invocation_end", observe_end) + before_context = context.get_current() + common_spec = importlib.util.spec_from_file_location( + "common", SRC_DIR / "common.py" + ) + assert common_spec is not None and common_spec.loader is not None + common = importlib.util.module_from_spec(common_spec) + monkeypatch.setitem(sys.modules, "common", common) + common_spec.loader.exec_module(common) + monkeypatch.setattr(common, "otel_plugin", lambda: plugin) + spec = importlib.util.spec_from_file_location( + "otel_24_invocation_retry_status", + SRC_DIR / "otel_24_invocation_retry_status.py", + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + calls: list[str] = [] + + def track(factory: Any, label: str) -> Any: + def make_step() -> Any: + step = factory() + + def counted(step_context: Any) -> str: + calls.append(label) + return step(step_context) + + return counted + + return make_step + + monkeypatch.setattr( + module, + "before_invocation_retry", + track(module.before_invocation_retry, "saved"), + ) + try: + with DurableFunctionTestRunner(handler=module.handler) as runner: + result = runner.run( + input=json.dumps({"scenario": "invocation-retry-status"}), timeout=15 + ) + assert result.status.value == "SUCCEEDED" + assert result.result is not None + assert json.loads(result.result) == "retry-complete" + assert calls == ["saved"] + assert statuses == [InvocationStatus.RETRY, InvocationStatus.SUCCEEDED] + spans = exporter.get_finished_spans() + invocations = [span for span in spans if span.name == "Invocation"] + assert len(invocations) == 2 + assert [span.status.status_code.name for span in invocations] == ["UNSET", "OK"] + workflows = [span for span in spans if span.name == "Workflow"] + assert len(workflows) == 1 + assert workflows[0].status.status_code.name == "OK" + assert context.get_current() == before_context + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py index 5908dd7d9..fec2b7326 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py @@ -50,6 +50,12 @@ ("Otel18ChainedInvokeFailure", "otel-invocation-18"), ("Otel19ExecutionFailure", "otel-invocation-19"), ("Otel20VirtualContext", "otel-invocation-20"), + ("Otel21CompletedStepReplay", "otel-invocation-21"), + ("Otel22UserFunctionContext", "otel-invocation-22"), + ("Otel23CallbackFunctionContext", "otel-invocation-23"), + ("Otel24InvocationRetryStatus", "otel-invocation-24"), + ("Otel25CallbackFailureWithoutError", "otel-invocation-25"), + ("Otel26ExternalCallbackReplay", "otel-invocation-26"), ("OtelExecution1Success", "otel-execution-1"), ("OtelExecution2WaitResume", "otel-execution-2"), ("OtelExecution3Retry", "otel-execution-3"), @@ -70,6 +76,12 @@ ("OtelExecution18ChainedInvokeFailure", "otel-execution-18"), ("OtelExecution19ExecutionFailure", "otel-execution-19"), ("OtelExecution20VirtualContext", "otel-execution-20"), + ("OtelExecution21CompletedStepReplay", "otel-execution-21"), + ("OtelExecution22UserFunctionContext", "otel-execution-22"), + ("OtelExecution23CallbackFunctionContext", "otel-execution-23"), + ("OtelExecution24InvocationRetryStatus", "otel-execution-24"), + ("OtelExecution25CallbackFailureWithoutError", "otel-execution-25"), + ("OtelExecution26ExternalCallbackReplay", "otel-execution-26"), ] EXPECTED_LONG_RUNNING_MAPPINGS: list[tuple[str, str]] = [ ("OtelLongRunning1Wait", "otel-long-running-1"), @@ -152,6 +164,12 @@ "otel_18_chained_invoke_failure", "otel_19_execution_failure", "otel_20_virtual_context", + "otel_21_completed_step_replay", + "otel_22_user_function_context", + "otel_23_callback_function_context", + "otel_24_invocation_retry_status", + "otel_25_caught_callback_failure_replay", + "otel_26_external_callback_completion_replay", "otel_long_running_1_wait", "otel_long_running_2_retry", "otel_long_running_3_callback", @@ -391,6 +409,25 @@ def test_every_registered_handler_exists( ) +def test_case25_uses_its_caught_replay_handler_without_retargeting_other_cases( + template: dict[str, Any], +) -> None: + resources = template["Resources"] + for prefix in ("Otel", "OtelExecution"): + assert resources[prefix + "25CallbackFailureWithoutError"]["Properties"][ + "Handler" + ] == ("otel_25_caught_callback_failure_replay.handler") + assert resources[prefix + "17WaitForCallbackFailure"]["Properties"][ + "Handler" + ] == ("otel_17_wait_for_callback_failure.handler") + assert resources[prefix + "24InvocationRetryStatus"]["Properties"][ + "Handler" + ] == ("otel_24_invocation_retry_status.handler") + assert resources[prefix + "26ExternalCallbackReplay"]["Properties"][ + "Handler" + ] == ("otel_26_external_callback_completion_replay.handler") + + def test_template_deploys_only_the_selected_view(template: dict[str, Any]) -> None: assert template["Parameters"]["OtelSuite"]["AllowedValues"] == [ "all", diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_user_function_context.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_user_function_context.py new file mode 100644 index 000000000..2b3cc26d6 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_user_function_context.py @@ -0,0 +1,258 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Exercise case 22's callback contexts through real concurrent work and resume.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +from collections import Counter +from contextlib import nullcontext +from pathlib import Path +from threading import Barrier +from typing import Any + +import pytest +from aws_durable_execution_sdk_python.execution import DurableExecutionInvocationInput +from aws_durable_execution_sdk_python.types import LambdaContext +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry import context, trace +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +SRC_DIR = Path(__file__).resolve().parents[1] / "src" +CALLBACK_PARENTS = { + "step": "otel-context-step attempt 1", + "child": "otel-context-child", + "child-step": "otel-context-child-step attempt 1", + "child-restored": "otel-context-child", + "parallel-a": "otel-context-branch-a", + "parallel-step-a": "otel-context-branch-step-a attempt 1", + "parallel-b": "otel-context-branch-b", + "parallel-step-b": "otel-context-branch-step-b attempt 1", + "map-0": "otel-context-iteration-0", + "map-step-0": "otel-context-map-step-0 attempt 1", + "map-1": "otel-context-iteration-1", + "map-step-1": "otel-context-map-step-1 attempt 1", +} +HANDLER_COUNTS = {"handler": 2, "handler-restored": 2, "handler-after-resume": 1} + + +@pytest.mark.parametrize("plugin_type", [ExecutionOtelPlugin, InvocationOtelPlugin]) +@pytest.mark.parametrize( + "ambient", [False, True], ids=["no-ambient", "unrelated-ambient"] +) +def test_user_function_probes_keep_sdk_context_across_resume( + monkeypatch: pytest.MonkeyPatch, + plugin_type: type[ExecutionOtelPlugin] | type[InvocationOtelPlugin], + ambient: bool, +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + monkeypatch.setattr(trace, "get_tracer_provider", lambda: provider) + plugin = plugin_type( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + before_context = context.get_current() + common_spec = importlib.util.spec_from_file_location( + "common", SRC_DIR / "common.py" + ) + assert common_spec is not None and common_spec.loader is not None + common = importlib.util.module_from_spec(common_spec) + monkeypatch.setitem(sys.modules, "common", common) + common_spec.loader.exec_module(common) + monkeypatch.setattr(common, "otel_plugin", lambda: plugin) + spec = importlib.util.spec_from_file_location( + "otel_22_user_function_context", SRC_DIR / "otel_22_user_function_context.py" + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + barriers: dict[str, Barrier] = {} + for pair in ( + ("parallel-a", "parallel-b"), + ("parallel-step-a", "parallel-step-b"), + ("map-0", "map-1"), + ("map-step-0", "map-step-1"), + ): + barrier = Barrier(2, timeout=5) + barriers.update(dict.fromkeys(pair, barrier)) + original_probe = module.probe + + def synchronized_probe(label: str) -> None: + callback_context = context.get_current() + # Both sibling callbacks must be active before either observes its + # parent. Synchronize only in the test; never supply a tracing context. + if label in barriers: + barriers[label].wait() + original_probe(label) + if label in barriers: + barriers[label].wait() + assert context.get_current() is callback_context + + monkeypatch.setattr(module, "probe", synchronized_probe) + invocation_contexts: list[tuple[context.Context, context.Context]] = [] + + def entry( + event: DurableExecutionInvocationInput, lambda_context: LambdaContext + ) -> dict[str, Any]: + # Simulate host instrumentation outside the unmodified durable handler. + # With no extracted parent, its unrelated trace must not become the + # parent of any SDK operation or handler probe. + host_context = context.get_current() + scope = ( + provider.get_tracer("host-instrumentation").start_as_current_span( + "ambient-invocation" + ) + if ambient + else nullcontext() + ) + try: + with scope: + invocation_context = context.get_current() + try: + return module.handler(event, lambda_context) + finally: + invocation_contexts.append( + (invocation_context, context.get_current()) + ) + finally: + assert context.get_current() is host_context + + try: + with DurableFunctionTestRunner(handler=entry, skip_time=False) as runner: + arn = runner.run_async( + input=json.dumps({"scenario": "user-function-context"}), + execution_timeout=30, + ) + result = runner.wait_for_result(arn, timeout=30) + history = runner.get_execution_history(arn) + assert result.status.value == "SUCCEEDED", result.error + assert result.result is not None + assert json.loads(result.result) == "context-complete" + assert [ + event.event_type + for event in history.events + if event.event_type.startswith("Wait") + ] == ["WaitStarted", "WaitSucceeded"] + + spans = exporter.get_finished_spans() + invocations = [span for span in spans if span.name == "Invocation"] + assert len(invocations) == 2 + assert [ + span.attributes["durable.invocation.status"] + for span in invocations + if span.attributes is not None + ] == ["PENDING", "SUCCEEDED"] + assert [ + span.attributes["durable.invocation.first"] + for span in invocations + if span.attributes is not None + ] == [True, False] + workflows = [span for span in spans if span.name == "Workflow"] + assert len(workflows) == 1 + assert workflows[0].context is not None + canonical_trace_id = workflows[0].context.trace_id + assert { + span.context.trace_id + for span in spans + if span.context is not None + and span.attributes is not None + and span.attributes.get("durable.execution.arn") == arn + } == {canonical_trace_id} + + probes = [span for span in spans if span.name.startswith("conformance.")] + # Count raw exports so replayed callback bodies and duplicate exports + # fail even if they reuse a span ID. + assert Counter(span.name for span in probes) == Counter( + { + f"conformance.{label}": count + for label, count in ( + {**dict.fromkeys(CALLBACK_PARENTS, 1), **HANDLER_COUNTS} + ).items() + } + ) + spans_by_id = { + (span.context.trace_id, span.context.span_id): span + for span in spans + if span.context is not None + } + handler_parent = ( + "Workflow" if plugin_type is ExecutionOtelPlugin else "Invocation" + ) + for probe_span in probes: + assert probe_span.context is not None + assert probe_span.context.trace_id == canonical_trace_id + assert probe_span.parent is not None + assert probe_span.parent.trace_id == canonical_trace_id + assert probe_span.attributes is not None + assert "durable.execution.arn" not in probe_span.attributes + label = str(probe_span.attributes["conformance.callback"]) + assert probe_span.name == f"conformance.{label}" + parent = spans_by_id[(canonical_trace_id, probe_span.parent.span_id)] + assert parent.name == CALLBACK_PARENTS.get(label, handler_parent) + assert parent.attributes is not None + assert parent.attributes["durable.execution.arn"] == arn + + for index, invocation in enumerate(invocations): + assert invocation.start_time is not None and invocation.end_time is not None + invocation_probes = [ + span + for span in probes + if span.start_time is not None + and invocation.start_time <= span.start_time <= invocation.end_time + ] + expected = ( + { + **dict.fromkeys(CALLBACK_PARENTS, 1), + "handler": 1, + "handler-restored": 1, + } + if index == 0 + else dict.fromkeys(HANDLER_COUNTS, 1) + ) + assert Counter(span.name for span in invocation_probes) == Counter( + {f"conformance.{label}": count for label, count in expected.items()} + ) + handler_probes = [ + span + for span in invocation_probes + if span.attributes is not None + and span.attributes["conformance.callback"] in HANDLER_COUNTS + ] + expected_parent = ( + workflows[0] if plugin_type is ExecutionOtelPlugin else invocation + ) + assert all( + span.parent == expected_parent.context for span in handler_probes + ) + + ambient_spans = [span for span in spans if span.name == "ambient-invocation"] + assert len(ambient_spans) == (2 if ambient else 0) + for span in ambient_spans: + assert span.context is not None + assert span.context.trace_id != canonical_trace_id + # Observe restoration on the runner's invocation threads, including + # PENDING, rather than checking only the pytest caller's context. + assert len(invocation_contexts) == 2 + assert all(before is after for before, after in invocation_contexts) + assert context.get_current() is before_context + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 187971ba5..a7a07a772 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -395,13 +395,85 @@ each durable span in the same invocation. `durable.invocation.status` uses `RETRYING` when the core plugin hook reports `InvocationStatus.RETRY`. The core enum remains unchanged. Operation spans use -OTel `OK` only for `SUCCEEDED`, `ERROR` when error details are delivered, and -`UNSET` for other outcomes without error details, including `FAILED`, -`CANCELLED`, `TIMED_OUT`, and `STOPPED`. The original durable operation status -remains in `durable.operation.status`. +OTel `ERROR` for explicit `FAILED`, `TIMED_OUT`, or `STOPPED` outcomes even when +error details are absent. Supplied error details continue to set `ERROR` and +record an exception event; a failure status alone does not invent that event. +Without error details, `SUCCEEDED` is `OK` and other statuses, including +`CANCELLED`, remain `UNSET`. The original durable operation status remains in +`durable.operation.status`. A handler that catches an operation failure and +returns successfully still has a successful invocation/workflow outcome; the +failed operation retains its own error status on initial completion and replay. + +Exception events use the same normalized end timestamp as their operation or +attempt span. This keeps events inside the exported span when a completion has +no durable start time or arrives later than its backend end time; the existing +view-specific span-boundary rules remain intact. + +When a concurrent branch's terminal completion arrives before replay has created +its parent span, Invocation view retains that completion until the actual parent +span is registered. The terminal segment is exported under that parent before +control returns to the branch's user code. No ancestor is invented and SDK +completion delivery is unchanged. If the parent never becomes active, normal +invocation cleanup and flushing still run before the missing parent is reported. + +### Invocation flushing + +After span and ContextVar cleanup, the plugins wait up to the OpenTelemetry +provider's existing 30-second default budget for `force_flush()`. Only provider +I/O runs on a background daemon worker; invocation Start/End hooks and token +cleanup stay on their original worker and Context. The I/O receives a snapshot +of that post-cleanup Context. + +There is at most one outstanding provider flush in the process and no pending +I/O queue. Concurrent requests, including requests using another provider, wait +for admission within the same absolute deadline used for completion. A timeout +does not cancel the exporter or free its slot. A request whose budget expires +before admission returns with a warning without starting more work. +A failed or skipped flush is not an export +acknowledgment, and late exports may continue. Applications still own exporter +transport timeouts and provider shutdown. This bound does not cover synchronous +span-processor `on_end` work that runs before flushing. + +Ordinary provider exceptions that finish within the budget retain the existing +plugin exception isolation; fatal exceptions are propagated on the lifecycle +owner. Exceptions after the caller has stopped waiting are logged and cannot be +raised retroactively into the completed invocation. ### Invocation context isolation +Tracing state belongs to each invocation's Context, including its trace IDs, +span registry, deferred completions, timestamps and scope tokens. A shared plugin +instance can therefore serve overlapping invocations without one invocation's +cleanup clearing another's trace. Provider configuration and registration remain +instance resources. SDK-owned worker tasks capture the selected invocation +Context; Start/End hooks and token reset retain their existing owner and order. + +Visible terminal input updates are delivered after all Start hooks. The service +omits descendants of completed contexts unless their ancestors retain children +through `ReplayChildren`; an `UpdatedOperationIds` entry alone does not provide +missing operation metadata. The SDK does not synthesize a hook or a span from +that missing record. Public input-delivery tests use a real large child result +to retain metadata, while still asserting one saved-step execution and terminal +export with the correct parent in each view. + +Known completed context metadata is retained from the full visible input +snapshot for descendant completions delivered by later checkpoint responses. +A child that finishes during the invocation retains its actual non-recording +parent context until cleanup. In Invocation view, completions wait for any +current replay segment, including a large-result ReplayChildren parent. Only +at End does a completion still lacking a current segment use a known completed +input parent's identity. Missing or nonterminal parents remain deferred; +unresolved parentage is reported after cleanup. No root parent or missing +metadata is invented, and queue draining does not re-enter a public hook. + +Completion metadata can be redelivered after a failed invocation or checkpoint +transport attempt that has no successful response acknowledgment. There is no +exporter acknowledgment in durable state, so this does not promise exactly-once +telemetry across failed/redelivered invocations. Successful replay preserves +checkpointed work; this limit is not a claim that already-checkpointed step +bodies run again. The plugin does not use a global exported-operation cache that +could discard a first delivery after a failed attempt. + With core 2.1+, invocation hooks and the handler run on one worker in an invocation-local Context initialized from the host's bindings. Successful Start bindings are visible to later hooks, the handler and resource cleanup. The host's diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_flush.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_flush.py new file mode 100644 index 000000000..26c1b46e8 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_flush.py @@ -0,0 +1,121 @@ +"""Bound caller waiting for provider I/O without moving plugin lifecycle hooks.""" + +from __future__ import annotations + +import contextvars +import logging +import threading +import time +from collections.abc import Callable +from dataclasses import dataclass, field + + +logger = logging.getLogger(__name__) + +# Match the OTel provider's existing default budget. Some processors ignore +# that cooperative budget, so the SDK bounds its own wait independently. +_FLUSH_TIMEOUT_SECONDS = 30.0 +_lock = threading.Condition() + + +@dataclass +class _Attempt: + done: threading.Event = field(default_factory=threading.Event) + result: object = None + error: BaseException | None = None + timed_out: bool = False + + +_active: _Attempt | None = None + + +def _run(flush: Callable[[], object], attempt: _Attempt) -> None: + global _active # noqa: PLW0603 + try: + try: + attempt.result = flush() + except BaseException as error: + # Fatal exceptions completed within the wait are re-raised on the + # lifecycle owner. A late error cannot be thrown into a returned call. + attempt.error = error + with _lock: + late = attempt.timed_out + if not late and _active is attempt: + _active = None + _lock.notify_all() + attempt.done.set() + if late and attempt.error is not None: + logger.error( + "Provider flush failed after the caller stopped waiting", + exc_info=( + type(attempt.error), + attempt.error, + attempt.error.__traceback__, + ), + ) + finally: + # Retain the slot through actual completion, even after timeout. There + # is no queue and no new worker per invocation for a stuck exporter. + with _lock: + if _active is attempt: + _active = None + _lock.notify_all() + + +def flush_provider(flush: Callable[[], object]) -> bool: + """Wait for provider I/O, with one outstanding attempt process-wide. + + Start/End and ContextVar-token cleanup stay on their existing owner. Only + the supplied provider method runs elsewhere, using the post-cleanup Context + snapshot. Admission and completion share one absolute deadline. Expired + requests return False without queueing more work; active I/O is never cancelled. + This does not bound synchronous span processors' earlier ``on_end`` calls, + application-owned provider shutdown, or exporter transport I/O itself. + """ + global _active # noqa: PLW0603 + deadline = time.monotonic() + _FLUSH_TIMEOUT_SECONDS + with _lock: + while _active is not None: + remaining = deadline - time.monotonic() + if remaining <= 0: + break + _lock.wait(remaining) + if _active is None and time.monotonic() < deadline: + attempt = _Attempt() + _active = attempt + else: + attempt = None + if attempt is None: + logger.warning("Provider flush admission exceeded the caller wait budget") + return False + + owner_context = contextvars.copy_context() + try: + threading.Thread( + target=owner_context.run, + args=(_run, flush, attempt), + name="durable-otel-flush", + daemon=True, + ).start() + except BaseException: + with _lock: + if _active is attempt: + _active = None + _lock.notify_all() + raise + + completed = attempt.done.wait(max(0.0, deadline - time.monotonic())) + if not completed: + with _lock: + completed = attempt.done.is_set() + if not completed: + attempt.timed_out = True + if not completed: + logger.warning("Provider flush exceeded the caller wait budget") + return False + if attempt.error is not None: + raise attempt.error + if attempt.result is False: + logger.warning("Provider reported an incomplete flush") + return False + return True diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_invocation_state.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_invocation_state.py new file mode 100644 index 000000000..b4212541b --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_invocation_state.py @@ -0,0 +1,130 @@ +"""Invocation-owned tracing data, propagated by the SDK's captured Contexts.""" + +from __future__ import annotations + +from collections.abc import Callable +from contextvars import ContextVar, Token +from dataclasses import dataclass, field +import datetime +import threading +from typing import Any, Generic, TypeVar + +from aws_durable_execution_sdk_python.plugin import ( + InvocationStartInfo, + OperationEndInfo, + OperationInfo, + OperationStatus, + OperationType, +) +from opentelemetry.sdk.trace.sampling import Sampler +from opentelemetry.trace import Span, Tracer, TracerProvider + +from aws_durable_execution_sdk_python_otel.context_extractors import ExtractedContext +from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( + DeterministicIdGenerator, +) +from aws_durable_execution_sdk_python_otel.durable_sampling import DurableSamplingIntent +from aws_durable_execution_sdk_python_otel.execution_trace_context import ( + ExecutionTraceContext, +) + + +_Key = TypeVar("_Key") + + +@dataclass +class _InvocationState(Generic[_Key]): + provider: TracerProvider + tracer: Tracer + id_generator: DeterministicIdGenerator + sampling_delegate: Sampler | None = None + execution_arn: str = "" + execution_trace_id: int | None = None + execution_start_time: datetime.datetime | None = None + extracted_context: ExtractedContext | None = None + execution_trace_context: ExecutionTraceContext | None = None + sampling_intent: DurableSamplingIntent | None = None + workflow_span: Span | None = None + operation_spans: dict[_Key, Span] = field(default_factory=dict) + context_tokens: dict[str, tuple[int, object]] = field(default_factory=dict) + tracing_enabled: bool = False + owner_token: Token[Any] | None = None + owner_thread: threading.Thread | None = None + + +@dataclass +class ExecutionState(_InvocationState[str]): + invocation_span: Span | None = None + checkpointed_context_ids: set[str] = field(default_factory=set) + ended_operation_ids: set[str] = field(default_factory=set) + pending_operation_ends: dict[str, list[OperationEndInfo]] = field( + default_factory=dict + ) + + +@dataclass +class InvocationState(_InvocationState[str | None]): + span_time_floor_ns: int | None = None + input_parent_refs: dict[str, OperationInfo] = field(default_factory=dict) + pending_operation_ends: dict[str, list[OperationEndInfo]] = field( + default_factory=dict + ) + context_operation_replays: dict[str, bool] = field(default_factory=dict) + incomplete_attempt_span_keys: set[str] = field(default_factory=set) + + +_State = TypeVar("_State", bound=_InvocationState[Any]) + + +def completed_input_parents(info: InvocationStartInfo) -> list[OperationInfo]: + """Known terminal contexts that may parent input or later checkpoint updates.""" + return [ + operation + for operation in info.operations.values() + if operation.operation_type is OperationType.CONTEXT + and operation.status + in ( + OperationStatus.SUCCEEDED, + OperationStatus.FAILED, + OperationStatus.TIMED_OUT, + OperationStatus.STOPPED, + OperationStatus.CANCELLED, + ) + ] + + +def begin_state( + variable: ContextVar[_State], + state: _State, + reset_previous: Callable[[], None], +) -> None: + """Start fresh data without clearing another invocation inherited in a copy. + + A direct warm-start caller can omit its previous End; retain that cleanup + behavior only when its token proves this is the actual owning Context. + Thread identity alone cannot distinguish a Context from a copy of it. + """ + previous = variable.get() + token = previous.owner_token + if token is not None and previous.owner_thread is threading.current_thread(): + try: + variable.reset(token) + except (ValueError, RuntimeError): + # An inherited Context (or an already-ended inherited token) does + # not own the previous invocation and must not mutate its data. + pass + else: + previous.owner_token = variable.set(previous) + reset_previous() + state.owner_thread = threading.current_thread() + state.owner_token = variable.set(state) + + +def finish_state(variable: ContextVar[_State]) -> None: + """Restore the binding on the Context that ran Start, after normal cleanup.""" + state = variable.get() + token = state.owner_token + if token is not None: + state.owner_token = None + state.owner_thread = None + variable.reset(token) diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py index 2008615e1..393ffb175 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py @@ -31,6 +31,7 @@ from __future__ import annotations +import contextvars import datetime import logging import threading @@ -64,12 +65,19 @@ TraceState, ) +from aws_durable_execution_sdk_python_otel._invocation_state import ( + ExecutionState, + begin_state, + completed_input_parents, + finish_state, +) from aws_durable_execution_sdk_python_otel.context_extractors import ( ContextExtractor, ExtractedContext, _ensure_extracted_context, xray_context_extractor, ) +from aws_durable_execution_sdk_python_otel._flush import flush_provider from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( DeterministicIdGenerator, derive_workflow_span_id, @@ -140,43 +148,167 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: self._workflow_span_name = self._config.workflow_span_name result = create_tracer_provider(self._config) - self._provider = result.tracer_provider self._uses_global_provider = result.uses_global_provider - - self._tracer: Tracer = self._provider.get_tracer(self._config.instrument_name) - self._id_generator = DeterministicIdGenerator() - self._sampling_delegate: Sampler | None = None - self._bind_sdk_tracer() - - # Per-invocation state. - self._execution_arn = "" - self._execution_trace_id: int | None = None - self._execution_start_time: datetime.datetime | None = None - self._extracted_context: ExtractedContext | None = None - self._execution_trace_context: ExecutionTraceContext | None = None - self._sampling_intent: DurableSamplingIntent | None = None - self._workflow_span: Span | None = None - self._invocation_span: Span | None = None - self._operation_spans: dict[str, Span] = {} - # CONTEXT operations that emitted a durable START hook this invocation. - # A context absent from this set is checkpointless (for example, a FLAT - # map/parallel branch) and needs a fresh per-invocation span identity. - self._checkpointed_context_ids: set[str] = set() - # Operations whose span was already exported this invocation, so a - # repeated on_operation_end does not export it twice. - self._ended_operation_ids: set[str] = set() - # Tokens returned by context.attach(), keyed by the span registry key, - # paired with the thread that attached them. Every attach the plugin - # owns is released through _detach_context so the plugin never leaves a - # scope on the context stack. - self._context_tokens: dict[str, tuple[int, object]] = {} + self._template = ExecutionState( + provider=result.tracer_provider, + tracer=result.tracer_provider.get_tracer(self._config.instrument_name), + id_generator=DeterministicIdGenerator(), + ) + self._state: contextvars.ContextVar[ExecutionState] = contextvars.ContextVar( + "durable_execution_otel_invocation", default=self._template + ) self._lock = threading.RLock() - self._tracing_enabled = False + self._bind_sdk_tracer() self._registration_accepted = False if self._config.enrich_logger: install_log_filter(self) + # Private accessors route existing helpers through the current invocation. + @property + def _provider(self) -> trace.TracerProvider: + return self._state.get().provider + + @_provider.setter + def _provider(self, value: trace.TracerProvider) -> None: + self._state.get().provider = value + + @property + def _tracer(self) -> Tracer: + return self._state.get().tracer + + @_tracer.setter + def _tracer(self, value: Tracer) -> None: + self._state.get().tracer = value + + @property + def _id_generator(self) -> DeterministicIdGenerator: + return self._state.get().id_generator + + @_id_generator.setter + def _id_generator(self, value: DeterministicIdGenerator) -> None: + self._state.get().id_generator = value + + @property + def _sampling_delegate(self) -> Sampler | None: + return self._state.get().sampling_delegate + + @_sampling_delegate.setter + def _sampling_delegate(self, value: Sampler | None) -> None: + self._state.get().sampling_delegate = value + + @property + def _execution_arn(self) -> str: + return self._state.get().execution_arn + + @_execution_arn.setter + def _execution_arn(self, value: str) -> None: + self._state.get().execution_arn = value + + @property + def _execution_trace_id(self) -> int | None: + return self._state.get().execution_trace_id + + @_execution_trace_id.setter + def _execution_trace_id(self, value: int | None) -> None: + self._state.get().execution_trace_id = value + + @property + def _execution_start_time(self) -> datetime.datetime | None: + return self._state.get().execution_start_time + + @_execution_start_time.setter + def _execution_start_time(self, value: datetime.datetime | None) -> None: + self._state.get().execution_start_time = value + + @property + def _extracted_context(self) -> ExtractedContext | None: + return self._state.get().extracted_context + + @_extracted_context.setter + def _extracted_context(self, value: ExtractedContext | None) -> None: + self._state.get().extracted_context = value + + @property + def _execution_trace_context(self) -> ExecutionTraceContext | None: + return self._state.get().execution_trace_context + + @_execution_trace_context.setter + def _execution_trace_context(self, value: ExecutionTraceContext | None) -> None: + self._state.get().execution_trace_context = value + + @property + def _sampling_intent(self) -> DurableSamplingIntent | None: + return self._state.get().sampling_intent + + @_sampling_intent.setter + def _sampling_intent(self, value: DurableSamplingIntent | None) -> None: + self._state.get().sampling_intent = value + + @property + def _workflow_span(self) -> Span | None: + return self._state.get().workflow_span + + @_workflow_span.setter + def _workflow_span(self, value: Span | None) -> None: + self._state.get().workflow_span = value + + @property + def _context_tokens(self) -> dict[str, tuple[int, object]]: + return self._state.get().context_tokens + + @_context_tokens.setter + def _context_tokens(self, value: dict[str, tuple[int, object]]) -> None: + self._state.get().context_tokens = value + + @property + def _tracing_enabled(self) -> bool: + return self._state.get().tracing_enabled + + @_tracing_enabled.setter + def _tracing_enabled(self, value: bool) -> None: + self._state.get().tracing_enabled = value + + @property + def _invocation_span(self) -> Span | None: + return self._state.get().invocation_span + + @_invocation_span.setter + def _invocation_span(self, value: Span | None) -> None: + self._state.get().invocation_span = value + + @property + def _operation_spans(self) -> dict[str, Span]: + return self._state.get().operation_spans + + @_operation_spans.setter + def _operation_spans(self, value: dict[str, Span]) -> None: + self._state.get().operation_spans = value + + @property + def _checkpointed_context_ids(self) -> set[str]: + return self._state.get().checkpointed_context_ids + + @_checkpointed_context_ids.setter + def _checkpointed_context_ids(self, value: set[str]) -> None: + self._state.get().checkpointed_context_ids = value + + @property + def _ended_operation_ids(self) -> set[str]: + return self._state.get().ended_operation_ids + + @_ended_operation_ids.setter + def _ended_operation_ids(self, value: set[str]) -> None: + self._state.get().ended_operation_ids = value + + @property + def _pending_operation_ends(self) -> dict[str, list[OperationEndInfo]]: + return self._state.get().pending_operation_ends + + @_pending_operation_ends.setter + def _pending_operation_ends(self, value: dict[str, list[OperationEndInfo]]) -> None: + self._state.get().pending_operation_ends = value + def on_registration_result(self, registered: bool) -> None: """Preserve accepted resources, releasing only a discarded constructor.""" if registered: @@ -210,6 +342,9 @@ def _bind_sdk_tracer(self) -> bool: def _set_span(self, key: str, span: Span) -> None: with self._lock: self._operation_spans[key] = span + pending = self._pending_operation_ends.pop(key, []) + for info in pending: + self._render_operation_end(info) def _get_span(self, key: str | None) -> Span | None: if key is None: @@ -433,7 +568,15 @@ def _with_sampling(self, parent_context: Context) -> Context: def on_invocation_start(self, info: InvocationStartInfo) -> None: logger.debug("Durable invocation started: %s", info) self._registration_accepted = True - self._reset_state() + begin_state( + self._state, + ExecutionState( + provider=self._template.provider, + tracer=self._template.tracer, + id_generator=self._template.id_generator, + ), + self._reset_state, + ) if info.execution_start_time is None: logger.warning( "ExecutionOtelPlugin requires InvocationStartInfo.execution_start_time " @@ -502,6 +645,10 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: self._start_workflow_span(info) # Keep the invocation on the shared execution trace. self._start_invocation_span(info) + # A completed child will not run another user-function Start hook. + # Reference its original logical identity without exporting it again. + for parent in completed_input_parents(info): + self._register_operation_placeholder(parent.operation_id, parent.start_time) # Make the Workflow span the active span so auto-instrumented spans # created during the invocation become its children. The token is @@ -627,6 +774,9 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: self._reset_state() return + with self._lock: + unresolved_parents = tuple(self._pending_operation_ends) + # End the invocation span regardless of terminal status. Record the # invocation status and map it to a span status: # SUCCEEDED/PENDING -> OK (this invocation did its work, whether it @@ -662,13 +812,18 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: if info.status in _TERMINAL_INVOCATION_STATUSES: self._export_workflow_span(info) + provider = self._provider self._reset_state() - if hasattr(self._provider, "force_flush"): + if hasattr(provider, "force_flush"): try: - self._provider.force_flush() + flush_provider(provider.force_flush) except Exception: # noqa: BLE001 logger.exception("force_flush failed at invocation end") + if unresolved_parents: + raise ValueError( + f"No parent span found for deferred operation ends: {unresolved_parents}" + ) def _reset_state(self) -> None: self._detach_remaining_contexts() @@ -684,7 +839,9 @@ def _reset_state(self) -> None: self._operation_spans = {} self._checkpointed_context_ids = set() self._ended_operation_ids = set() + self._pending_operation_ends = {} self._tracing_enabled = False + finish_state(self._state) # ------------------------------------------------------------------ # Operation lifecycle @@ -703,6 +860,10 @@ def on_operation_start(self, info: OperationStartInfo) -> None: self._note_parent_start(info.parent_id, info.start_time) def on_operation_end(self, info: OperationEndInfo) -> None: + self._render_operation_end(info) + + def _render_operation_end(self, info: OperationEndInfo) -> None: + """Render or defer one delivery without re-entering a public hook.""" logger.debug("Durable operation ended: %s", info) if not self._tracing_enabled: return @@ -711,15 +872,22 @@ def on_operation_end(self, info: OperationEndInfo) -> None: # replay-only and must not re-export the deterministic logical span. if info.is_replayed: return + with self._lock: + if ( + info.parent_id is not None + and info.parent_id not in self._operation_spans + ): + self._pending_operation_ends.setdefault(info.parent_id, []).append(info) + return # Export the span only on the first end for this operation. with self._lock: if info.operation_id in self._ended_operation_ids: return self._ended_operation_ids.add(info.operation_id) - # An open operation is held as a non-recording placeholder; drop it and - # create the single recording span for the operation now. + # Keep the parent context registered while materializing its span. + # _start_span replaces the placeholder atomically under the registry lock, + # so a concurrent child end cannot be stranded by a temporary absence. placeholder = self._get_span(info.operation_id) - self._pop_span(info.operation_id) parent = self._resolve_parent(info.parent_id) start_time = info.start_time end_time = info.end_time @@ -758,13 +926,27 @@ def on_operation_end(self, info: OperationEndInfo) -> None: if info.error: span.set_status(StatusCode.ERROR, info.error.message or "") span.record_exception( - Exception(info.error.message or info.error.type or "Unknown error") + Exception(info.error.message or info.error.type or "Unknown error"), + timestamp=_to_otel_timestamp(end_time), ) + elif info.status in ( + OperationStatus.FAILED, + OperationStatus.TIMED_OUT, + OperationStatus.STOPPED, + ): + span.set_status(StatusCode.ERROR) elif info.status is OperationStatus.SUCCEEDED: span.set_status(StatusCode.OK) self._note_parent_end(info.parent_id, end_time) - popped = self._pop_span(info.operation_id) + with self._lock: + popped = self._operation_spans.pop(info.operation_id, None) + if popped is not None and info.operation_type is OperationType.CONTEXT: + # A callback created by this child can finish in a later root + # checkpoint. Retain only its actual parent context until End. + self._operation_spans[info.operation_id] = DurableParentSpan( + popped.get_span_context(), start_time=start_time + ) if popped is not None: popped.end(end_time=_to_otel_timestamp(end_time)) @@ -895,7 +1077,8 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: (info.error.message or info.error.type) if info.error else "Unknown error" - ) + ), + timestamp=_to_otel_timestamp(end_time), ) else: span.set_status(StatusCode.OK) diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py index 4fc2bf02b..ad919106f 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py @@ -2,6 +2,7 @@ from __future__ import annotations +import contextvars import datetime import logging import threading @@ -34,12 +35,19 @@ TraceState, ) +from aws_durable_execution_sdk_python_otel._invocation_state import ( + InvocationState, + begin_state, + completed_input_parents, + finish_state, +) from aws_durable_execution_sdk_python_otel.context_extractors import ( ContextExtractor, ExtractedContext, _ensure_extracted_context, xray_context_extractor, ) +from aws_durable_execution_sdk_python_otel._flush import flush_provider from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( DeterministicIdGenerator, derive_workflow_span_id, @@ -140,42 +148,167 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: self._enrich_logger = self._config.enrich_logger result = create_tracer_provider(self._config) - self._provider = result.tracer_provider self._uses_global_provider = result.uses_global_provider - self._tracer: Tracer = self._provider.get_tracer(self._config.instrument_name) - self._id_generator = DeterministicIdGenerator() - self._sampling_delegate: Sampler | None = None - self._bind_sdk_tracer() - - # per invocation status: - self._execution_arn = "" - self._execution_trace_id: int | None = None - self._execution_start_time: datetime.datetime | None = None - self._extracted_context: ExtractedContext | None = None - self._execution_trace_context: ExecutionTraceContext | None = None - self._sampling_intent: DurableSamplingIntent | None = None - self._workflow_span: Span | None = None - self._span_time_floor_ns: int | None = None - # Maps operation ID (None for root) to the active span. - self._operation_spans: dict[str | None, Span] = {} - # Replay state supplied by CONTEXT operation START hooks. Missing - # entries identify checkpointless contexts such as FLAT branches. - self._context_operation_replays: dict[str, bool] = {} - # STEP attempt spans whose user function explicitly ended INCOMPLETE. - # Only these spans are marked truncated during invocation cleanup. - self._incomplete_attempt_span_keys: set[str] = set() - # Tokens returned by context.attach(), keyed by the span registry key, - # paired with the thread that attached them. Every attach the plugin - # owns is released through _detach_context so the plugin never leaves a - # scope on the context stack. - self._context_tokens: dict[str, tuple[int, object]] = {} + self._template = InvocationState( + provider=result.tracer_provider, + tracer=result.tracer_provider.get_tracer(self._config.instrument_name), + id_generator=DeterministicIdGenerator(), + ) + self._state: contextvars.ContextVar[InvocationState] = contextvars.ContextVar( + "durable_invocation_otel_invocation", default=self._template + ) self._operation_spans_lock = threading.RLock() - self._tracing_enabled = False + self._bind_sdk_tracer() self._registration_accepted = False if self._enrich_logger: install_log_filter(self) + # Private accessors route existing helpers through the current invocation. + @property + def _provider(self) -> trace.TracerProvider: + return self._state.get().provider + + @_provider.setter + def _provider(self, value: trace.TracerProvider) -> None: + self._state.get().provider = value + + @property + def _tracer(self) -> Tracer: + return self._state.get().tracer + + @_tracer.setter + def _tracer(self, value: Tracer) -> None: + self._state.get().tracer = value + + @property + def _id_generator(self) -> DeterministicIdGenerator: + return self._state.get().id_generator + + @_id_generator.setter + def _id_generator(self, value: DeterministicIdGenerator) -> None: + self._state.get().id_generator = value + + @property + def _sampling_delegate(self) -> Sampler | None: + return self._state.get().sampling_delegate + + @_sampling_delegate.setter + def _sampling_delegate(self, value: Sampler | None) -> None: + self._state.get().sampling_delegate = value + + @property + def _execution_arn(self) -> str: + return self._state.get().execution_arn + + @_execution_arn.setter + def _execution_arn(self, value: str) -> None: + self._state.get().execution_arn = value + + @property + def _execution_trace_id(self) -> int | None: + return self._state.get().execution_trace_id + + @_execution_trace_id.setter + def _execution_trace_id(self, value: int | None) -> None: + self._state.get().execution_trace_id = value + + @property + def _execution_start_time(self) -> datetime.datetime | None: + return self._state.get().execution_start_time + + @_execution_start_time.setter + def _execution_start_time(self, value: datetime.datetime | None) -> None: + self._state.get().execution_start_time = value + + @property + def _extracted_context(self) -> ExtractedContext | None: + return self._state.get().extracted_context + + @_extracted_context.setter + def _extracted_context(self, value: ExtractedContext | None) -> None: + self._state.get().extracted_context = value + + @property + def _execution_trace_context(self) -> ExecutionTraceContext | None: + return self._state.get().execution_trace_context + + @_execution_trace_context.setter + def _execution_trace_context(self, value: ExecutionTraceContext | None) -> None: + self._state.get().execution_trace_context = value + + @property + def _sampling_intent(self) -> DurableSamplingIntent | None: + return self._state.get().sampling_intent + + @_sampling_intent.setter + def _sampling_intent(self, value: DurableSamplingIntent | None) -> None: + self._state.get().sampling_intent = value + + @property + def _workflow_span(self) -> Span | None: + return self._state.get().workflow_span + + @_workflow_span.setter + def _workflow_span(self, value: Span | None) -> None: + self._state.get().workflow_span = value + + @property + def _context_tokens(self) -> dict[str, tuple[int, object]]: + return self._state.get().context_tokens + + @_context_tokens.setter + def _context_tokens(self, value: dict[str, tuple[int, object]]) -> None: + self._state.get().context_tokens = value + + @property + def _tracing_enabled(self) -> bool: + return self._state.get().tracing_enabled + + @_tracing_enabled.setter + def _tracing_enabled(self, value: bool) -> None: + self._state.get().tracing_enabled = value + + @property + def _operation_spans(self) -> dict[str | None, Span]: + return self._state.get().operation_spans + + @_operation_spans.setter + def _operation_spans(self, value: dict[str | None, Span]) -> None: + self._state.get().operation_spans = value + + @property + def _span_time_floor_ns(self) -> int | None: + return self._state.get().span_time_floor_ns + + @_span_time_floor_ns.setter + def _span_time_floor_ns(self, value: int | None) -> None: + self._state.get().span_time_floor_ns = value + + @property + def _pending_operation_ends(self) -> dict[str, list[OperationEndInfo]]: + return self._state.get().pending_operation_ends + + @_pending_operation_ends.setter + def _pending_operation_ends(self, value: dict[str, list[OperationEndInfo]]) -> None: + self._state.get().pending_operation_ends = value + + @property + def _context_operation_replays(self) -> dict[str, bool]: + return self._state.get().context_operation_replays + + @_context_operation_replays.setter + def _context_operation_replays(self, value: dict[str, bool]) -> None: + self._state.get().context_operation_replays = value + + @property + def _incomplete_attempt_span_keys(self) -> set[str]: + return self._state.get().incomplete_attempt_span_keys + + @_incomplete_attempt_span_keys.setter + def _incomplete_attempt_span_keys(self, value: set[str]) -> None: + self._state.get().incomplete_attempt_span_keys = value + def on_registration_result(self, registered: bool) -> None: """Preserve accepted resources, releasing only a discarded constructor.""" if registered: @@ -493,14 +626,27 @@ def _start_span( links=links, ) self._operation_spans[registry_key] = span + pending_ends = ( + self._pending_operation_ends.pop(registry_key, []) + if registry_key is not None + else [] + ) if operation_id is None: self._span_time_floor_ns = span_start_time logger.debug("Started OTel span: %s", span) + # Registration and dequeue share the lock, so an arriving completion + # either sees its parent or is drained here. Do not hold it in callbacks. + for pending_end in pending_ends: + self._render_operation_end(pending_end) return span def _end_span( - self, operation_id: str | None, end_timestamp: datetime.datetime | None = None + self, + operation_id: str | None, + end_timestamp: datetime.datetime | None = None, + *, + exception: Exception | None = None, ): """End and unregister the active span for an operation ID. @@ -509,6 +655,8 @@ def _end_span( span. end_timestamp: Optional durable end timestamp to use as the span end time. When omitted, OpenTelemetry uses the current time. + exception: Existing failure details to record at the same normalized + timestamp as the span end. """ logger.debug("Ending OTel span: operation_id=%s", operation_id) with self._operation_spans_lock: @@ -516,6 +664,8 @@ def _end_span( if span: # the span is not going to be populated if it has the same end_time and start_time end_time = self._next_ordered_timestamp(end_timestamp) + if exception is not None: + span.record_exception(exception, timestamp=end_time) span.end(end_time=end_time) logger.debug("Ended OTel span: %s", span) @@ -526,7 +676,15 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: """Called at the start of each invocation. Creates the invocation span.""" logger.debug("Durable invocation started: %s", info) self._registration_accepted = True - self._reset_state() + begin_state( + self._state, + InvocationState( + provider=self._template.provider, + tracer=self._template.tracer, + id_generator=self._template.id_generator, + ), + self._reset_state, + ) if info.execution_start_time is None: logger.warning( "InvocationOtelPlugin requires InvocationStartInfo.execution_start_time " @@ -599,6 +757,11 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: name="Invocation", attributes=self._extract_attributes(info), ) + # A terminal checkpoint can still use ReplayChildren. Give the handler + # a chance to register its current segment before using an old identity. + self._state.get().input_parent_refs = { + parent.operation_id: parent for parent in completed_input_parents(info) + } # Start and End share the invocation worker and token-owning Context. # Retain a valid same-trace ambient parent; otherwise bind this invocation. @@ -691,6 +854,24 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: self._reset_state() return + # Actual replay parent segments drain their children on registration. + # Only now can an input completion with no such segment use the known + # completed checkpoint's identity. Missing/nonterminal parents still fail. + pending_ends = [] + with self._operation_spans_lock: + for parent_id in tuple(self._pending_operation_ends): + parent = self._state.get().input_parent_refs.get(parent_id) + if parent is None or parent_id in self._operation_spans: + continue + parent_context = self._operation_link_context(parent_id) + if parent_context is not None: + self._operation_spans[parent_id] = DurableParentSpan( + parent_context, start_time=parent.start_time + ) + pending_ends.extend(self._pending_operation_ends.pop(parent_id)) + for pending_end in pending_ends: + self._render_operation_end(pending_end) + # User work and checkpoint cleanup have finished. Release Start bindings # in the Context that created their tokens before closing the spans. self._detach_remaining_contexts() @@ -700,6 +881,7 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: with self._operation_spans_lock: operation_ids = list(reversed(self._operation_spans)) incomplete_attempt_span_keys = set(self._incomplete_attempt_span_keys) + unresolved_parents = tuple(self._pending_operation_ends) for operation_id in operation_ids: if operation_id: span = self._get_span(operation_id) @@ -743,11 +925,18 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: if info.status in _TERMINAL_INVOCATION_STATUSES: self._export_workflow_span(info) + provider = self._provider self._reset_state() # Flush before Lambda freeze - if hasattr(self._provider, "force_flush"): - self._provider.force_flush() + if hasattr(provider, "force_flush"): + flush_provider(provider.force_flush) + if unresolved_parents: + # Keep the previous missing-parent failure observable, but finish + # normal token/span cleanup and flushing before reporting it. + raise ValueError( + f"No parent span found for deferred operation ends: {unresolved_parents}" + ) def _reset_state(self) -> None: """Clear per-invocation state for warm Lambda environment reuse.""" @@ -762,9 +951,11 @@ def _reset_state(self) -> None: self._span_time_floor_ns = None with self._operation_spans_lock: self._operation_spans = {} + self._pending_operation_ends = {} self._context_operation_replays = {} self._incomplete_attempt_span_keys = set() self._tracing_enabled = False + finish_state(self._state) def on_operation_start(self, info: OperationStartInfo) -> None: """Called when an operation begins. Creates a span for the operation.""" @@ -797,6 +988,10 @@ def on_operation_end(self, info: OperationEndInfo) -> None: invocation is completing an operation that began earlier, so a short continuation span with a fresh span ID is created and ended. """ + self._render_operation_end(info) + + def _render_operation_end(self, info: OperationEndInfo) -> None: + """Render or defer one delivery without re-entering a public hook.""" logger.debug("Durable operation ended: %s", info) if not self._tracing_enabled: return @@ -805,7 +1000,15 @@ def on_operation_end(self, info: OperationEndInfo) -> None: # The operation started in a prior invocation. Create a new # correlated segment and link it to the deterministic logical # operation context shared across invocations. - parent_span = self._resolve_parent_span(info.parent_id) + with self._operation_spans_lock: + parent_span = self._operation_spans.get(info.parent_id) + if parent_span is None and info.parent_id is not None: + self._pending_operation_ends.setdefault(info.parent_id, []).append( + info + ) + return + if parent_span is None: + parent_span = self._resolve_parent_span(info.parent_id) attributes = self._extract_attributes(info) span = self._start_span( operation_id=info.operation_id, @@ -818,15 +1021,22 @@ def on_operation_end(self, info: OperationEndInfo) -> None: else: span.set_attributes(self._extract_attributes(info)) + exception: Exception | None = None if info.error: span.set_status(StatusCode.ERROR, info.error.message or "") - span.record_exception( - Exception(info.error.message or info.error.type or "Unknown error") + exception = Exception( + info.error.message or info.error.type or "Unknown error" ) + elif info.status in ( + OperationStatus.FAILED, + OperationStatus.TIMED_OUT, + OperationStatus.STOPPED, + ): + span.set_status(StatusCode.ERROR) elif info.status is OperationStatus.SUCCEEDED: span.set_status(StatusCode.OK) - self._end_span(info.operation_id, info.end_time) + self._end_span(info.operation_id, info.end_time, exception=exception) def on_user_function_start(self, info: UserFunctionStartInfo) -> None: """Called when a context or step operation starts user code. @@ -938,16 +1148,15 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: and info.outcome is not UserFunctionOutcome.INCOMPLETE ): span.set_attributes(self._extract_attributes(info)) + exception: Exception | None = None if info.outcome is UserFunctionOutcome.FAILED: span.set_status( StatusCode.ERROR, info.error.message if info.error else "" ) - span.record_exception( - Exception( - (info.error.message or info.error.type) - if info.error - else "Unknown error" - ) + exception = Exception( + (info.error.message or info.error.type) + if info.error + else "Unknown error" ) else: span.set_status(StatusCode.OK) @@ -955,7 +1164,7 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: end_timestamp = info.end_time if end_timestamp is not None and end_timestamp == info.start_time: end_timestamp += datetime.timedelta(microseconds=1) - self._end_span(span_key, end_timestamp) + self._end_span(span_key, end_timestamp, exception=exception) # Restore the enclosing context by releasing the scope this user # function attached. Code that runs after this operation (e.g. between # steps in a child context) correlates to its enclosing operation diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_cached_child_input_delivery_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_cached_child_input_delivery_int.py new file mode 100644 index 000000000..dfc1c0254 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_cached_child_input_delivery_int.py @@ -0,0 +1,252 @@ +"""Visible terminal input updates are delivered once with real retained children.""" + +import json +import logging +import os +from pathlib import Path +import time + +import pytest +from aws_durable_execution_sdk_python import durable_execution, durable_step +from aws_durable_execution_sdk_python.config import CallbackConfig +from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin +from aws_durable_execution_sdk_python.serdes import JsonSerDes +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from aws_durable_execution_sdk_python_testing.stores.filesystem import ( + FileSystemExecutionStore, +) +from aws_durable_execution_sdk_python_testing.stores.memory import ( + InMemoryExecutionStore, +) +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +@pytest.mark.parametrize( + ("view", "filesystem", "failed_delivery"), + [ + (ExecutionOtelPlugin, False, False), + (InvocationOtelPlugin, False, False), + (ExecutionOtelPlugin, True, False), + (InvocationOtelPlugin, True, False), + (ExecutionOtelPlugin, False, True), + ], +) +def test_retained_child_callback_notifies_before_input_ack( + view, filesystem, failed_delivery, tmp_path, monkeypatch, caplog +): + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + store = ( + FileSystemExecutionStore.create(tmp_path / "store") + if filesystem + else InMemoryExecutionStore() + ) + starts, ends, notifications, child_calls, marker_calls = [], [], [], [], [] + attempts = [] + saved_child_calls = [] + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + + class Observe(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + starts.append(info) + + def on_invocation_end(self, info): + ends.append(info) + + def on_operation_end(self, info): + if info.name in ("a", "b"): + notifications.append((len(starts), info)) + + def saved_child(_step): + saved_child_calls.append(1) + return "saved" + + def child(context): + child_calls.append(1) + callback_id = context.create_callback( + name="b", config=CallbackConfig(serdes=JsonSerDes()) + ).callback_id + assert context.step(saved_child, name="saved-child") == "saved" + # The real service hides descendants of a completed ordinary context. + # A real large result makes this context use ReplayChildren, retaining + # B's metadata without injecting input state or hook notifications. + return {"id": callback_id, "large": "x" * (300 * 1024)} + + @durable_step + def marker(_context, result): + marker_calls.append(1) + return result + + def handler(_event, context): + a = context.create_callback( + name="a", config=CallbackConfig(serdes=JsonSerDes()) + ) + b_id = context.run_in_child_context(child, name="cached-child")["id"] + result = context.step(marker(a.result()), name="after-a") + return {"a": result, "b_id": b_id} + + wrapped = durable_execution(handler, plugins=[Observe(), plugin]) + + def entrypoint(event, context): + attempts.append(1) + if failed_delivery and len(attempts) == 2: + raise RuntimeError("invocation failed before SDK input delivery") + return wrapped(event, context) + + def wait_for_completed(runner, arn, count): + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + events = runner.get_execution_history( + arn, include_execution_data=True + ).events + if ( + len(ends) >= count + and sum(e.event_type == "InvocationCompleted" for e in events) >= count + ): + return events + time.sleep(0.01) + raise AssertionError("Expected invocation did not finish") + + try: + with DurableFunctionTestRunner( + handler=entrypoint, + store=store, + skip_time=False, + poll_interval=0.01, + execution_timeout=25, + ) as runner: + arn = runner.run_async(input="{}") + a_id = runner.wait_for_callback(arn, name="a", timeout=5) + b_id = runner.wait_for_callback(arn, name="b", timeout=5) + wait_for_completed(runner, arn, 1) + runner.send_callback_success(b_id, result=b'"B"') + # A pre-SDK failure is a physical attempt, not an SDK End hook. + middle_history = wait_for_completed(runner, arn, 2) + middle_state = store.load(arn) + middle_updated_ids = list(middle_state.updated_operation_ids) + runner.send_callback_success(a_id, result=b'"A"') + result = runner.wait_for_result(arn, timeout=10) + final_history = runner.get_execution_history( + arn, include_execution_data=True + ).events + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == {"a": "A", "b_id": b_id} + assert [i.status.value for i in ends] == ["PENDING", "PENDING", "SUCCEEDED"] + assert child_calls == [1, 1, 1] and saved_child_calls == [1] + assert marker_calls == [1] + assert len(attempts) == (4 if failed_delivery else 3) + b_info = next( + info for info in starts[1].updated_operations.values() if info.name == "b" + ) + parent = starts[1].operations[b_info.parent_id] + assert parent.name == "cached-child" and parent.status.value == "SUCCEEDED" + b_ends = [ + (invocation, info) for invocation, info in notifications if info.name == "b" + ] + terminal_spans = [ + s + for s in exporter.get_finished_spans() + if s.attributes.get("durable.operation.name") == "b" + and s.attributes.get("durable.operation.status") == "SUCCEEDED" + ] + parent_spans = [ + s + for s in exporter.get_finished_spans() + if s.attributes.get("durable.operation.id") == b_info.parent_id + and s.attributes.get("durable.operation.type") == "CONTEXT" + and s.attributes.get("durable.operation.status") == "SUCCEEDED" + ] + hook_errors = [ + r.getMessage() + for r in caplog.records + if r.levelno >= logging.ERROR and ("plugin" in r.name or "otel" in r.name) + ] + proof_dir = os.environ.get("CACHED_CHILD_PROOF_DIR") + if proof_dir: + destination = Path(proof_dir) + destination.mkdir(parents=True, exist_ok=True) + ( + destination + / f"{view.__name__}-{'fs' if filesystem else 'memory'}-failed{failed_delivery}.json" + ).write_text( + json.dumps( + { + "business_succeeded": True, + "result_b_matches_issued": True, + "invocation_statuses": [i.status.value for i in ends], + "child_calls": len(child_calls), + "saved_child_calls": len(saved_child_calls), + "retention_via_public_large_result": True, + "marker_calls": len(marker_calls), + "physical_attempts": len(attempts), + "failed_before_SDK": failed_delivery, + "input_updated_b": b_info.operation_id, + "input_parent_id": b_info.parent_id, + "parent_status": parent.status.value, + "middle_updated_ids": middle_updated_ids, + "b_notifications": [ + (n, i.operation_id, i.parent_id, i.is_replayed) + for n, i in b_ends + ], + "b_terminal_spans": len(terminal_spans), + "hook_errors": hook_errors, + "parent_segment_span_ids": [ + s.context.span_id for s in parent_spans + ], + "b_terminal_parent_ids": [ + s.parent.span_id if s.parent else None + for s in terminal_spans + ], + "middle_history_events": [ + (e.event_id, e.event_type, e.name) for e in middle_history + ], + "final_history_events": [ + (e.event_id, e.event_type, e.name) for e in final_history + ], + }, + indent=2, + ) + + "\n" + ) + assert len(b_ends) == 1 + assert b_ends[0][0] == 2 + assert b_ends[0][1].operation_id == b_info.operation_id + assert b_ends[0][1].parent_id == b_info.parent_id + assert not b_ends[0][1].is_replayed + assert len(terminal_spans) == 1 + assert terminal_spans[0].parent is not None + if view is ExecutionOtelPlugin: + assert len(parent_spans) == 1 + assert terminal_spans[0].parent.span_id == parent_spans[0].context.span_id + else: + assert len(parent_spans) == 3 + assert len({span.context.span_id for span in parent_spans}) == 3 + invocations = [ + span + for span in exporter.get_finished_spans() + if span.name == "Invocation" + ] + assert len(invocations) == 3 + assert [span.parent.span_id for span in parent_spans] == [ + span.context.span_id for span in invocations + ] + assert terminal_spans[0].parent.span_id == parent_spans[1].context.span_id + assert not hook_errors + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_callback_input_delivery_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_callback_input_delivery_int.py new file mode 100644 index 000000000..917716607 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_callback_input_delivery_int.py @@ -0,0 +1,198 @@ +"""Input-only callback resumes export each terminal outcome once.""" + +import json +import logging +import threading +import time + +import pytest +from aws_durable_execution_sdk_python import ( + DurableContext, + durable_execution, + durable_step, + durable_with_child_context, +) +from aws_durable_execution_sdk_python.config import CallbackConfig +from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin +from aws_durable_execution_sdk_python.serdes import JsonSerDes +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +@pytest.mark.parametrize( + "view,nested,page_bytes,fresh_executor,concurrent_completion", + [ + (ExecutionOtelPlugin, False, None, False, False), + (InvocationOtelPlugin, False, None, False, False), + (ExecutionOtelPlugin, True, None, False, False), + (InvocationOtelPlugin, True, None, False, False), + (InvocationOtelPlugin, True, 1, False, False), + (ExecutionOtelPlugin, False, None, True, False), + (InvocationOtelPlugin, True, None, False, True), + ], + ids=[ + "execution", + "invocation", + "execution-nested", + "invocation-nested", + "paged", + "fresh-executor", + "concurrent", + ], +) +def test_two_callbacks_do_not_repeat_terminal_spans_after_input_only_resume( + view, nested, page_bytes, fresh_executor, concurrent_completion, caplog, monkeypatch +): + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + starts = [] + ends = [] + marker_calls = [] + second_end = threading.Event() + release_end = threading.Event() + + class Observe(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + starts.append(info) + + def on_operation_end(self, info): + if info.name in ("a", "b"): + ends.append((len(starts), info.name, info.status.value)) + + def on_invocation_end(self, info): + if concurrent_completion and len(starts) == 2: + second_end.set() + assert release_end.wait(6) + + observer = Observe() + + def make_plugin(): + return view( + OtelPluginConfig( + tracer_provider=provider, + enrich_logger=False, + context_extractor=lambda _: None, + ) + ) + + @durable_step + def marker(context, value): + marker_calls.append(1) + return value + + def callbacks(context): + a = context.create_callback( + name="a", config=CallbackConfig(serdes=JsonSerDes()) + ) + b = context.create_callback( + name="b", config=CallbackConfig(serdes=JsonSerDes()) + ) + return [a.result(), b.result()] + + @durable_with_child_context + def child(context): + return callbacks(context) + + def body(event, context: DurableContext): + result = ( + context.run_in_child_context(child(), name="group") + if nested + else callbacks(context) + ) + return context.step(marker(result), name="after-b") + + wrapped = durable_execution(body, plugins=[make_plugin(), observer]) + + def entrypoint(event, context): + target = ( + durable_execution(body, plugins=[make_plugin(), observer]) + if fresh_executor + else wrapped + ) + return target(event, context) + + try: + with DurableFunctionTestRunner( + handler=entrypoint, + skip_time=False, + poll_interval=0.01, + execution_timeout=25, + max_invocation_page_bytes=page_bytes, + ) as runner: + arn = runner.run_async(input="{}") + a_id = runner.wait_for_callback(arn, name="a", timeout=5) + b_id = runner.wait_for_callback(arn, name="b", timeout=5) + runner.send_callback_success(a_id, result=b'"A"') + if concurrent_completion: + assert second_end.wait(5) + runner.send_callback_success(b_id, result=b'"B"') + release_end.set() + else: + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + events = runner.get_execution_history( + arn, include_execution_data=True + ).events + if ( + sum( + event.event_type == "InvocationCompleted" + for event in events + ) + >= 2 + ): + break + time.sleep(0.01) + else: + raise AssertionError("The second invocation did not suspend") + runner.send_callback_success(b_id, result=b'"B"') + result = runner.wait_for_result(arn, timeout=10) + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == ["A", "B"] + assert len(starts) == 3 + assert marker_calls == [1] + assert ends == [(2, "a", "SUCCEEDED"), (3, "b", "SUCCEEDED")] + spans = exporter.get_finished_spans() + callback_spans = [ + span + for span in spans + if (span.attributes or {}).get("durable.operation.type") == "CALLBACK" + ] + # Retain the nonterminal segments as well as the terminal segments; + # filtering down to successful spans would hide the invocation story. + expected = [("a", "SUCCEEDED"), ("b", "SUCCEEDED")] + if view is InvocationOtelPlugin: + expected = [ + ("b", "STARTED"), + ("a", "STARTED"), + ("a", "SUCCEEDED"), + ("b", "STARTED"), + ("b", "SUCCEEDED"), + ] + assert [ + (span.name, span.attributes["durable.operation.status"]) + for span in callback_spans + ] == expected + exported_ids = {span.context.span_id for span in spans} + assert all( + span.parent is not None and span.parent.span_id in exported_ids + for span in callback_spans + ) + assert not [ + record + for record in caplog.records + if record.levelno >= logging.ERROR + and ("plugin" in record.name or "otel" in record.name) + ] + finally: + release_end.set() + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_exception_event_bounds_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_exception_event_bounds_int.py new file mode 100644 index 000000000..ed72ee651 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_exception_event_bounds_int.py @@ -0,0 +1,78 @@ +"""Real checkpointless child failures keep diagnostic events inside their span.""" + +import json + +import pytest +from aws_durable_execution_sdk_python import ( + durable_execution, + durable_with_child_context, +) +from aws_durable_execution_sdk_python.config import ChildConfig +from aws_durable_execution_sdk_python.exceptions import ChildContextError +from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +@pytest.mark.parametrize("view", [ExecutionOtelPlugin, InvocationOtelPlugin]) +def test_caught_virtual_child_failure_has_bounded_event_time(view, monkeypatch): + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + enrich_logger=False, + context_extractor=lambda _: None, + ) + ) + hooks = [] + + class Observe(DurableInstrumentationPlugin): + def on_operation_end(self, info): + if info.name == "virtual-error": + hooks.append(info) + + @durable_with_child_context + def child(_context): + raise ValueError("actual child failure") + + def body(_event, context): + try: + context.run_in_child_context( + child(), name="virtual-error", config=ChildConfig(is_virtual=True) + ) + except ChildContextError: + return "caught" + raise AssertionError("Expected the stored child failure") + + try: + with DurableFunctionTestRunner( + handler=durable_execution(body, plugins=[plugin, Observe()]), + poll_interval=0.01, + ) as runner: + result = runner.run(input="{}") + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == "caught" + (info,) = hooks + assert info.start_time is None + assert info.status.value == "FAILED" + assert info.error.message == "actual child failure" + (span,) = [ + s for s in exporter.get_finished_spans() if s.name == "virtual-error" + ] + assert span.status.status_code.name == "ERROR" + assert len(span.events) == 1 + assert span.events[0].attributes["exception.message"] == "actual child failure" + assert span.start_time <= span.events[0].timestamp <= span.end_time + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_flush_lifecycle_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_flush_lifecycle_int.py new file mode 100644 index 000000000..6fa26b08a --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_flush_lifecycle_int.py @@ -0,0 +1,215 @@ +"""Only provider I/O may outlive the existing handler-worker End scope.""" + +import contextvars +import json +import threading +import time +from datetime import UTC, datetime +from unittest.mock import Mock, patch + +import pytest +from aws_durable_execution_sdk_python import durable_execution +from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin +from aws_durable_execution_sdk_python.lambda_service import ( + ExecutionDetails, + Operation, + OperationStatus, + OperationType, +) +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import BatchSpanProcessor +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, + _flush, +) + + +@pytest.mark.parametrize("view", [ExecutionOtelPlugin, InvocationOtelPlugin]) +def test_public_handler_finishes_while_one_provider_flush_is_blocked(view, monkeypatch): + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + monkeypatch.setattr(_flush, "_FLUSH_TIMEOUT_SECONDS", 0.05) + ambient = contextvars.ContextVar("flush-owner", default="caller") + release = threading.Event() + exporting = threading.Event() + owners = [] + flush_calls = [] + exporter_contexts = [] + + class Exporter(InMemorySpanExporter): + def export(self, spans): + exporting.set() + exporter_contexts.append(ambient.get()) + assert release.wait(10) + return super().export(spans) + + class Provider(TracerProvider): + def force_flush(self, timeout_millis=30000): + flush_calls.append(threading.get_ident()) + return super().force_flush(timeout_millis) + + class Owner(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + self.token = ambient.set("handler") + self.thread = threading.get_ident() + + def on_invocation_end(self, info): + assert ambient.get() == "handler" + assert threading.get_ident() == self.thread + # A successful token reset proves ownership of the actual Context, + # rather than only comparing an OS thread ID that can be reused. + ambient.reset(self.token) + owners.append(self.thread) + + exporter = Exporter() + provider = Provider() + provider.add_span_processor( + BatchSpanProcessor(exporter, schedule_delay_millis=60000) + ) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + enrich_logger=False, + context_extractor=lambda _: None, + ) + ) + + def body(_event, _context): + assert ambient.get() == "handler" + return "done" + + handler = durable_execution(body, plugins=[Owner(), plugin]) + + def invoke(): + started = time.monotonic() + with DurableFunctionTestRunner(handler=handler, poll_interval=0.01) as runner: + result = runner.run(input="{}") + assert time.monotonic() - started < 2 + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == "done" + assert ambient.get() == "caller" + + try: + invoke() + assert exporting.wait(1) + invoke() + assert len(owners) == 2 + assert len(flush_calls) == 1 + assert all(worker != flush_calls[0] for worker in owners) + assert exporter_contexts == ["caller"] + finally: + release.set() + deadline = time.monotonic() + 2 + while time.monotonic() < deadline: + with _flush._lock: + idle = _flush._active is None + if idle: + break + time.sleep(0.001) + assert idle + try: + # Actual completion restores admission; no accumulated late error or + # stale invocation Context is delivered to this next request. + monkeypatch.setattr(_flush, "_FLUSH_TIMEOUT_SECONDS", 1.0) + invoke() + assert len(flush_calls) == 2 + assert len(owners) == 3 + assert all(value == "caller" for value in exporter_contexts) + finally: + provider.shutdown() + + +@pytest.mark.parametrize("view", [ExecutionOtelPlugin, InvocationOtelPlugin]) +@pytest.mark.parametrize("fatal", [False, True]) +def test_public_fast_flush_failure_preserves_owner_cleanup_and_reuse( + view, fatal, monkeypatch +): + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + monkeypatch.setattr(_flush, "_FLUSH_TIMEOUT_SECONDS", 1.0) + binding = contextvars.ContextVar("flush-failure-owner", default="caller") + ended = [] + error = SystemExit("provider fatal") if fatal else RuntimeError("provider ordinary") + + class Owner(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + self.token = binding.set("handler") + self.thread = threading.get_ident() + + def on_invocation_end(self, info): + assert threading.get_ident() == self.thread + binding.reset(self.token) + ended.append(self.thread) + + class Provider(TracerProvider): + failure: BaseException | None = error + + def force_flush(self, timeout_millis=30000): + if self.failure is not None: + raise self.failure + return super().force_flush(timeout_millis) + + provider = Provider() + provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter())) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + enrich_logger=False, + context_extractor=lambda _: None, + ) + ) + + def body(_event, _context): + assert binding.get() == "handler" + return "done" + + handler = durable_execution(body, plugins=[Owner(), plugin]) + operation = Operation( + operation_id="flush-execution", + operation_type=OperationType.EXECUTION, + status=OperationStatus.STARTED, + start_timestamp=datetime.now(UTC), + execution_details=ExecutionDetails(input_payload="{}"), + ) + event = { + "DurableExecutionArn": "test-arn/flush-execution", + "CheckpointToken": "token", + "InitialExecutionState": { + "Operations": [operation.to_json_dict()], + "NextMarker": "", + }, + "LocalRunner": True, + } + context = Mock() + context.aws_request_id = "request" + context.invoked_function_arn = "test-arn" + context.client_context = None + context.identity = None + context.tenant_id = None + context._epoch_deadline_time_in_ms = 0 + try: + with patch("aws_durable_execution_sdk_python.execution.LambdaClient"): + if fatal: + with pytest.raises(SystemExit) as caught: + handler(event, context) + assert caught.value is error + else: + assert handler(event, context) == { + "Status": "SUCCEEDED", + "Result": '"done"', + } + assert len(ended) == 1 + assert binding.get() == "caller" + provider.failure = None + assert handler(event, context) == { + "Status": "SUCCEEDED", + "Result": '"done"', + } + assert len(ended) == 2 + assert binding.get() == "caller" + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py index fa7036920..7477a779d 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py @@ -15,10 +15,13 @@ InvocationStatus, durable_execution, ) +from aws_durable_execution_sdk_python.exceptions import InvokeError from aws_durable_execution_sdk_python.lambda_service import ( + ChainedInvokeDetails, CheckpointOutput, CheckpointUpdatedExecutionState, ExecutionDetails, + ErrorObject, Operation, OperationAction, OperationStatus, @@ -142,6 +145,128 @@ def checkpoint( return checkpoint, operations +@pytest.mark.parametrize("plugin_type", [ExecutionOtelPlugin, InvocationOtelPlugin]) +@pytest.mark.parametrize("with_error", [False, True]) +def test_caught_invoke_failure_without_details_is_error_on_live_and_replay( + plugin_type, with_error +) -> None: + """Use real SDK replay with service-shaped failed-invoke checkpoint responses. + + The local child runner synthesizes a diagnostic for FAILED children, so the + absent-error case belongs at the LambdaClient transport boundary. No plugin + lifecycle/status callback is fabricated here. + """ + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = plugin_type( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + error = ( + ErrorObject( + message="supplied invoke failure", + type="External", + data=None, + stack_trace=None, + ) + if with_error + else None + ) + caught: list[str] = [] + saved_calls: list[str] = [] + + @durable_step + def saved(_context) -> str: + saved_calls.append("saved") + return "caught" + + def body(_event, context: DurableContext) -> str: + try: + context.invoke("child", {}, name="failed-invoke") + except InvokeError as observed: + caught.append(type(observed).__name__) + value = context.step(saved(), name="saved-result") + context.wait(Duration.from_seconds(1), name="resume-wait") + return value + + handler = durable_execution(body, plugins=[plugin]) + + def failed_invoke_store(initial): + checkpoint, operations = _checkpoint_store(initial) + + def send(*args, **kwargs): + output = checkpoint(*args, **kwargs) + for key, operation in operations.items(): + if operation.name == "failed-invoke": + operations[key] = replace( + operation, + status=OperationStatus.FAILED, + sub_type=OperationSubType.CHAINED_INVOKE, + end_timestamp=datetime.now(UTC), + chained_invoke_details=ChainedInvokeDetails(error=error), + ) + return replace( + output, + new_execution_state=CheckpointUpdatedExecutionState( + operations=list(operations.values()) + ), + ) + + return send, operations + + initial = [_execution_operation()] + checkpoint, operations = failed_invoke_store(initial) + try: + with patch( + "aws_durable_execution_sdk_python.execution.LambdaClient" + ) as client_class: + client = Mock() + client.checkpoint = checkpoint + client_class.initialize_client.return_value = client + first = handler(_event(initial), _lambda_context()) + assert first["Status"] == "PENDING" + first_spans = [ + s for s in exporter.get_finished_spans() if s.name == "failed-invoke" + ] + resumed = [ + replace( + op, status=OperationStatus.SUCCEEDED, end_timestamp=datetime.now(UTC) + ) + if op.name == "resume-wait" + else op + for op in operations.values() + ] + wait_id = next(op.operation_id for op in resumed if op.name == "resume-wait") + checkpoint, _ = failed_invoke_store(resumed) + with patch( + "aws_durable_execution_sdk_python.execution.LambdaClient" + ) as client_class: + client = Mock() + client.checkpoint = checkpoint + client_class.initialize_client.return_value = client + second = handler(_event(resumed, [wait_id]), _lambda_context()) + assert second == {"Status": "SUCCEEDED", "Result": '"caught"'} + assert caught == ["InvokeError", "InvokeError"] + assert saved_calls == ["saved"] + spans = [s for s in exporter.get_finished_spans() if s.name == "failed-invoke"] + assert spans == first_spans + (span,) = spans + assert span.attributes is not None + assert span.attributes["durable.operation.status"] == "FAILED" + assert span.status.status_code is trace.StatusCode.ERROR + assert len(span.events) == int(with_error) + workflow = next( + s for s in exporter.get_finished_spans() if s.name == "Workflow" + ) + assert workflow.status.status_code is trace.StatusCode.OK + finally: + provider.shutdown() + + @pytest.mark.parametrize( "plugin_type", [InvocationOtelPlugin, ExecutionOtelPlugin], diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_late_child_parent_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_late_child_parent_int.py new file mode 100644 index 000000000..6845f2884 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_late_child_parent_int.py @@ -0,0 +1,236 @@ +"""Late callback completions keep the actual parent for their invocation.""" + +import json +import logging +import os +import queue +from pathlib import Path +import threading +import time + +import pytest +from aws_durable_execution_sdk_python import durable_execution +from aws_durable_execution_sdk_python.config import CallbackConfig +from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin +from aws_durable_execution_sdk_python.serdes import JsonSerDes +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +@pytest.mark.parametrize("view", [ExecutionOtelPlugin, InvocationOtelPlugin]) +@pytest.mark.parametrize( + ("large_result", "resume_before_completion"), + [(False, False), (True, False), (False, True)], +) +def test_late_callback_keeps_completed_or_replayed_parent( + view, large_result, resume_before_completion, monkeypatch, caplog +): + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + starts, ends, child_calls, saved_calls, root_calls = [], [], [], [], [] + root_step_started, release_root_step = threading.Event(), threading.Event() + callback_ids = queue.Queue(maxsize=1) + + class Observe(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + starts.append(info) + + def on_invocation_end(self, info): + ends.append(info) + + def saved(_step): + saved_calls.append(1) + return "saved" + + def child(ctx): + child_calls.append(1) + cb = ctx.create_callback(name="b", config=CallbackConfig(serdes=JsonSerDes())) + assert ctx.step(saved, name="save-child") == "saved" + return dict( + b_id=cb.callback_id, large="x" * (300 * 1024) if large_result else "small" + ) + + def root_step(a_id, b_id): + callback_ids.put((a_id, b_id)) + root_calls.append(1) + root_step_started.set() + assert release_root_step.wait(10) + return "root-saved" + + def handler(_event, ctx): + a = ctx.create_callback(name="a", config=CallbackConfig(serdes=JsonSerDes())) + result = ctx.run_in_child_context(child, name="returned-child") + if resume_before_completion: + gate = ctx.create_callback( + name="resume-gate", config=CallbackConfig(serdes=JsonSerDes()) + ) + assert gate.result() == "resume" + if not large_result: + assert ( + ctx.step( + lambda _: root_step(a.callback_id, result["b_id"]), + name="root-checkpoint", + ) + == "root-saved" + ) + return dict(a=a.result(), b_id=result["b_id"], length=len(result["large"])) + + wrapped = durable_execution(handler, plugins=[Observe(), plugin]) + + def completed(runner, arn, count): + deadline = time.monotonic() + 12 + while time.monotonic() < deadline: + history = runner.get_execution_history( + arn, include_execution_data=True + ).events + if ( + len(ends) >= count + and sum(e.event_type == "InvocationCompleted" for e in history) >= count + ): + return history + time.sleep(0.01) + raise AssertionError("invocation completion not observed") + + def spans(name, status=None): + return [ + s + for s in exporter.get_finished_spans() + if s.attributes.get("durable.operation.name") == name + and ( + status is None or s.attributes.get("durable.operation.status") == status + ) + ] + + try: + with DurableFunctionTestRunner( + handler=wrapped, skip_time=False, poll_interval=0.01, execution_timeout=35 + ) as runner: + arn = runner.run_async(input="{}") + if large_result: + a_id = runner.wait_for_callback(arn, name="a", timeout=5) + b_id = runner.wait_for_callback(arn, name="b", timeout=5) + completed(runner, arn, 1) + else: + if resume_before_completion: + gate_id = runner.wait_for_callback( + arn, name="resume-gate", timeout=5 + ) + completed(runner, arn, 1) + runner.send_callback_success(gate_id, result=b'"resume"') + # The runner wait_for_callback helper waits for suspension; + # this case deliberately completes during the active root step. + a_id, b_id = callback_ids.get(timeout=5) + assert root_step_started.wait(5) + runner.send_callback_success(b_id, result=b'"B"') + release_root_step.set() + completed(runner, arn, 2 if large_result or resume_before_completion else 1) + middle_callbacks = spans("b", "SUCCEEDED") + middle_parents = spans("returned-child", "SUCCEEDED") + invocation_spans = [ + s for s in exporter.get_finished_spans() if s.name == "Invocation" + ] + runner.send_callback_success(a_id, result=b'"A"') + result = runner.wait_for_result(arn, timeout=10) + history = runner.get_execution_history( + arn, include_execution_data=True + ).events + expected_invocations = 3 if large_result or resume_before_completion else 2 + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == dict( + a="A", b_id=b_id, length=300 * 1024 if large_result else 5 + ) + assert len(starts) == expected_invocations and [ + x.status.value for x in ends + ] == ["PENDING"] * (expected_invocations - 1) + ["SUCCEEDED"] + assert saved_calls == [1] and root_calls == ([] if large_result else [1]) + assert len(child_calls) == (3 if large_result else 1) + if resume_before_completion: + assert all( + info.name != "b" for info in starts[1].updated_operations.values() + ) + # Service input prunes B even while it is STARTED. Its later + # checkpoint response still carries the actual completion unchanged. + assert all(info.name != "b" for info in starts[1].operations.values()) + completed_parent = next( + info + for info in starts[1].operations.values() + if info.name == "returned-child" + ) + assert completed_parent.status.value == "SUCCEEDED" + hook_errors = [ + r.getMessage() + for r in caplog.records + if r.levelno >= logging.ERROR and ("plugin" in r.name or "otel" in r.name) + ] + proof = dict( + view=view.__name__, + large_result=large_result, + resume_before_completion=resume_before_completion, + statuses=[x.status.value for x in ends], + saved_step_calls=len(saved_calls), + root_step_calls=len(root_calls), + child_calls=len(child_calls), + middle_input_updated_names=[ + info.name for info in starts[1].updated_operations.values() + ] + if resume_before_completion + else None, + middle_callback_spans=[ + dict(span_id=s.context.span_id, parent_id=s.parent.span_id) + for s in middle_callbacks + ], + middle_parent_spans=[ + dict(span_id=s.context.span_id, parent_id=s.parent.span_id) + for s in middle_parents + ], + middle_invocation_ids=[s.context.span_id for s in invocation_spans], + final_callback_spans=len(spans("b", "SUCCEEDED")), + hook_errors=hook_errors, + history=[(e.event_id, e.event_type, e.name) for e in history], + ) + if os.environ.get("LATE_PARENT_PROOF_DIR"): + out = Path(os.environ["LATE_PARENT_PROOF_DIR"]) + out.mkdir(parents=True, exist_ok=True) + ( + out + / f"{view.__name__}-large{large_result}-resume{resume_before_completion}.json" + ).write_text(json.dumps(proof, indent=2) + "\n") + assert len(middle_callbacks) == 1 + assert len(spans("b", "SUCCEEDED")) == 1 + if large_result and view is InvocationOtelPlugin: + assert len(middle_parents) == 2 + assert ( + middle_parents[-1].parent.span_id + == invocation_spans[-1].context.span_id + ) + assert ( + middle_callbacks[0].parent.span_id == middle_parents[-1].context.span_id + ) + assert ( + middle_callbacks[0].parent.span_id != middle_parents[0].context.span_id + ) + else: + assert ( + middle_callbacks[0].parent.span_id == middle_parents[0].context.span_id + ) + assert not hook_errors + finally: + release_root_step.set() + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_parent_creation_order_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_parent_creation_order_int.py new file mode 100644 index 000000000..16490c3f8 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_parent_creation_order_int.py @@ -0,0 +1,198 @@ +"""A real completion can precede its replayed sibling's parent span.""" + +from __future__ import annotations + +import json +import threading +import time +from typing import Any + +import pytest +from aws_durable_execution_sdk_python import DurableContext, durable_execution +from aws_durable_execution_sdk_python.config import ( + CallbackConfig, + Duration, + ParallelConfig, +) +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationStartInfo, + OperationEndInfo, + OperationType, + UserFunctionStartInfo, +) +from aws_durable_execution_sdk_python.serdes import JsonSerDes +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry import trace +from opentelemetry.sdk.trace import ReadableSpan, TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import InvocationOtelPlugin, OtelPluginConfig + + +def _span_start(span: ReadableSpan) -> int: + assert span.start_time is not None + return span.start_time + + +def test_early_sibling_callback_exports_under_real_parent_before_user_observation( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = InvocationOtelPlugin( + OtelPluginConfig( + tracer_provider=provider, + enrich_logger=False, + context_extractor=lambda _: None, + ) + ) + second_start = threading.Event() + both_completed = threading.Event() + end_observed = threading.Event() + parent_entered = threading.Event() + generations = [0] + observed: list[tuple[int, int]] = [] + hook_order: list[str] = [] + + class Gate(DurableInstrumentationPlugin): + def on_invocation_start(self, _info: InvocationStartInfo) -> None: + generations[0] += 1 + if generations[0] == 2: + second_start.set() + assert both_completed.wait(10) + + def on_user_function_start(self, info: UserFunctionStartInfo) -> None: + if generations[0] == 2 and info.name == "parallel-branch-1": + hook_order.append("parent-gated") + assert end_observed.wait(10) + parent_entered.set() + + class ObserveEnd(DurableInstrumentationPlugin): + def on_operation_end(self, info: OperationEndInfo) -> None: + if ( + generations[0] == 2 + and info.name == "target-1" + and info.operation_type is OperationType.CALLBACK + ): + # This observer follows OTel in the actual SDK dispatch order. + assert not parent_entered.is_set() + hook_order.append("early-completion-observed") + end_observed.set() + + def branch(child: DurableContext, index: int) -> Any: + value = child.create_callback( + name=f"target-{index}", config=CallbackConfig(serdes=JsonSerDes()) + ).result() + if index == 1: + finished = [ + span + for span in exporter.get_finished_spans() + if span.name == "target-1" + ] + observed.append( + (len(finished), trace.get_current_span().get_span_context().span_id) + ) + with provider.get_tracer("customer").start_as_current_span( + "after-target-1" + ): + pass + return value + + def handler(_event: Any, durable: DurableContext) -> list[Any]: + result = durable.parallel( + [lambda child: branch(child, 0), lambda child: branch(child, 1)], + name="parallel", + config=ParallelConfig(max_concurrency=2), + ) + durable.wait(Duration.from_seconds(1), name="later-one") + durable.wait(Duration.from_seconds(1), name="later-two") + return result.get_results() + + wrapped = durable_execution(handler, plugins=[Gate(), plugin, ObserveEnd()]) + try: + with DurableFunctionTestRunner( + handler=wrapped, poll_interval=0.01, skip_time=False, execution_timeout=25 + ) as runner: + arn = runner.run_async(input="{}") + deadline = time.monotonic() + 10 + callbacks: dict[str, str] = {} + while time.monotonic() < deadline: + history = runner.get_execution_history( + arn, include_execution_data=True + ).events + callbacks = {} + for event in history: + if event.event_type == "CallbackStarted": + assert event.name is not None + assert event.callback_started_details is not None + assert event.callback_started_details.callback_id is not None + callbacks[event.name] = ( + event.callback_started_details.callback_id + ) + if len(callbacks) == 2 and any( + event.event_type == "InvocationCompleted" for event in history + ): + break + time.sleep(0.01) + assert set(callbacks) == {"target-0", "target-1"} + runner.send_callback_success( + callbacks["target-0"], result=json.dumps("left").encode() + ) + assert second_start.wait(10) + # The second callback completes after the invocation input snapshot, + # so another branch's checkpoint returns this genuine new completion. + runner.send_callback_success( + callbacks["target-1"], result=json.dumps("right").encode() + ) + both_completed.set() + result = runner.wait_for_result(arn, timeout=15) + history = runner.get_execution_history( + arn, include_execution_data=True + ).events + assert result.status.value == "SUCCEEDED" + assert result.result is not None + assert json.loads(result.result) == ["left", "right"] + assert end_observed.is_set() and parent_entered.is_set() + assert len(observed) == 1 and observed[0][0] == 2 + spans = exporter.get_finished_spans() + # Keep every span: Invocation view exports a nonterminal segment in + # invocation one and a distinct terminal continuation in invocation two. + # Neither segment may stand in for the other or reappear on later replays. + for index in (0, 1): + segments = sorted( + [span for span in spans if span.name == f"target-{index}"], + key=_span_start, + ) + assert len(segments) == 2 + initial, terminal = segments + assert initial.attributes is not None and terminal.attributes is not None + assert initial.attributes["durable.operation.status"] == "STARTED" + assert terminal.attributes["durable.operation.status"] == "SUCCEEDED" + assert initial.context.span_id != terminal.context.span_id + parents = sorted( + [span for span in spans if span.name == f"parallel-branch-{index}"], + key=_span_start, + ) + assert len(parents) == 2 + assert initial.parent is not None + assert initial.parent.span_id == parents[0].context.span_id + assert terminal.parent is not None + assert terminal.parent.span_id == parents[1].context.span_id + assert initial.end_time is not None and terminal.start_time is not None + assert initial.end_time <= terminal.start_time + if index == 1: + assert terminal.parent.span_id == observed[0][1] + (marker,) = [span for span in spans if span.name == "after-target-1"] + assert terminal.end_time is not None and marker.start_time is not None + assert terminal.end_time <= marker.start_time + assert sum(event.event_type == "InvocationCompleted" for event in history) == 4 + assert not [record for record in caplog.records if record.exc_info] + assert plugin._context_tokens == {} + finally: + both_completed.set() + end_observed.set() + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_parent_registration_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_parent_registration_int.py new file mode 100644 index 000000000..9d590f2bf --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_parent_registration_int.py @@ -0,0 +1,182 @@ +"""A closing batch remains the parent of a concurrent virtual branch end.""" + +import json +import logging +import threading + +import pytest +from aws_durable_execution_sdk_python import durable_execution +from aws_durable_execution_sdk_python.config import ( + CompletionConfig, + MapConfig, + NestingType, + ParallelBranch, + ParallelConfig, +) +from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ExecutionOtelPlugin, OtelPluginConfig + + +@pytest.mark.parametrize("operation", ["map", "parallel"]) +@pytest.mark.parametrize("interleave", [False, True]) +def test_closing_batch_keeps_parent_registration( + operation, interleave, monkeypatch, caplog, tmp_path +): + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = ExecutionOtelPlugin( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + gates = {} + saved = [] + hook_ends = [] + invocation_ends = [] + + class Observe(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + gates.clear() + gates.update( + slow_ready=threading.Event(), + parent_materializing=threading.Event(), + slow_end=threading.Event(), + ended=threading.Event(), + parent_waited=False, + arn=info.execution_arn, + ) + + def on_user_function_end(self, info): + if info.name == "slow-branch": + gates["slow_ready"].set() + if interleave: + assert gates["parent_materializing"].wait(5) + elif info.name == "fast-branch": + assert gates["slow_ready"].wait(5) + + def on_operation_end(self, info): + hook_ends.append((gates["arn"], info.name, info.operation_id)) + if info.name == "slow-branch": + gates["slow_end"].set() + + def on_invocation_end(self, info): + invocation_ends.append(info.status.value) + gates["ended"].set() + + original_start_span = plugin._start_span + + def gated_start_span(**kwargs): + # Only pause scheduling. The real SDK delivers every hook, and the + # original renderer/provider still creates and exports every span. + if interleave and kwargs["name"] == "race-parent": + gates["parent_materializing"].set() + assert gates["slow_end"].wait(5) + gates["parent_waited"] = True + return original_start_span(**kwargs) + + monkeypatch.setattr(plugin, "_start_span", gated_start_span) + + def save(_step): + saved.append(1) + return "complete" + + def handler(_event, ctx): + completion = ( + CompletionConfig.first_successful() + if interleave + else CompletionConfig.all_successful() + ) + if operation == "map": + ctx.map( + [0, 1], + lambda _ctx, value, _index, _items: value, + name="race-parent", + config=MapConfig( + max_concurrency=2, + nesting_type=NestingType.FLAT, + completion_config=completion, + item_namer=lambda _value, index: ("fast-branch", "slow-branch")[ + index + ], + ), + ) + else: + ctx.parallel( + [ + ParallelBranch(func=lambda _ctx: 0, name="fast-branch"), + ParallelBranch(func=lambda _ctx: 1, name="slow-branch"), + ], + name="race-parent", + config=ParallelConfig( + max_concurrency=2, + nesting_type=NestingType.FLAT, + completion_config=completion, + ), + ) + return ctx.step(save, name="root-save") + + wrapped = durable_execution(handler, plugins=[plugin, Observe()]) + try: + with ( + caplog.at_level(logging.ERROR), + DurableFunctionTestRunner( + handler=wrapped, + skip_time=False, + poll_interval=0.01, + execution_timeout=20, + ) as runner, + ): + for execution in range(2): + arn = runner.run_async(input="{}") + result = runner.wait_for_result(arn, timeout=15) + assert gates["ended"].wait(5) + spans = [ + span + for span in exporter.get_finished_spans() + if span.attributes.get("durable.execution.arn") == arn + ] + rows = [ + dict( + name=span.name, + operation=span.attributes.get("durable.operation.name"), + operation_id=span.attributes.get("durable.operation.id"), + operation_status=span.attributes.get( + "durable.operation.status" + ), + trace_id=format(span.context.trace_id, "032x"), + span_id=format(span.context.span_id, "016x"), + parent_id=format(span.parent.span_id, "016x") + if span.parent + else None, + ) + for span in spans + ] + (tmp_path / f"execution-{execution}.json").write_text( + json.dumps(rows, indent=2) + ) + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == "complete" + assert gates["slow_ready"].is_set() and gates["slow_end"].is_set() + assert gates["parent_waited"] is interleave + parents = [r for r in rows if r["operation"] == "race-parent"] + slow = [r for r in rows if r["operation"] == "slow-branch"] + fast = [r for r in rows if r["operation"] == "fast-branch"] + assert len(parents) == len(fast) == len(slow) == 1, rows + assert slow[0]["parent_id"] == parents[0]["span_id"] + assert fast[0]["parent_id"] == parents[0]["span_id"] + assert slow[0]["operation_status"] == "SUCCEEDED" + assert len([r for r in hook_ends if r[:2] == (arn, "slow-branch")]) == 1 + assert saved == [1] * (execution + 1) + assert invocation_ends == ["SUCCEEDED"] * (execution + 1) + assert not [r.message for r in caplog.records if r.levelno >= logging.ERROR] + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_shared_invocation_state_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_shared_invocation_state_int.py new file mode 100644 index 000000000..1a0bbd625 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_shared_invocation_state_int.py @@ -0,0 +1,113 @@ +"""A shared plugin routes worker/checkpoint events to their real invocation.""" + +import json +import logging +import threading + +import pytest +from aws_durable_execution_sdk_python import durable_execution, durable_step +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +@pytest.mark.parametrize("view", [ExecutionOtelPlugin, InvocationOtelPlugin]) +def test_same_handler_shared_plugin_keeps_two_execution_traces( + view, monkeypatch, caplog +): + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + ready = threading.Event() + release = threading.Event() + side_effects = [] + + @durable_step + def save(_context, label, stage): + side_effects.append((label, stage)) + return label + stage + + def body(event, context): + label = event["label"] + first = context.step(save(label, "first"), name="first") + if label == "a": + ready.set() + assert release.wait(5) + second = context.step(save(label, "second"), name="second") + return [first, second] + + handler = durable_execution(body, plugins=[plugin]) + try: + with ( + DurableFunctionTestRunner(handler=handler, poll_interval=0.01) as ra, + DurableFunctionTestRunner(handler=handler, poll_interval=0.01) as rb, + ): + a = ra.run_async(input='{"label":"a"}') + assert ready.wait(5) + try: + b = rb.run_async(input='{"label":"b"}') + result_b = rb.wait_for_result(b, timeout=5) + finally: + release.set() + result_a = ra.wait_for_result(a, timeout=5) + assert result_a.status.value == result_b.status.value == "SUCCEEDED" + assert json.loads(result_a.result) == ["afirst", "asecond"] + assert json.loads(result_b.result) == ["bfirst", "bsecond"] + assert sorted(side_effects) == [ + ("a", "first"), + ("a", "second"), + ("b", "first"), + ("b", "second"), + ] + spans = exporter.get_finished_spans() + workflows = [s for s in spans if s.name == "Workflow"] + assert len(workflows) == 2 + assert {s.attributes["durable.execution.arn"] for s in workflows} == {a, b} + trace_for = { + s.attributes["durable.execution.arn"]: s.context.trace_id for s in workflows + } + assert len(set(trace_for.values())) == 2 + for arn in (a, b): + own = [s for s in spans if s.attributes.get("durable.execution.arn") == arn] + assert all(s.context.trace_id == trace_for[arn] for s in own) + steps = [ + s for s in own if s.attributes.get("durable.operation.type") == "STEP" + ] + assert len(steps) == 4 + assert len({s.context.span_id for s in steps}) == 4 + terminals = [ + s + for s in steps + if s.attributes.get("durable.operation.status") == "SUCCEEDED" + ] + assert len(terminals) == 2 + assert {s.attributes["durable.operation.name"] for s in terminals} == { + "first", + "second", + } + assert all(s.parent is not None for s in terminals) + own_ids = {s.context.span_id for s in own} + assert all(s.parent.span_id in own_ids for s in steps) + assert not [ + r + for r in caplog.records + if r.levelno >= logging.ERROR and ("plugin" in r.name or "otel" in r.name) + ] + finally: + release.set() + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_terminal_status_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_terminal_status_int.py new file mode 100644 index 000000000..6e0dbb692 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_terminal_status_int.py @@ -0,0 +1,169 @@ +"""A caught durable failure remains an error span without invented diagnostics.""" + +import json +import threading + +import pytest +from aws_durable_execution_sdk_python import durable_execution, durable_step +from aws_durable_execution_sdk_python.exceptions import CallbackError +from aws_durable_execution_sdk_python.lambda_service import ErrorObject +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationStatus, +) +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import StatusCode + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +@pytest.mark.parametrize("view", [ExecutionOtelPlugin, InvocationOtelPlugin]) +@pytest.mark.parametrize("completion", ["initial", "resume"]) +@pytest.mark.parametrize("outcome", ["empty-failure", "rich-failure", "success"]) +def test_caught_callback_status_survives_two_later_replays( + view, completion, outcome, monkeypatch +): + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + created = threading.Event() + completed = threading.Event() + suspended = threading.Event() + finished = threading.Condition() + completed_invocations = [] + invocations = [] + target_ends = [] + side_effects = [] + first_exports = [] + caught = [] + callback_ids = [] + + class Observer(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + invocations.append(info) + + def on_operation_end(self, info): + if info.name == "target": + target_ends.append((len(invocations), info)) + + def on_invocation_end(self, info): + if info.status is InvocationStatus.PENDING: + suspended.set() + with finished: + completed_invocations.append(info) + finished.notify_all() + + @durable_step + def saved(_context, name, value): + side_effects.append(name) + if name == "observed": + first_exports.extend( + span for span in exporter.get_finished_spans() if span.name == "target" + ) + return value + + def body(_event, context): + callback = context.create_callback(name="target") + context.step(saved("before", "saved"), name="before") + if completion == "initial" and len(invocations) == 1: + callback_ids.append(callback.callback_id) + created.set() + assert completed.wait(5), "Driver did not complete the callback" + # A real checkpoint refreshes an external completion during the active + # invocation; on resume this completed step is consumed from history. + context.step(saved("refresh", "saved"), name="refresh") + try: + callback.result() + result = "success" + except CallbackError as error: + result = type(error).__name__ + caught.append(result) + result = context.step(saved("observed", result), name="observed") + for name in ("barrier-1", "barrier-2"): + context.create_callback(name=name).result() + return result + + handler = durable_execution(body, plugins=[plugin, Observer()]) + error = ( + ErrorObject( + message="supplied detail", type="External", data=None, stack_trace=None + ) + if outcome == "rich-failure" + else None + ) + try: + with DurableFunctionTestRunner( + handler=handler, skip_time=False, poll_interval=0.01, execution_timeout=25 + ) as runner: + arn = runner.run_async(input="{}") + if completion == "initial": + assert created.wait(5) + callback_id = callback_ids[0] + else: + callback_id = runner.wait_for_callback(arn, name="target", timeout=5) + assert suspended.wait(5) + if outcome == "success": + runner.send_callback_success(callback_id) + else: + runner.send_callback_failure(callback_id, error=error) + completed.set() + first_completion = 1 if completion == "initial" else 2 + for index, name in enumerate(("barrier-1", "barrier-2")): + callback_id = runner.wait_for_callback(arn, name=name, timeout=5) + with finished: + assert finished.wait_for( + lambda: len(completed_invocations) >= first_completion + index, + timeout=5, + ) + runner.send_callback_success(callback_id) + result = runner.wait_for_result(arn, timeout=10) + assert result.status.value == "SUCCEEDED" + expected_result = "success" if outcome == "success" else "CallbackExternalError" + assert json.loads(result.result) == expected_result + assert side_effects == ["before", "refresh", "observed"] + assert len(invocations) == (3 if completion == "initial" else 4) + assert caught == ([] if outcome == "success" else [expected_result] * 3) + assert len(target_ends) == 1 + invocation_number, info = target_ends[0] + assert invocation_number == (1 if completion == "initial" else 2) + assert info.status.value == ("SUCCEEDED" if outcome == "success" else "FAILED") + assert info.error == error + spans = [ + span for span in exporter.get_finished_spans() if span.name == "target" + ] + expected_statuses = [info.status.value] + if view is InvocationOtelPlugin and completion == "resume": + expected_statuses.insert(0, "STARTED") + assert [ + span.attributes["durable.operation.status"] for span in spans + ] == expected_statuses + assert first_exports == spans + terminal = spans[-1] + assert terminal.status.status_code is ( + StatusCode.OK if outcome == "success" else StatusCode.ERROR + ) + assert len(terminal.events) == (1 if error is not None else 0) + if error is not None: + assert terminal.events[0].attributes["exception.message"] == error.message + workflow = next( + s for s in exporter.get_finished_spans() if s.name == "Workflow" + ) + assert workflow.status.status_code is StatusCode.OK + finally: + completed.set() + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py index a3d79c8aa..d24232741 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py @@ -4,6 +4,7 @@ import threading from concurrent.futures import ThreadPoolExecutor +from contextvars import copy_context from datetime import UTC, datetime, timedelta from types import SimpleNamespace @@ -1300,7 +1301,7 @@ def test_invocation_end_restores_context_from_before_invocation_start(status): ) def test_step_scope_is_released_at_user_function_end(outcome): """Verify a finished step restores the context that enclosed it.""" - plugin, _ = _create_plugin() + plugin, exporter = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) enclosing_context = otel_context.get_current() @@ -1320,6 +1321,12 @@ def test_step_scope_is_released_at_user_function_end(outcome): plugin.on_invocation_end(_invocation_end_info()) assert plugin._context_tokens == {} + events = [ + (span, event) for span in exporter.get_finished_spans() for event in span.events + ] + assert len(events) == (1 if outcome is UserFunctionOutcome.FAILED else 0) + for span, event in events: + assert span.start_time <= event.timestamp <= span.end_time def test_user_function_start_preserves_baggage_in_current_context(): @@ -1446,7 +1453,9 @@ def test_detach_ignores_token_attached_on_another_thread(): with ThreadPoolExecutor(max_workers=1) as executor: executor.submit( - plugin.on_user_function_start, _step_start_info("step-1") + copy_context().run, + plugin.on_user_function_start, + _step_start_info("step-1"), ).result() plugin._detach_context("step-1:attempt:1") @@ -1539,20 +1548,23 @@ def test_suspension_releases_the_scope_on_the_originating_worker(): with ThreadPoolExecutor(max_workers=1) as worker: def suspend_on_worker() -> tuple[int, bool]: + worker_context = otel_context.get_current() plugin.on_user_function_start(_step_start_info("step-1")) attached_span_id = trace.get_current_span().get_span_context().span_id plugin.on_user_function_end(_step_incomplete_info("step-1")) return ( attached_span_id, - trace.get_current_span().get_span_context().is_valid, + otel_context.get_current() == worker_context, ) - attached_span_id, span_still_current = worker.submit(suspend_on_worker).result() + attached_span_id, context_restored = worker.submit( + copy_context().run, suspend_on_worker + ).result() suspended_span = plugin._get_span(span_key) # The scope was released on the worker, and its span is left open. assert attached_span_id != 0 - assert span_still_current is False + assert context_restored is True assert span_key not in plugin._context_tokens assert suspended_span is not None assert {s.name for s in exporter.get_finished_spans()} == { diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_flush.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_flush.py new file mode 100644 index 000000000..ec67d7725 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_flush.py @@ -0,0 +1,231 @@ +"""Bound provider I/O without spawning replacement workers for a stuck call.""" + +import contextvars +import threading +import time +from concurrent.futures import ThreadPoolExecutor + +import pytest +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import BatchSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import _flush + + +def _wait_idle() -> None: + deadline = time.monotonic() + 2 + while time.monotonic() < deadline: + with _flush._lock: + if _flush._active is None: + return + time.sleep(0.001) + raise AssertionError("Provider worker did not actually complete") + + +@pytest.fixture(autouse=True) +def idle_flush(monkeypatch): + _wait_idle() + monkeypatch.setattr(_flush, "_FLUSH_TIMEOUT_SECONDS", 1.0) + yield + _wait_idle() + + +def test_timeout_retains_one_worker_and_rejects_new_work(monkeypatch): + monkeypatch.setattr(_flush, "_FLUSH_TIMEOUT_SECONDS", 0.02) + started = threading.Event() + release = threading.Event() + calls = [] + + def blocked(): + calls.append(threading.get_ident()) + started.set() + assert release.wait(2) + return True + + try: + began = time.monotonic() + assert _flush.flush_provider(blocked) is False + assert time.monotonic() - began < 1 + assert started.wait(1) + for _ in range(5): + assert _flush.flush_provider(lambda: calls.append("unexpected")) is False + assert len(calls) == 1 + workers = [t for t in threading.enumerate() if t.name == "durable-otel-flush"] + assert len(workers) == 1 + assert workers[0].daemon + finally: + release.set() + _wait_idle() + monkeypatch.setattr(_flush, "_FLUSH_TIMEOUT_SECONDS", 1.0) + assert _flush.flush_provider(lambda: True) is True + + +def test_healthy_distinct_providers_wait_for_admission(): + exporting = threading.Event() + release = threading.Event() + second_called = threading.Event() + + class GatedExporter(InMemorySpanExporter): + def export(self, spans): + exporting.set() + assert release.wait(3) + return super().export(spans) + + first_exporter = GatedExporter() + second_exporter = InMemorySpanExporter() + first_provider = TracerProvider() + second_provider = TracerProvider() + first_provider.add_span_processor( + BatchSpanProcessor(first_exporter, schedule_delay_millis=60000) + ) + second_provider.add_span_processor( + BatchSpanProcessor(second_exporter, schedule_delay_millis=60000) + ) + with first_provider.get_tracer("first").start_as_current_span("first"): + pass + with second_provider.get_tracer("second").start_as_current_span("second"): + pass + + def second_flush(): + second_called.set() + return _flush.flush_provider(second_provider.force_flush) + + try: + with ThreadPoolExecutor(max_workers=2) as callers: + first = callers.submit(_flush.flush_provider, first_provider.force_flush) + try: + assert exporting.wait(1) + second = callers.submit(second_flush) + assert second_called.wait(1) + time.sleep(0.03) + assert not second.done(), ( + "A healthy caller was skipped instead of waiting" + ) + finally: + release.set() + assert first.result(timeout=2) is True + assert second.result(timeout=2) is True + assert [span.name for span in first_exporter.get_finished_spans()] == ["first"] + assert [span.name for span in second_exporter.get_finished_spans()] == [ + "second" + ] + finally: + release.set() + first_provider.shutdown() + second_provider.shutdown() + + +def test_admission_and_execution_share_the_same_deadline(): + first_entered = threading.Event() + first_release = threading.Event() + second_called = threading.Event() + second_entered = threading.Event() + second_release = threading.Event() + + def first_io(): + first_entered.set() + assert first_release.wait(3) + return True + + def second_io(): + second_entered.set() + assert second_release.wait(3) + return True + + def second_call(): + second_called.set() + return _flush.flush_provider(second_io) + + with ThreadPoolExecutor(max_workers=2) as callers: + first = callers.submit(_flush.flush_provider, first_io) + try: + assert first_entered.wait(1) + second = callers.submit(second_call) + assert second_called.wait(1) + # Spend most of the existing one-second budget on admission. + time.sleep(0.6) + first_release.set() + assert first.result(timeout=1) is True + assert second_entered.wait(1) + # A reset one-second I/O budget would still be blocked here. + assert second.result(timeout=0.65) is False + finally: + first_release.set() + second_release.set() + _wait_idle() + + +@pytest.mark.parametrize("result", [True, False, None]) +def test_completed_provider_result(result): + assert _flush.flush_provider(lambda: result) is (result is not False) + + +@pytest.mark.parametrize( + "error", [RuntimeError("ordinary"), SystemExit("fatal"), KeyboardInterrupt()] +) +def test_fast_error_preserves_identity_on_the_caller(error): + owner = threading.get_ident() + worker = [] + + def fail(): + worker.append(threading.get_ident()) + raise error + + with pytest.raises(type(error)) as caught: + _flush.flush_provider(fail) + assert caught.value is error + assert worker[0] != owner + assert threading.get_ident() == owner + + +@pytest.mark.parametrize( + "error", [RuntimeError("late ordinary"), SystemExit("late fatal")] +) +def test_late_error_is_reported_without_poisoning_next_call(error, monkeypatch, caplog): + monkeypatch.setattr(_flush, "_FLUSH_TIMEOUT_SECONDS", 0.02) + release = threading.Event() + + def fail_late(): + assert release.wait(2) + raise error + + try: + assert _flush.flush_provider(fail_late) is False + finally: + release.set() + _wait_idle() + messages = [ + r for r in caplog.records if "after the caller stopped waiting" in r.message + ] + assert len(messages) == 1 + assert messages[0].exc_info[1] is error + monkeypatch.setattr(_flush, "_FLUSH_TIMEOUT_SECONDS", 1.0) + assert _flush.flush_provider(lambda: True) is True + + +def test_start_failure_releases_capacity(monkeypatch): + def fail_start(_thread): + raise RuntimeError("start failed") + + with monkeypatch.context() as patch: + patch.setattr(threading.Thread, "start", fail_start) + with pytest.raises(RuntimeError, match="start failed"): + _flush.flush_provider(lambda: True) + assert _flush.flush_provider(lambda: True) is True + + +def test_only_provider_io_gets_post_cleanup_context_snapshot(): + value = contextvars.ContextVar("flush-test", default="unset") + token = value.set("owner-after-cleanup") + try: + + def flush(): + assert value.get() == "owner-after-cleanup" + value.set("provider-only") + return True + + assert _flush.flush_provider(flush) + assert value.get() == "owner-after-cleanup" + finally: + value.reset(token) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_owned_state.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_owned_state.py new file mode 100644 index 000000000..e56c52fea --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_owned_state.py @@ -0,0 +1,118 @@ +"""Invocation-owned state must not be reset through an inherited Context.""" + +import contextvars +from datetime import UTC, datetime + +import pytest +from aws_durable_execution_sdk_python.plugin import ( + InvocationEndInfo, + InvocationStartInfo, + InvocationStatus, +) +from opentelemetry import context as otel_context +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +@pytest.mark.parametrize("view", [ExecutionOtelPlugin, InvocationOtelPlugin]) +def test_copied_context_on_same_thread_does_not_clear_its_parent_invocation(view): + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + a = contextvars.Context() + info_a = dict( + request_id="a", + execution_arn="execution/a", + execution_start_time=datetime.now(UTC), + is_first_invocation=True, + ) + info_b = dict( + request_id="b", + execution_arn="execution/b", + execution_start_time=datetime.now(UTC), + is_first_invocation=True, + ) + try: + original = a.run(otel_context.get_current) + a.run(plugin.on_invocation_start, InvocationStartInfo(**info_a)) + state_a = a.run(plugin._state.get) + span_a = a.run(plugin.get_current_span_context) + b = a.copy() + b.run(plugin.on_invocation_start, InvocationStartInfo(**info_b)) + state_b = b.run(plugin._state.get) + assert state_a is not state_b + b.run( + plugin.on_invocation_end, + InvocationEndInfo(**info_b, status=InvocationStatus.SUCCEEDED), + ) + assert a.run(plugin._state.get) is state_a + assert a.run(plugin.get_current_span_context) == span_a + assert state_a.execution_arn == "execution/a" and state_a.tracing_enabled + a.run( + plugin.on_invocation_end, + InvocationEndInfo(**info_a, status=InvocationStatus.SUCCEEDED), + ) + assert a.run(otel_context.get_current) == original + assert plugin._state not in a + workflows = [s for s in exporter.get_finished_spans() if s.name == "Workflow"] + assert len(workflows) == 2 + assert {s.attributes["durable.execution.arn"] for s in workflows} == { + "execution/a", + "execution/b", + } + assert len({s.context.trace_id for s in workflows}) == 2 + finally: + provider.shutdown() + + +@pytest.mark.parametrize("view", [ExecutionOtelPlugin, InvocationOtelPlugin]) +def test_same_owner_warm_restart_releases_abandoned_scope(view): + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter())) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + before = otel_context.get_current() + first = dict( + request_id="first", + execution_arn="execution/first", + execution_start_time=datetime.now(UTC), + is_first_invocation=True, + ) + second = dict( + request_id="second", + execution_arn="execution/second", + execution_start_time=datetime.now(UTC), + is_first_invocation=True, + ) + try: + plugin.on_invocation_start(InvocationStartInfo(**first)) + abandoned = plugin._state.get() + plugin.on_invocation_start(InvocationStartInfo(**second)) + assert abandoned.context_tokens == {} + assert abandoned.owner_token is None + plugin.on_invocation_end( + InvocationEndInfo(**second, status=InvocationStatus.SUCCEEDED) + ) + assert otel_context.get_current() == before + assert plugin._state not in contextvars.copy_context() + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py index b22d438c3..dcab099b6 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py @@ -6,6 +6,7 @@ import threading import time from concurrent.futures import ThreadPoolExecutor +from contextvars import copy_context from datetime import UTC, datetime from types import SimpleNamespace @@ -1304,7 +1305,7 @@ def test_user_function_start_preserves_baggage_in_current_context(): def test_user_function_end_restores_enclosing_context_on_failure(): """Verify the enclosing context is restored even when the step fails.""" - plugin, _ = _create_plugin() + plugin, exporter = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) enclosing_context = otel_context.get_current() @@ -1317,6 +1318,12 @@ def test_user_function_end_restores_enclosing_context_on_failure(): assert otel_context.get_current() == enclosing_context assert set(plugin._context_tokens) == {"__invocation_context__"} plugin.on_invocation_end(_invocation_end_info()) + events = [ + (span, event) for span in exporter.get_finished_spans() for event in span.events + ] + assert len(events) == 1 + for span, event in events: + assert span.start_time <= event.timestamp <= span.end_time def test_user_function_end_restores_enclosing_context_across_multiple_steps(): @@ -2009,7 +2016,9 @@ def test_detach_ignores_token_attached_on_another_thread(): with ThreadPoolExecutor(max_workers=1) as executor: executor.submit( - plugin.on_user_function_start, _user_function_start_info("step-1") + copy_context().run, + plugin.on_user_function_start, + _user_function_start_info("step-1"), ).result() plugin._detach_context("step-1:attempt:1") @@ -2127,20 +2136,23 @@ def test_suspension_releases_the_scope_on_the_originating_worker(): with ThreadPoolExecutor(max_workers=1) as worker: def suspend_on_worker() -> tuple[int, bool]: + worker_context = otel_context.get_current() plugin.on_user_function_start(_user_function_start_info(operation_id)) attached_span_id = trace.get_current_span().get_span_context().span_id plugin.on_user_function_end(_user_function_incomplete_info(operation_id)) return ( attached_span_id, - trace.get_current_span().get_span_context().is_valid, + otel_context.get_current() == worker_context, ) - attached_span_id, span_still_current = worker.submit(suspend_on_worker).result() + attached_span_id, context_restored = worker.submit( + copy_context().run, suspend_on_worker + ).result() suspended_span = plugin._get_span(span_key) # The scope was released on the worker, and its span is left open. assert attached_span_id != 0 - assert span_still_current is False + assert context_restored is True assert span_key not in plugin._context_tokens assert suspended_span is not None assert {s.name for s in exporter.get_finished_spans()} == { @@ -2300,3 +2312,142 @@ def body() -> None: finally: otel_context.detach(token) provider.shutdown() + + +@pytest.mark.parametrize( + ("status", "error", "expected"), + [ + (OperationStatus.SUCCEEDED, None, StatusCode.OK), + (OperationStatus.FAILED, None, StatusCode.ERROR), + (OperationStatus.CANCELLED, None, StatusCode.UNSET), + (OperationStatus.TIMED_OUT, None, StatusCode.ERROR), + (OperationStatus.STOPPED, None, StatusCode.ERROR), + ( + OperationStatus.FAILED, + ErrorObject(message="failure", type="Example", data=None, stack_trace=None), + StatusCode.ERROR, + ), + ], +) +def test_deferred_end_retains_real_parent_and_status_mapping( + status: OperationStatus, error: ErrorObject | None, expected: StatusCode +) -> None: + plugin, exporter = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + end = OperationEndInfo( + operation_id="child", + operation_type=OperationType.WAIT, + sub_type=None, + name="deferred-child", + parent_id="parent", + start_time=START_TIME, + is_replayed=False, + status=status, + end_time=END_TIME, + error=error, + ) + plugin.on_operation_end(end) + (root,) = exporter.get_finished_spans() + assert root.name == "DurableExecutionRoot" + assert root.start_time == root.end_time + assert plugin._pending_operation_ends == {"parent": [end]} + plugin.on_user_function_start( + _user_function_start_info("parent", operation_type=OperationType.CONTEXT) + ) + parent = plugin._get_span("parent") + assert parent is not None + same_root, child = exporter.get_finished_spans() + assert same_root is root + assert child.parent is not None and child.attributes is not None + assert child.parent.span_id == parent.get_span_context().span_id + assert child.status.status_code is expected + assert child.attributes["durable.operation.status"] == status.value + assert len(child.events) == (1 if error is not None else 0) + assert plugin._pending_operation_ends == {} + plugin.on_user_function_end( + _user_function_end_info("parent", operation_type=OperationType.CONTEXT) + ) + plugin.on_invocation_end(_invocation_end_info()) + + +def test_deferred_descendants_drain_from_real_parent_completion_in_child_first_order() -> ( + None +): + plugin, exporter = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + for operation_id, parent_id, kind in [ + ("child", "inner", OperationType.WAIT), + ("inner", "outer", OperationType.CONTEXT), + ]: + plugin.on_operation_end( + OperationEndInfo( + operation_id=operation_id, + operation_type=kind, + sub_type=None, + name=operation_id, + parent_id=parent_id, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.SUCCEEDED, + end_time=END_TIME, + ) + ) + (root,) = exporter.get_finished_spans() + assert root.name == "DurableExecutionRoot" + assert root.start_time == root.end_time + plugin.on_user_function_start( + _user_function_start_info("outer", operation_type=OperationType.CONTEXT) + ) + same_root, child, inner = exporter.get_finished_spans() + assert same_root is root + assert [child.name, inner.name] == ["child", "inner"] + assert child.parent is not None and inner.parent is not None + assert child.parent.span_id == inner.context.span_id + outer = plugin._get_span("outer") + assert outer is not None + assert inner.parent.span_id == outer.get_span_context().span_id + assert child.end_time is not None and inner.end_time is not None + assert child.end_time <= inner.end_time + assert plugin._pending_operation_ends == {} + plugin.on_user_function_end( + _user_function_end_info("outer", operation_type=OperationType.CONTEXT) + ) + plugin.on_invocation_end(_invocation_end_info()) + + +def test_unresolved_parent_is_reported_after_cleanup_and_not_carried_into_reuse() -> ( + None +): + plugin, exporter = _create_plugin() + before = otel_context.get_current() + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_operation_end( + OperationEndInfo( + operation_id="missing-child", + operation_type=OperationType.WAIT, + sub_type=None, + name="missing-child", + parent_id="missing-parent", + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.SUCCEEDED, + end_time=END_TIME, + ) + ) + with pytest.raises(ValueError, match="No parent span found for deferred"): + plugin.on_invocation_end(_invocation_end_info()) + assert otel_context.get_current() == before + assert plugin._operation_spans == {} + assert plugin._pending_operation_ends == {} + assert plugin._context_tokens == {} + assert {span.name for span in exporter.get_finished_spans()} == { + "DurableExecutionRoot", + "Invocation", + "Workflow", + } + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_invocation_end(_invocation_end_info(InvocationStatus.PENDING)) + assert not [ + span for span in exporter.get_finished_spans() if span.name == "missing-child" + ] + assert plugin._pending_operation_ends == {} diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_status_mapping.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_status_mapping.py index d1224a305..a97df2e1c 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_status_mapping.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_status_mapping.py @@ -79,7 +79,7 @@ def test_operation_status_and_exception( with_error: bool, started_here: bool, ) -> None: - """A delivered terminal hook must not label an unsuccessful operation OK.""" + """Explicit failures are errors even when diagnostic details are absent.""" plugin, exporter = instrumented_view invocation: dict[str, Any] = dict( request_id="first", @@ -123,6 +123,8 @@ def test_operation_status_and_exception( expected = ( StatusCode.ERROR if with_error + or status + in (OperationStatus.FAILED, OperationStatus.TIMED_OUT, OperationStatus.STOPPED) else ( StatusCode.OK if status is OperationStatus.SUCCEEDED else StatusCode.UNSET ) @@ -134,10 +136,50 @@ def test_operation_status_and_exception( assert span.events[0].name == "exception" assert span.events[0].attributes is not None assert span.events[0].attributes["exception.message"] == "deadline exceeded" + assert span.start_time is not None and span.end_time is not None + assert span.start_time <= span.events[0].timestamp <= span.end_time else: assert not span.events +def test_error_event_is_inside_completion_without_a_start( + instrumented_view: InstrumentedView, +) -> None: + plugin, exporter = instrumented_view + invocation: dict[str, Any] = dict( + request_id="missing-start", + execution_arn=ARN, + execution_start_time=START, + is_first_invocation=True, + ) + plugin.on_invocation_start(InvocationStartInfo(**invocation)) + plugin.on_operation_end( + OperationEndInfo( + operation_id="missing-start", + name="missing-start", + parent_id=None, + operation_type=OperationType.WAIT, + sub_type=OperationSubType.WAIT, + start_time=None, + end_time=END, + is_replayed=False, + status=OperationStatus.FAILED, + error=ErrorObject( + message="real failure", type="Failure", data=None, stack_trace=None + ), + ) + ) + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.SUCCEEDED) + ) + span = next(s for s in exporter.get_finished_spans() if s.name == "missing-start") + assert len(span.events) == 1 + assert span.events[0].attributes is not None + assert span.events[0].attributes["exception.message"] == "real failure" + assert span.start_time is not None and span.end_time is not None + assert span.start_time <= span.events[0].timestamp <= span.end_time + + def test_retry_then_resume_exports_retrying( instrumented_view: InstrumentedView, ) -> None: diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_terminal_input_parent_refs.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_terminal_input_parent_refs.py new file mode 100644 index 000000000..b69284966 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_terminal_input_parent_refs.py @@ -0,0 +1,450 @@ +"""Missing parent metadata is deferred rather than guessed or reparented.""" + +import contextvars +import threading +from concurrent.futures import ThreadPoolExecutor +from datetime import UTC, datetime, timedelta + +import pytest +from aws_durable_execution_sdk_python.plugin import ( + InvocationEndInfo, + InvocationStartInfo, + InvocationStatus, + OperationEndInfo, + OperationInfo, + OperationStartInfo, + OperationStatus, + OperationType, + UserFunctionEndInfo, + UserFunctionStartInfo, +) +from opentelemetry import context as otel_context +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +def _info(cls, now, **changes): + fields = dict( + operation_id="b", + operation_type=OperationType.CALLBACK, + sub_type=None, + name="b", + parent_id="parent", + start_time=now, + is_replayed=False, + status=OperationStatus.SUCCEEDED, + end_time=now + timedelta(seconds=1), + ) + fields.update(changes) + return cls(**fields) + + +@pytest.mark.parametrize("view", [ExecutionOtelPlugin, InvocationOtelPlugin]) +@pytest.mark.parametrize("parent_visible", [False, True]) +def test_nonterminal_or_missing_parent_waits_for_its_actual_span(view, parent_visible): + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + now = datetime.now(UTC) + invocation = dict( + request_id="request", + execution_arn="execution/parent", + is_first_invocation=False, + execution_start_time=now, + ) + b = _info(OperationInfo, now) + parent = _info( + OperationInfo, + now, + operation_id="parent", + operation_type=OperationType.CONTEXT, + name="parent", + parent_id=None, + status=OperationStatus.STARTED, + end_time=None, + ) + operations = {"b": b} + if parent_visible: + operations["parent"] = parent + + def invoke(): + before = otel_context.get_current() + plugin.on_invocation_start( + InvocationStartInfo( + **invocation, operations=operations, updated_operations={"b": b} + ) + ) + plugin.on_operation_end(_info(OperationEndInfo, now)) + assert not [ + s + for s in exporter.get_finished_spans() + if s.attributes.get("durable.operation.id") == "b" + ] + parent_fields = dict( + operation_id="parent", + operation_type=OperationType.CONTEXT, + name="parent", + parent_id=None, + status=OperationStatus.STARTED, + is_replayed=True, + end_time=None, + ) + plugin.on_operation_start(_info(OperationStartInfo, now, **parent_fields)) + start = _info(UserFunctionStartInfo, now, **parent_fields) + plugin.on_user_function_start(start) + actual_parent = plugin._get_span("parent").get_span_context() + plugin.on_user_function_end(UserFunctionEndInfo.from_start_info(start, None)) + plugin.on_operation_end( + _info( + OperationEndInfo, + now, + operation_id="parent", + operation_type=OperationType.CONTEXT, + name="parent", + parent_id=None, + ) + ) + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.SUCCEEDED) + ) + spans = [ + s + for s in exporter.get_finished_spans() + if s.attributes.get("durable.operation.id") == "b" + ] + assert len(spans) == 1 and spans[0].parent.span_id == actual_parent.span_id + assert otel_context.get_current() == before + assert plugin._state not in contextvars.copy_context() + + try: + contextvars.Context().run(invoke) + finally: + provider.shutdown() + + +@pytest.mark.parametrize("view", [ExecutionOtelPlugin, InvocationOtelPlugin]) +def test_unavailable_parent_is_reported_after_context_cleanup(view): + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + now = datetime.now(UTC) + invocation = dict( + request_id="request", + execution_arn="execution/unknown-parent", + is_first_invocation=False, + execution_start_time=now, + ) + b = _info(OperationInfo, now) + + def invoke(): + before = otel_context.get_current() + plugin.on_invocation_start( + InvocationStartInfo( + **invocation, operations={"b": b}, updated_operations={"b": b} + ) + ) + plugin.on_operation_end(_info(OperationEndInfo, now)) + with pytest.raises(ValueError, match="No parent span found for deferred"): + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.PENDING) + ) + assert not [ + s + for s in exporter.get_finished_spans() + if s.attributes.get("durable.operation.id") == "b" + ] + assert otel_context.get_current() == before + assert plugin._state not in contextvars.copy_context() + + try: + contextvars.Context().run(invoke) + finally: + provider.shutdown() + + +def test_completed_context_remains_parent_for_late_execution_checkpoint(): + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = ExecutionOtelPlugin( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + now = datetime.now(UTC) + invocation = dict( + request_id="request", + execution_arn="execution/late-parent", + is_first_invocation=True, + execution_start_time=now, + ) + parent_fields = dict( + operation_id="parent", + operation_type=OperationType.CONTEXT, + name="parent", + parent_id=None, + status=OperationStatus.STARTED, + end_time=None, + ) + before = otel_context.get_current() + try: + plugin.on_invocation_start(InvocationStartInfo(**invocation)) + plugin.on_operation_start(_info(OperationStartInfo, now, **parent_fields)) + start = _info(UserFunctionStartInfo, now, **parent_fields) + plugin.on_user_function_start(start) + plugin.on_operation_start( + _info( + OperationStartInfo, now, status=OperationStatus.STARTED, end_time=None + ) + ) + plugin.on_user_function_end(UserFunctionEndInfo.from_start_info(start, None)) + plugin.on_operation_end( + _info( + OperationEndInfo, + now, + operation_id="parent", + operation_type=OperationType.CONTEXT, + name="parent", + parent_id=None, + ) + ) + parents = [s for s in exporter.get_finished_spans() if s.name == "parent"] + assert len(parents) == 1 + plugin.on_operation_end(_info(OperationEndInfo, now)) + children = [s for s in exporter.get_finished_spans() if s.name == "b"] + assert ( + len(children) == 1 + and children[0].parent.span_id == parents[0].context.span_id + ) + assert ( + len([s for s in exporter.get_finished_spans() if s.name == "parent"]) == 1 + ) + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.SUCCEEDED) + ) + assert plugin._operation_spans == {} and otel_context.get_current() == before + finally: + provider.shutdown() + + +def test_completed_input_parent_waits_for_current_invocation_segment(): + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = InvocationOtelPlugin( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + now = datetime.now(UTC) + invocation = dict( + request_id="request", + execution_arn="execution/replayed-parent", + is_first_invocation=False, + execution_start_time=now, + ) + parent_fields = dict( + operation_id="parent", + operation_type=OperationType.CONTEXT, + name="parent", + parent_id=None, + is_replayed=True, + ) + b = _info(OperationInfo, now) + parent = _info(OperationInfo, now, **parent_fields) + before = otel_context.get_current() + try: + plugin.on_invocation_start( + InvocationStartInfo( + **invocation, + operations={"parent": parent, "b": b}, + updated_operations={"b": b}, + ) + ) + plugin.on_operation_end(_info(OperationEndInfo, now)) + assert not [s for s in exporter.get_finished_spans() if s.name == "b"] + start = _info( + UserFunctionStartInfo, now, **parent_fields, is_replay_children=True + ) + plugin.on_user_function_start(start) + actual_parent = plugin._get_span("parent").get_span_context() + children = [s for s in exporter.get_finished_spans() if s.name == "b"] + assert ( + len(children) == 1 and children[0].parent.span_id == actual_parent.span_id + ) + plugin.on_user_function_end(UserFunctionEndInfo.from_start_info(start, None)) + plugin.on_operation_end(_info(OperationEndInfo, now, **parent_fields)) + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.PENDING) + ) + assert plugin._operation_spans == {} and otel_context.get_current() == before + assert plugin._state.get().input_parent_refs == {} + finally: + provider.shutdown() + + +def test_completed_parent_refs_cover_later_updates_without_guessing(): + from aws_durable_execution_sdk_python_otel._invocation_state import ( + completed_input_parents, + ) + + now = datetime.now(UTC) + parent = _info( + OperationInfo, + now, + operation_id="parent", + operation_type=OperationType.CONTEXT, + name="parent", + parent_id=None, + ) + unfinished = _info( + OperationInfo, + now, + operation_id="unfinished", + operation_type=OperationType.CONTEXT, + status=OperationStatus.STARTED, + ) + step = _info( + OperationInfo, now, operation_id="step", operation_type=OperationType.STEP + ) + invocation = dict( + request_id="request", + execution_arn="execution/full-parent", + is_first_invocation=False, + execution_start_time=now, + ) + assert completed_input_parents( + InvocationStartInfo( + **invocation, + operations={"parent": parent, "unfinished": unfinished, "step": step}, + updated_operations={}, + ) + ) == [parent] + missing_parent_child = _info(OperationInfo, now) + assert ( + completed_input_parents( + InvocationStartInfo( + **invocation, + operations={"b": missing_parent_child}, + updated_operations={"b": missing_parent_child}, + ) + ) + == [] + ) + + +def test_parent_registration_survives_concurrent_terminal_materialization(monkeypatch): + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = ExecutionOtelPlugin( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + now = datetime.now(UTC) + invocation = dict( + request_id="atomic-parent", + execution_arn="execution/atomic-parent", + is_first_invocation=True, + execution_start_time=now, + ) + parent_entered, allow_parent = threading.Event(), threading.Event() + original_start = plugin._start_span + + def start_span(**kwargs): + if kwargs["operation_id"] == "parent": + parent_entered.set() + assert allow_parent.wait(5) + return original_start(**kwargs) + + monkeypatch.setattr(plugin, "_start_span", start_span) + + def invoke(): + before = otel_context.get_current() + plugin.on_invocation_start(InvocationStartInfo(**invocation)) + parent_fields = dict( + operation_id="parent", + operation_type=OperationType.CONTEXT, + name="parent", + parent_id=None, + status=OperationStatus.STARTED, + end_time=None, + ) + plugin.on_operation_start(_info(OperationStartInfo, now, **parent_fields)) + start = _info(UserFunctionStartInfo, now, **parent_fields) + plugin.on_user_function_start(start) + plugin.on_user_function_end(UserFunctionEndInfo.from_start_info(start, None)) + parent_fields.update( + status=OperationStatus.SUCCEEDED, end_time=now + timedelta(seconds=2) + ) + parent_context, child_context = ( + contextvars.copy_context(), + contextvars.copy_context(), + ) + with ThreadPoolExecutor(max_workers=2) as pool: + parent = pool.submit( + parent_context.run, + plugin.on_operation_end, + _info(OperationEndInfo, now, **parent_fields), + ) + try: + assert parent_entered.wait(5) + child = pool.submit( + child_context.run, + plugin.on_operation_end, + _info(OperationEndInfo, now), + ) + child.result(timeout=5) + finally: + allow_parent.set() + parent.result(timeout=5) + close_error = None + try: + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.SUCCEEDED) + ) + except ValueError as exc: + close_error = exc + spans = exporter.get_finished_spans() + parents = [ + s for s in spans if s.attributes.get("durable.operation.id") == "parent" + ] + children = [s for s in spans if s.attributes.get("durable.operation.id") == "b"] + assert len(parents) == len(children) == 1 + assert children[0].parent.span_id == parents[0].context.span_id + assert children[0].context.trace_id == parents[0].context.trace_id + assert close_error is None + assert plugin._operation_spans == {} and otel_context.get_current() == before + + try: + contextvars.Context().run(invoke) + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-testing/README.md b/packages/aws-durable-execution-sdk-python-testing/README.md index 58b21e793..fc31dbefc 100644 --- a/packages/aws-durable-execution-sdk-python-testing/README.md +++ b/packages/aws-durable-execution-sdk-python-testing/README.md @@ -170,6 +170,31 @@ flow, pagination, and a map of the code. ## Documentation +### Operation state on replay + +Invocation input and `GetDurableExecutionState` follow the service's retained +operation view. An operation is omitted when any ancestor is a `CONTEXT` that +has `SUCCEEDED` or `FAILED` without `ReplayChildren: true`. The walk includes +running intermediate contexts; a nested context's own retention flag cannot +override a completed ancestor that prunes it. Roots and operations whose parent +is absent are kept, and ordering does not change the rule. + +`ContextOptions.ReplayChildren` on START is ignored. A closing update records +true only when it explicitly supplies true; otherwise it records false. Testing +service responses include that closing false, while Core's public model and +generic serialization contract remain unchanged. + +Checkpoint responses still contain all updated operations. `UpdatedOperationIds` +can include an ID whose operation record is hidden from invocation input and +GetState. Raw stores, history, results and nested operation inspection retain the +complete execution. An updated ID alone does not restore missing metadata for a +plugin. A real large child result retains its children through ReplayChildren +when a test needs that public replay behavior. + +This corrects local service parity; it does not fix duplicate execution-view +spans for virtual FLAT map iterations while their map is still running. + + ### Error Handling The testing framework implements AWS-compliant error responses that match the exact format expected by boto3 and AWS services. For detailed information about error response formats, exception types, and troubleshooting, see: diff --git a/packages/aws-durable-execution-sdk-python-testing/docs/architecture.md b/packages/aws-durable-execution-sdk-python-testing/docs/architecture.md index c46f076a1..e82f81337 100644 --- a/packages/aws-durable-execution-sdk-python-testing/docs/architecture.md +++ b/packages/aws-durable-execution-sdk-python-testing/docs/architecture.md @@ -234,10 +234,18 @@ list in a checkpoint response. The delta for the next response is the operations whose `operation_last_touched_seq` is greater than `handler_seen_seq`. -The watermark advances only when a checkpoint response is built. It never -advances when invocation input is built. That is what makes retries -correct. A failed invocation never moves the watermark. So the retry -delivers the same state again. +The watermark advances when a checkpoint response is built or an invocation +response passes validation, including a `PENDING` response. Building invocation +input does not acknowledge it: input updates remain pending until delivery is +confirmed. A failed call or rejected response does not consume that input, so a +retry can receive it again. + +The invocation acknowledgment uses the sequence captured when its input was +built, on the execution's worker. It never moves a higher checkpoint watermark +backward or consumes operations that changed after the captured snapshot. This +also prevents an input-only callback resume from receiving the same terminal +update again at a later checkpoint. The acknowledgment concerns runner state +delivery, not telemetry export. ``` seq_counter handler_seen_seq delta next checkpoint diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/processor.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/processor.py index 42ea9333b..f1b1d6da6 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/processor.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/processor.py @@ -21,6 +21,9 @@ CheckpointRequestDispatcher, ) from aws_durable_execution_sdk_python_testing.clock import RealClock +from aws_durable_execution_sdk_python_testing.operation_state import ( + prune_replay_children, +) from aws_durable_execution_sdk_python_testing.exceptions import ( InvalidParameterValueException, ) @@ -153,7 +156,7 @@ def get_execution_state( token = CheckpointToken.from_str(checkpoint_token) execution = self._store.load(token.execution_arn) return StateOutput( - operations=execution.get_navigable_operations(), + operations=prune_replay_children(execution.get_navigable_operations()), next_marker=None, ) diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/processors/base.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/processors/base.py index d2eaa4ec4..9ea3cb256 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/processors/base.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/processors/base.py @@ -100,9 +100,11 @@ def _create_context_details(self, update: OperationUpdate) -> ContextDetails | N ContextDetails( result=update.payload, error=update.error, - replay_children=update.context_options.replay_children - if update.context_options - else False, + replay_children=( + update.action in (OperationAction.SUCCEED, OperationAction.FAIL) + and update.context_options is not None + and update.context_options.replay_children is True + ), ) if update.operation_type == OperationType.CONTEXT else None diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py index 82880a1ba..833081dcd 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py @@ -24,6 +24,9 @@ ) from aws_durable_execution_sdk_python_testing.clock import real_now +from aws_durable_execution_sdk_python_testing.operation_state import ( + prune_replay_children, +) from aws_durable_execution_sdk_python_testing.exceptions import ( IllegalStateException, InvalidParameterValueException, @@ -434,15 +437,11 @@ def record_invocation_completion( ) def mark_state_delivered(self) -> None: - """Reset the list of changed operations once an invocation's input - carries them. - - The service reports, on each invocation, the operations that - changed after the last state the handler observed. The runner - delivers state in the invocation input. So the list is reset when - that input is built, not when the invocation completes: an - operation that completes while the handler is still running is - reported on the next invocation. + """Explicitly clear all pending operation IDs after their delivery. + + The live executor uses versioned ``advance_handler_seen`` acknowledgments + so failed invocations retain their input updates and concurrent changes + newer than the acknowledged input snapshot remain pending. """ self.updated_operation_ids = [] @@ -647,6 +646,7 @@ def complete_callback_success( end_timestamp=now if now is not None else real_now(), callback_details=updated_callback_details, ) + self._record_updated_operation(operation.operation_id) return self.operations[index] def complete_callback_failure( @@ -663,8 +663,11 @@ def complete_callback_failure( self.touch_operation(operation.operation_id) updated_callback_details = None if operation.callback_details: + # Match CallbackDetails.from_dict without depending on a store + # serialization round trip: an empty wire Error has no details. updated_callback_details = replace( - operation.callback_details, error=error + operation.callback_details, + error=error if error is not None and error.to_dict() else None, ) self.operations[index] = replace( @@ -673,6 +676,7 @@ def complete_callback_failure( end_timestamp=now if now is not None else real_now(), callback_details=updated_callback_details, ) + self._record_updated_operation(operation.operation_id) return self.operations[index] def complete_callback_timeout( @@ -699,6 +703,7 @@ def complete_callback_timeout( end_timestamp=now if now is not None else real_now(), callback_details=updated_callback_details, ) + self._record_updated_operation(operation.operation_id) return self.operations[index] def complete_chained_invoke( @@ -852,10 +857,9 @@ def page( ``max_size_bytes`` and, when given, ``max_items``. Second element of the tuple is a marker for the next page, or ``None`` when the page fits everything.""" - start_idx = self._resolve_marker(marker) - return self._walk_page( - self.snapshot_operations, start_idx, max_size_bytes, max_items - ) + operations = prune_replay_children(self.snapshot_operations) + start_idx = self._resolve_marker(marker, len(operations)) + return self._walk_page(operations, start_idx, max_size_bytes, max_items) def unseen_operations(self) -> list[Operation]: """Operations whose ``operation_last_touched_seq`` is strictly @@ -875,11 +879,20 @@ def unseen_operations(self) -> list[Operation]: ] def advance_handler_seen(self, seq: int) -> None: - """Advance the wrapped execution's ``handler_seen_seq`` to - ``seq`` if that represents forward progress. Monotonic: - smaller or equal values are ignored.""" + """Acknowledge delivered updates without retreating the watermark. + + Equal acknowledgments also consume old-format, unversioned update IDs. + Only versions covered by this delivery may be removed. + """ if seq > self.execution.handler_seen_seq: self.execution.handler_seen_seq = seq + # A checkpoint response or completed invocation delivered this state. + # Retain changes newer than that acknowledged snapshot. + self.execution.updated_operation_ids = [ + operation_id + for operation_id in self.execution.updated_operation_ids + if self.execution.operation_last_touched_seq.get(operation_id, 0) > seq + ] # --- internals ------------------------------------------------- @@ -908,14 +921,21 @@ def _size_for(self, op: Operation) -> int: # floor at 1 to guarantee pagination always advances. return max(self.execution.operation_size_bytes.get(op.operation_id, 0), 1) - def _resolve_marker(self, marker: str | None) -> int: + def _resolve_marker( + self, marker: str | None, operation_count: int | None = None + ) -> int: if marker is None: return 0 seq, idx = self._decode_marker(marker) if seq != self.pinned_token_sequence: msg = "Invalid marker" raise InvalidParameterValueException(msg) - if idx < 0 or idx > len(self.snapshot_operations): + count = ( + len(self.snapshot_operations) + if operation_count is None + else operation_count + ) + if idx < 0 or idx > count: msg = "Invalid marker" raise InvalidParameterValueException(msg) return idx diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/executor.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/executor.py index 44a01bc01..b4033d976 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/executor.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/executor.py @@ -1526,7 +1526,8 @@ def _begin_invocation( self._set_invocation_gate(execution_arn, InvocationState.INVOKING) execution.begin_new_invocation() invocation_input = self._invoker.create_invocation_input(execution=execution) - execution.mark_state_delivered() + # Keep input updates until a valid response acknowledges this snapshot. + # A failed invocation must not consume their delivery before a retry. self._store.save(execution) return execution, invocation_input, execution.seq_counter @@ -1584,6 +1585,13 @@ def _finish_invocation( # A clean invocation resets the retry counter. reset_target = self._store.load(execution_arn) reset_target.consecutive_failed_invocation_attempts = 0 + if invocation_seq is not None: + # A valid response acknowledges the input snapshot, even when this + # invocation suspended without another checkpoint. Newer changes + # must remain available for the next invocation or checkpoint. + OperationPaginatorState.pin(reset_target).advance_handler_seen( + invocation_seq + ) self._store.save(reset_target) # Release the gate: terminal -> COMPLETED, otherwise PRE_INVOKE diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/invoker.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/invoker.py index 47e7ac5ed..69070cbb6 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/invoker.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/invoker.py @@ -25,6 +25,9 @@ ResourceNotFoundException, ) from aws_durable_execution_sdk_python_testing.execution import OperationPaginatorState +from aws_durable_execution_sdk_python_testing.operation_state import ( + operation_state_to_dict, +) from aws_durable_execution_sdk_python_testing.model import ( LambdaContext, executed_version, @@ -472,10 +475,15 @@ def invoke( input.durable_execution_arn, endpoint_url, region_name ) + invocation_payload = input.to_json_dict() + invocation_payload["InitialExecutionState"]["Operations"] = [ + operation_state_to_dict(operation, json_timestamps=True) + for operation in input.initial_execution_state.operations + ] invoke_kwargs: dict[str, Any] = { "FunctionName": function_name, "InvocationType": "RequestResponse", # Synchronous invocation - "Payload": json.dumps(input.to_json_dict()), + "Payload": json.dumps(invocation_payload), } if tenant_id is not None: invoke_kwargs["TenantId"] = tenant_id diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/model.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/model.py index 43026ec60..617a40089 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/model.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/model.py @@ -1,6 +1,9 @@ """Model classes for the web API.""" from __future__ import annotations +from aws_durable_execution_sdk_python_testing.operation_state import ( + operation_state_to_dict, +) import datetime @@ -572,7 +575,7 @@ def from_dict(cls, data: dict) -> GetDurableExecutionStateResponse: def to_dict(self) -> dict[str, Any]: result: dict[str, Any] = { - "Operations": [op.to_dict() for op in self.operations] + "Operations": [operation_state_to_dict(op) for op in self.operations] } if self.next_marker is not None: result["NextMarker"] = self.next_marker @@ -677,7 +680,8 @@ class EventError: @classmethod def from_dict(cls, data: dict) -> EventError: payload = None - if payload_data := data.get("Payload"): + payload_data = data.get("Payload") + if payload_data is not None: payload = ErrorObject.from_dict(payload_data) return cls( @@ -2239,6 +2243,14 @@ def create_callback_event_failed(cls, context: EventCreationContext) -> Event: event_error: EventError | None = ( EventError.from_details(callback_details) if callback_details else None ) + if ( + context.include_execution_data + and callback_details is not None + and callback_details.error is None + ): + # Detailed service history retains an empty Error.Payload object. + # This projection must not turn the SDK-facing absent error into one. + event_error = EventError(payload=ErrorObject.from_dict({}), truncated=False) return cls( event_type=EventType.CALLBACK_FAILED.value, event_timestamp=context.end_timestamp, @@ -2592,6 +2604,11 @@ class HistoryEventTypeConfig: } +def _operation_error(error: ErrorObject | None) -> ErrorObject | None: + # History retains present Payload={}, while SDK state treats Error={} as absent. + return error if error is not None and error.to_dict() else None + + def events_to_operations(events: list[Event]) -> list[Operation]: """Convert a list of history events into operations. @@ -2744,7 +2761,7 @@ def events_to_operations(events: list[Event]) -> list[Operation]: callback_details=CallbackDetails( callback_id=callback_id, result=result, - error=error, + error=_operation_error(error), ), ) @@ -2785,7 +2802,7 @@ def events_to_operations(events: list[Event]) -> list[Operation]: operation, step_details=StepDetails( result=result_val, - error=error_val, + error=_operation_error(error_val), attempt=attempt, next_attempt_timestamp=next_attempt_ts, ), @@ -2818,7 +2835,9 @@ def events_to_operations(events: list[Event]) -> list[Operation]: operation, context_details=ContextDetails( result=None, - error=event.context_failed_details.error.payload, + error=_operation_error( + event.context_failed_details.error.payload + ), ), ) @@ -2843,7 +2862,9 @@ def events_to_operations(events: list[Event]) -> list[Operation]: operation, chained_invoke_details=ChainedInvokeDetails( result=None, - error=event.chained_invoke_failed_details.error.payload, + error=_operation_error( + event.chained_invoke_failed_details.error.payload + ), ), ) @@ -3139,7 +3160,7 @@ def from_dict(cls, data: dict) -> CheckpointUpdatedExecutionState: def to_dict(self) -> dict[str, Any]: result: dict[str, Any] = { - "Operations": [op.to_dict() for op in self.operations] + "Operations": [operation_state_to_dict(op) for op in self.operations] } if self.next_marker is not None: result["NextMarker"] = self.next_marker diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/operation_state.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/operation_state.py new file mode 100644 index 000000000..850eea00b --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/operation_state.py @@ -0,0 +1,64 @@ +"""Service views of retained operation state; raw execution storage stays complete.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +from aws_durable_execution_sdk_python.lambda_service import ( + Operation, + OperationStatus, + OperationType, +) + + +_CLOSED_CONTEXT = (OperationStatus.SUCCEEDED, OperationStatus.FAILED) + + +def prune_replay_children(operations: Sequence[Operation]) -> list[Operation]: + """Omit a descendant when any known completed context prunes its children.""" + by_id = {operation.operation_id: operation for operation in operations} + visible = [] + for operation in operations: + parent_id = operation.parent_id + visited = {operation.operation_id} + hidden = False + while parent_id is not None and parent_id not in visited: + parent = by_id.get(parent_id) + if parent is None: + break + if ( + parent.operation_type is OperationType.CONTEXT + and parent.status in _CLOSED_CONTEXT + and not ( + parent.context_details + and parent.context_details.replay_children is True + ) + ): + hidden = True + break + visited.add(parent_id) + parent_id = parent.parent_id + if not hidden: + visible.append(operation) + return visible + + +def operation_state_to_dict( + operation: Operation, *, json_timestamps: bool = False +) -> dict[str, Any]: + """Project service ReplayChildren presence without changing Core's model API.""" + data = dict(operation.to_json_dict() if json_timestamps else operation.to_dict()) + if ( + operation.operation_type is OperationType.CONTEXT + and operation.context_details is not None + ): + details = dict(data.get("ContextDetails", {})) + if operation.status is OperationStatus.STARTED: + details.pop("ReplayChildren", None) + elif operation.status in _CLOSED_CONTEXT: + details["ReplayChildren"] = ( + operation.context_details.replay_children is True + ) + data["ContextDetails"] = details + return data diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/e2e/callback_updated_operations_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/callback_updated_operations_test.py new file mode 100644 index 000000000..cae94c294 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/callback_updated_operations_test.py @@ -0,0 +1,253 @@ +"""Real callback completions are delivered once before resumed user code.""" + +from __future__ import annotations + +import json +import time +from pathlib import Path +from queue import Queue +from threading import Event +from typing import Any + +import pytest +from aws_durable_execution_sdk_python import ( + DurableContext, + StepContext, + durable_execution, + durable_step, +) +from aws_durable_execution_sdk_python.config import ( + CallbackConfig, + Duration, + WaitForCallbackConfig, +) +from aws_durable_execution_sdk_python.exceptions import CallbackError +from aws_durable_execution_sdk_python.lambda_service import ErrorObject +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationStartInfo, + OperationEndInfo, + UserFunctionStartInfo, +) +from aws_durable_execution_sdk_python.serdes import JsonSerDes +from aws_durable_execution_sdk_python.types import WaitForCallbackContext + +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from aws_durable_execution_sdk_python_testing.stores.filesystem import ( + FileSystemExecutionStore, +) + + +class CallbackObserver(DurableInstrumentationPlugin): + def __init__(self) -> None: + self.invocations: list[InvocationStartInfo] = [] + self.ends: list[OperationEndInfo] = [] + self.order: list[str] = [] + + def on_invocation_start(self, info: InvocationStartInfo) -> None: + self.invocations.append(info) + + def on_operation_end(self, info: OperationEndInfo) -> None: + if info.name == "target": + self.ends.append(info) + self.order.append("target-end") + + def on_user_function_start(self, info: UserFunctionStartInfo) -> None: + if info.name == "observed": + self.order.append("observed-start") + + +def _submit(_callback_id: str, _context: WaitForCallbackContext) -> None: + return None + + +def _suspended_callback(runner: DurableFunctionTestRunner, arn: str, name: str) -> str: + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + events = runner.get_execution_history(arn, include_execution_data=True).events + starts = [ + event + for event in events + if event.event_type == "CallbackStarted" and event.name == name + ] + if starts and any( + event.event_type == "InvocationCompleted" + and event.event_id > starts[0].event_id + for event in events + ): + details = starts[0].callback_started_details + assert details is not None + assert details.callback_id is not None + return details.callback_id + time.sleep(0.01) + raise AssertionError(f"Callback {name} did not suspend") + + +@pytest.mark.parametrize("outcome", ["success", "failure", "timeout"]) +@pytest.mark.parametrize("filesystem", [False, True]) +def test_callback_update_is_consumed_before_user_code( + outcome: str, filesystem: bool, tmp_path: Path +) -> None: + observer = CallbackObserver() + marker_calls: list[str] = [] + + @durable_step + def observed(_context: StepContext, value: str) -> str: + marker_calls.append(value) + return value + + def handler(_event: Any, context: DurableContext) -> str: + config = CallbackConfig( + timeout=Duration.from_seconds(1 if outcome == "timeout" else 30), + serdes=JsonSerDes(), + ) + try: + target = context.create_callback(name="target", config=config).result() + except CallbackError: + target = outcome + assert isinstance(target, str) + saved = context.step(observed(target), name="observed") + callback_config = WaitForCallbackConfig(serdes=JsonSerDes()) + one = context.wait_for_callback(_submit, name="one", config=callback_config) + two = context.wait_for_callback(_submit, name="two", config=callback_config) + return "/".join((saved, one, two)) + + wrapped = durable_execution(handler, plugins=[observer]) + store = FileSystemExecutionStore.create(tmp_path) if filesystem else None + with DurableFunctionTestRunner( + handler=wrapped, + store=store, + skip_time=False, + poll_interval=0.01, + execution_timeout=25, + ) as runner: + arn = runner.run_async(input="{}") + callback_id = _suspended_callback(runner, arn, "target") + if outcome == "success": + runner.send_callback_success( + callback_id, result=json.dumps("target").encode() + ) + elif outcome == "failure": + runner.send_callback_failure( + callback_id, error=ErrorObject.from_message("explicit callback failure") + ) + # The timeout case uses the runner's real scheduled callback deadline. + for name in ["one", "two"]: + callback_id = _suspended_callback(runner, arn, name + " create callback id") + runner.send_callback_success(callback_id, result=json.dumps(name).encode()) + result = runner.wait_for_result(arn, timeout=10) + + expected = "target" if outcome == "success" else outcome + assert result.status.value == "SUCCEEDED" + assert result.result is not None + assert json.loads(result.result) == expected + "/one/two" + assert marker_calls == [expected] + assert len(observer.invocations) == 4 + assert len(observer.ends) == 1 + terminal = observer.ends[0] + assert ( + terminal.status.value + == {"success": "SUCCEEDED", "failure": "FAILED", "timeout": "TIMED_OUT"}[ + outcome + ] + ) + assert terminal.is_replayed is False + assert set(observer.invocations[1].updated_operations) == {terminal.operation_id} + assert all( + terminal.operation_id not in invocation.updated_operations + for invocation in observer.invocations[2:] + ) + assert observer.order == ["target-end", "observed-start"] + if outcome == "failure": + assert terminal.error is not None + assert terminal.error.message == "explicit callback failure" + + +@pytest.mark.parametrize("outcome", ["success", "failure", "timeout"]) +@pytest.mark.parametrize("filesystem", [False, True]) +def test_callback_delivered_by_checkpoint_is_not_updated_on_later_resume( + outcome: str, filesystem: bool, tmp_path: Path +) -> None: + observer = CallbackObserver() + submitter_entered = Event() + release_submitter = Event() + submission_calls: list[str] = [] + published_callbacks: Queue[str] = Queue() + + @durable_step + def submit(_context: StepContext, callback_id: str) -> str: + submission_calls.append(callback_id) + published_callbacks.put(callback_id) + submitter_entered.set() + assert release_submitter.wait(10) + return "submitted" + + def handler(_event: Any, context: DurableContext) -> str: + callback = context.create_callback( + name="target", + config=CallbackConfig( + timeout=Duration.from_seconds(1 if outcome == "timeout" else 30), + serdes=JsonSerDes(), + ), + ) + context.step(submit(callback.callback_id), name="submit") + try: + result = callback.result() + except CallbackError: + result = outcome + context.wait(Duration.from_seconds(1), name="first-replay") + context.wait(Duration.from_seconds(1), name="second-replay") + assert isinstance(result, str) + return result + + store = FileSystemExecutionStore.create(tmp_path) if filesystem else None + with DurableFunctionTestRunner( + handler=durable_execution(handler, plugins=[observer]), + store=store, + skip_time=False, + poll_interval=0.01, + execution_timeout=25, + ) as runner: + arn = runner.run_async(input="{}") + try: + assert submitter_entered.wait(5) + callback_id = published_callbacks.get(timeout=5) + if outcome == "success": + runner.send_callback_success(callback_id, result=b'"target"') + elif outcome == "failure": + runner.send_callback_failure( + callback_id, error=ErrorObject.from_message("early failure") + ) + expected_event = { + "success": "CallbackSucceeded", + "failure": "CallbackFailed", + "timeout": "CallbackTimedOut", + }[outcome] + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + events = runner.get_execution_history( + arn, include_execution_data=True + ).events + if any(event.event_type == expected_event for event in events): + break + time.sleep(0.01) + else: + raise AssertionError("Callback did not complete during submission") + assert not any( + event.event_type == "InvocationCompleted" for event in events + ) + finally: + release_submitter.set() + result = runner.wait_for_result(arn, timeout=15) + + assert result.status.value == "SUCCEEDED" + assert result.result is not None + assert json.loads(result.result) == ("target" if outcome == "success" else outcome) + assert len(submission_calls) == 1 + assert len(observer.invocations) >= 3 + assert len(observer.ends) == 1 + target_id = observer.ends[0].operation_id + assert all( + target_id not in invocation.updated_operations + for invocation in observer.invocations[1:] + ) diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/e2e/invocation_input_delivery_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/invocation_input_delivery_test.py new file mode 100644 index 000000000..dfc64e7aa --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/invocation_input_delivery_test.py @@ -0,0 +1,144 @@ +"""An input-only callback resume must not redeliver its terminal update later.""" + +import json +import time + +import pytest +from aws_durable_execution_sdk_python import ( + DurableContext, + durable_execution, + durable_step, +) +from aws_durable_execution_sdk_python.config import CallbackConfig +from aws_durable_execution_sdk_python.serdes import JsonSerDes +from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin +from aws_durable_execution_sdk_python_testing.execution import Execution +from aws_durable_execution_sdk_python_testing.executor import Executor, InvocationState +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from aws_durable_execution_sdk_python_testing.stores.filesystem import ( + FileSystemExecutionStore, +) +from aws_durable_execution_sdk_python_testing.stores.memory import ( + InMemoryExecutionStore, +) + + +@pytest.mark.parametrize("filesystem", [False, True]) +@pytest.mark.parametrize("failed_delivery", [False, True]) +@pytest.mark.parametrize("legacy_format", [False, True]) +def test_callback_input_acknowledgment_survives_checkpoint_free_resume( + filesystem, failed_delivery, legacy_format, tmp_path, monkeypatch +): + store = ( + FileSystemExecutionStore.create(tmp_path / "store") + if filesystem + else InMemoryExecutionStore() + ) + starts = [] + invocation_ends = [] + operation_ends = [] + attempts = [] + marker_calls = [] + legacy_restores = [] + + original_begin = Executor._begin_invocation # noqa: SLF001 + + def begin_with_old_store(self, arn): + # Rehydrate a real completed callback using the persisted format that + # predates sequence accounting. No SDK hook or invocation input is faked. + if ( + legacy_format + and not legacy_restores + and self._invocation_gate(arn) is InvocationState.PRE_INVOKE # noqa: SLF001 + ): + execution = store.load(arn) + if execution.updated_operation_ids: + snapshot = execution.to_json_dict() + for field in ( + "SeqCounter", + "HandlerSeenSeq", + "OperationLastTouchedSeq", + "OperationSizeBytes", + ): + snapshot.pop(field) + restored = Execution.from_json_dict(snapshot) + assert restored.seq_counter == restored.handler_seen_seq == 0 + assert restored.updated_operation_ids == execution.updated_operation_ids + store.save(restored) + legacy_restores.append(snapshot) + return original_begin(self, arn) + + monkeypatch.setattr(Executor, "_begin_invocation", begin_with_old_store) + + class Observe(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + starts.append(info) + + def on_invocation_end(self, info): + invocation_ends.append(info) + + def on_operation_end(self, info): + if info.name in ("a", "b"): + operation_ends.append((len(starts), info.name, info.is_replayed)) + + @durable_step + def marker(context, value): + marker_calls.append(1) + return value + + @durable_execution(plugins=[Observe()]) + def handler(event, context: DurableContext): + a = context.create_callback( + name="a", config=CallbackConfig(serdes=JsonSerDes()) + ) + b = context.create_callback( + name="b", config=CallbackConfig(serdes=JsonSerDes()) + ) + result = [a.result(), b.result()] + return context.step(marker(result), name="after-b") + + def entrypoint(event, context): + attempts.append(1) + if failed_delivery and len(attempts) == 2: + raise RuntimeError("invocation failed before the SDK received its input") + return handler(event, context) + + with DurableFunctionTestRunner( + handler=entrypoint, + store=store, + skip_time=False, + poll_interval=0.01, + execution_timeout=25, + ) as runner: + arn = runner.run_async(input="{}") + a_id = runner.wait_for_callback(arn, name="a", timeout=5) + b_id = runner.wait_for_callback(arn, name="b", timeout=5) + runner.send_callback_success(a_id, result=b'"A"') + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + events = runner.get_execution_history( + arn, include_execution_data=True + ).events + if ( + len(invocation_ends) >= 2 + and sum(event.event_type == "InvocationCompleted" for event in events) + >= 2 + ): + break + time.sleep(0.01) + else: + raise AssertionError("The second SDK invocation did not suspend") + runner.send_callback_success(b_id, result=b'"B"') + result = runner.wait_for_result(arn, timeout=10) + + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == ["A", "B"] + assert [info.status.value for info in invocation_ends] == [ + "PENDING", + "PENDING", + "SUCCEEDED", + ] + assert operation_ends == [(2, "a", False), (3, "b", False)] + assert marker_calls == [1] + assert len(attempts) == (4 if failed_delivery else 3) + assert len(legacy_restores) == int(legacy_format) diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/e2e/replay_children_service_parity_int_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/replay_children_service_parity_int_test.py new file mode 100644 index 000000000..3c31a9412 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/replay_children_service_parity_int_test.py @@ -0,0 +1,256 @@ +"""Real public handlers replay through the service's retained-operation view.""" + +import json +import os +from pathlib import Path + +import pytest +from aws_durable_execution_sdk_python import durable_execution +from aws_durable_execution_sdk_python.config import CallbackConfig +from aws_durable_execution_sdk_python.exceptions import ChildContextError +from aws_durable_execution_sdk_python.serdes import JsonSerDes +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from aws_durable_execution_sdk_python_testing.stores.memory import ( + InMemoryExecutionStore, +) +from aws_durable_execution_sdk_python_testing.stores.filesystem import ( + FileSystemExecutionStore, +) + + +@pytest.mark.parametrize("filesystem", [False, True]) +@pytest.mark.parametrize("mode", ["completed", "retained", "nested-retained", "failed"]) +def test_public_resume_uses_service_state_but_inspection_remains_complete( + filesystem, mode, tmp_path +): + store = ( + FileSystemExecutionStore.create(tmp_path / "state") + if filesystem + else InMemoryExecutionStore() + ) + snapshots = [] + child_calls = [] + step_calls = [] + + def save(_step): + step_calls.append(1) + return "saved" + + def child(ctx): + child_calls.append(1) + saved = ctx.step(save, name="saved-child-step") + if mode == "failed": + raise ValueError("stored-child-failure") + return { + "saved": saved, + "large": "x" * (300 * 1024) + if mode in ["retained", "nested-retained"] + else "", + } + + def outer(ctx): + value = ctx.run_in_child_context(child, name="inner-retained") + return {"saved": value["saved"], "large": ""} + + def handler(_event, ctx): + try: + value = ctx.run_in_child_context( + outer if mode == "nested-retained" else child, name="outer" + ) + result = {"saved": value["saved"], "length": len(value["large"])} + except ChildContextError: + assert mode == "failed" + result = {"failed": True} + gate = ctx.create_callback( + name="resume", config=CallbackConfig(serdes=JsonSerDes()) + ) + assert gate.result() == "resume" + return result + + wrapped = durable_execution(handler) + + def entrypoint(event, context): + initial = event.initial_execution_state.operations + state = event.service_client.get_execution_state( + event.durable_execution_arn, event.checkpoint_token, "" + ) + snapshots.append( + { + "input_names": [op.name for op in initial], + "state_names": [op.name for op in state.operations], + "input_ids": [op.operation_id for op in initial], + "state_ids": [op.operation_id for op in state.operations], + "updated_ids": list(event.updated_operation_ids), + } + ) + return wrapped(event, context) + + with DurableFunctionTestRunner( + handler=entrypoint, + store=store, + skip_time=False, + poll_interval=0.01, + execution_timeout=20, + ) as runner: + arn = runner.run_async(input="{}") + callback_id = runner.wait_for_callback(arn, name="resume", timeout=5) + runner.send_callback_success(callback_id, result=b'"resume"') + result = runner.wait_for_result(arn, timeout=10) + history = runner.get_execution_history(arn, include_execution_data=True).events + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == ( + {"failed": True} + if mode == "failed" + else {"saved": "saved", "length": 300 * 1024 if mode == "retained" else 0} + ) + assert ( + len(snapshots) == 2 and snapshots[0]["input_ids"] == snapshots[0]["state_ids"] + ) + assert snapshots[1]["input_ids"] == snapshots[1]["state_ids"] + visible = "saved-child-step" in snapshots[1]["input_names"] + assert visible is (mode == "retained") + if mode == "nested-retained": + assert "inner-retained" not in snapshots[1]["input_names"] + assert "outer" in snapshots[1]["input_names"] + assert step_calls == [1] and len(child_calls) == (2 if mode == "retained" else 1) + raw = store.load(arn) + assert any(op.name == "saved-child-step" for op in raw.operations) + inspected = result.get_operation_by_name("outer") + if mode == "nested-retained": + inspected = inspected.get_operation_by_name("inner-retained") + assert ( + inspected.get_operation_by_name("saved-child-step").name == "saved-child-step" + ) + assert ( + sum( + e.event_type == "StepSucceeded" and e.name == "saved-child-step" + for e in history + ) + == 1 + ) + assert any(e.name == "saved-child-step" for e in history) + if os.environ.get("SERVICE_PRUNING_PROOF_DIR"): + out = Path(os.environ["SERVICE_PRUNING_PROOF_DIR"]) + out.mkdir(parents=True, exist_ok=True) + (out / f"{mode}-fs{filesystem}.json").write_text( + json.dumps( + dict( + mode=mode, + filesystem=filesystem, + snapshots=snapshots, + result=json.loads(result.result), + step_calls=len(step_calls), + child_calls=len(child_calls), + history_names=[e.name for e in history], + ), + indent=2, + ) + + "\n" + ) + + +@pytest.mark.parametrize("filesystem", [False, True]) +def test_hidden_callback_updated_id_does_not_make_its_metadata_visible( + filesystem, tmp_path +): + import time + + store = ( + FileSystemExecutionStore.create(tmp_path / "hidden-state") + if filesystem + else InMemoryExecutionStore() + ) + snapshots = [] + child_calls = [] + saved_calls = [] + + def saved(_step): + saved_calls.append(1) + return "saved" + + def child(ctx): + child_calls.append(1) + callback = ctx.create_callback( + name="hidden-b", config=CallbackConfig(serdes=JsonSerDes()) + ) + assert ctx.step(saved, name="hidden-saved") == "saved" + return callback.callback_id + + def handler(_event, ctx): + a = ctx.create_callback( + name="root-a", config=CallbackConfig(serdes=JsonSerDes()) + ) + b_id = ctx.run_in_child_context(child, name="closed-parent") + return {"a": a.result(), "b_id": b_id} + + wrapped = durable_execution(handler) + + def entrypoint(event, context): + state = event.service_client.get_execution_state( + event.durable_execution_arn, event.checkpoint_token, "" + ) + snapshots.append( + { + "input_ids": [ + op.operation_id for op in event.initial_execution_state.operations + ], + "state_ids": [op.operation_id for op in state.operations], + "updated_ids": list(event.updated_operation_ids), + } + ) + return wrapped(event, context) + + def completed(runner, arn, count): + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + events = runner.get_execution_history( + arn, include_execution_data=True + ).events + if sum(e.event_type == "InvocationCompleted" for e in events) >= count: + return events + time.sleep(0.01) + raise AssertionError("invocation completion not observed") + + with DurableFunctionTestRunner( + handler=entrypoint, + store=store, + skip_time=False, + poll_interval=0.01, + execution_timeout=25, + ) as runner: + arn = runner.run_async(input="{}") + a_id = runner.wait_for_callback(arn, name="root-a", timeout=5) + b_id = runner.wait_for_callback(arn, name="hidden-b", timeout=5) + first = completed(runner, arn, 1) + b_operation_id = next( + e.operation_id + for e in first + if e.name == "hidden-b" and e.event_type == "CallbackStarted" + ) + runner.send_callback_success(b_id, result=b'"B"') + completed(runner, arn, 2) + runner.send_callback_success(a_id, result=b'"A"') + result = runner.wait_for_result(arn, timeout=10) + history = runner.get_execution_history(arn, include_execution_data=True).events + assert result.status.value == "SUCCEEDED" and json.loads(result.result) == { + "a": "A", + "b_id": b_id, + } + assert len(snapshots) == 3 and child_calls == [1] and saved_calls == [1] + assert b_operation_id in snapshots[1]["updated_ids"] + assert b_operation_id not in snapshots[1]["input_ids"] + assert b_operation_id not in snapshots[1]["state_ids"] + assert any(op.operation_id == b_operation_id for op in store.load(arn).operations) + assert ( + result.get_operation_by_name("closed-parent") + .get_operation_by_name("hidden-b") + .result + == '"B"' + ) + assert ( + sum( + e.event_type == "CallbackSucceeded" and e.operation_id == b_operation_id + for e in history + ) + == 1 + ) diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/event_factory_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/event_factory_test.py index cd2fd8743..04334b6db 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/event_factory_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/event_factory_test.py @@ -9,8 +9,10 @@ import pytest from aws_durable_execution_sdk_python.lambda_service import ( + CallbackDetails, ChainedInvokeOptions, ErrorObject, + Operation as ServiceOperation, OperationStatus, OperationType, StepDetails, @@ -849,6 +851,65 @@ def test_create_callback_failed(): assert event.callback_failed_details.error.payload.message == "Callback failed" +@pytest.mark.parametrize("include_data", [False, True]) +@pytest.mark.parametrize("has_error", [False, True]) +def test_callback_failure_history_projection_preserves_sdk_state( + include_data, has_error +): + error = ErrorObject.from_message("details") if has_error else None + operation = ServiceOperation( + operation_id="callback", + operation_type=OperationType.CALLBACK, + status=OperationStatus.FAILED, + callback_details=CallbackDetails(callback_id="callback-id", error=error), + ) + context = EventCreationContext.create( + operation=operation, + event_id=3, + durable_execution_arn="arn:test", + start_input=StartDurableExecutionInput( + account_id="123", + function_name="test", + function_qualifier="$LATEST", + execution_name="test", + execution_timeout_seconds=300, + execution_retention_period_days=7, + ), + include_execution_data=include_data, + ) + event = Event.create_callback_event(context) + expected = {"Truncated": not (include_data and not has_error)} + if has_error or include_data: + expected["Payload"] = error.to_dict() if error else {} + assert event.callback_failed_details.error.to_dict() == expected + assert operation.callback_details.error is error + + +def test_callback_timeout_history_projection_is_unchanged(): + operation = ServiceOperation( + operation_id="callback", + operation_type=OperationType.CALLBACK, + status=OperationStatus.TIMED_OUT, + callback_details=CallbackDetails(callback_id="callback-id", error=None), + ) + context = EventCreationContext.create( + operation=operation, + event_id=3, + durable_execution_arn="arn:test", + start_input=StartDurableExecutionInput( + account_id="123", + function_name="test", + function_qualifier="$LATEST", + execution_name="test", + execution_timeout_seconds=300, + execution_retention_period_days=7, + ), + include_execution_data=True, + ) + event = Event.create_callback_event(context) + assert event.callback_timed_out_details.error.to_dict() == {"Truncated": True} + + def test_create_callback_timed_out(): operation = create_mock_operation("callback-1", status=OperationStatus.TIMED_OUT) error_obj = ErrorObject.from_message("Callback timed out") diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/execution_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/execution_test.py index bde5e3d4b..c9457a6dd 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/execution_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/execution_test.py @@ -3,6 +3,7 @@ import json from dataclasses import replace from datetime import datetime, timezone +from threading import Event, Thread from unittest.mock import patch, Mock import pytest @@ -955,6 +956,38 @@ def test_from_dict_with_none_result(): # region callback +@pytest.mark.parametrize("outcome", ["success", "failure", "timeout"]) +def test_callback_completion_records_only_successful_state_changes(outcome): + """Preserve payloads, token versions and consumed metadata on rejection.""" + operation = Operation( + operation_id="callback", + operation_type=OperationType.CALLBACK, + status=OperationStatus.STARTED, + callback_details=CallbackDetails(callback_id="callback-id"), + ) + execution = Execution("test-arn", _make_start_input(), [operation]) + complete = getattr(execution, f"complete_callback_{outcome}") + payload = b'"result"' if outcome == "success" else ErrorObject.from_message("error") + token_version = execution.token_sequence + result = complete("callback-id", payload) + + assert execution.updated_operation_ids == ["callback"] + assert execution.token_sequence == token_version + assert execution.seq_counter == 1 + assert result.callback_details.result == ( + '"result"' if outcome == "success" else None + ) + assert result.callback_details.error == (None if outcome == "success" else payload) + restored = Execution.from_json_dict(execution.to_json_dict()) + assert restored.updated_operation_ids == ["callback"] + execution.mark_state_delivered() + with pytest.raises(IllegalStateException, match="not in STARTED state"): + complete("callback-id", payload) + assert execution.updated_operation_ids == [] + assert execution.seq_counter == 1 + assert execution.token_sequence == token_version + + def test_find_callback_operation_not_found(): """Test find_callback_operation raises exception when callback not found.""" execution = Execution("test-arn", Mock(), []) @@ -1690,6 +1723,87 @@ def test_record_invocation_completion_keeps_updated_operation_ids(): assert execution.updated_operation_ids == [] +@pytest.mark.parametrize("complete_before_advance", [False, True]) +def test_checkpoint_consumes_only_updates_covered_by_its_watermark( + complete_before_advance, +): + """A late update survives reads and retries of an older state delivery.""" + execution = Execution( + "test-arn", + _make_start_input(), + [ + Operation( + operation_id=name, + operation_type=OperationType.CALLBACK, + status=OperationStatus.STARTED, + callback_details=CallbackDetails(callback_id=name), + ) + for name in ["delivered", "later"] + ], + ) + execution.complete_callback_success("delivered", b"first") + response = OperationPaginatorState.pin(execution) + ready = Event() + finished = Event() + + def complete_later(): + assert ready.wait(5) + execution.complete_callback_failure("later", ErrorObject.from_message("later")) + finished.set() + + worker = Thread(target=complete_later) + worker.start() + if complete_before_advance: + ready.set() + assert finished.wait(5) + expected = ["delivered", "later"] if complete_before_advance else ["delivered"] + assert execution.updated_operation_ids == expected + response.page(None, max_size_bytes=1024 * 1024) + assert execution.updated_operation_ids == expected + + response.advance_handler_seen(1) + ready.set() + assert finished.wait(5) + worker.join(timeout=5) + assert not worker.is_alive() + assert execution.updated_operation_ids == ["later"] + assert execution.handler_seen_seq == 1 + assert execution.token_sequence == 0 + assert execution.seq_counter == 2 + # An idempotent/older delivery cannot consume the later completion. + response.advance_handler_seen(1) + response.advance_handler_seen(0) + assert execution.updated_operation_ids == ["later"] + OperationPaginatorState.pin(execution).advance_handler_seen(2) + assert execution.updated_operation_ids == [] + + +@pytest.mark.parametrize("seen,ack", [(0, 0), (5, 5), (5, 3)]) +def test_acknowledged_updates_are_pruned_without_watermark_progress(seen, ack): + execution = Execution("test-arn", _make_start_input(), []) + snapshot = execution.to_json_dict() + for field in ( + "SeqCounter", + "HandlerSeenSeq", + "OperationLastTouchedSeq", + "OperationSizeBytes", + ): + snapshot.pop(field) + snapshot["UpdatedOperationIds"] = ["legacy", "newer"] + restored = Execution.from_json_dict(snapshot) + restored.handler_seen_seq = seen + restored.operation_last_touched_seq["newer"] = ack + 1 + paginator = OperationPaginatorState.pin(restored) + + paginator.advance_handler_seen(ack) + + assert restored.updated_operation_ids == ["newer"] + assert restored.handler_seen_seq == seen + assert restored.token_sequence == restored.seq_counter == 0 + paginator.advance_handler_seen(ack) + assert restored.updated_operation_ids == ["newer"] + + def test_function_arn_is_qualified_with_the_executed_version(): execution = Execution.new(_make_start_input()) execution.region = "ap-southeast-2" diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/executor_checkpoint_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/executor_checkpoint_test.py index c789dcd74..fb21cf68f 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/executor_checkpoint_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/executor_checkpoint_test.py @@ -17,6 +17,7 @@ import pytest from aws_durable_execution_sdk_python.execution import InvocationStatus from aws_durable_execution_sdk_python.lambda_service import ( + CallbackDetails, ErrorObject, Operation as SvcOperation, OperationAction, @@ -126,6 +127,49 @@ def test_empty_poll_returns_empty_operations_and_advances_token(): assert CheckpointToken.from_str(response.checkpoint_token).token_sequence == 1 +@pytest.mark.parametrize("rejection", ["token", "operation"]) +def test_rejected_checkpoint_retains_undelivered_callback_update(rejection): + executor, store, execution, token = _make_executor_with_started_execution() + execution.operations.append( + SvcOperation( + operation_id="callback", + operation_type=OperationType.CALLBACK, + status=OperationStatus.STARTED, + callback_details=CallbackDetails(callback_id="callback-id"), + ) + ) + execution.complete_callback_success("callback-id", b"result") + store.save(execution) + before = ( + execution.token_sequence, + execution.handler_seen_seq, + execution.seq_counter, + ) + assert execution.updated_operation_ids == ["callback"] + with pytest.raises(InvalidParameterValueException): + executor.checkpoint_execution( + execution_arn=execution.durable_execution_arn, + checkpoint_token="invalid" if rejection == "token" else token, + updates=( + [ + OperationUpdate( + operation_id="callback", + operation_type=OperationType.CALLBACK, + action=OperationAction.SUCCEED, + ) + ] + if rejection == "operation" + else [] + ), + ) + assert execution.updated_operation_ids == ["callback"] + assert ( + execution.token_sequence, + execution.handler_seen_seq, + execution.seq_counter, + ) == before + + # endregion # region: Non-empty checkpoint returns only the delta diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/executor_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/executor_test.py index bb6b71f4f..8bd384aea 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/executor_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/executor_test.py @@ -698,6 +698,10 @@ def test_pending_response_is_valid_when_an_operation_completed_after_the_handler mock_execution.consecutive_failed_invocation_attempts = 0 mock_execution.seq_counter = 4 mock_execution.handler_seen_seq = 2 + mock_execution.token_sequence = 0 + mock_execution.operations = [] + mock_execution.updated_operation_ids = [] + mock_execution.operation_last_touched_seq = {} mock_execution.has_pending_operations.return_value = False mock_execution.has_changes_after.side_effect = lambda seq: seq == 4 mock_execution.has_unseen_changes.return_value = True @@ -1237,6 +1241,12 @@ def test_should_retry_invocation_when_under_limit_through_public_api( mock_execution.is_complete = False mock_execution.start_input = start_input mock_execution.consecutive_failed_invocation_attempts = 3 # Under limit (5 is max) + mock_execution.seq_counter = 0 + mock_execution.handler_seen_seq = 0 + mock_execution.token_sequence = 0 + mock_execution.operations = [] + mock_execution.updated_operation_ids = [] + mock_execution.operation_last_touched_seq = {} # Configure invoker to fail initially with validation error, then succeed on retry mock_invocation_input = Mock() diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/invocation_input_delivery_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/invocation_input_delivery_test.py new file mode 100644 index 000000000..42c64a35b --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/tests/invocation_input_delivery_test.py @@ -0,0 +1,98 @@ +"""Invocation acknowledgments preserve newer checkpoint and completion state.""" + +from datetime import UTC, datetime +from unittest.mock import Mock, patch + +import pytest + +from aws_durable_execution_sdk_python.execution import ( + DurableExecutionInvocationOutput, + InvocationStatus, +) +from aws_durable_execution_sdk_python.lambda_service import ( + CallbackDetails, + Operation, + OperationStatus, + OperationType, +) +from aws_durable_execution_sdk_python_testing.execution import ( + Execution, + OperationPaginatorState, +) +from aws_durable_execution_sdk_python_testing.executor import Executor +from aws_durable_execution_sdk_python_testing.invoker import InvokeResponse +from aws_durable_execution_sdk_python_testing.model import StartDurableExecutionInput +from aws_durable_execution_sdk_python_testing.stores.memory import ( + InMemoryExecutionStore, +) +from aws_durable_execution_sdk_python_testing.token import CallbackToken + + +@pytest.mark.parametrize("checkpoint_during_invocation", [False, True]) +def test_input_ack_preserves_newer_completion_and_checkpoint_watermark( + checkpoint_during_invocation, +): + store = InMemoryExecutionStore() + execution = Execution.new( + StartDurableExecutionInput( + account_id="123456789012", + function_name="test", + function_qualifier="$LATEST", + execution_name="test", + execution_timeout_seconds=30, + execution_retention_period_days=1, + invocation_id="invocation", + ) + ) + execution.start() + arn = execution.durable_execution_arn + tokens = {} + for name in ("a", "b", "c"): + tokens[name] = CallbackToken(execution_arn=arn, operation_id=name).to_str() + execution.operations.append( + Operation( + operation_id=name, + operation_type=OperationType.CALLBACK, + status=OperationStatus.STARTED, + callback_details=CallbackDetails(callback_id=tokens[name]), + ) + ) + execution.touch_operation(name) + OperationPaginatorState.pin(execution).advance_handler_seen(execution.seq_counter) + execution.complete_callback_success(tokens["a"], b'"A"') + input_sequence = execution.seq_counter + if checkpoint_during_invocation: + execution.complete_callback_success(tokens["c"], b'"C"') + OperationPaginatorState.pin(execution).advance_handler_seen( + execution.seq_counter + ) + already_seen = execution.handler_seen_seq + execution.complete_callback_success(tokens["b"], b'"B"') + execution.needs_reinvoke = True + store.save(execution) + executor = Executor(store, Mock(), Mock(), Mock()) + try: + with patch.object(executor, "_invoke_execution") as reinvoke: + executor._finish_invocation( # noqa: SLF001 + arn, + InvokeResponse( + invocation_output=DurableExecutionInvocationOutput( + status=InvocationStatus.PENDING + ), + request_id="request", + ), + datetime.now(UTC), + datetime.now(UTC), + input_sequence, + ) + saved = store.load(arn) + assert saved.handler_seen_seq == max(already_seen, input_sequence) + assert saved.updated_operation_ids == ["b"] + assert [ + op.operation_id + for op in OperationPaginatorState.pin(saved).unseen_operations() + ] == ["b"] + reinvoke.assert_called_once_with(arn) + finally: + executor.shutdown() + executor._registry.shutdown() # noqa: SLF001 diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/model_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/model_test.py index 746f083d6..9885003a5 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/model_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/model_test.py @@ -7,8 +7,11 @@ import pytest from aws_durable_execution_sdk_python.lambda_service import ( + ChainedInvokeDetails as ServiceChainedInvokeDetails, + ContextDetails as ServiceContextDetails, OperationStatus, OperationType, + StepDetails as ServiceStepDetails, ) from aws_durable_execution_sdk_python_testing.exceptions import ( @@ -1461,6 +1464,117 @@ def test_event_error_with_payload_only(): } +@pytest.mark.parametrize( + "payload", + [ + {}, + {"ErrorMessage": ""}, + {"ErrorType": ""}, + {"ErrorData": ""}, + {"StackTrace": []}, + ], +) +def test_event_error_roundtrip_preserves_present_payload(payload): + wire = {"Payload": payload, "Truncated": False} + assert EventError.from_dict(wire).to_dict() == wire + + +@pytest.mark.parametrize("truncated", [False, True]) +@pytest.mark.parametrize("payload_field", [{}, {"Payload": None}]) +def test_event_error_absent_or_null_payload_retains_existing_meaning( + payload_field, truncated +): + parsed = EventError.from_dict({**payload_field, "Truncated": truncated}) + assert parsed.payload is None + assert parsed.to_dict() == {"Truncated": truncated} + + +def test_event_error_empty_payload_retains_truncation_flag(): + wire = {"Payload": {}, "Truncated": True} + assert EventError.from_dict(wire).to_dict() == wire + + +@pytest.mark.parametrize( + "event_type,details_attribute,service_details", + [ + ("StepFailed", "step_details", ServiceStepDetails), + ("ContextFailed", "context_details", ServiceContextDetails), + ("ChainedInvokeFailed", "chained_invoke_details", ServiceChainedInvokeDetails), + ], +) +@pytest.mark.parametrize( + "wire_error", + [ + {"Truncated": False}, + {"Payload": None, "Truncated": False}, + {"Payload": {}, "Truncated": False}, + {"Payload": {}, "Truncated": True}, + {"Truncated": True}, + {"Payload": {"ErrorMessage": ""}, "Truncated": False}, + {"Payload": {"ErrorType": ""}, "Truncated": False}, + {"Payload": {"ErrorData": ""}, "Truncated": False}, + {"Payload": {"StackTrace": []}, "Truncated": False}, + { + "Payload": { + "ErrorMessage": "failure", + "ErrorType": "RuntimeError", + "ErrorData": "details", + "StackTrace": [], + }, + "Truncated": False, + }, + ], +) +def test_failed_history_preserves_payload_and_matches_service_error_semantics( + event_type, details_attribute, service_details, wire_error +): + from aws_durable_execution_sdk_python_testing.runner import ( + DurableFunctionTestResult, + ) + + history = GetDurableExecutionHistoryResponse.from_dict( + { + "Events": [ + { + "EventType": event_type, + "EventId": 1, + "Id": "operation", + "Name": "failed-operation", + "EventTimestamp": TIMESTAMP_2023_01_01_00_01, + event_type + "Details": {"Error": wire_error}, + } + ] + } + ) + history_wire = history.to_dict() + expected_history_error = dict(wire_error) + if expected_history_error.get("Payload") is None: + expected_history_error.pop("Payload", None) + assert history_wire["Events"][0][event_type + "Details"]["Error"] == ( + expected_history_error + ) + assert GetDurableExecutionHistoryResponse.from_dict(history_wire).to_dict() == ( + history_wire + ) + + expected = service_details.from_dict({"Error": wire_error.get("Payload")}).error + (operation,) = events_to_operations(history.events) + assert operation.status is OperationStatus.FAILED + + execution = GetDurableExecutionResponse( + durable_execution_arn="arn:aws:lambda:us-west-2:123456789012:function:test:execution:one", + durable_execution_name="one", + function_arn="arn:aws:lambda:us-west-2:123456789012:function:test", + status="FAILED", + start_timestamp=TIMESTAMP_2023_01_01_00_00, + ) + result = DurableFunctionTestResult.from_execution_history(execution, history) + assert result.get_operation_by_name("failed-operation").error == expected + assert getattr(operation, details_attribute).error == expected + # Reconstructing SDK-facing state must not mutate detailed history payloads. + assert history.to_dict() == history_wire + + # Tests for RetryDetails def test_retry_details_serialization(): """Test RetryDetails from_dict/to_dict round-trip.""" diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/replay_children_service_parity_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/replay_children_service_parity_test.py new file mode 100644 index 000000000..65302dbec --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/tests/replay_children_service_parity_test.py @@ -0,0 +1,239 @@ +"""The service's retained-operation tree applies only to invocation/state reads.""" + +from datetime import UTC, datetime +from unittest.mock import Mock + +import pytest +from aws_durable_execution_sdk_python.lambda_service import ( + ContextDetails, + ContextOptions, + Operation, + OperationAction, + OperationStatus, + OperationType, + OperationUpdate, +) +from aws_durable_execution_sdk_python_testing.execution import ( + Execution, + OperationPaginatorState, +) +from aws_durable_execution_sdk_python_testing.invoker import ( + InProcessInvoker, + LambdaInvoker, +) +from aws_durable_execution_sdk_python_testing.model import ( + StartDurableExecutionInput, + GetDurableExecutionStateResponse, + CheckpointUpdatedExecutionState, +) +from aws_durable_execution_sdk_python_testing.checkpoint.processor import ( + CheckpointProcessor, +) +from aws_durable_execution_sdk_python_testing.checkpoint.processors.context import ( + ContextProcessor, +) +from aws_durable_execution_sdk_python_testing.stores.memory import ( + InMemoryExecutionStore, +) +from aws_durable_execution_sdk_python_testing.stores.filesystem import ( + FileSystemExecutionStore, +) + + +def step(identifier, parent=None, status=OperationStatus.SUCCEEDED): + return Operation( + operation_id=identifier, + parent_id=parent, + operation_type=OperationType.STEP, + status=status, + ) + + +def ctx(identifier, parent=None, status=OperationStatus.SUCCEEDED, replay=None): + return Operation( + operation_id=identifier, + parent_id=parent, + operation_type=OperationType.CONTEXT, + status=status, + context_details=ContextDetails.from_dict( + {} if replay is None else {"ReplayChildren": replay} + ), + ) + + +TREE = [ + step("1-step"), + ctx("2-context", status=OperationStatus.STARTED), + step("2.1-step", "2-context", OperationStatus.PENDING), + step("3-step", "missing"), + ctx("4-context", status=OperationStatus.FAILED), + ctx("4.1-context", "4-context", OperationStatus.FAILED, True), + step("4.1.1-step", "4.1-context", OperationStatus.FAILED), + step("4.1.2-step", "4.1-context", OperationStatus.FAILED), + ctx("5-context", replay=True), + ctx("5.1-context", "5-context", replay=False), + step("5.1.1-step", "5.1-context"), + step("5.1-step", "5-context"), + ctx("map"), + ctx("running-branch", "map", OperationStatus.STARTED), + step("pending-branch-step", "running-branch", OperationStatus.PENDING), + ctx("retained-inside-pruned", "running-branch", replay=True), + step("nested-step", "retained-inside-pruned"), +] +KEPT = [ + "1-step", + "2-context", + "2.1-step", + "3-step", + "4-context", + "5-context", + "5.1-context", + "5.1-step", + "map", +] + + +def execution_with(operations): + start = StartDurableExecutionInput( + account_id="123456789012", + function_name="parity", + function_qualifier="$LATEST", + execution_name="parity", + execution_timeout_seconds=60, + execution_retention_period_days=1, + invocation_id="invocation", + ) + execution = Execution( + "arn:aws:lambda:us-west-2:123456789012:function:parity:$LATEST/durable-execution/parity/id", + start, + list(operations), + ) + for op in operations: + execution.touch_operation(op.operation_id) + execution.updated_operation_ids = ["4.1.1-step", "pending-branch-step"] + return execution + + +@pytest.mark.parametrize("filesystem", [False, True]) +@pytest.mark.parametrize("reverse", [False, True]) +def test_service_tree_paging_inputs_and_get_state_keep_raw_inspection( + filesystem, reverse, tmp_path +): + operations = list(reversed(TREE)) if reverse else list(TREE) + expected = list(reversed(KEPT)) if reverse else KEPT + execution = execution_with(operations) + store = ( + FileSystemExecutionStore.create(tmp_path / "state") + if filesystem + else InMemoryExecutionStore() + ) + store.save(execution) + restored = store.load(execution.durable_execution_arn) + paginator = OperationPaginatorState.pin(restored) + paged = [] + marker = None + while True: + page, marker = paginator.page(marker, 100000, max_items=2) + paged.extend(op.operation_id for op in page) + if marker is None: + break + assert paged == expected + for invoker in [InProcessInvoker(Mock(), Mock()), LambdaInvoker(Mock())]: + invocation = invoker.create_invocation_input(restored) + assert [ + o.operation_id for o in invocation.initial_execution_state.operations + ] == expected + assert invocation.updated_operation_ids == ["4.1.1-step", "pending-branch-step"] + state = CheckpointProcessor(store, Mock()).get_execution_state( + restored.get_new_checkpoint_token(), "" + ) + assert [o.operation_id for o in state.operations] == expected + # Checkpoint deltas, raw store, inspection, and hidden updated IDs remain full. + assert paginator.unseen_operations() == operations + assert ( + restored.operations == operations + and restored.get_assertable_operations() == operations + ) + assert restored.updated_operation_ids == ["4.1.1-step", "pending-branch-step"] + + +@pytest.mark.parametrize("action", [OperationAction.SUCCEED, OperationAction.FAIL]) +@pytest.mark.parametrize( + "close_options", + [ + None, + ContextOptions(replay_children=True), + ContextOptions(replay_children=False), + ContextOptions(), + ], +) +def test_replay_children_is_recorded_only_by_closing_update(action, close_options): + processor = ContextProcessor() + now = datetime.now(UTC) + notifier = Mock() + start = OperationUpdate( + operation_id="ctx", + operation_type=OperationType.CONTEXT, + action=OperationAction.START, + context_options=ContextOptions(replay_children=True), + ) + opened = processor.process(start, None, notifier, "execution", now) + assert ( + opened.context_details is not None + and opened.context_details.replay_children is False + ) + assert "ReplayChildren" not in opened.to_json_dict()["ContextDetails"] + close = OperationUpdate( + operation_id="ctx", + operation_type=OperationType.CONTEXT, + action=action, + context_options=close_options, + payload="saved", + ) + completed = processor.process(close, opened, notifier, "execution", now) + expected = close_options is not None and close_options.replay_children is True + assert completed.context_details.replay_children is expected + assert ( + GetDurableExecutionStateResponse([completed]).to_dict()["Operations"][0][ + "ContextDetails" + ]["ReplayChildren"] + is expected + ) + assert ( + CheckpointUpdatedExecutionState([completed]).to_dict()["Operations"][0][ + "ContextDetails" + ]["ReplayChildren"] + is expected + ) + assert completed.context_details.result == "saved" + + +@pytest.mark.parametrize( + "operations", [[step("a", "a")], [step("a", "b"), step("b", "a")]] +) +def test_invalid_cycles_terminate_without_dropping_unrelated_operations(operations): + execution = execution_with(operations) + assert OperationPaginatorState.pin(execution).page(None, 100000)[0] == operations + + +def test_lambda_invocation_projects_close_false_and_start_absence(): + from io import BytesIO + import json + + execution = execution_with(TREE) + client = Mock() + client.invoke.return_value = { + "StatusCode": 200, + "Payload": BytesIO(b'{"Status":"SUCCEEDED","Result":"null"}'), + } + invoker = LambdaInvoker(client) + invocation = invoker.create_invocation_input(execution) + invoker.invoke("target", invocation) + payload = json.loads(client.invoke.call_args.kwargs["Payload"]) + operations = {op["Id"]: op for op in payload["InitialExecutionState"]["Operations"]} + assert list(operations) == KEPT + assert "ReplayChildren" not in operations["2-context"]["ContextDetails"] + assert operations["4-context"]["ContextDetails"]["ReplayChildren"] is False + assert operations["5-context"]["ContextDetails"]["ReplayChildren"] is True + assert operations["5.1-context"]["ContextDetails"]["ReplayChildren"] is False + assert payload["UpdatedOperationIds"] == ["4.1.1-step", "pending-branch-step"] diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/runner_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/runner_test.py index ab75f0d24..0662c07f1 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/runner_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/runner_test.py @@ -1116,6 +1116,82 @@ def test_durable_child_context_test_runner_init_with_args( # Tests for DurableFunctionCloudTestRunner and from_execution_history +@pytest.mark.parametrize("decode_wire", [False, True]) +@pytest.mark.parametrize( + "payload", + [ + {}, + {"ErrorMessage": ""}, + {"ErrorType": ""}, + {"ErrorData": ""}, + {"StackTrace": []}, + {"ErrorMessage": "failure"}, + ], +) +def test_callback_history_result_preserves_wire_and_canonical_error( + payload, decode_wire +): + from aws_durable_execution_sdk_python.lambda_service import ErrorObject + from aws_durable_execution_sdk_python_testing.model import ( + CallbackFailedDetails, + CallbackStartedDetails, + Event, + EventError, + GetDurableExecutionResponse, + ) + + timestamp = datetime.datetime(2026, 10, 7, tzinfo=datetime.UTC) + caller_error = ErrorObject.from_message("Callback failed") + execution = GetDurableExecutionResponse( + durable_execution_arn="arn:execution", + durable_execution_name="execution", + function_arn="arn:function", + status="FAILED", + start_timestamp=timestamp, + error=caller_error, + ) + history = GetDurableExecutionHistoryResponse( + events=[ + Event( + event_type="CallbackStarted", + event_timestamp=timestamp, + event_id=1, + operation_id="callback", + name="callback", + callback_started_details=CallbackStartedDetails( + callback_id="callback-id" + ), + ), + Event( + event_type="CallbackFailed", + event_timestamp=timestamp, + event_id=2, + operation_id="callback", + name="callback", + callback_failed_details=CallbackFailedDetails( + error=EventError( + payload=ErrorObject.from_dict(payload), truncated=False + ) + ), + ), + ] + ) + if decode_wire: + history = GetDurableExecutionHistoryResponse.from_dict(history.to_dict()) + assert history.events[-1].callback_failed_details.error.to_dict() == { + "Payload": payload, + "Truncated": False, + } + result = DurableFunctionTestResult.from_execution_history(execution, history) + callback = result.get_callback("callback") + assert callback.status is OperationStatus.FAILED + assert (callback.error.to_dict() if callback.error is not None else None) == ( + payload if payload else None + ) + assert result.status is InvocationStatus.FAILED + assert result.error is caller_error + + def test_durable_function_test_result_from_execution_history(): """Test DurableFunctionTestResult.from_execution_history factory method.""" import datetime diff --git a/packages/aws-durable-execution-sdk-python/README.md b/packages/aws-durable-execution-sdk-python/README.md index 4b2463e52..a2b9486ea 100644 --- a/packages/aws-durable-execution-sdk-python/README.md +++ b/packages/aws-durable-execution-sdk-python/README.md @@ -122,6 +122,14 @@ checkpointing is available, stops and waits for checkpoint processing, then call End before returning the outcome. The caller shuts down the handler executor. There is no separate handler-context hook or context-manager plugin API. +After all Start hooks, visible terminal operations named by the invocation +input's `UpdatedOperationIds` receive their existing `on_operation_end` callback. +This includes descendants that a cached completed child will not revisit. +Input delivery and later checkpoint/replay observations share the per-invocation +notification ledger; an operation/status is not reported twice in that invocation. +Missing operation metadata is not fabricated, and failed invocation delivery can +be redelivered on a later attempt. This is not an exporter acknowledgment. + When plugins are registered, the worker begins with a copy of the caller's context-variable bindings. Successful Start bindings flow into later Start hooks, the handler, serialization and resource cleanup. A later successful Start may diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py index e59820007..71d07605e 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py @@ -190,9 +190,9 @@ def durable_execution( logger.debug("Starting durable execution handler...") configured_plugins = load_configured_plugins(plugins) - plugin_executor = PluginExecutor(configured_plugins) def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: + plugin_executor = PluginExecutor(configured_plugins) invocation_input: DurableExecutionInvocationInput service_client: DurableServiceClient @@ -464,9 +464,20 @@ def wait_for_checkpoint() -> None: updated_operation_ids=invocation_input.updated_operation_ids, ) # No checkpoint work starts until all Start hooks finish. - checkpoint_future = executor.submit( - execution_state.checkpoint_batches_forever - ) + if configured_plugins: + checkpoint_context = ( + plugin_executor._run_in_invocation_context( + contextvars.copy_context + ) + ) + checkpoint_future = executor.submit( + checkpoint_context.run, + execution_state.checkpoint_batches_forever, + ) + else: + checkpoint_future = executor.submit( + execution_state.checkpoint_batches_forever + ) output = plugin_executor._run_in_invocation_context( invoke_handler ) diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py index 856fc68c6..195588296 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py @@ -9,6 +9,7 @@ from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from enum import Enum +from threading import Lock from typing import Any, Callable, cast from aws_durable_execution_sdk_python.identifier import OperationIdentifier @@ -415,9 +416,11 @@ def on_operation_start(self, info: OperationStartInfo) -> None: def on_operation_end(self, info: OperationEndInfo) -> None: """ - Called when an operation reaches a terminal status. Terminal operations - are not emitted again during replay. Child contexts without a terminal - checkpoint may emit this from the thread that runs the operation. + Called when an operation reaches a terminal status. Visible terminal + input updates are reported after all invocation-start hooks, even if + replay skips their completed child context. Duplicate deliveries of an + operation/status within one invocation are suppressed. Child contexts + without a terminal checkpoint may emit this from their operation thread. Args: info: Information about the operation. @@ -470,9 +473,13 @@ def __init__(self, plugins: list[DurableInstrumentationPlugin] | None): # work use its pre-hook snapshot. End still uses each hook's token owner. self._startup_context: contextvars.Context | None = None self._invocation_contexts: list[contextvars.Context | None] = [] + self._reported_terminal_updates: set[tuple[str, OperationStatus]] = set() + self._terminal_updates_lock = Lock() @contextlib.contextmanager def run(self): + with self._terminal_updates_lock: + self._reported_terminal_updates.clear() if self._plugins: self._executor = ThreadPoolExecutor( max_workers=1, @@ -488,6 +495,8 @@ def run(self): # Shut down the thread pool, waiting for pending tasks to complete. if self._executor: self._executor.shutdown(wait=True) + with self._terminal_updates_lock: + self._reported_terminal_updates.clear() @staticmethod def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> bool: @@ -553,7 +562,9 @@ def execute_plugins(self, info, sync): self._dispatch_plugin(plugin, info) else: # this is called asynchronously, so plugins cannot manipulate thread local objects - self._executor.submit(self._dispatch_plugin, plugin, info) + self._executor.submit( + contextvars.copy_context().run, self._dispatch_plugin, plugin, info + ) def _run_in_invocation_context(self, invoke: Callable[[], Any]) -> Any: """Continue in the pre-hook Context only after a failed Start hook.""" @@ -564,6 +575,8 @@ def _run_in_invocation_context(self, invoke: Callable[[], Any]) -> Any: def _snapshot_operation_infos( self, operations_provider: Callable[[], Mapping[str, Operation]] | None, + *, + raw_snapshot: dict[str, Operation] | None = None, ) -> dict[str, OperationInfo]: """Build the plugin ``OperationInfo`` view of the current operation map. @@ -584,7 +597,11 @@ def _snapshot_operation_infos( if not self._plugins or operations_provider is None: return {} try: - return _to_operation_info_map(operations_provider()) + source = dict(operations_provider()) + snapshot = _to_operation_info_map(source) + if raw_snapshot is not None: + raw_snapshot.update(source) + return snapshot except Exception: # A plugin-facing view must never break the execution. logger.exception("Failed to snapshot operations for plugin hook") @@ -615,7 +632,10 @@ def on_invocation_start( """ aws_request_id = lambda_context.aws_request_id if lambda_context else None self._operations_provider = operations_provider if self._plugins else None - operations = self._snapshot_operation_infos(operations_provider) + raw_operations: dict[str, Operation] = {} + operations = self._snapshot_operation_infos( + operations_provider, raw_snapshot=raw_operations + ) self._invocation_status = InvocationStartInfo( execution_arn=execution_arn, request_id=aws_request_id, @@ -629,7 +649,38 @@ def on_invocation_start( if operation_id in operations }, ) + # Capture delivery data before exposing the mutable snapshot maps to + # Start hooks. A cached child may never visit its updated descendant. + input_ends = [ + OperationEndInfo( + operation_id=operation.operation_id, + operation_type=_to_operation_type(operation.operation_type), + sub_type=operation.sub_type, + name=operation.name, + parent_id=operation.parent_id, + start_time=operation.start_timestamp, + end_time=operation.end_timestamp, + result=_extract_result(operation), + error=self._extract_error(operation), + attempt=( + operation.step_details.attempt if operation.step_details else None + ), + is_replayed=False, + status=operation.status, + ) + for operation_id in (updated_operation_ids or []) + if operation_id in raw_operations + for operation in (raw_operations[operation_id],) + if self._is_terminal_status(operation.status) + ] self.execute_plugins(self._invocation_status, sync=True) + if input_ends: + + def report_input_updates() -> None: + for operation_end in input_ends: + self._report_terminal_update(operation_end) + + self._run_in_invocation_context(report_input_updates) def _snapshot_execution_input(self, execution_input: Any) -> Any: """Deep-copy the execution input so the plugin view is isolated. @@ -815,7 +866,10 @@ def on_operation_update( ) for operation in updated_operations: if self._is_terminal_status(operation.status): - self.execute_plugins( + # Replay delivery and checkpoint responses can report the same + # completion. Deduplicate actual notifications, not state that + # may have arrived without an UpdatedOperationIds notification. + self._report_terminal_update( OperationEndInfo( operation_id=operation.operation_id, operation_type=_to_operation_type(operation.operation_type), @@ -834,7 +888,6 @@ def on_operation_update( ), is_replayed=False, ), - sync=True, ) if ( @@ -865,6 +918,16 @@ def on_operation_update( sync=True, ) + def _report_terminal_update(self, info: OperationEndInfo) -> None: + """Claim actual input/checkpoint delivery once in this invocation.""" + if self._plugins: + key = (info.operation_id, info.status) + with self._terminal_updates_lock: + if key in self._reported_terminal_updates: + return + self._reported_terminal_updates.add(key) + self.execute_plugins(info, sync=True) + @staticmethod def _extract_error(operation: Operation): return _extract_error(operation) diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_overlap_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_overlap_int_test.py new file mode 100644 index 000000000..6073b9fce --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_overlap_int_test.py @@ -0,0 +1,65 @@ +"""Two local executions may overlap without sharing invocation-hook state.""" + +import contextvars +import json +import threading + +from aws_durable_execution_sdk_python import durable_execution +from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner + + +def test_shared_handler_reports_both_invocation_ends(monkeypatch): + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + ready = threading.Event() + release = threading.Event() + binding: contextvars.ContextVar[str | None] = contextvars.ContextVar( + "integration-plugin-owner", default=None + ) + tokens: contextvars.ContextVar[contextvars.Token[str | None] | None] = ( + contextvars.ContextVar("integration-plugin-token", default=None) + ) + starts = [] + ends = [] + + class SafePlugin(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + tokens.set(binding.set(info.execution_arn)) + starts.append(info.execution_arn) + + def on_invocation_end(self, info): + assert binding.get() == info.execution_arn + token = tokens.get() + assert token is not None + binding.reset(token) + ends.append(info.execution_arn) + + def body(event, _context): + if event["label"] == "a": + ready.set() + assert release.wait(5) + return event["label"] + + wrapped = durable_execution(body, plugins=[SafePlugin()]) + try: + with ( + DurableFunctionTestRunner(handler=wrapped, poll_interval=0.01) as first, + DurableFunctionTestRunner(handler=wrapped, poll_interval=0.01) as second, + ): + first_arn = first.run_async(input='{"label":"a"}') + assert ready.wait(5) + try: + second_result = second.run(input='{"label":"b"}') + finally: + release.set() + first_result = first.wait_for_result(first_arn, timeout=5) + assert first_result.status.value == second_result.status.value == "SUCCEEDED" + assert [json.loads(first_result.result), json.loads(second_result.result)] == [ + "a", + "b", + ] + assert len(starts) == 2 and len(set(starts)) == 2 + assert ends == list(reversed(starts)) + assert binding.get() is None + finally: + release.set() diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_operation_type_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_operation_type_int_test.py index f4fca018c..c1ee85d63 100644 --- a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_operation_type_int_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_operation_type_int_test.py @@ -120,10 +120,16 @@ def handler(_event: Any, context: DurableContext) -> str: assert {info.sub_type for info in plugin.operation_starts} == set(expected) assert {info.sub_type for info in plugin.operation_ends} == set(expected) - # The final replay starts with every earlier operation checkpointed, - # including the execution root, whose service record has no subtype. + # Live hooks above cover every category. The service prunes descendants of + # the completed callback/map/parallel contexts from the replay input; the + # surviving records, including the subtype-less execution root, keep type. replay_operations = plugin.invocation_starts[-1].operations.values() - assert {info.sub_type for info in replay_operations} == {*expected, None} + retained_subtypes = set(expected) - { + OperationSubType.CALLBACK, + OperationSubType.MAP_ITERATION, + OperationSubType.PARALLEL_BRANCH, + } + assert {info.sub_type for info in replay_operations} == {*retained_subtypes, None} all_infos: list[OperationInfo] = [ *plugin.operation_starts, *plugin.operation_ends, diff --git a/packages/aws-durable-execution-sdk-python/tests/execution_test.py b/packages/aws-durable-execution-sdk-python/tests/execution_test.py index 85ea72a97..3a6c9c1c1 100644 --- a/packages/aws-durable-execution-sdk-python/tests/execution_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/execution_test.py @@ -4,7 +4,11 @@ import json import time import warnings +import contextvars +import threading from collections.abc import Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import replace from copy import deepcopy from typing import Any from unittest.mock import Mock, patch @@ -3503,6 +3507,75 @@ def test_handler(event: Any, context: DurableContext) -> dict: assert callable(test_handler) +def test_overlapping_handler_calls_keep_each_plugin_end_and_payload(): + """One wrapper must not let another call replace its invocation-hook state.""" + entered = threading.Event() + release = threading.Event() + binding: contextvars.ContextVar[str | None] = contextvars.ContextVar( + "unit-plugin-owner", default=None + ) + tokens: contextvars.ContextVar[contextvars.Token[str | None] | None] = ( + contextvars.ContextVar("unit-plugin-token", default=None) + ) + starts = [] + ends = [] + + class SafePlugin(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + tokens.set(binding.set(info.request_id)) + starts.append(info.request_id) + + def on_invocation_end(self, info): + assert binding.get() == info.request_id + token = tokens.get() + assert token is not None + binding.reset(token) + ends.append( + ( + info.request_id, + info.execution_input["label"], + json.loads(info.execution_result), + ) + ) + + def body(event, _context): + if event["label"] == "a": + entered.set() + assert release.wait(5) + return event["label"] + + wrapped = durable_execution(body, plugins=[SafePlugin()]) + client = Mock(spec=DurableServiceClient) + client.checkpoint.return_value = CheckpointOutput( + checkpoint_token="new_token", + new_execution_state=CheckpointUpdatedExecutionState(), + ) + + def call(label): + incoming = replace( + _make_invocation_input(client, input_payload=json.dumps({"label": label})), + durable_execution_arn=f"arn:test:execution-{label}/exec1", + ) + context = _make_lambda_context() + context.aws_request_id = label + return wrapped(incoming, context) + + with ThreadPoolExecutor(max_workers=2) as callers: + first = callers.submit(call, "a") + try: + assert entered.wait(5) + second = callers.submit(call, "b") + second_result = second.result(timeout=5) + finally: + release.set() + first_result = first.result(timeout=5) + assert first_result == {"Status": "SUCCEEDED", "Result": '"a"'} + assert second_result == {"Status": "SUCCEEDED", "Result": '"b"'} + assert starts == ["a", "b"] + assert ends == [("b", "b", "b"), ("a", "a", "a")] + assert binding.get() is None + + def test_durable_execution_with_plugins_success(): """Test that plugins receive invocation start/end and execution end on success.""" mock_client = Mock(spec=DurableServiceClient) diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_terminal_input_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_terminal_input_test.py new file mode 100644 index 000000000..363a984c4 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_terminal_input_test.py @@ -0,0 +1,132 @@ +"""Visible terminal input updates notify plugins before cached user paths skip them.""" + +import contextvars +from datetime import UTC, datetime + +from aws_durable_execution_sdk_python.execution import ( + DurableExecutionInvocationOutput, + InvocationStatus, +) +from aws_durable_execution_sdk_python.lambda_service import ( + CallbackDetails, + ErrorObject, + Operation, + OperationStatus, + OperationType, +) +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + PluginExecutor, +) + + +def _callback(operation_id, status=OperationStatus.SUCCEEDED, error=None): + return Operation( + operation_id=operation_id, + operation_type=OperationType.CALLBACK, + status=status, + name=operation_id, + callback_details=CallbackDetails( + callback_id=operation_id, result="saved", error=error + ), + ) + + +def _start(executor, operations, updated_ids): + executor.on_invocation_start( + execution_arn="execution/input-delivery", + is_first_invocation=False, + execution_start_time=datetime.now(UTC), + lambda_context=None, + operations_provider=lambda: operations, + updated_operation_ids=updated_ids, + ) + + +def test_input_delivery_follows_all_starts_and_shares_actual_notification_ledger(): + marker = contextvars.ContextVar("input-owner", default="host") + events = [] + + class Observe(DurableInstrumentationPlugin): + def __init__(self, name): + self.name = name + + def on_invocation_start(self, info): + marker.set(self.name) + events.append(("start", self.name)) + + def on_operation_end(self, info): + events.append( + ("end", self.name, info.operation_id, marker.get(), info.result) + ) + + b, active = _callback("b"), _callback("active", OperationStatus.STARTED) + operations = {"b": b, "active": active} + executor = PluginExecutor([Observe("first"), Observe("second")]) + + def invoke(): + with executor.run(): + _start(executor, operations, ["b", "active", "not-visible"]) + assert events == [ + ("start", "first"), + ("start", "second"), + ("end", "first", "b", "second", "saved"), + ("end", "second", "b", "second", "saved"), + ] + executor.on_operation_update(b) + assert len(events) == 4 + # Missing input metadata did not claim an event that arrives later. + executor.on_operation_update(_callback("not-visible")) + assert [event[2] for event in events if event[0] == "end"] == [ + "b", + "b", + "not-visible", + "not-visible", + ] + executor.on_invocation_end( + DurableExecutionInvocationOutput(status=InvocationStatus.PENDING) + ) + + contextvars.Context().run(invoke) + assert marker.get() == "host" + events.clear() + # A new invocation can receive the same explicit update after redelivery; + # the ledger is not a cross-invocation exporter acknowledgment. + contextvars.Context().run(invoke) + + +def test_failed_start_and_snapshot_mutation_cannot_lose_input_notification(): + marker = contextvars.ContextVar("input-failed-start", default="host") + error = ErrorObject( + message="saved error", type="TestError", data=None, stack_trace=["original"] + ) + operation = _callback("b", OperationStatus.FAILED, error) + seen = [] + + class MutatingFailure(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + marker.set("poison") + info.updated_operations["b"].error.stack_trace.append("mutated") + info.updated_operations.clear() + raise RuntimeError("intentional Start failure") + + class Observe(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + assert marker.get() == "host" + marker.set("survivor") + + def on_operation_end(self, info): + seen.append((info.operation_id, marker.get(), list(info.error.stack_trace))) + + def invoke(): + executor = PluginExecutor([MutatingFailure(), Observe()]) + with executor.run(): + _start(executor, {"b": operation}, ["b"]) + executor.on_invocation_end( + DurableExecutionInvocationOutput(status=InvocationStatus.PENDING) + ) + + contextvars.Context().run(invoke) + assert seen == [("b", "survivor", ["original"])] + assert operation.callback_details.error.stack_trace == ["original"] + assert marker.get() == "host" diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py index 70d266e3a..e1745b9dc 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -2,6 +2,7 @@ import logging import pickle import unittest +from concurrent.futures import ThreadPoolExecutor from copy import deepcopy from dataclasses import asdict, fields from unittest.mock import MagicMock, patch @@ -1539,6 +1540,65 @@ def test_terminal_status_without_step_details_fires_operation_only(self): self.assertIn("operation_end:op-1", self.plugin.calls) + def test_checkpoint_does_not_repeat_an_observed_terminal_update(self): + """A resumed external result can also appear in the next checkpoint.""" + for status in ( + OperationStatus.SUCCEEDED, + OperationStatus.FAILED, + OperationStatus.CANCELLED, + OperationStatus.TIMED_OUT, + OperationStatus.STOPPED, + ): + with self.subTest(status=status): + plugin = _TrackingPlugin() + executor = PluginExecutor(plugins=[plugin]) + operation = self._make_operation(status=status) + with executor.run(): + # First completion delivered through UpdatedOperationIds. + executor.on_operation_update(operation) + # The next response carries the already observed state. + executor.on_operation_update( + [operation], + operations={operation.operation_id: operation}, + previous_operations={operation.operation_id: operation}, + ) + self.assertEqual(plugin.calls.count("operation_end:op-1"), 1) + + def test_checkpoint_terminal_transition_still_emits(self): + previous = self._make_operation(status=OperationStatus.STARTED) + operation = self._make_operation(status=OperationStatus.SUCCEEDED) + with self.executor.run(): + self.executor.on_operation_update( + [operation], + operations={operation.operation_id: operation}, + previous_operations={previous.operation_id: previous}, + ) + self.assertEqual(self.plugin.calls.count("operation_end:op-1"), 1) + + def test_checkpoint_preserves_first_terminal_notification(self): + """State can predate notification, e.g. without UpdatedOperationIds.""" + operation = self._make_operation(status=OperationStatus.SUCCEEDED) + with self.executor.run(): + self.executor.on_operation_update( + [operation], + operations={operation.operation_id: operation}, + previous_operations={operation.operation_id: operation}, + ) + self.assertEqual(self.plugin.calls.count("operation_end:op-1"), 1) + + def test_terminal_notification_resets_between_invocations(self): + operation = self._make_operation(status=OperationStatus.SUCCEEDED) + for _ in range(2): + with self.executor.run(): + self.executor.on_operation_update(operation) + self.assertEqual(self.plugin.calls.count("operation_end:op-1"), 2) + + def test_concurrent_completion_notifications_emit_once(self): + operation = self._make_operation(status=OperationStatus.SUCCEEDED) + with self.executor.run(), ThreadPoolExecutor(max_workers=8) as workers: + list(workers.map(self.executor.on_operation_update, [operation] * 16)) + self.assertEqual(self.plugin.calls.count("operation_end:op-1"), 1) + def test_non_terminal_status_without_step_details_fires_nothing(self): op = self._make_operation(status=OperationStatus.STARTED, step_details=None)