From ae272004a3a122b110b272b579684d0ad2b76894 Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Thu, 22 Oct 2026 16:20:17 -0700 Subject: [PATCH 1/3] fix(otel): preserve completion lifecycle Keep completion snapshots and export deduplication in each OTel invocation. Complete reentered replay segments, preserve missing-parent completions, and keep failed span starts retryable. Isolate plugin error snapshots and flush attempts by provider. Restore main-only CI filters; retain testing-wheel support downstream. Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- .../README.md | 97 +++- .../_flush.py | 128 +++++ .../_invocation_state.py | 45 ++ .../execution_plugin.py | 259 +++++++++- .../invocation_plugin.py | 319 +++++++++++- .../e2e/test_completion_export_retry_int.py | 159 ++++++ .../tests/e2e/test_completion_snapshot_int.py | 194 +++++++ .../e2e/test_context_segment_reentry_int.py | 162 ++++++ .../e2e/test_exception_event_bounds_int.py | 84 +++ .../tests/e2e/test_flush_lifecycle_int.py | 227 +++++++++ .../e2e/test_invocation_wait_resume_int.py | 127 +++++ .../tests/e2e/test_parent_registration_int.py | 190 +++++++ .../e2e/test_shared_invocation_state_int.py | 115 +++++ .../tests/test_completion_export_retry.py | 477 +++++++++++++++++ .../tests/test_completion_replay_parent.py | 132 +++++ .../tests/test_completion_snapshots.py | 215 ++++++++ .../tests/test_context_segment_reentry.py | 250 +++++++++ .../tests/test_execution_plugin.py | 22 +- .../tests/test_flush.py | 317 ++++++++++++ .../tests/test_invocation_owned_state.py | 128 +++++ .../tests/test_invocation_plugin.py | 165 +++++- .../tests/test_status_mapping.py | 44 +- .../tests/test_terminal_input_parent_refs.py | 479 ++++++++++++++++++ .../README.md | 8 + .../plugin.py | 30 +- .../e2e/plugin_invocation_overlap_int_test.py | 69 +++ .../tests/execution_test.py | 76 +++ .../tests/plugin_test.py | 178 +++++++ 28 files changed, 4616 insertions(+), 80 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_flush.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_invocation_state.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_completion_export_retry_int.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_completion_snapshot_int.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_context_segment_reentry_int.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_exception_event_bounds_int.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_flush_lifecycle_int.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_parent_registration_int.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_shared_invocation_state_int.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/test_completion_export_retry.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/test_completion_replay_parent.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/test_completion_snapshots.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/test_context_segment_reentry.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/test_flush.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_owned_state.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/test_terminal_input_parent_refs.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_overlap_int_test.py diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 6693855c3..b1d34cc2e 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -405,13 +405,102 @@ 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 terminal completion arrives before replay creates its parent span, the +plugin retains it until that span is registered. Parent registration exports the +pending child inside the actual parent segment. Completions whose parents remain +unavailable are exported at invocation End: Invocation view uses the invocation +span, and Execution view derives a logical parent context from the operation ID. +Each fallback is logged. A derived context is not a reconstructed recording span +or proof of an unavailable virtual context's original span identity. + +### 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 flush per provider and no pending I/O queue. +Distinct providers can flush independently. Concurrent requests for the same +provider wait for admission within the same absolute deadline used for completion, +then perform a fresh flush so spans added during the earlier attempt are covered. +A timeout does not cancel the exporter or free that provider's 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 +Factories create a separate plugin instance for each invocation. Each instance +owns its trace IDs, span registry, deferred completions, timestamps and scope +tokens, so one invocation's cleanup cannot clear another's trace. Factories share +provider configuration across instances. SDK user-code workers capture the +invocation Context; Start/End hooks and token reset retain their owner and order. + +Each OTel plugin retains visible terminal operations from +`InvocationStartInfo.updated_operations` and checkpoint `on_operation_change` +snapshots. Natural `on_operation_end` hooks export those operations and remove +them from the pending dictionary. Pending children are also indexed by parent, +so replay registration can export them under the actual parent before it closes. +The plugin exports any unvisited completions before ending the invocation span. +An invocation-local set suppresses repeated input, replay and checkpoint exports; +core does not synthesize additional completion hooks for other plugins. + +A completion enters that set only after its span ends successfully. An in-progress +claim prevents concurrent attempts, and a failed attempt leaves the completion +pending for a later notification or invocation-End retry. If a processor delivers +a span and then raises, retrying can repeat that delivery; the plugin has no +exporter acknowledgment that would distinguish partial delivery from failure. + +Invocation view can open another context segment within the same invocation +when a timed branch resumes and re-enters a completed child, or when a virtual +context runs again. Each new segment receives its own terminal outcome. Repeated +durable checkpoint observations still do not close or re-export that segment. + +The service omits descendants of completed contexts unless their ancestors retain +children through `ReplayChildren`. An `UpdatedOperationIds` entry alone does not +provide missing operation metadata, and the plugin does not synthesize an +operation span from that missing record. + +Known completed context metadata is retained from the full visible input snapshot +for descendants delivered by later checkpoint responses. Both views retain actual +completed parent contexts until invocation cleanup. Invocation view prefers its +current replay segment, including a large-result ReplayChildren parent; only at +End can a completion without such a segment use a known completed input parent's +identity. Remaining unknown parents use the logged fallback described above. +Queue draining does not re-enter public hooks. + +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 3.x, invocation hooks and the handler run on one worker in an invocation-local Context initialized from the host's bindings. Start bindings remain visible to later hooks, the handler and resource cleanup even when Start 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..87fbe9379 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_flush.py @@ -0,0 +1,128 @@ +"""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 dataclasses import dataclass, field + +from opentelemetry.trace import TracerProvider + + +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(threading.Lock()) + + +@dataclass +class _Attempt: + done: threading.Event = field(default_factory=threading.Event) + result: object = None + error: BaseException | None = None + timed_out: bool = False + + +# Identity keys support custom providers with value equality or no hash. The +# worker retains its provider until removal, so an active key cannot be reused. +_active: dict[int, _Attempt] = {} + + +def _run(provider: TracerProvider, attempt: _Attempt) -> None: + provider_id = id(provider) + try: + try: + # force_flush is an optional extension to the API provider. The + # plugin checks for it before invoking this helper. + attempt.result = getattr(provider, "force_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.get(provider_id) is attempt: + del _active[provider_id] + _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.get(provider_id) is attempt: + del _active[provider_id] + _lock.notify_all() + + +def flush_provider(provider: TracerProvider) -> bool: + """Wait for provider I/O, with one outstanding attempt per provider. + + 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. + Each admitted caller starts a new flush after prior I/O completes, since spans + queued during a prior attempt may not have been included in that attempt. + This does not bound synchronous span processors' earlier ``on_end`` calls, + application-owned provider shutdown, or exporter transport I/O itself. + """ + deadline = time.monotonic() + _FLUSH_TIMEOUT_SECONDS + provider_id = id(provider) + with _lock: + while provider_id in _active: + remaining = deadline - time.monotonic() + if remaining <= 0: + break + _lock.wait(remaining) + if provider_id not in _active and time.monotonic() < deadline: + attempt = _Attempt() + _active[provider_id] = 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, provider, attempt), + name="durable-otel-flush", + daemon=True, + ).start() + except BaseException: + with _lock: + if _active.get(provider_id) is attempt: + del _active[provider_id] + _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..52b46e231 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_invocation_state.py @@ -0,0 +1,45 @@ +"""Visible completed context metadata for terminal-update parent references.""" + +from collections.abc import Mapping + +from aws_durable_execution_sdk_python.plugin import ( + InvocationStartInfo, + OperationInfo, + OperationStatus, + OperationType, +) + + +def terminal_updates( + operations: Mapping[str, OperationInfo], +) -> dict[str, OperationInfo]: + """Retain visible terminal updates until their operation hook or invocation End.""" + return { + operation_id: operation + for operation_id, operation in operations.items() + if operation.status + in ( + OperationStatus.SUCCEEDED, + OperationStatus.FAILED, + OperationStatus.TIMED_OUT, + OperationStatus.STOPPED, + OperationStatus.CANCELLED, + ) + } + + +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, + ) + ] 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 64f55caf3..57bc13657 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 @@ -34,6 +34,7 @@ import datetime import logging import threading +import time from typing import Any, TYPE_CHECKING, cast from aws_durable_execution_sdk_python.plugin import ( @@ -41,7 +42,9 @@ InvocationEndInfo, InvocationStatus, InvocationStartInfo, + OperationChangeInfo, OperationEndInfo, + OperationInfo, OperationStartInfo, OperationStatus, OperationType, @@ -52,6 +55,7 @@ from opentelemetry import context as otel_context from opentelemetry import trace from opentelemetry.context import Context +from opentelemetry.sdk.trace import Span as SdkSpan from opentelemetry.sdk.trace import Tracer as SdkTracer from opentelemetry.sdk.trace.sampling import Sampler from opentelemetry.trace import ( @@ -64,12 +68,17 @@ TraceState, ) +from aws_durable_execution_sdk_python_otel._invocation_state import ( + completed_input_parents, + terminal_updates, +) 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, @@ -192,11 +201,16 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: # 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() + self._ending_operation_ids: set[str] = set() + self._exception_recorded_spans: dict[str, Span] = {} + self._completion_end_timestamps: dict[str, int] = {} # 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._pending_completions: dict[str, OperationInfo] = {} + self._pending_operation_ends: dict[str, dict[str, OperationInfo]] = {} self._lock = threading.RLock() self._tracing_enabled = False @@ -232,6 +246,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.values(): + self._render_operation_end(info) def _get_span(self, key: str | None) -> Span | None: if key is None: @@ -593,6 +610,11 @@ 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) + self._remember_completions(info.updated_operations) # Make the Workflow span the active span so auto-instrumented spans # created during the invocation become its children. The token is @@ -736,6 +758,52 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: def _end_invocation_spans(self, info: InvocationEndInfo) -> None: """End this invocation's open spans and export the Workflow span.""" + with self._lock: + pending = tuple(self._pending_completions.values()) + for operation in pending: + self._render_operation_end(operation) + # No more parent scopes will start after invocation End. Preserve the + # completion under a derived logical parent when its actual span is + # unavailable, without manufacturing a recording parent span. + while True: + with self._lock: + if not self._pending_operation_ends: + break + # Materialize pending ancestors before their descendants. A + # cycle has no outermost entry, so break it with one fallback. + pending_ids = { + operation_id + for group in self._pending_operation_ends.values() + for operation_id in group + } + parent_id = next( + ( + key + for key in self._pending_operation_ends + if key not in pending_ids + ), + next(iter(self._pending_operation_ends)), + ) + operations = self._pending_operation_ends.pop(parent_id) + parent_context = self._operation_span_context(parent_id) + parent = ( + self._get_span(parent_id) + or (DurableParentSpan(parent_context) if parent_context else None) + or self._workflow_span + ) + if parent is None: + raise ValueError( + "No execution trace context available for completion fallback" + ) + for operation in operations.values(): + logger.warning( + "Using fallback parent context for operation %s: parent %s " + "has no available span", + operation.operation_id, + parent_id, + ) + self._render_operation_end(operation, fallback_parent=parent) + # 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 @@ -782,7 +850,7 @@ def _force_flush(self) -> None: if not hasattr(self._provider, "force_flush"): return try: - self._provider.force_flush() + flush_provider(self._provider) except Exception: # noqa: BLE001 logger.exception("force_flush failed at invocation end") @@ -812,6 +880,12 @@ def _release_invocation_scope(self) -> None: self._detach_remaining_contexts() unbind_invocation(self) self._tracing_enabled = False + with self._lock: + self._pending_operation_ends.clear() + self._pending_completions.clear() + self._ending_operation_ids.clear() + self._exception_recorded_spans.clear() + self._completion_end_timestamps.clear() # ------------------------------------------------------------------ # Operation lifecycle @@ -837,6 +911,31 @@ 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 on_operation_change(self, info: OperationChangeInfo) -> None: + """Remember checkpoint completions not already exported by their End hook.""" + if not self._tracing_enabled: + return + self._remember_completions(info.updated_operations) + + def _remember_completions(self, updates: dict[str, OperationInfo]) -> None: + with self._lock: + for operation_id, operation in terminal_updates(updates).items(): + if operation_id in self._ended_operation_ids: + continue + self._pending_completions[operation_id] = operation + if operation.parent_id is not None: + # Index by parent so replay registration drains children + # inside that actual segment, without scanning other input. + self._pending_operation_ends.setdefault(operation.parent_id, {})[ + operation_id + ] = operation + + def _render_operation_end( + self, info: OperationInfo, *, fallback_parent: Span | None = None + ) -> 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 @@ -845,16 +944,51 @@ def on_operation_end(self, info: OperationEndInfo) -> None: # replay-only and must not re-export the deterministic logical span. if info.is_replayed: return - # Export the span only on the first end for this operation. with self._lock: - if info.operation_id in self._ended_operation_ids: + if ( + info.operation_id in self._ended_operation_ids + or info.operation_id in self._ending_operation_ids + ): + return + # Retain the delivery until materialization and span.end both + # succeed. A later notification or invocation End can retry it. + self._pending_completions[info.operation_id] = info + if ( + info.parent_id is not None + and info.parent_id not in self._operation_spans + and fallback_parent is None + ): + self._pending_operation_ends.setdefault(info.parent_id, {})[ + info.operation_id + ] = info 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. + self._ending_operation_ids.add(info.operation_id) + try: + self._finish_operation_end(info, fallback_parent) + with self._lock: + self._ended_operation_ids.add(info.operation_id) + self._exception_recorded_spans.pop(info.operation_id, None) + self._completion_end_timestamps.pop(info.operation_id, None) + self._pending_completions.pop(info.operation_id, None) + if info.parent_id is not None: + siblings = self._pending_operation_ends.get(info.parent_id) + if siblings is not None: + siblings.pop(info.operation_id, None) + if not siblings: + del self._pending_operation_ends[info.parent_id] + finally: + with self._lock: + self._ending_operation_ids.discard(info.operation_id) + + def _finish_operation_end( + self, info: OperationInfo, fallback_parent: Span | None + ) -> None: + """Materialize and end one claimed completion without losing retry state.""" + # 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) + parent = fallback_parent or self._resolve_parent(info.parent_id) start_time = info.start_time end_time = info.end_time if isinstance(placeholder, DurableParentSpan): @@ -869,10 +1003,13 @@ def on_operation_end(self, info: OperationEndInfo) -> None: # the tracer to choose a later wall-clock start. start_time = end_time end_time = ensure_end_after_start(start_time, end_time) + end_timestamp = ( + placeholder.end_time + if isinstance(placeholder, SdkSpan) and placeholder.end_time is not None + else _to_otel_timestamp(end_time) + ) placeholder_span_id = ( - placeholder.get_span_context().span_id - if isinstance(placeholder, DurableParentSpan) - else None + placeholder.get_span_context().span_id if placeholder is not None else None ) checkpointless_context = ( info.operation_type is OperationType.CONTEXT and info.start_time is None @@ -880,27 +1017,93 @@ def on_operation_end(self, info: OperationEndInfo) -> None: span_id_override = placeholder_span_id if span_id_override is None and checkpointless_context: span_id_override = self._id_generator.generate_span_id() - span = self._start_span( - operation_id=info.operation_id, - name=info.name or info.operation_id, - info=info, - parent=parent, - start_time=start_time, - span_id_override=span_id_override, - ) + if isinstance(placeholder, SdkSpan) and placeholder.end_time is not None: + # OTel marks a span ended before running processors. Such a span + # cannot be ended again, so recreate its original context and data. + span = self._retry_ended_span(info.operation_id, placeholder) + elif placeholder is not None and placeholder.is_recording(): + # A failed end that did not reach the provider can retry the same + # recording span rather than replacing and prematurely ending it. + span = placeholder + else: + span = self._start_span( + operation_id=info.operation_id, + name=info.name or info.operation_id, + info=info, + parent=parent, + start_time=start_time, + span_id_override=span_id_override, + ) + + with self._lock: + # Preserve the first attempt's normalized end even if later spans + # finish before this completion is retried. + end_timestamp = self._completion_end_timestamps.setdefault( + info.operation_id, + end_timestamp if end_timestamp is not None else time.time_ns(), + ) 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") - ) + if self._exception_recorded_spans.get(info.operation_id) is not span: + span.record_exception( + Exception(info.error.message or info.error.type or "Unknown error"), + timestamp=end_timestamp, + ) + self._exception_recorded_spans[info.operation_id] = span + 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) - if popped is not None: - popped.end(end_time=_to_otel_timestamp(end_time)) + # Leave the actual parent context registered through every processor + # callback, including a failing one. Only success releases its span. + span.end(end_time=end_timestamp) + with self._lock: + if self._operation_spans.get(info.operation_id) is not span: + return + del self._operation_spans[info.operation_id] + if 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( + span.get_span_context(), start_time=start_time + ) + + def _retry_ended_span(self, operation_id: str, failed: SdkSpan) -> Span: + """Recreate a span whose processor failed after OTel marked it ended.""" + parent_context = self._with_sampling(Context()) + if failed.parent is not None: + parent_context = self._with_sampling( + trace.set_span_in_context(DurableParentSpan(failed.parent), Context()) + ) + failed_context = failed.get_span_context() + with self._id_generator.use_ids( + trace_id=failed_context.trace_id, span_id=failed_context.span_id + ): + span = self._tracer.start_span( + name=failed.name, + kind=failed.kind, + attributes=dict(failed.attributes or {}), + start_time=failed.start_time, + context=parent_context, + links=failed.links, + ) + for event in failed.events: + span.add_event( + event.name, attributes=event.attributes, timestamp=event.timestamp + ) + span.set_status(failed.status) + with self._lock: + self._operation_spans[operation_id] = span + if self._exception_recorded_spans.get(operation_id) is failed: + self._exception_recorded_spans[operation_id] = span + return span def _start_span( self, @@ -950,6 +1153,9 @@ def _start_span( links=links, ) self._operation_spans[key] = span + pending = self._pending_operation_ends.pop(key, {}) + for operation in pending.values(): + self._render_operation_end(operation) return span # ------------------------------------------------------------------ @@ -1032,7 +1238,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 a861c76b6..f333ea297 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 @@ -12,7 +12,9 @@ InvocationEndInfo, InvocationStatus, InvocationStartInfo, + OperationChangeInfo, OperationEndInfo, + OperationInfo, OperationStartInfo, OperationStatus, OperationType, @@ -22,6 +24,7 @@ ) from opentelemetry import context, trace from opentelemetry.context import Context +from opentelemetry.sdk.trace import Span as SdkSpan from opentelemetry.sdk.trace import Tracer as SdkTracer from opentelemetry.sdk.trace.sampling import Sampler from opentelemetry.trace import ( @@ -34,12 +37,17 @@ TraceState, ) +from aws_durable_execution_sdk_python_otel._invocation_state import ( + completed_input_parents, + terminal_updates, +) 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, @@ -195,6 +203,15 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: # 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._pending_completions: dict[str, OperationInfo] = {} + self._ended_operation_ids: set[str] = set() + self._ending_operation_ids: set[str] = set() + self._failed_operation_ids: set[str] = set() + self._exception_recorded_spans: dict[str, Span] = {} + self._completion_end_timestamps: dict[str, int] = {} + self._pending_operation_ends: dict[str, dict[str, OperationInfo]] = {} + self._input_parent_refs: dict[str, OperationInfo] = {} + self._completed_parent_spans: dict[str, Span] = {} self._operation_spans_lock = threading.RLock() self._tracing_enabled = False @@ -405,6 +422,13 @@ def _resolve_parent_span(self, parent_id: str | None = None) -> Span: raise ValueError("No parent span found") + def _completion_parent_span(self, parent_id: str | None) -> Span | None: + """Prefer the current segment, retaining its identity after context End.""" + with self._operation_spans_lock: + return self._operation_spans.get(parent_id) or ( + self._completed_parent_spans.get(parent_id) if parent_id else None + ) + def _invocation_parent_context(self) -> Context: """Return same-trace ambient context, else execution ancestor context.""" execution_trace_context = self._execution_trace_context @@ -553,14 +577,28 @@ 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.values(): + 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, + operation_completion: bool = False, ): """End and unregister the active span for an operation ID. @@ -569,14 +607,38 @@ 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. + operation_completion: Deduplicate the terminal exception event when + retrying a completion whose prior end failed. """ logger.debug("Ending OTel span: operation_id=%s", operation_id) with self._operation_spans_lock: - span = self._operation_spans.pop(operation_id, None) + span = self._operation_spans.get(operation_id) 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) + with self._operation_spans_lock: + end_time = ( + self._completion_end_timestamps.get(operation_id) + if operation_completion and operation_id is not None + else None + ) + if end_time is None: + end_time = self._next_ordered_timestamp(end_timestamp) + if operation_completion and operation_id is not None: + self._completion_end_timestamps[operation_id] = end_time + if exception is not None and ( + not operation_completion + or operation_id is None + or self._exception_recorded_spans.get(operation_id) is not span + ): + span.record_exception(exception, timestamp=end_time) + if operation_completion and operation_id is not None: + self._exception_recorded_spans[operation_id] = span span.end(end_time=end_time) + with self._operation_spans_lock: + if self._operation_spans.get(operation_id) is span: + del self._operation_spans[operation_id] logger.debug("Ended OTel span: %s", span) # ------------------------------------------------------------------ @@ -697,6 +759,12 @@ 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._input_parent_refs = { + parent.operation_id: parent for parent in completed_input_parents(info) + } + self._remember_completions(info.updated_operations) # Last, so that everything this hook makes current is accounted for. # Start and End now run in the same worker Context. Attach the @@ -811,6 +879,62 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: def _end_invocation_spans(self, info: InvocationEndInfo) -> None: """End this invocation's open spans and export the Workflow span.""" + with self._operation_spans_lock: + pending = tuple(self._pending_completions.values()) + for operation in pending: + self._render_operation_end(operation) + # 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. + pending_ends: list[OperationInfo] = [] + with self._operation_spans_lock: + for parent_id in tuple(self._pending_operation_ends): + parent = self._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).values() + ) + for pending_end in pending_ends: + self._render_operation_end(pending_end) + + invocation_span = self._get_span(None) + if invocation_span is None and self._pending_operation_ends: + raise ValueError("No invocation span available for completion fallback") + while True: + with self._operation_spans_lock: + if not self._pending_operation_ends: + break + # Materialize pending ancestors before their descendants. A + # cycle has no outermost entry, so break it with one fallback. + pending_ids = { + operation_id + for group in self._pending_operation_ends.values() + for operation_id in group + } + parent_id = next( + ( + key + for key in self._pending_operation_ends + if key not in pending_ids + ), + next(iter(self._pending_operation_ends)), + ) + operations = self._pending_operation_ends.pop(parent_id) + for operation in operations.values(): + logger.warning( + "Using invocation fallback parent for operation %s: parent %s " + "has no available span", + operation.operation_id, + parent_id, + ) + self._render_operation_end(operation, fallback_parent=invocation_span) + # Spans are registered parent-first, so close pending spans in reverse # order to keep every child contained within its parent. with self._operation_spans_lock: @@ -870,7 +994,7 @@ def _force_flush(self) -> None: if not hasattr(self._provider, "force_flush"): return try: - self._provider.force_flush() + flush_provider(self._provider) except Exception: # noqa: BLE001 logger.exception("force_flush failed at invocation end") @@ -900,6 +1024,15 @@ def _release_invocation_scope(self) -> None: self._detach_remaining_contexts() unbind_invocation(self) self._tracing_enabled = False + with self._operation_spans_lock: + self._pending_operation_ends.clear() + self._pending_completions.clear() + self._ending_operation_ids.clear() + self._failed_operation_ids.clear() + self._exception_recorded_spans.clear() + self._completion_end_timestamps.clear() + self._input_parent_refs.clear() + self._completed_parent_spans.clear() def on_operation_start(self, info: OperationStartInfo) -> None: """Called when an operation begins. Creates a span for the operation.""" @@ -939,15 +1072,110 @@ 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 on_operation_change(self, info: OperationChangeInfo) -> None: + """Remember checkpoint completions not already exported by their End hook.""" + if not self._tracing_enabled: + return + self._remember_completions(info.updated_operations) + + def _remember_completions(self, updates: dict[str, OperationInfo]) -> None: + with self._operation_spans_lock: + for operation_id, operation in terminal_updates(updates).items(): + if operation_id in self._ended_operation_ids: + continue + self._pending_completions[operation_id] = operation + if operation.parent_id is not None: + # Index by parent so replay registration drains children + # inside that actual segment, without scanning other input. + self._pending_operation_ends.setdefault(operation.parent_id, {})[ + operation_id + ] = operation + + def _render_operation_end( + self, info: OperationInfo, *, fallback_parent: Span | None = None + ) -> 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 - span = self._get_span(info.operation_id) + with self._operation_spans_lock: + if info.operation_id in self._ending_operation_ids: + return + span = self._operation_spans.get(info.operation_id) + # ReplayChildren and virtual contexts can open another segment in + # this invocation after their logical operation already completed. + # Their checkpointless End closes that segment; a repeated durable + # checkpoint must not close it while its user function is running. + ends_active_context = ( + info.operation_type is OperationType.CONTEXT + and info.start_time is None + and ( + info.is_replayed + or info.operation_id not in self._context_operation_replays + ) + and span is not None + and span.is_recording() + ) + if ( + info.operation_id in self._ended_operation_ids + and not ends_active_context + and self._pending_completions.get(info.operation_id) != info + ): + return + # An already-exported logical context may have a new segment whose + # finalization failed. Keep that exact End retryable, without letting + # an old checkpoint consume the segment in its place. + self._pending_completions[info.operation_id] = info + if ( + span is None + and info.parent_id is not None + and self._completion_parent_span(info.parent_id) is None + and fallback_parent is None + ): + self._pending_operation_ends.setdefault(info.parent_id, {})[ + info.operation_id + ] = info + return + self._ending_operation_ids.add(info.operation_id) + try: + self._finish_operation_end(info, span, fallback_parent) + with self._operation_spans_lock: + self._ended_operation_ids.add(info.operation_id) + self._failed_operation_ids.discard(info.operation_id) + self._pending_completions.pop(info.operation_id, None) + self._exception_recorded_spans.pop(info.operation_id, None) + self._completion_end_timestamps.pop(info.operation_id, None) + if info.parent_id is not None: + siblings = self._pending_operation_ends.get(info.parent_id) + if siblings is not None: + siblings.pop(info.operation_id, None) + if not siblings: + del self._pending_operation_ends[info.parent_id] + except BaseException: + with self._operation_spans_lock: + self._failed_operation_ids.add(info.operation_id) + raise + finally: + with self._operation_spans_lock: + self._ending_operation_ids.discard(info.operation_id) + + def _finish_operation_end( + self, info: OperationInfo, span: Span | None, fallback_parent: Span | None + ) -> None: + """Finish the claimed segment without allowing another End to overtake it.""" + if isinstance(span, SdkSpan) and span.end_time is not None: + span = self._retry_ended_span(info.operation_id, span) if span is 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) + parent_span = fallback_parent or self._completion_parent_span( + info.parent_id + ) + 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, @@ -960,15 +1188,62 @@ 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, + operation_completion=True, + ) + if info.operation_type is OperationType.CONTEXT: + with self._operation_spans_lock: + self._completed_parent_spans[info.operation_id] = DurableParentSpan( + span.get_span_context(), start_time=info.start_time + ) + + def _retry_ended_span(self, operation_id: str, failed: SdkSpan) -> Span: + """Recreate a span whose processor failed after OTel marked it ended.""" + parent_context = self._with_sampling(Context()) + if failed.parent is not None: + parent_context = self._with_sampling( + trace.set_span_in_context(DurableParentSpan(failed.parent), Context()) + ) + failed_context = failed.get_span_context() + with self._id_generator.use_ids( + trace_id=failed_context.trace_id, span_id=failed_context.span_id + ): + span = self._tracer.start_span( + name=failed.name, + kind=failed.kind, + attributes=dict(failed.attributes or {}), + start_time=failed.start_time, + context=parent_context, + links=failed.links, + ) + for event in failed.events: + span.add_event( + event.name, attributes=event.attributes, timestamp=event.timestamp + ) + span.set_status(failed.status) + with self._operation_spans_lock: + self._operation_spans[operation_id] = span + if self._exception_recorded_spans.get(operation_id) is failed: + self._exception_recorded_spans[operation_id] = span + return span def on_user_function_start(self, info: UserFunctionStartInfo) -> None: """Called when a context or step operation starts user code. @@ -993,6 +1268,19 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None: raise RuntimeError( "on_user_function_start should only be called for CONTEXT and STEP operations" ) + if info.operation_type is OperationType.CONTEXT: + with self._operation_spans_lock: + failed_completion = ( + self._pending_completions.get(info.operation_id) + if info.operation_id in self._failed_operation_ids + else None + ) + # A failed completed segment must not be mistaken for an incomplete + # segment to continue or replaced by a fresh replay segment. Only + # explicitly failed Ends are retried here; input snapshots still + # wait for their normal operation/parent hooks or invocation End. + if failed_completion is not None: + self._render_operation_end(failed_completion) if info.operation_type is OperationType.STEP: parent_span = self._get_span( info.operation_id @@ -1083,16 +1371,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) @@ -1100,7 +1387,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_completion_export_retry_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_completion_export_retry_int.py new file mode 100644 index 000000000..4aa692b36 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_completion_export_retry_int.py @@ -0,0 +1,159 @@ +"""Public handlers retry a failed terminal export before invocation cleanup.""" + +import json +import logging +from threading import Event +from typing import Any + +import pytest +from aws_durable_execution_sdk_python import ( + DurableContext, + StepContext, + durable_execution, +) +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationEndInfo, + InvocationStartInfo, + InvocationStatus, +) +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry.context import Context +from opentelemetry.sdk.trace import ReadableSpan, Span, 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 ( + ExecutionOtelPluginFactory, + InvocationOtelPluginFactory, + OtelPluginConfig, +) + + +class _FailTargetOnceProcessor(SimpleSpanProcessor): + def __init__( + self, exporter: InMemorySpanExporter, invocation_ending: Event + ) -> None: + super().__init__(exporter) + self.invocation_ending = invocation_ending + self.target_starts: list[Span] = [] + self.target_ends: list[tuple[bool, ReadableSpan]] = [] + self.failed_after_close = False + self.error = RuntimeError("target span processor failed before export") + + def on_start(self, span: Span, parent_context: Context | None = None) -> None: + if span.name == "target": + self.target_starts.append(span) + super().on_start(span, parent_context) + + def on_end(self, span: ReadableSpan) -> None: + if span.name == "target": + self.target_ends.append((self.invocation_ending.is_set(), span)) + if len(self.target_ends) == 1: + original = self.target_starts[0] + self.failed_after_close = ( + not original.is_recording() and original.end_time is not None + ) + raise self.error + super().on_end(span) + + +@pytest.mark.parametrize( + "factory_type", [ExecutionOtelPluginFactory, InvocationOtelPluginFactory] +) +def test_handler_retries_closed_terminal_span_during_invocation_end( + factory_type: type[ExecutionOtelPluginFactory] | type[InvocationOtelPluginFactory], + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + caplog.set_level(logging.ERROR, logger="aws_durable_execution_sdk_python.plugin") + provider = TracerProvider() + exporter = InMemorySpanExporter() + invocation_ending = Event() + processor = _FailTargetOnceProcessor(exporter, invocation_ending) + provider.add_span_processor(processor) + invocation_statuses: list[InvocationStatus] = [] + exports_before_drain: list[ReadableSpan] = [] + step_calls: list[str] = [] + + class Observe(DurableInstrumentationPlugin): + def on_invocation_end(self, info: InvocationEndInfo) -> None: + invocation_statuses.append(info.status) + exports_before_drain.extend( + span for span in exporter.get_finished_spans() if span.name == "target" + ) + invocation_ending.set() + + class ObserveFactory: + def create_plugin(self, info: InvocationStartInfo) -> Observe: + return Observe() + + factory = factory_type( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + + def saved(_context: StepContext) -> str: + step_calls.append("target") + return "saved result" + + def body(_event: Any, context: DurableContext) -> str: + return context.step(saved, name="target") + + handler = durable_execution(body, plugins=[ObserveFactory(), factory]) + try: + with DurableFunctionTestRunner( + handler=handler, poll_interval=0.01, execution_timeout=10 + ) as runner: + result = runner.run(input="{}") + + assert result.status.value == "SUCCEEDED" + assert result.result is not None + assert json.loads(result.result) == "saved result" + assert step_calls == ["target"] + assert invocation_statuses == [InvocationStatus.SUCCEEDED] + assert processor.failed_after_close + assert [during_end for during_end, _ in processor.target_ends] == [False, True] + assert exports_before_drain == [] + (failed_attempt, successful_attempt) = [ + span for _, span in processor.target_ends + ] + (target,) = [ + span for span in exporter.get_finished_spans() if span.name == "target" + ] + assert target is successful_attempt + assert target.context is not None + assert target.context == failed_attempt.context + assert target.parent == failed_attempt.parent + assert target.start_time == failed_attempt.start_time + assert target.end_time == failed_attempt.end_time + assert target.status.status_code is StatusCode.OK + assert target.attributes is not None + assert target.attributes["durable.operation.status"] == "SUCCEEDED" + assert target.attributes["durable.operation.type"] == "STEP" + assert target.events == () + parent_name = ( + "Workflow" if factory_type is ExecutionOtelPluginFactory else "Invocation" + ) + (parent,) = [ + span for span in exporter.get_finished_spans() if span.name == parent_name + ] + assert parent.context is not None + assert target.parent == parent.context + assert target.context.trace_id == parent.context.trace_id + failures = [ + record + for record in caplog.records + if record.name == "aws_durable_execution_sdk_python.plugin" + and record.exc_info is not None + and record.exc_info[1] is processor.error + ] + assert len(failures) == 1 + assert "exception ignored" in failures[0].getMessage() + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_completion_snapshot_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_completion_snapshot_int.py new file mode 100644 index 000000000..b65ef8bda --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_completion_snapshot_int.py @@ -0,0 +1,194 @@ +"""Unvisited callback updates survive the public handler's invocation lifecycle.""" + +from datetime import UTC, datetime, timedelta +import logging +from typing import Any +from unittest.mock import Mock, patch + +import pytest +from aws_durable_execution_sdk_python import DurableContext, durable_execution +from aws_durable_execution_sdk_python.lambda_service import ( + CallbackDetails, + ErrorObject, + ExecutionDetails, + Operation, + OperationStatus, + OperationType, +) +from opentelemetry import context as otel_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 ( + ExecutionOtelPluginFactory, + InvocationOtelPluginFactory, + OtelPluginConfig, +) +from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( + operation_id_to_span_id, +) + + +EXECUTION_ARN = ( + "arn:aws:lambda:us-west-2:123456789012:function:completion:1" + "/durable-execution/unvisited-callback/execution-id" +) +CALLBACK_ID = "unvisited-callback" + + +@pytest.mark.parametrize( + "factory_type", [ExecutionOtelPluginFactory, InvocationOtelPluginFactory] +) +@pytest.mark.parametrize("status", [OperationStatus.SUCCEEDED, OperationStatus.FAILED]) +@pytest.mark.parametrize("missing_parent", [False, True]) +def test_unvisited_input_completion_exports_before_invocation_closes( + factory_type: type[ExecutionOtelPluginFactory] | type[InvocationOtelPluginFactory], + status: OperationStatus, + missing_parent: bool, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + """Use service-shaped resume input without fabricating plugin callbacks. + + A handler can finish without revisiting an updated callback, such as a + callback in a sibling branch after another branch determines the result. + A later invocation may still contain its terminal checkpoint without + naming it as newly updated. Only the first invocation should export it. + """ + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + caplog.set_level(logging.WARNING) + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + factory = factory_type( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + start = datetime.now(UTC) - timedelta(seconds=10) + error = ( + ErrorObject( + type="CallbackFailed", + message="callback failed while suspended", + data=None, + stack_trace=None, + ) + if status is OperationStatus.FAILED + else None + ) + operations = [ + Operation( + operation_id="execution-id", + operation_type=OperationType.EXECUTION, + status=OperationStatus.STARTED, + start_timestamp=start, + execution_details=ExecutionDetails(input_payload='{"result": "winner"}'), + ), + Operation( + operation_id=CALLBACK_ID, + operation_type=OperationType.CALLBACK, + status=status, + name="unvisited callback", + parent_id="missing-context" if missing_parent else None, + start_timestamp=start + timedelta(seconds=1), + end_timestamp=start + timedelta(seconds=2), + callback_details=CallbackDetails( + callback_id="service-callback-id", + result='"callback-result"' if error is None else None, + error=error, + ), + ), + ] + invocation: dict[str, Any] = { + "DurableExecutionArn": EXECUTION_ARN, + "CheckpointToken": "resume-token", + "InitialExecutionState": { + "Operations": [operation.to_json_dict() for operation in operations], + "NextMarker": "", + }, + "UpdatedOperationIds": [CALLBACK_ID], + } + lambda_context = Mock( + aws_request_id="first-request", + client_context=None, + identity=None, + _epoch_deadline_time_in_ms=0, + invoked_function_arn="completion:1", + tenant_id=None, + ) + counts_during_handler: list[int] = [] + + def body(event: dict[str, str], _context: DurableContext) -> str: + counts_during_handler.append( + sum( + (span.attributes or {}).get("durable.operation.id") == CALLBACK_ID + for span in exporter.get_finished_spans() + ) + ) + return event["result"] + + handler = durable_execution(body, plugins=[factory]) + service = Mock() + before = otel_context.get_current() + try: + with patch( + "aws_durable_execution_sdk_python.execution.LambdaClient.initialize_client", + return_value=service, + ): + result = handler(invocation, lambda_context) + assert result == {"Status": "SUCCEEDED", "Result": '"winner"'} + spans = exporter.get_finished_spans() + (callback_span,) = [ + span + for span in spans + if (span.attributes or {}).get("durable.operation.id") == CALLBACK_ID + ] + (invocation_span,) = [span for span in spans if span.name == "Invocation"] + assert spans.index(callback_span) < spans.index(invocation_span) + assert callback_span.attributes is not None + assert callback_span.attributes["durable.operation.status"] == status.value + assert callback_span.status.status_code is ( + trace.StatusCode.ERROR if error else trace.StatusCode.OK + ) + assert len(callback_span.events) == int(error is not None) + assert callback_span.context is not None + assert callback_span.parent is not None + assert invocation_span.context is not None + assert callback_span.context.trace_id == invocation_span.context.trace_id + + if factory_type is InvocationOtelPluginFactory: + expected_parent_id = invocation_span.context.span_id + elif missing_parent: + expected_parent_id = operation_id_to_span_id( + EXECUTION_ARN, "missing-context" + ) + else: + (workflow_span,) = [span for span in spans if span.name == "Workflow"] + assert workflow_span.context is not None + expected_parent_id = workflow_span.context.span_id + assert callback_span.parent.span_id == expected_parent_id + + invocation["UpdatedOperationIds"] = [] + lambda_context.aws_request_id = "second-request" + assert handler(invocation, lambda_context) == result + + service.checkpoint.assert_not_called() + assert counts_during_handler == [0, 1] + assert [ + span + for span in exporter.get_finished_spans() + if (span.attributes or {}).get("durable.operation.id") == CALLBACK_ID + ] == [callback_span] + fallback_warnings = [ + record + for record in caplog.records + if "fallback parent" in record.getMessage() + and CALLBACK_ID in record.getMessage() + ] + assert len(fallback_warnings) == int(missing_parent) + assert otel_context.get_current() == before + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_context_segment_reentry_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_context_segment_reentry_int.py new file mode 100644 index 000000000..92ec5aca9 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_context_segment_reentry_int.py @@ -0,0 +1,162 @@ +"""A completed child can run again during one invocation's timed branch resume.""" + +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 ( + Duration, + 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 opentelemetry.trace import StatusCode + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPluginFactory, + InvocationOtelPluginFactory, + OtelPluginConfig, +) + + +@pytest.mark.parametrize( + "view", [InvocationOtelPluginFactory, ExecutionOtelPluginFactory] +) +def test_timed_branch_resume_ends_each_reentered_child_segment( + view, monkeypatch, caplog +): + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + sibling_running = threading.Event() + resumed = threading.Event() + child_calls = [] + saved_calls = [] + invocations = [] + child_starts = [] + child_ends = [] + replay_exports = [] + + class Observer(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + invocations.append(info.request_id) + + def on_user_function_start(self, info): + if info.name == "large-child": + child_starts.append(info) + + def on_operation_end(self, info): + if info.name == "large-child": + child_ends.append(info) + if info.is_replayed: + replay_exports.append( + len( + [ + span + for span in exporter.get_finished_spans() + if span.name == "large-child" + ] + ) + ) + + class ObserverFactory: + def create_plugin(self, _info): + return Observer() + + def saved(_step): + saved_calls.append(1) + return "saved" + + def large_child(context): + child_calls.append(1) + assert context.step(saved, name="saved-inside-child") == "saved" + # This public large-result path checkpoints ReplayChildren rather than + # storing the result. It reruns the child body on the timed resume. + return "x" * (300 * 1024) + + def resuming_branch(context): + assert sibling_running.wait(10) + value = context.run_in_child_context(large_child, name="large-child") + context.wait(Duration.from_seconds(1), name="resume-branch") + resumed.set() + return len(value) + + def running_sibling(_context): + sibling_running.set() + assert resumed.wait(10) + return "sibling" + + def handler(_event, context): + return context.parallel( + [ + ParallelBranch(func=resuming_branch, name="resuming-branch"), + ParallelBranch(func=running_sibling, name="running-sibling"), + ], + name="parallel-parent", + config=ParallelConfig(max_concurrency=2), + ).get_results() + + wrapped = durable_execution( + handler, + plugins=[ + view( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ), + ObserverFactory(), + ], + ) + try: + with DurableFunctionTestRunner( + handler=wrapped, skip_time=False, poll_interval=0.01, execution_timeout=20 + ) as runner: + result = runner.run(input="{}") + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == [300 * 1024, "sibling"] + assert len(invocations) == 1 + assert child_calls == [1, 1] + assert saved_calls == [1] + assert [info.is_replay_children for info in child_starts] == [False, True] + assert any(not info.is_replayed for info in child_ends) + assert any(info.is_replayed for info in child_ends) + + spans = exporter.get_finished_spans() + children = [span for span in spans if span.name == "large-child"] + expected_children = 2 if view is InvocationOtelPluginFactory else 1 + assert len(children) == expected_children + assert replay_exports == [expected_children] + assert all( + span.attributes["durable.operation.status"] == "SUCCEEDED" + and span.status.status_code is StatusCode.OK + for span in children + ), [(span.attributes, span.status.status_code) for span in children] + parents = [span for span in spans if span.name == "resuming-branch"] + assert len(parents) == 1 + for child in children: + assert child.parent is not None + assert child.parent.span_id == parents[0].context.span_id + assert child.context.trace_id == parents[0].context.trace_id + assert parents[0].start_time <= child.start_time < child.end_time + assert child.end_time <= parents[0].end_time + if len(children) == 2: + assert children[0].context.span_id != children[1].context.span_id + assert children[0].end_time <= children[1].start_time + assert not [ + record + for record in caplog.records + if record.levelno >= logging.ERROR + and ("plugin" in record.name or "otel" in record.name) + ] + finally: + resumed.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..0510a1328 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_exception_event_bounds_int.py @@ -0,0 +1,84 @@ +"""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 ( + ExecutionOtelPluginFactory, + InvocationOtelPluginFactory, + OtelPluginConfig, +) + + +@pytest.mark.parametrize( + "view", [ExecutionOtelPluginFactory, InvocationOtelPluginFactory] +) +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) + + class ObserveFactory: + def create_plugin(self, info): + return Observe() + + @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, ObserveFactory()]), + 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..dda0276b7 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_flush_lifecycle_int.py @@ -0,0 +1,227 @@ +"""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 ( + ExecutionOtelPluginFactory, + InvocationOtelPluginFactory, + OtelPluginConfig, + _flush, +) + + +@pytest.mark.parametrize( + "view", [ExecutionOtelPluginFactory, InvocationOtelPluginFactory] +) +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" + + class OwnerFactory: + def create_plugin(self, info): + return Owner() + + handler = durable_execution(body, plugins=[OwnerFactory(), 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 = not _flush._active + 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", [ExecutionOtelPluginFactory, InvocationOtelPluginFactory] +) +@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" + + class OwnerFactory: + def create_plugin(self, info): + return Owner() + + handler = durable_execution(body, plugins=[OwnerFactory(), 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 bf1365fb9..99ea3c16e 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 @@ -16,10 +16,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, @@ -152,6 +155,130 @@ def checkpoint( return checkpoint, operations +@pytest.mark.parametrize( + "plugin_type", [ExecutionOtelPluginFactory, InvocationOtelPluginFactory] +) +@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( "factory_type", [InvocationOtelPluginFactory, ExecutionOtelPluginFactory], 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..fdf93b895 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_parent_registration_int.py @@ -0,0 +1,190 @@ +"""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)) + 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() + + class OtelFactory: + def create_plugin(self, info): + plugin = ExecutionOtelPlugin( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + 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) + + return plugin + + class ObserveFactory: + def create_plugin(self, info): + return Observe() + + 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=[OtelFactory(), ObserveFactory()]) + 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..af7452b2d --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_shared_invocation_state_int.py @@ -0,0 +1,115 @@ +"""A shared factory 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 ( + ExecutionOtelPluginFactory, + InvocationOtelPluginFactory, + OtelPluginConfig, +) + + +@pytest.mark.parametrize( + "view", [ExecutionOtelPluginFactory, InvocationOtelPluginFactory] +) +def test_same_handler_shared_factory_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/test_completion_export_retry.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_completion_export_retry.py new file mode 100644 index 000000000..a3f65ab33 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_completion_export_retry.py @@ -0,0 +1,477 @@ +"""A failed completion remains retryable until its span successfully ends.""" + +import threading +from concurrent.futures import ThreadPoolExecutor +from datetime import UTC, datetime, timedelta + +import pytest +from aws_durable_execution_sdk_python.lambda_service import ErrorObject +from aws_durable_execution_sdk_python.plugin import ( + InvocationEndInfo, + InvocationStartInfo, + InvocationStatus, + OperationChangeInfo, + OperationEndInfo, + OperationStartInfo, + OperationStatus, + OperationType, + UserFunctionEndInfo, + UserFunctionStartInfo, +) +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, +) + + +class _FailOnceProcessor(SimpleSpanProcessor): + def __init__(self, exporter, operation_id): + super().__init__(exporter) + self.operation_id = operation_id + self.error = RuntimeError("one-time processor failure") + self.failed = False + + def on_end(self, span): + if ( + span.attributes.get("durable.operation.id") == self.operation_id + and not self.failed + ): + self.failed = True + raise self.error + super().on_end(span) + + +def _setup(view, failing_operation=None): + provider = TracerProvider() + exporter = InMemorySpanExporter() + processor = ( + _FailOnceProcessor(exporter, failing_operation) + if failing_operation is not None + else SimpleSpanProcessor(exporter) + ) + provider.add_span_processor(processor) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + invocation = dict( + request_id="retry-completion", + execution_arn="execution/retry-completion", + execution_start_time=datetime.now(UTC), + is_first_invocation=False, + ) + plugin.on_invocation_start(InvocationStartInfo(**invocation)) + return plugin, provider, exporter, invocation, processor + + +def _completion( + operation_id="operation", + parent_id=None, + operation_type=OperationType.CALLBACK, + error=None, + end_time=None, +): + now = end_time or datetime.now(UTC) + return OperationEndInfo( + operation_id=operation_id, + operation_type=operation_type, + sub_type=None, + name=operation_id, + parent_id=parent_id, + start_time=now - timedelta(seconds=1), + end_time=now, + is_replayed=False, + status=OperationStatus.FAILED if error else OperationStatus.SUCCEEDED, + error=error, + ) + + +def _spans(exporter, operation_id="operation"): + return [ + span + for span in exporter.get_finished_spans() + if span.attributes.get("durable.operation.id") == operation_id + ] + + +@pytest.mark.parametrize("view", [InvocationOtelPlugin, ExecutionOtelPlugin]) +@pytest.mark.parametrize("failure", ["start", "end", "processor"]) +@pytest.mark.parametrize("retry", ["notification", "invocation-end"]) +def test_failed_completion_retries_without_poisoning_success_ledger( + view, failure, retry, monkeypatch +): + plugin, provider, exporter, invocation, processor = _setup( + view, "operation" if failure == "processor" else None + ) + error = ( + processor.error + if failure == "processor" + else RuntimeError("one-time span failure") + ) + info = _completion( + error=ErrorObject( + type="ValueError", + message="operation failed", + data=None, + stack_trace=["frame"], + ) + ) + original_start = plugin._start_span + injected = False + attempted_end_times = [] + + def fail_once_start(*args, **kwargs): + nonlocal injected + if failure == "start" and not injected: + injected = True + raise error + span = original_start(*args, **kwargs) + if failure == "end": + original_end = span.end + + def fail_once_end(*end_args, **end_kwargs): + nonlocal injected + attempted_end_times.append(end_kwargs["end_time"]) + if not injected: + injected = True + raise error + return original_end(*end_args, **end_kwargs) + + monkeypatch.setattr(span, "end", fail_once_end) + return span + + monkeypatch.setattr(plugin, "_start_span", fail_once_start) + try: + with pytest.raises(RuntimeError) as caught: + plugin.on_operation_end(info) + assert caught.value is error + assert "operation" not in plugin._ended_operation_ids + assert "operation" not in plugin._ending_operation_ids + assert plugin._pending_completions["operation"] == info + assert _spans(exporter) == [] + original_end_time = ( + attempted_end_times[0] + if failure == "end" + else plugin._get_span("operation").end_time + if failure == "processor" + else None + ) + if original_end_time is not None: + # An unrelated completion advances invocation-local ordering. It + # must not inflate the failed operation's retained duration. + plugin.on_operation_end( + _completion("unrelated", end_time=info.end_time + timedelta(seconds=5)) + ) + if retry == "notification": + plugin.on_operation_end(info) + assert "operation" in plugin._ended_operation_ids + assert "operation" not in plugin._pending_completions + plugin.on_operation_end(info) + plugin.on_operation_change( + OperationChangeInfo( + execution_arn=invocation["execution_arn"], + operations={"operation": info}, + updated_operations={"operation": info}, + ) + ) + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.SUCCEEDED) + ) + spans = _spans(exporter) + assert len(spans) == 1 + if original_end_time is not None: + assert spans[0].end_time == original_end_time + assert spans[0].status.status_code is StatusCode.ERROR + assert spans[0].attributes["durable.operation.status"] == "FAILED" + exceptions = [event for event in spans[0].events if event.name == "exception"] + assert len(exceptions) == 1 + assert exceptions[0].attributes["exception.message"] == "operation failed" + assert spans[0].start_time <= exceptions[0].timestamp <= spans[0].end_time + assert not plugin._pending_completions + assert not plugin._ending_operation_ids + assert not plugin._completion_end_timestamps + finally: + if plugin._tracing_enabled: + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.SUCCEEDED) + ) + provider.shutdown() + + +@pytest.mark.parametrize("view", [InvocationOtelPlugin, ExecutionOtelPlugin]) +def test_concurrent_completion_cannot_overtake_failed_attempt_then_retry( + view, monkeypatch +): + plugin, provider, exporter, invocation, _ = _setup(view) + entered, release = threading.Event(), threading.Event() + original_start = plugin._start_span + attempts = [] + error = RuntimeError("first attempt failed") + + def gated_start(*args, **kwargs): + attempts.append(1) + if len(attempts) == 1: + entered.set() + assert release.wait(5) + raise error + return original_start(*args, **kwargs) + + monkeypatch.setattr(plugin, "_start_span", gated_start) + info = _completion() + try: + with ThreadPoolExecutor(max_workers=2) as callers: + first = callers.submit(plugin.on_operation_end, info) + try: + assert entered.wait(5) + duplicate = callers.submit(plugin.on_operation_end, info) + duplicate.result(timeout=5) + assert attempts == [1] + finally: + release.set() + with pytest.raises(RuntimeError) as caught: + first.result(timeout=5) + assert caught.value is error + assert "operation" not in plugin._ended_operation_ids + assert not plugin._ending_operation_ids + plugin.on_operation_end(info) + plugin.on_operation_end(info) + assert attempts == [1, 1] + assert len(_spans(exporter)) == 1 + finally: + release.set() + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.SUCCEEDED) + ) + provider.shutdown() + + +@pytest.mark.parametrize("view", [InvocationOtelPlugin, ExecutionOtelPlugin]) +def test_parent_processor_retry_preserves_context_for_already_exported_child(view): + plugin, provider, exporter, invocation, _ = _setup(view, "parent") + parent_end = _completion("parent", operation_type=OperationType.CONTEXT) + start_fields = dict( + operation_id="parent", + operation_type=OperationType.CONTEXT, + sub_type=None, + name="parent", + parent_id=None, + start_time=parent_end.start_time, + status=OperationStatus.STARTED, + is_replayed=False, + ) + try: + plugin.on_operation_start(OperationStartInfo(**start_fields)) + start = UserFunctionStartInfo(**start_fields) + plugin.on_user_function_start(start) + original_context = plugin._get_span("parent").get_span_context() + plugin.on_user_function_end(UserFunctionEndInfo.from_start_info(start, None)) + with pytest.raises(RuntimeError, match="processor failure"): + plugin.on_operation_end(parent_end) + failed_parent = plugin._get_span("parent") + assert failed_parent is not None + original_start_ns = failed_parent.start_time + plugin.on_operation_end(_completion(parent_id="parent")) + assert len(_spans(exporter)) == 1 + plugin.on_operation_end(parent_end) + parent = _spans(exporter, "parent") + child = _spans(exporter)[0] + assert len(parent) == 1 + assert parent[0].context.span_id == original_context.span_id + assert child.parent.span_id == original_context.span_id + assert parent[0].start_time == original_start_ns + plugin.on_operation_end(parent_end) + assert len(_spans(exporter, "parent")) == 1 + finally: + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.SUCCEEDED) + ) + provider.shutdown() + + +@pytest.mark.parametrize("failure", ["end", "processor"]) +@pytest.mark.parametrize("retry", ["notification", "invocation-end"]) +def test_failed_replay_segment_remains_retryable_after_logical_export( + failure, retry, monkeypatch +): + plugin, provider, exporter, invocation, processor = _setup( + InvocationOtelPlugin, "operation" if failure == "processor" else None + ) + if failure == "processor": + processor.failed = True + first_end = _completion(operation_type=OperationType.CONTEXT) + start_fields = dict( + operation_id="operation", + operation_type=OperationType.CONTEXT, + sub_type=None, + name="operation", + parent_id=None, + start_time=first_end.start_time, + status=OperationStatus.STARTED, + is_replayed=False, + ) + try: + plugin.on_operation_start(OperationStartInfo(**start_fields)) + start = UserFunctionStartInfo(**start_fields) + plugin.on_user_function_start(start) + plugin.on_user_function_end(UserFunctionEndInfo.from_start_info(start, None)) + plugin.on_operation_end(first_end) + assert len(_spans(exporter)) == 1 + + replay_start = UserFunctionStartInfo(**start_fields, is_replay_children=True) + plugin.on_user_function_start(replay_start) + failed_span = plugin._get_span("operation") + if failure == "end": + original_end = failed_span.end + injected = False + + def fail_once_end(*args, **kwargs): + nonlocal injected + if not injected: + injected = True + raise RuntimeError("replay end failure") + return original_end(*args, **kwargs) + + monkeypatch.setattr(failed_span, "end", fail_once_end) + else: + processor.failed = False + plugin.on_user_function_end( + UserFunctionEndInfo.from_start_info(replay_start, None) + ) + replay_end = OperationEndInfo( + **{ + **start_fields, + "start_time": None, + "is_replayed": True, + "status": OperationStatus.SUCCEEDED, + }, + end_time=datetime.now(UTC), + ) + with pytest.raises(RuntimeError): + plugin.on_operation_end(replay_end) + assert "operation" in plugin._ended_operation_ids + assert plugin._pending_completions["operation"] == replay_end + assert not plugin._ending_operation_ids + + # The old terminal checkpoint must neither finalize nor overwrite the + # failed fresh segment's retained replay End. + plugin.on_operation_end(first_end) + plugin.on_operation_change( + OperationChangeInfo( + execution_arn=invocation["execution_arn"], + operations={"operation": first_end}, + updated_operations={"operation": first_end}, + ) + ) + assert plugin._pending_completions["operation"] == replay_end + assert len(_spans(exporter)) == 1 + if retry == "notification": + plugin.on_operation_end(replay_end) + plugin.on_operation_end(replay_end) + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.SUCCEEDED) + ) + spans = _spans(exporter) + assert len(spans) == 2 + assert spans[1].context.span_id == failed_span.get_span_context().span_id + assert spans[1].status.status_code is StatusCode.OK + assert spans[1].attributes["durable.operation.status"] == "SUCCEEDED" + finally: + if plugin._tracing_enabled: + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.SUCCEEDED) + ) + provider.shutdown() + + +@pytest.mark.parametrize("failure", ["end", "processor"]) +def test_failed_first_segment_finishes_before_next_context_segment_starts( + failure, monkeypatch +): + plugin, provider, exporter, invocation, _ = _setup( + InvocationOtelPlugin, "operation" if failure == "processor" else None + ) + first_end = _completion(operation_type=OperationType.CONTEXT) + fields = dict( + operation_id="operation", + operation_type=OperationType.CONTEXT, + sub_type=None, + name="operation", + parent_id=None, + start_time=first_end.start_time, + status=OperationStatus.STARTED, + is_replayed=False, + ) + try: + plugin.on_operation_start(OperationStartInfo(**fields)) + first_start = UserFunctionStartInfo(**fields) + plugin.on_user_function_start(first_start) + first_span = plugin._get_span("operation") + if failure == "end": + original_end = first_span.end + injected = False + + def fail_once_end(*args, **kwargs): + nonlocal injected + if not injected: + injected = True + raise RuntimeError("first segment end failure") + return original_end(*args, **kwargs) + + monkeypatch.setattr(first_span, "end", fail_once_end) + plugin.on_user_function_end( + UserFunctionEndInfo.from_start_info(first_start, None) + ) + with pytest.raises(RuntimeError): + plugin.on_operation_end(first_end) + assert _spans(exporter) == [] + + # No new End notification arrives for the failed first segment. The + # next user-function Start must finish it before replacing its registry + # entry or treating it as an incomplete segment to continue. + second_start = UserFunctionStartInfo(**fields, is_replay_children=True) + plugin.on_user_function_start(second_start) + assert len(_spans(exporter)) == 1 + assert ( + _spans(exporter)[0].context.span_id == first_span.get_span_context().span_id + ) + second_span = plugin._get_span("operation") + assert ( + second_span.get_span_context().span_id + != first_span.get_span_context().span_id + ) + plugin.on_user_function_end( + UserFunctionEndInfo.from_start_info(second_start, None) + ) + second_end = OperationEndInfo( + **{ + **fields, + "start_time": None, + "is_replayed": True, + "status": OperationStatus.SUCCEEDED, + }, + end_time=datetime.now(UTC), + ) + plugin.on_operation_end(second_end) + plugin.on_operation_end(first_end) + assert len(_spans(exporter)) == 2 + assert ( + _spans(exporter)[1].context.span_id + == second_span.get_span_context().span_id + ) + assert all( + span.status.status_code is StatusCode.OK for span in _spans(exporter) + ) + finally: + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.SUCCEEDED) + ) + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_completion_replay_parent.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_completion_replay_parent.py new file mode 100644 index 000000000..e12ade477 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_completion_replay_parent.py @@ -0,0 +1,132 @@ +"""Input-only completions retain actual parent segments after replay finishes.""" + +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.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("replay_children", [False, True]) +def test_input_only_completion_uses_parent_segment_that_finished_during_replay( + view, replay_children +): + now = datetime.now(UTC) + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + invocation = dict( + request_id="input-only-completion", + execution_arn="execution/replayed-parent", + is_first_invocation=False, + execution_start_time=now, + ) + parent_fields = dict( + operation_id="parent", + operation_type=OperationType.CONTEXT, + sub_type=None, + name="parent", + parent_id=None, + start_time=now, + is_replayed=True, + status=OperationStatus.STARTED, + ) + child = OperationInfo( + operation_id="child", + operation_type=OperationType.CALLBACK, + sub_type=None, + name="child", + parent_id="parent", + start_time=now, + end_time=now + timedelta(seconds=1), + is_replayed=False, + status=OperationStatus.SUCCEEDED, + ) + parent_input = OperationInfo( + **{ + **parent_fields, + "status": OperationStatus.SUCCEEDED + if replay_children + else OperationStatus.STARTED, + "is_replayed": False, + }, + end_time=now + timedelta(seconds=2) if replay_children else None, + ) + try: + plugin.on_invocation_start( + InvocationStartInfo( + **invocation, + operations={"parent": parent_input, "child": child}, + updated_operations={"child": child}, + ) + ) + if not replay_children: + plugin.on_operation_start(OperationStartInfo(**parent_fields)) + start = UserFunctionStartInfo( + **parent_fields, is_replay_children=replay_children + ) + 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( + OperationEndInfo( + **{ + **parent_fields, + "status": OperationStatus.SUCCEEDED, + "is_replayed": replay_children, + }, + end_time=now + timedelta(seconds=2), + ) + ) + # The application skips this callback on replay. Its updated input + # snapshot is the only completion delivery; no child End hook fires. + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.PENDING) + ) + + completions = [ + span + for span in exporter.get_finished_spans() + if span.attributes.get("durable.operation.id") == "child" + ] + assert len(completions) == 1 + assert completions[0].parent is not None + assert completions[0].parent.span_id == actual_parent.span_id + assert completions[0].context.trace_id == actual_parent.trace_id + if view is InvocationOtelPlugin: + parent_span = next( + span + for span in exporter.get_finished_spans() + if span.context.span_id == actual_parent.span_id + ) + assert parent_span.start_time <= completions[0].start_time + assert completions[0].end_time <= parent_span.end_time + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_completion_snapshots.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_completion_snapshots.py new file mode 100644 index 000000000..62dc0b73c --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_completion_snapshots.py @@ -0,0 +1,215 @@ +"""Input, replay and checkpoint observations share an invocation's export ledger.""" + +from concurrent.futures import ThreadPoolExecutor +from datetime import UTC, datetime, timedelta +from threading import Event + +import pytest +from aws_durable_execution_sdk_python.plugin import ( + InvocationEndInfo, + InvocationStartInfo, + InvocationStatus, + OperationChangeInfo, + OperationEndInfo, + OperationInfo, + OperationStatus, + OperationType, +) +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.fixture(params=[ExecutionOtelPlugin, InvocationOtelPlugin]) +def trace_case(request): + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = request.param( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + invocation = dict( + request_id="request", + execution_arn="execution/snapshot-delivery", + execution_start_time=datetime.now(UTC) - timedelta(seconds=2), + is_first_invocation=False, + ) + try: + yield plugin, exporter, invocation + finally: + if plugin._tracing_enabled: + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.PENDING) + ) + provider.shutdown() + + +def completion(operation_id="b", parent_id=None, operation_type=OperationType.CALLBACK): + now = datetime.now(UTC) + return OperationEndInfo( + operation_id=operation_id, + operation_type=operation_type, + sub_type=None, + name=operation_id, + parent_id=parent_id, + start_time=now - timedelta(seconds=1), + end_time=now, + is_replayed=False, + status=OperationStatus.SUCCEEDED, + result="saved", + ) + + +@pytest.mark.parametrize("input_update", [False, True]) +@pytest.mark.parametrize("visited", [False, True]) +def test_snapshot_and_natural_notifications_export_once( + trace_case, input_update, visited +): + plugin, exporter, invocation = trace_case + end = completion() + active = OperationInfo( + operation_id="active", + operation_type=OperationType.CALLBACK, + sub_type=None, + name="active", + parent_id=None, + start_time=end.start_time, + is_replayed=False, + status=OperationStatus.STARTED, + ) + updates = {"b": end, "active": active} + plugin.on_invocation_start( + InvocationStartInfo( + **invocation, + operations=updates, + updated_operations=updates if input_update else {}, + ) + ) + change = OperationChangeInfo( + execution_arn=invocation["execution_arn"], + operations=updates, + updated_operations=updates, + ) + plugin.on_operation_change(change) + if visited: + plugin.on_operation_end(end) + plugin.on_operation_end(end) + plugin.on_operation_change(change) + before_end = [span for span in exporter.get_finished_spans() if span.name == "b"] + assert len(before_end) == int(visited) + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.PENDING) + ) + spans = exporter.get_finished_spans() + assert len([span for span in spans if span.name == "b"]) == 1 + assert not [span for span in spans if span.name == "active"] + assert plugin._pending_completions == plugin._pending_operation_ends == {} + + +def test_concurrent_end_and_checkpoint_do_not_duplicate_export(trace_case, monkeypatch): + plugin, exporter, invocation = trace_case + end = completion() + plugin.on_invocation_start( + InvocationStartInfo(**invocation, updated_operations={"b": end}) + ) + entered, release = Event(), Event() + original = plugin._start_span + + def block_start(*args, **kwargs): + if kwargs.get("operation_id") == "b": + entered.set() + assert release.wait(5) + return original(*args, **kwargs) + + monkeypatch.setattr(plugin, "_start_span", block_start) + with ThreadPoolExecutor(max_workers=2) as pool: + first = pool.submit(plugin.on_operation_end, end) + try: + assert entered.wait(5) + second = pool.submit(plugin.on_operation_end, end) + second.result(timeout=5) + plugin.on_operation_change( + OperationChangeInfo( + execution_arn=invocation["execution_arn"], + operations={"b": end}, + updated_operations={"b": end}, + ) + ) + finally: + release.set() + first.result(timeout=5) + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.PENDING) + ) + assert ( + len([span for span in exporter.get_finished_spans() if span.name == "b"]) == 1 + ) + + +def test_missing_ancestor_fallback_preserves_available_descendant_hierarchy( + trace_case, caplog +): + plugin, exporter, invocation = trace_case + # Deliberately child-first: fallback must register its pending parent first. + child = completion(parent_id="parent") + parent = completion("parent", "missing-ancestor", OperationType.CONTEXT) + plugin.on_invocation_start( + InvocationStartInfo( + **invocation, updated_operations={"b": child, "parent": parent} + ) + ) + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.PENDING) + ) + spans = {span.name: span for span in exporter.get_finished_spans()} + assert spans["b"].parent.span_id == spans["parent"].context.span_id + assert spans["b"].context.trace_id == spans["parent"].context.trace_id + assert ( + len( + [record for record in caplog.records if "fallback parent" in record.message] + ) + == 1 + ) + assert plugin._pending_completions == plugin._pending_operation_ends == {} + + +def test_failed_invocation_start_cannot_requeue_fallback_forever( + trace_case, monkeypatch +): + plugin, _, invocation = trace_case + + def fail_extraction(_info): + raise RuntimeError("context extraction failed") + + monkeypatch.setattr(plugin, "_context_extractor", fail_extraction) + with pytest.raises(RuntimeError, match="context extraction failed"): + plugin.on_invocation_start(InvocationStartInfo(**invocation)) + plugin.on_operation_end(completion(parent_id="missing-parent")) + + # Bound the regression itself if a future change retries a fallback without + # an initialized trace or invocation span. + attempts = [] + original = plugin._render_operation_end + + def bounded_render(*args, **kwargs): + attempts.append(1) + assert len(attempts) <= 1, "completion fallback kept requeueing" + return original(*args, **kwargs) + + monkeypatch.setattr(plugin, "_render_operation_end", bounded_render) + with pytest.raises(ValueError, match="available for completion fallback"): + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.PENDING) + ) + assert not plugin._tracing_enabled + assert plugin._pending_completions == plugin._pending_operation_ends == {} diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_context_segment_reentry.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_context_segment_reentry.py new file mode 100644 index 000000000..0648ed5aa --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_context_segment_reentry.py @@ -0,0 +1,250 @@ +"""Context replay segments have their own End within a single invocation.""" + +import threading +from concurrent.futures import ThreadPoolExecutor +from datetime import UTC, datetime + +import pytest +from aws_durable_execution_sdk_python.plugin import ( + InvocationEndInfo, + InvocationStartInfo, + InvocationStatus, + OperationChangeInfo, + OperationEndInfo, + OperationStartInfo, + OperationStatus, + OperationType, + UserFunctionEndInfo, + UserFunctionStartInfo, +) +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", [InvocationOtelPlugin, ExecutionOtelPlugin]) +@pytest.mark.parametrize( + "checkpointed", [True, False], ids=["replay-children", "virtual"] +) +def test_replay_context_end_closes_new_segment_without_repeating_checkpoint( + view, checkpointed +): + 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="same-invocation", + execution_arn="execution/context-segment-reentry", + execution_start_time=now, + is_first_invocation=True, + ) + fields = dict( + operation_id="child", + operation_type=OperationType.CONTEXT, + sub_type=None, + name="child", + parent_id=None, + start_time=now, + is_replayed=False, + status=OperationStatus.STARTED, + ) + + def child_spans(): + return [span for span in exporter.get_finished_spans() if span.name == "child"] + + plugin.on_invocation_start(InvocationStartInfo(**invocation)) + try: + if checkpointed: + plugin.on_operation_start(OperationStartInfo(**fields)) + first_start = UserFunctionStartInfo(**fields) + plugin.on_user_function_start(first_start) + plugin.on_user_function_end( + UserFunctionEndInfo.from_start_info(first_start, None) + ) + first_end = OperationEndInfo( + **{ + **fields, + "start_time": now if checkpointed else None, + "status": OperationStatus.SUCCEEDED, + }, + end_time=datetime.now(UTC), + ) + plugin.on_operation_end(first_end) + assert len(child_spans()) == 1 + first_span = child_spans()[0] + + replay_start = UserFunctionStartInfo( + **{**fields, "start_time": datetime.now(UTC)}, + is_replay_children=checkpointed, + ) + plugin.on_user_function_start(replay_start) + replay_span = plugin._get_span("child") + assert replay_span is not None + if view is InvocationOtelPlugin: + assert replay_span.is_recording() + assert replay_span.get_span_context().span_id != first_span.context.span_id + + # An old durable checkpoint may be delivered again while the new user + # function is running. It must not close that invocation-local segment. + if checkpointed: + plugin.on_operation_end(first_end) + plugin.on_operation_change( + OperationChangeInfo( + execution_arn=invocation["execution_arn"], + updated_operations={"child": first_end}, + operations={"child": first_end}, + ) + ) + assert len(child_spans()) == 1 + if view is InvocationOtelPlugin: + assert replay_span.is_recording() + + plugin.on_user_function_end( + UserFunctionEndInfo.from_start_info(replay_start, None) + ) + replay_end = OperationEndInfo( + **{ + **fields, + "start_time": None, + "is_replayed": checkpointed, + "status": OperationStatus.SUCCEEDED, + }, + end_time=datetime.now(UTC), + ) + plugin.on_operation_end(replay_end) + expected_count = 2 if view is InvocationOtelPlugin else 1 + assert len(child_spans()) == expected_count + assert not replay_span.is_recording() + assert all( + span.attributes["durable.operation.status"] == "SUCCEEDED" + and span.status.status_code is StatusCode.OK + for span in child_spans() + ) + if view is InvocationOtelPlugin: + first, second = child_spans() + assert first.context.trace_id == second.context.trace_id + assert first.parent.span_id == second.parent.span_id + assert ( + first.start_time < first.end_time <= second.start_time < second.end_time + ) + + # Retrying either notification without a fresh segment must export + # nothing else, including at invocation End. + plugin.on_operation_end(replay_end) + plugin.on_operation_end(first_end) + assert len(child_spans()) == expected_count + finally: + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.SUCCEEDED) + ) + provider.shutdown() + assert len(child_spans()) == expected_count + + +@pytest.mark.parametrize("replayed", [False, True], ids=["virtual", "replay-children"]) +def test_concurrent_context_ends_leave_finalization_with_first_caller( + monkeypatch, replayed +): + 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="concurrent-context-end", + execution_arn="execution/concurrent-context-end", + execution_start_time=now, + is_first_invocation=False, + ) + fields = dict( + operation_id="child", + operation_type=OperationType.CONTEXT, + sub_type=None, + name="child", + parent_id=None, + start_time=now, + is_replayed=False, + status=OperationStatus.STARTED, + ) + finalizing = threading.Event() + release = threading.Event() + original_finish = plugin._finish_operation_end + + def paused_finish(info, span, fallback_parent): + if not finalizing.is_set(): + finalizing.set() + assert release.wait(5) + original_finish(info, span, fallback_parent) + + monkeypatch.setattr(plugin, "_finish_operation_end", paused_finish) + plugin.on_invocation_start(InvocationStartInfo(**invocation)) + try: + start = UserFunctionStartInfo(**fields, is_replay_children=replayed) + plugin.on_user_function_start(start) + active_span = plugin._get_span("child") + assert active_span is not None and active_span.is_recording() + plugin.on_user_function_end(UserFunctionEndInfo.from_start_info(start, None)) + end = OperationEndInfo( + **{ + **fields, + "start_time": None, + "is_replayed": replayed, + "status": OperationStatus.SUCCEEDED, + }, + end_time=datetime.now(UTC), + ) + with ThreadPoolExecutor(max_workers=2) as callers: + first = callers.submit(plugin.on_operation_end, end) + try: + assert finalizing.wait(5) + duplicate = callers.submit(plugin.on_operation_end, end) + duplicate.result(timeout=5) + # The duplicate must not finalize the span while the winning + # caller still owns it outside the registry lock. + assert active_span.is_recording() + assert not [ + span + for span in exporter.get_finished_spans() + if span.name == "child" + ] + finally: + release.set() + first.result(timeout=5) + assert not active_span.is_recording() + plugin.on_operation_end(end) + children = [ + span for span in exporter.get_finished_spans() if span.name == "child" + ] + assert len(children) == 1 + assert children[0].status.status_code is StatusCode.OK + assert children[0].attributes["durable.operation.status"] == "SUCCEEDED" + assert children[0].context.span_id == active_span.get_span_context().span_id + assert children[0].start_time < children[0].end_time + assert not plugin._ending_operation_ids + finally: + release.set() + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.SUCCEEDED) + ) + 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 ffb9ee1c8..b37b2e185 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 @@ -5,6 +5,7 @@ import logging import threading from concurrent.futures import ThreadPoolExecutor +from contextvars import copy_context from datetime import UTC, datetime, timedelta from types import SimpleNamespace @@ -1321,7 +1322,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() @@ -1341,6 +1342,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(): @@ -1473,7 +1480,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") @@ -1566,20 +1575,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..b030a0422 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_flush.py @@ -0,0 +1,317 @@ +"""Bound provider I/O without spawning replacement workers for a stuck call.""" + +import contextvars +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from unittest.mock import Mock + +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 not _flush._active: + 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 + + provider = Mock(spec=TracerProvider) + provider.force_flush.side_effect = blocked + try: + began = time.monotonic() + assert _flush.flush_provider(provider) is False + assert time.monotonic() - began < 1 + assert started.wait(1) + for _ in range(5): + assert _flush.flush_provider(provider) 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) + provider.force_flush.side_effect = None + provider.force_flush.return_value = True + assert _flush.flush_provider(provider) is True + + +@pytest.mark.parametrize("first_timed_out", [False, True]) +def test_distinct_provider_flushes_while_first_provider_is_blocked( + first_timed_out, monkeypatch +): + exporting = threading.Event() + release = threading.Event() + + class GatedExporter(InMemorySpanExporter): + def export(self, spans): + exporting.set() + assert release.wait(3) + return super().export(spans) + + class UnhashableProvider(TracerProvider): + def __eq__(self, other): + return isinstance(other, UnhashableProvider) + + first_exporter = GatedExporter() + second_exporter = InMemorySpanExporter() + first_provider = UnhashableProvider() + second_provider = UnhashableProvider() + 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 + + try: + with ThreadPoolExecutor(max_workers=2) as callers: + monkeypatch.setattr( + _flush, "_FLUSH_TIMEOUT_SECONDS", 0.02 if first_timed_out else 3.0 + ) + first = callers.submit(_flush.flush_provider, first_provider) + try: + assert exporting.wait(1) + if first_timed_out: + assert first.result(timeout=1) is False + monkeypatch.setattr(_flush, "_FLUSH_TIMEOUT_SECONDS", 1.0) + # Complete B's real export while A's exporter is still blocked, + # including after A's caller has exhausted its wait budget. + second = callers.submit(_flush.flush_provider, second_provider) + assert second.result(timeout=1.5) is True + assert not first_exporter.get_finished_spans() + assert [span.name for span in second_exporter.get_finished_spans()] == [ + "second" + ] + finally: + release.set() + assert first.result(timeout=2) is (not first_timed_out) + _wait_idle() + assert [span.name for span in first_exporter.get_finished_spans()] == ["first"] + finally: + release.set() + first_provider.shutdown() + second_provider.shutdown() + + +def test_same_provider_flushes_spans_queued_after_active_attempt(monkeypatch): + first_flushed = threading.Event() + release_first = threading.Event() + waiting_for_provider = threading.Event() + calls = [] + original_wait = _flush._lock.wait + + def admission_wait(timeout=None): + waiting_for_provider.set() + return original_wait(timeout) + + monkeypatch.setattr(_flush._lock, "wait", admission_wait) + + class GatedProvider(TracerProvider): + def force_flush(self, timeout_millis=30000): + calls.append(threading.get_ident()) + result = super().force_flush(timeout_millis) + if len(calls) == 1: + first_flushed.set() + assert release_first.wait(3) + return result + + exporter = InMemorySpanExporter() + provider = GatedProvider() + provider.add_span_processor( + BatchSpanProcessor(exporter, schedule_delay_millis=60000) + ) + tracer = provider.get_tracer("shared") + with tracer.start_as_current_span("first"): + pass + + try: + with ThreadPoolExecutor(max_workers=2) as callers: + first = callers.submit(_flush.flush_provider, provider) + try: + assert first_flushed.wait(1) + with tracer.start_as_current_span("second"): + pass + second = callers.submit(_flush.flush_provider, provider) + assert waiting_for_provider.wait(1) + assert len(calls) == 1 + assert [span.name for span in exporter.get_finished_spans()] == [ + "first" + ] + finally: + release_first.set() + assert first.result(timeout=2) is True + assert second.result(timeout=2) is True + assert len(calls) == 2 + assert [span.name for span in exporter.get_finished_spans()] == [ + "first", + "second", + ] + finally: + release_first.set() + 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(provider) + + provider = Mock(spec=TracerProvider) + provider.force_flush.side_effect = lambda: ( + second_io() if first_entered.is_set() else first_io() + ) + with ThreadPoolExecutor(max_workers=2) as callers: + first = callers.submit(_flush.flush_provider, provider) + 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): + provider = Mock(spec=TracerProvider) + provider.force_flush.return_value = result + assert _flush.flush_provider(provider) 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 + + provider = Mock(spec=TracerProvider) + provider.force_flush.side_effect = fail + with pytest.raises(type(error)) as caught: + _flush.flush_provider(provider) + 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 + + provider = Mock(spec=TracerProvider) + provider.force_flush.side_effect = fail_late + try: + assert _flush.flush_provider(provider) 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) + provider.force_flush.side_effect = None + provider.force_flush.return_value = True + assert _flush.flush_provider(provider) is True + + +def test_start_failure_releases_capacity(monkeypatch): + def fail_start(_thread): + raise RuntimeError("start failed") + + provider = Mock(spec=TracerProvider) + provider.force_flush.return_value = True + with monkeypatch.context() as patch: + patch.setattr(threading.Thread, "start", fail_start) + with pytest.raises(RuntimeError, match="start failed"): + _flush.flush_provider(provider) + assert _flush.flush_provider(provider) 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 + + provider = Mock(spec=TracerProvider) + provider.force_flush.side_effect = flush + assert _flush.flush_provider(provider) + 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..b237325eb --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_owned_state.py @@ -0,0 +1,128 @@ +"""Factory-created invocations stay isolated through inherited Contexts.""" + +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 ( + ExecutionOtelPluginFactory, + InvocationOtelPluginFactory, + OtelPluginConfig, +) + + +@pytest.mark.parametrize( + "view", [ExecutionOtelPluginFactory, InvocationOtelPluginFactory] +) +def test_copied_context_on_same_thread_does_not_clear_its_parent_invocation(view): + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + factory = 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, + ) + plugin_a = factory.create_plugin(InvocationStartInfo(**info_a)) + plugin_b = factory.create_plugin(InvocationStartInfo(**info_b)) + assert plugin_a is not plugin_b + try: + original = a.run(otel_context.get_current) + a.run(plugin_a.on_invocation_start, InvocationStartInfo(**info_a)) + span_a = a.run(plugin_a.get_current_span_context) + b = a.copy() + b.run(plugin_b.on_invocation_start, InvocationStartInfo(**info_b)) + b.run( + plugin_b.on_invocation_end, + InvocationEndInfo(**info_b, status=InvocationStatus.SUCCEEDED), + ) + assert a.run(plugin_a.get_current_span_context) == span_a + assert plugin_a._execution_arn == "execution/a" and plugin_a._tracing_enabled + a.run( + plugin_a.on_invocation_end, + InvocationEndInfo(**info_a, status=InvocationStatus.SUCCEEDED), + ) + assert a.run(otel_context.get_current) == original + assert not plugin_a._tracing_enabled and not plugin_b._tracing_enabled + 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", [ExecutionOtelPluginFactory, InvocationOtelPluginFactory] +) +def test_factory_creates_fresh_state_after_previous_invocation_end(view): + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter())) + factory = 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: + first_plugin = factory.create_plugin(InvocationStartInfo(**first)) + first_plugin.on_invocation_start(InvocationStartInfo(**first)) + first_plugin.on_invocation_end( + InvocationEndInfo(**first, status=InvocationStatus.PENDING) + ) + assert first_plugin._context_tokens == {} + assert not first_plugin._tracing_enabled + plugin = factory.create_plugin(InvocationStartInfo(**second)) + assert plugin is not first_plugin + assert plugin._pending_operation_ends == {} + plugin.on_invocation_start(InvocationStartInfo(**second)) + plugin.on_invocation_end( + InvocationEndInfo(**second, status=InvocationStatus.SUCCEEDED) + ) + assert otel_context.get_current() == before + assert plugin._context_tokens == {} + assert not plugin._tracing_enabled + 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 324d31f06..4c9f1e644 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 @@ -1466,7 +1467,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() @@ -1479,6 +1480,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(): @@ -2179,7 +2186,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") @@ -2297,20 +2306,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()} == { @@ -2470,3 +2482,146 @@ 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": {"child": 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_uses_invocation_fallback_and_is_not_carried_into_next_instance() -> ( + 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, + ) + ) + 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", + "missing-child", + } + spans = {span.name: span for span in exporter.get_finished_spans()} + assert spans["missing-child"].parent is not None + assert spans["missing-child"].parent.span_id == spans["Invocation"].context.span_id + plugin, exporter = _create_plugin() + 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..28d16a3b2 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_terminal_input_parent_refs.py @@ -0,0 +1,479 @@ +"""Completion parentage prefers actual spans and falls back only at invocation End.""" + +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, +) +from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( + operation_id_to_span_id, +) + + +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 not plugin._tracing_enabled and plugin._context_tokens == {} + + try: + contextvars.Context().run(invoke) + finally: + provider.shutdown() + + +@pytest.mark.parametrize("view", [ExecutionOtelPlugin, InvocationOtelPlugin]) +def test_unavailable_parent_exports_once_with_fallback_before_context_cleanup( + view, caplog +): + 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)) + plugin.on_operation_end(_info(OperationEndInfo, now)) + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.PENDING) + ) + children = [ + s + for s in exporter.get_finished_spans() + if s.attributes.get("durable.operation.id") == "b" + ] + assert len(children) == 1 + invocation_span = next( + span for span in exporter.get_finished_spans() if span.name == "Invocation" + ) + assert children[0].parent is not None + assert children[0].context.trace_id == invocation_span.context.trace_id + assert children[0].parent.span_id == ( + invocation_span.context.span_id + if view is InvocationOtelPlugin + else operation_id_to_span_id(invocation["execution_arn"], "parent") + ) + exported = exporter.get_finished_spans() + assert exported.index(children[0]) < exported.index(invocation_span) + assert ( + len( + [ + record + for record in caplog.records + if "fallback parent" in record.message + ] + ) + == 1 + ) + assert plugin._pending_completions == plugin._pending_operation_ends == {} + assert otel_context.get_current() == before + assert not plugin._tracing_enabled and plugin._context_tokens == {} + + 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 not plugin._tracing_enabled 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 not plugin._tracing_enabled and otel_context.get_current() == before + assert plugin._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 not plugin._tracing_enabled and otel_context.get_current() == before + + try: + contextvars.Context().run(invoke) + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python/README.md b/packages/aws-durable-execution-sdk-python/README.md index f05e878ee..335b32549 100644 --- a/packages/aws-durable-execution-sdk-python/README.md +++ b/packages/aws-durable-execution-sdk-python/README.md @@ -167,6 +167,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. +Invocation Start exposes the input's visible `UpdatedOperationIds` through +`InvocationStartInfo.updated_operations`. The operation hooks follow ordinary +execution and replay. Repeated observations of a terminal checkpoint deliver one +operation-end hook per invocation. Cached parent scopes can skip their descendant +hooks when the completion is present only in the invocation input. Plugins that +require completion delivery for those unvisited operations use the invocation +snapshots and lifecycle hooks. + When plugins are registered, the worker begins with a copy of the caller's context-variable bindings. Start hooks, the handler, serialization, resource cleanup and End hooks use that same Context. Changes made by a Start hook remain 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 fd72dae91..bf84f60e0 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 @@ -10,7 +10,7 @@ from collections.abc import Collection, Iterator, Mapping, Sequence from dataclasses import dataclass, field from enum import Enum -from typing import Any, Callable, MutableMapping, Protocol, cast +from typing import Any, Callable, MutableMapping, Protocol, Self, cast from aws_durable_execution_sdk_python.identifier import OperationIdentifier from aws_durable_execution_sdk_python.lambda_service import ( @@ -108,13 +108,14 @@ class OperationInfo: """The operation error, when available.""" attempt: int | None = field(default=None, kw_only=True) - @staticmethod + @classmethod def from_operation( + cls, operation: Operation, *, is_replayed: bool = False, - ) -> OperationInfo: - return OperationInfo( + ) -> Self: + return cls( operation_id=operation.operation_id, operation_type=_to_operation_type(operation.operation_type), sub_type=operation.sub_type, @@ -137,7 +138,7 @@ def _copy_error(error: ErrorObject | None) -> ErrorObject | None: The checkpointed ``ErrorObject`` is handed straight to user code on replay, and its ``stack_trace`` is a mutable list. Without a copy a plugin reading - ``info.operations`` could append to (or clear) that list and change the error + operation hooks or ``info.operations`` could change that list and the error the execution later raises. Only the list needs cloning -- the other fields are immutable strings -- so this is cheaper than a full deep copy. """ @@ -1629,8 +1630,6 @@ def on_invocation_start( if operation_id in operations }, ) - # Build this invocation's plugin instances from the very info their first - # hook receives, and before that hook is dispatched. self._create_plugins(self._invocation_status) self.execute_plugins(self._invocation_status) @@ -1853,22 +1852,7 @@ def on_operation_update( self._submit_change(change_info) def _end_info(self, operation: Operation) -> OperationEndInfo: - return 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), - status=operation.status, - error=self._extract_error(operation), - attempt=( - operation.step_details.attempt if operation.step_details else None - ), - is_replayed=False, - ) + return OperationEndInfo.from_operation(operation) def _change_info( self, 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..f78e38566 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_overlap_int_test.py @@ -0,0 +1,69 @@ +"""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"] + + class Factory: + def create_plugin(self, info): + return SafePlugin() + + wrapped = durable_execution(body, plugins=[Factory()]) + 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/execution_test.py b/packages/aws-durable-execution-sdk-python/tests/execution_test.py index 0cf360f34..4bf8c4d89 100644 --- a/packages/aws-durable-execution-sdk-python/tests/execution_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/execution_test.py @@ -7,7 +7,10 @@ import threading import time import warnings +import contextvars from collections.abc import Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import replace from copy import deepcopy from typing import Any, cast from unittest.mock import Mock, patch @@ -3937,6 +3940,79 @@ 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"] + + class Factory: + def create_plugin(self, info): + return SafePlugin() + + wrapped = durable_execution(body, plugins=[Factory()]) + 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_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py index acad8b9db..79cffa27f 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -9,10 +9,14 @@ from collections.abc import Iterator from copy import deepcopy from dataclasses import asdict, fields +from itertools import product from unittest.mock import MagicMock, patch from aws_durable_execution_sdk_python.identifier import OperationIdentifier from aws_durable_execution_sdk_python.lambda_service import ( + CallbackDetails, + ChainedInvokeDetails, + ContextDetails, DurableExecutionInvocationOutput, ErrorObject, InvocationStatus as ServiceInvocationStatus, @@ -391,6 +395,25 @@ def test_operation_info_converts_service_operation_type(self): self.assertIs(info.operation_type, OperationType.WAIT) + def test_operation_info_factory_preserves_subclass_and_replay_flag(self) -> None: + operation = Operation( + operation_id="wait-1", + operation_type=ServiceOperationType.WAIT, + status=OperationStatus.SUCCEEDED, + start_timestamp=START_TS, + end_timestamp=END_TS, + ) + + for info_type in (OperationInfo, OperationStartInfo, OperationEndInfo): + with self.subTest(info_type=info_type): + info = info_type.from_operation(operation, is_replayed=True) + + self.assertIs(type(info), info_type) + self.assertIs(info.operation_type, OperationType.WAIT) + self.assertEqual(info.start_time, START_TS) + self.assertEqual(info.end_time, END_TS) + self.assertTrue(info.is_replayed) + def test_invocation_end_info_converts_service_invocation_status(self): output = DurableExecutionInvocationOutput( status=ServiceInvocationStatus.PENDING, @@ -2260,6 +2283,161 @@ def test_terminal_status_without_step_details_fires_operation_only(self): self.assertIn("operation_end:op-1", self.plugin.calls) + def test_repeated_checkpoint_observations_reach_plugins_once(self): + """Replay and checkpoint observations share one terminal hook per invocation.""" + 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_factory(plugin)]) + operation = self._make_operation(status=status) + with _invocation(executor, plugin): + # Natural replay delivery of an UpdatedOperationIds result. + 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 _invocation(self.executor, self.plugin): + 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 _invocation(self.executor, self.plugin): + 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): + executor = PluginExecutor([plugin_factory(self.plugin)]) + with _invocation(executor, self.plugin): + executor.on_operation_update(operation) + self.assertEqual(self.plugin.calls.count("operation_end:op-1"), 1) + + def test_plugin_error_mutations_do_not_change_checkpoint_or_later_hooks( + self, + ) -> None: + """Snapshots and terminal hooks must own their mutable error stack traces.""" + for operation_type, checkpoint_delivery in product( + ( + ServiceOperationType.STEP, + ServiceOperationType.CALLBACK, + ServiceOperationType.CHAINED_INVOKE, + ServiceOperationType.CONTEXT, + ), + (False, True), + ): + with self.subTest( + operation_type=operation_type, checkpoint_delivery=checkpoint_delivery + ): + error = ErrorObject( + message="saved error", + type="ExternalError", + data="details", + stack_trace=["original frame"], + ) + operation = Operation( + operation_id="op-1", + operation_type=operation_type, + status=OperationStatus.FAILED, + step_details=( + StepDetails(attempt=2, error=error) + if operation_type is ServiceOperationType.STEP + else None + ), + callback_details=( + CallbackDetails(callback_id="callback-1", error=error) + if operation_type is ServiceOperationType.CALLBACK + else None + ), + chained_invoke_details=( + ChainedInvokeDetails(error=error) + if operation_type is ServiceOperationType.CHAINED_INVOKE + else None + ), + context_details=( + ContextDetails(error=error) + if operation_type is ServiceOperationType.CONTEXT + else None + ), + ) + seen: list[list[str]] = [] + + class MutatingPlugin(DurableInstrumentationPlugin): + def on_invocation_start(self, info: InvocationStartInfo) -> None: + snapshot_error = info.updated_operations["op-1"].error + assert snapshot_error is not None + assert snapshot_error.stack_trace is not None + snapshot_error.stack_trace.clear() + + def on_operation_end(self, info: OperationEndInfo) -> None: + assert info.error is not None + assert info.error.stack_trace is not None + seen.append(list(info.error.stack_trace)) + info.error.stack_trace.clear() + + def on_invocation_end(self, info: InvocationEndInfo) -> None: + snapshot_error = info.operations["op-1"].error + assert snapshot_error is not None + assert snapshot_error.stack_trace is not None + seen.append(list(snapshot_error.stack_trace)) + + executor = PluginExecutor([plugin_factory(MutatingPlugin())]) + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=False, + operations_provider=lambda: {"op-1": operation}, + updated_operation_ids=["op-1"], + ) + self.assertEqual(error.stack_trace, ["original frame"]) + self.assertEqual(seen, []) + # Both delivery paths must snapshot errors before hooks run. + if checkpoint_delivery: + executor.record_end_hooks([operation]) + executor.drain_end_hooks() + else: + executor.on_operation_update(operation) + # A later observation must not repeat the terminal hook. + executor.on_operation_update( + [operation], + operations={"op-1": operation}, + previous_operations={"op-1": operation}, + ) + executor.on_invocation_end( + DurableExecutionInvocationOutput( + status=ServiceInvocationStatus.SUCCEEDED, + ) + ) + + self.assertEqual(seen, [["original frame"]] * 2) + self.assertEqual(error.stack_trace, ["original frame"]) + def test_non_terminal_status_without_step_details_fires_nothing(self): op = self._make_operation(status=OperationStatus.STARTED, step_details=None) From c3d7ae1e1fdcfd77f975e290837b3e741a21a92e Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Thu, 22 Oct 2026 16:58:57 -0700 Subject: [PATCH 2/3] fix(otel): isolate invocation finalization failures --- .../README.md | 6 + .../_invocation_state.py | 51 ++- .../execution_plugin.py | 84 +++-- .../invocation_plugin.py | 91 ++++-- ...st_invocation_finalization_failures_int.py | 266 +++++++++++++++ .../tests/test_finalization_parent_start.py | 163 ++++++++++ .../test_invocation_finalization_failures.py | 306 ++++++++++++++++++ .../tests/test_log_filter.py | 15 +- 8 files changed, 909 insertions(+), 73 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_finalization_failures_int.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/test_finalization_parent_start.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_finalization_failures.py diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index b1d34cc2e..fd99e5cb4 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -475,6 +475,12 @@ pending for a later notification or invocation-End retry. If a processor deliver a span and then raises, retrying can repeat that delivery; the plugin has no exporter acknowledgment that would distinguish partial delivery from failure. +During invocation End, an ordinary completion-export failure receives at most one +retry. Failures do not prevent other completions or the Invocation/Workflow spans +from finishing. Errors are reported after scope cleanup and flushing, including +recovered transient failures. Control exceptions are not retried and propagate +after cleanup; they are not converted into ordinary export failures. + Invocation view can open another context segment within the same invocation when a timed branch resumes and re-enters a completed child, or when a virtual context runs again. Each new segment receives its own terminal outcome. Repeated 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 index 52b46e231..abeca2dcd 100644 --- 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 @@ -1,6 +1,6 @@ -"""Visible completed context metadata for terminal-update parent references.""" +"""Invocation-owned completion metadata and bounded finalization.""" -from collections.abc import Mapping +from collections.abc import Callable, Mapping from aws_durable_execution_sdk_python.plugin import ( InvocationStartInfo, @@ -10,6 +10,53 @@ ) +class InvocationFinalization: + """Defer failures until invocation cleanup, with one ordinary export retry. + + Operation claims and failures use the plugin's registry lock. Finalization + actions and error reporting run on the invocation owner. Control exceptions + are held only to finish cleanup; they are never retried or swallowed. + """ + + def __init__(self) -> None: + self._attempts: dict[str, int] = {} + self._stopped_operations: set[str] = set() + self._retryable_operations: set[str] = set() + self._errors: list[BaseException] = [] + + def claim_operation(self, operation_id: str) -> bool: + attempts = self._attempts.get(operation_id, 0) + if attempts >= 2 or operation_id in self._stopped_operations: + return False + self._attempts[operation_id] = attempts + 1 + self._retryable_operations.discard(operation_id) + return True + + def operation_failed(self, operation_id: str, error: BaseException) -> None: + self._errors.append(error) + if not isinstance(error, Exception): + self._stopped_operations.add(operation_id) + elif self._attempts[operation_id] < 2: + self._retryable_operations.add(operation_id) + + def can_retry(self, operation_id: str) -> bool: + return operation_id in self._retryable_operations + + def run(self, action: Callable[[], None]) -> None: + try: + action() + except BaseException as error: + self._errors.append(error) + + def raise_errors(self) -> None: + if len(self._errors) == 1: + raise self._errors[0] + if self._errors: + raise BaseExceptionGroup( + "OpenTelemetry invocation finalization failed", self._errors + ) + + def terminal_updates( operations: Mapping[str, OperationInfo], ) -> dict[str, OperationInfo]: 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 57bc13657..c04d0ddef 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 @@ -69,6 +69,7 @@ ) from aws_durable_execution_sdk_python_otel._invocation_state import ( + InvocationFinalization, completed_input_parents, terminal_updates, ) @@ -202,6 +203,7 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: # repeated on_operation_end does not export it twice. self._ended_operation_ids: set[str] = set() self._ending_operation_ids: set[str] = set() + self._finalization: InvocationFinalization | None = None self._exception_recorded_spans: dict[str, Span] = {} self._completion_end_timestamps: dict[str, int] = {} # Tokens returned by context.attach(), keyed by the span registry key, @@ -710,18 +712,21 @@ def _end_open_recording_spans(self) -> None: """ with self._lock: keys = list(reversed(self._operation_spans)) + finalization = self._finalization + assert finalization is not None for key in keys: - if key == _INVOCATION_KEY: + if key == _INVOCATION_KEY or key in self._pending_completions: continue - span = self._get_span(key) - if span is None or not span.is_recording(): - continue - popped = self._pop_span(key) - if popped is not None: - popped.set_attribute( - "durable.span.truncated_at_invocation_boundary", True - ) - popped.end() + finalization.run(lambda: self._finish_open_span(key)) + + def _finish_open_span(self, key: str) -> None: + span = self._get_span(key) + if span is None or not span.is_recording(): + return + popped = self._pop_span(key) + if popped is not None: + popped.set_attribute("durable.span.truncated_at_invocation_boundary", True) + popped.end() def _start_invocation_span(self, info: InvocationStartInfo) -> None: self._invocation_span = self._tracer.start_span( @@ -750,14 +755,30 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: self._release_invocation_scope() return + finalization = InvocationFinalization() + self._finalization = finalization try: - self._end_invocation_spans(info) + finalization.run(lambda: self._end_invocation_spans(info)) finally: - self._release_invocation_scope() - self._force_flush() + try: + finalization.run(self._release_invocation_scope) + finalization.run(self._force_flush) + finally: + self._finalization = None + finalization.raise_errors() def _end_invocation_spans(self, info: InvocationEndInfo) -> None: - """End this invocation's open spans and export the Workflow span.""" + """Finish each phase even when another completion or root export fails.""" + finalization = self._finalization + assert finalization is not None + finalization.run(self._drain_pending_completions) + finalization.run(lambda: self._finish_invocation_span(info)) + finalization.run(self._end_open_recording_spans) + if info.status in _TERMINAL_INVOCATION_STATUSES: + finalization.run(lambda: self._export_workflow_span(info)) + + def _drain_pending_completions(self) -> None: + """Drain visible completions; the finalization ledger bounds retries.""" with self._lock: pending = tuple(self._pending_completions.values()) for operation in pending: @@ -804,15 +825,7 @@ def _end_invocation_spans(self, info: InvocationEndInfo) -> None: ) self._render_operation_end(operation, fallback_parent=parent) - # 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 - # completed the execution or cleanly suspended) - # FAILED -> ERROR - # RETRY -> UNSET - # RETRY is left UNSET because the plugin interface cannot tell whether the - # execution/workflow was STOPPED or TIMED_OUT: a RETRY invocation is not a - # definitive failure of the execution, so we avoid marking the span ERROR. + def _finish_invocation_span(self, info: InvocationEndInfo) -> None: if self._invocation_span is not None: self._invocation_span.set_attribute( "durable.invocation.status", @@ -828,17 +841,6 @@ def _end_invocation_spans(self, info: InvocationEndInfo) -> None: ) self._invocation_span.end() - # End recording user-function spans left open by a suspended operation. - self._end_open_recording_spans() - - # The Workflow span (execution view) is a non-recording placeholder - # during the invocation, so only a terminal status materializes and ends - # the recording span. Its span status reflects the execution outcome: - # SUCCEEDED -> OK, FAILED -> ERROR (RETRY/PENDING are non-terminal and - # leave the Workflow span unexported until a later terminal invocation). - if info.status in _TERMINAL_INVOCATION_STATUSES: - self._export_workflow_span(info) - def _force_flush(self) -> None: """Flush pending spans, containing any error the flush raises. @@ -962,6 +964,11 @@ def _render_operation_end( info.operation_id ] = info return + finalization = self._finalization + if finalization is not None and not finalization.claim_operation( + info.operation_id + ): + return self._ending_operation_ids.add(info.operation_id) try: self._finish_operation_end(info, fallback_parent) @@ -976,9 +983,18 @@ def _render_operation_end( siblings.pop(info.operation_id, None) if not siblings: del self._pending_operation_ends[info.parent_id] + except BaseException as error: + if finalization is None: + raise + with self._lock: + finalization.operation_failed(info.operation_id, error) finally: with self._lock: self._ending_operation_ids.discard(info.operation_id) + if finalization is not None and finalization.can_retry(info.operation_id): + # Retry before returning to a parent registration that may close + # its span. Reuse the same fallback and enforce the shared bound. + self._render_operation_end(info, fallback_parent=fallback_parent) def _finish_operation_end( self, info: OperationInfo, fallback_parent: Span | None 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 f333ea297..40ae63a4e 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 @@ -38,6 +38,7 @@ ) from aws_durable_execution_sdk_python_otel._invocation_state import ( + InvocationFinalization, completed_input_parents, terminal_updates, ) @@ -206,6 +207,7 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: self._pending_completions: dict[str, OperationInfo] = {} self._ended_operation_ids: set[str] = set() self._ending_operation_ids: set[str] = set() + self._finalization: InvocationFinalization | None = None self._failed_operation_ids: set[str] = set() self._exception_recorded_spans: dict[str, Span] = {} self._completion_end_timestamps: dict[str, int] = {} @@ -867,18 +869,49 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: self._release_invocation_scope() return + finalization = InvocationFinalization() + self._finalization = finalization try: - # Start and End share the owner Context. Release its bindings before - # span processors run while closing this invocation's spans. - self._detach_remaining_contexts() - self._end_invocation_spans(info) + # Release owner bindings before invoking customer span processors. + finalization.run(self._detach_remaining_contexts) + finalization.run(lambda: self._end_invocation_spans(info)) finally: - self._release_invocation_scope() - # Flush before Lambda freeze. - self._force_flush() + try: + finalization.run(self._release_invocation_scope) + finalization.run(self._force_flush) + finally: + self._finalization = None + finalization.raise_errors() def _end_invocation_spans(self, info: InvocationEndInfo) -> None: - """End this invocation's open spans and export the Workflow span.""" + """Finish each phase even when another completion or root export fails.""" + finalization = self._finalization + assert finalization is not None + finalization.run(self._drain_pending_completions) + # Failed completions already exhausted their bounded attempts. Do not + # retry them a third time as ordinary open-span cleanup. + with self._operation_spans_lock: + operation_ids = list(reversed(self._operation_spans)) + incomplete_attempt_span_keys = set(self._incomplete_attempt_span_keys) + for operation_id in operation_ids: + if operation_id and operation_id not in self._pending_completions: + finalization.run( + lambda: self._finish_open_span( + operation_id, operation_id in incomplete_attempt_span_keys + ) + ) + finalization.run(lambda: self._finish_invocation_span(info)) + if info.status in _TERMINAL_INVOCATION_STATUSES: + finalization.run(lambda: self._export_workflow_span(info)) + + def _finish_open_span(self, operation_id: str, incomplete_attempt: bool) -> None: + span = self._get_span(operation_id) + if incomplete_attempt and span is not None and span.is_recording(): + span.set_attribute("durable.span.truncated_at_invocation_boundary", True) + self._end_span(operation_id) + + def _drain_pending_completions(self) -> None: + """Drain visible completions; the finalization ledger bounds retries.""" with self._operation_spans_lock: pending = tuple(self._pending_completions.values()) for operation in pending: @@ -935,24 +968,7 @@ def _end_invocation_spans(self, info: InvocationEndInfo) -> None: ) self._render_operation_end(operation, fallback_parent=invocation_span) - # Spans are registered parent-first, so close pending spans in reverse - # order to keep every child contained within its parent. - with self._operation_spans_lock: - operation_ids = list(reversed(self._operation_spans)) - incomplete_attempt_span_keys = set(self._incomplete_attempt_span_keys) - for operation_id in operation_ids: - if operation_id: - span = self._get_span(operation_id) - if ( - operation_id in incomplete_attempt_span_keys - and span is not None - and span.is_recording() - ): - span.set_attribute( - "durable.span.truncated_at_invocation_boundary", True - ) - self._end_span(operation_id) - + def _finish_invocation_span(self, info: InvocationEndInfo) -> None: invocation_span = self._get_span(None) if invocation_span: invocation_span.set_attribute( @@ -976,13 +992,6 @@ def _end_invocation_spans(self, info: InvocationEndInfo) -> None: # end the invocation span self._end_span(None) - # The Workflow span (execution view) is a non-recording placeholder - # during the invocation, so only a terminal status materializes and ends - # the recording span. SUCCEEDED -> OK, FAILED -> ERROR; RETRY/PENDING are - # non-terminal and leave it unexported until a later terminal invocation. - if info.status in _TERMINAL_INVOCATION_STATUSES: - self._export_workflow_span(info) - def _force_flush(self) -> None: """Flush pending spans, containing any error the flush raises. @@ -1138,6 +1147,11 @@ def _render_operation_end( info.operation_id ] = info return + finalization = self._finalization + if finalization is not None and not finalization.claim_operation( + info.operation_id + ): + return self._ending_operation_ids.add(info.operation_id) try: self._finish_operation_end(info, span, fallback_parent) @@ -1153,13 +1167,20 @@ def _render_operation_end( siblings.pop(info.operation_id, None) if not siblings: del self._pending_operation_ends[info.parent_id] - except BaseException: + except BaseException as error: with self._operation_spans_lock: self._failed_operation_ids.add(info.operation_id) - raise + if finalization is not None: + finalization.operation_failed(info.operation_id, error) + if finalization is None: + raise finally: with self._operation_spans_lock: self._ending_operation_ids.discard(info.operation_id) + if finalization is not None and finalization.can_retry(info.operation_id): + # Retry before falling back to a completed-parent reference, so + # descendants keep this invocation's actual parent segment. + self._render_operation_end(info, fallback_parent=fallback_parent) def _finish_operation_end( self, info: OperationInfo, span: Span | None, fallback_parent: Span | None diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_finalization_failures_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_finalization_failures_int.py new file mode 100644 index 000000000..0d0302a8e --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_finalization_failures_int.py @@ -0,0 +1,266 @@ +"""A deferred completion failure must not discard unrelated invocation telemetry.""" + +from copy import deepcopy +from datetime import UTC, datetime, timedelta +import logging +from threading import Event +from typing import Any +from unittest.mock import Mock, patch + +import pytest +from aws_durable_execution_sdk_python import DurableContext, durable_execution +from aws_durable_execution_sdk_python.lambda_service import ( + CallbackDetails, + ErrorObject, + ExecutionDetails, + Operation, + OperationStatus, + OperationType, +) +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationEndInfo, + InvocationStartInfo, + InvocationStatus, +) +from opentelemetry import context as otel_context +from opentelemetry.context import Context +from opentelemetry.sdk.trace import ReadableSpan, Span, 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 ( + ExecutionOtelPluginFactory, + InvocationOtelPluginFactory, + OtelPluginConfig, +) + + +class _FailCallbackProcessor(SimpleSpanProcessor): + def __init__( + self, + exporter: InMemorySpanExporter, + invocation_ending: Event, + permanent: bool, + ) -> None: + super().__init__(exporter) + self.invocation_ending = invocation_ending + self.permanent = permanent + self.error = RuntimeError("callback A processor failed before export") + self.a_starts: list[Span] = [] + self.a_attempts: list[ReadableSpan] = [] + self.failed_after_close: list[bool] = [] + self.during_invocation_end: list[bool] = [] + + def on_start(self, span: Span, parent_context: Context | None = None) -> None: + if span.name == "a": + self.a_starts.append(span) + super().on_start(span, parent_context) + + def on_end(self, span: ReadableSpan) -> None: + if span.name == "a": + self.a_attempts.append(span) + self.during_invocation_end.append(self.invocation_ending.is_set()) + if self.permanent or len(self.a_attempts) == 1: + current = self.a_starts[-1] + self.failed_after_close.append( + not current.is_recording() and current.end_time is not None + ) + raise self.error + super().on_end(span) + + +def _contains_error(error: BaseException, expected: BaseException) -> bool: + if error is expected: + return True + return isinstance(error, BaseExceptionGroup) and any( + _contains_error(nested, expected) for nested in error.exceptions + ) + + +@pytest.mark.parametrize( + "factory_type", [ExecutionOtelPluginFactory, InvocationOtelPluginFactory] +) +@pytest.mark.parametrize("permanent", [False, True]) +def test_deferred_export_failure_preserves_other_completions_and_roots( + factory_type: type[ExecutionOtelPluginFactory] | type[InvocationOtelPluginFactory], + permanent: bool, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + caplog.set_level(logging.ERROR) + provider = TracerProvider() + exporter = InMemorySpanExporter() + invocation_ending = Event() + processor = _FailCallbackProcessor(exporter, invocation_ending, permanent) + provider.add_span_processor(processor) + invocation_ends: list[InvocationEndInfo] = [] + + class Observe(DurableInstrumentationPlugin): + def on_invocation_end(self, info: InvocationEndInfo) -> None: + invocation_ends.append(info) + invocation_ending.set() + + class ObserveFactory: + def create_plugin(self, info: InvocationStartInfo) -> Observe: + return Observe() + + factory = factory_type( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + start = datetime.now(UTC) - timedelta(seconds=10) + callback_error = ErrorObject( + type="RemoteCallbackError", + message="callback A failed", + data="callback details", + stack_trace=["callback frame"], + ) + operations = [ + Operation( + operation_id="execution-id", + operation_type=OperationType.EXECUTION, + status=OperationStatus.STARTED, + start_timestamp=start, + execution_details=ExecutionDetails(input_payload='{"result":"winner"}'), + ), + Operation( + operation_id="a", + operation_type=OperationType.CALLBACK, + name="a", + status=OperationStatus.FAILED, + start_timestamp=start + timedelta(seconds=1), + end_timestamp=start + timedelta(seconds=2), + callback_details=CallbackDetails( + callback_id="callback-a", error=callback_error + ), + ), + Operation( + operation_id="b", + operation_type=OperationType.CALLBACK, + name="b", + status=OperationStatus.SUCCEEDED, + start_timestamp=start + timedelta(seconds=1), + end_timestamp=start + timedelta(seconds=2), + callback_details=CallbackDetails(callback_id="callback-b", result='"B"'), + ), + ] + invocation: dict[str, Any] = { + "DurableExecutionArn": ( + "arn:aws:lambda:us-west-2:123456789012:function:completion:1" + "/durable-execution/finalization/execution-id" + ), + "CheckpointToken": "resume-token", + "InitialExecutionState": { + "Operations": [operation.to_json_dict() for operation in operations], + "NextMarker": "", + }, + "UpdatedOperationIds": ["a", "b"], + } + original_invocation = deepcopy(invocation) + lambda_context = Mock( + aws_request_id="request", + client_context=None, + identity=None, + _epoch_deadline_time_in_ms=0, + invoked_function_arn="completion:1", + tenant_id=None, + ) + handler_calls: list[dict[str, str]] = [] + + def body(event: dict[str, str], _context: DurableContext) -> str: + handler_calls.append(dict(event)) + assert processor.a_attempts == [] + assert not any( + span.name in ("a", "b") for span in exporter.get_finished_spans() + ) + return event["result"] + + handler = durable_execution(body, plugins=[ObserveFactory(), factory]) + service = Mock() + before = otel_context.get_current() + try: + with patch( + "aws_durable_execution_sdk_python.execution.LambdaClient.initialize_client", + return_value=service, + ): + result = handler(invocation, lambda_context) + + assert result == {"Status": "SUCCEEDED", "Result": '"winner"'} + assert handler_calls == [{"result": "winner"}] + assert invocation == original_invocation + service.checkpoint.assert_not_called() + (end_info,) = invocation_ends + assert end_info.status is InvocationStatus.SUCCEEDED + assert end_info.operations["a"].error == callback_error + assert end_info.operations["b"].result == '"B"' + spans = exporter.get_finished_spans() + names = [span.name for span in spans] + assert ( + names.count("b") + == names.count("Invocation") + == names.count("Workflow") + == 1 + ) + assert names.count("a") == int(not permanent) + assert len(processor.a_attempts) == 2 + assert processor.during_invocation_end == [True, True] + assert processor.failed_after_close == [True] * (2 if permanent else 1) + + first_a, retried_a = processor.a_attempts + assert retried_a.context == first_a.context + assert retried_a.parent == first_a.parent + assert retried_a.start_time == first_a.start_time + assert retried_a.end_time == first_a.end_time + assert retried_a.attributes == first_a.attributes + for attempt in processor.a_attempts: + assert attempt.status.status_code is StatusCode.ERROR + assert attempt.attributes is not None + assert attempt.attributes["durable.operation.status"] == "FAILED" + assert len(attempt.events) == 1 + assert attempt.events[0].attributes is not None + assert ( + attempt.events[0].attributes["exception.message"] + == callback_error.message + ) + if not permanent: + assert [span for span in spans if span.name == "a"] == [retried_a] + + (b,) = [span for span in spans if span.name == "b"] + (invocation_span,) = [span for span in spans if span.name == "Invocation"] + (workflow_span,) = [span for span in spans if span.name == "Workflow"] + assert b.attributes is not None + assert b.attributes["durable.operation.id"] == "b" + assert b.attributes["durable.operation.status"] == "SUCCEEDED" + assert b.status.status_code is StatusCode.OK + assert b.events == () + assert invocation_span.attributes is not None + assert invocation_span.attributes["durable.invocation.status"] == "SUCCEEDED" + assert workflow_span.attributes is not None + assert workflow_span.attributes["durable.execution.status"] == "SUCCEEDED" + assert invocation_span.status.status_code is StatusCode.OK + assert workflow_span.status.status_code is StatusCode.OK + expected_parent = ( + workflow_span + if factory_type is ExecutionOtelPluginFactory + else invocation_span + ) + assert expected_parent.context is not None + for callback in (b, retried_a): + assert callback.parent == expected_parent.context + assert callback.context is not None + assert callback.context.trace_id == expected_parent.context.trace_id + assert any( + record.exc_info is not None + and record.exc_info[1] is not None + and _contains_error(record.exc_info[1], processor.error) + for record in caplog.records + ) + assert otel_context.get_current() == before + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_finalization_parent_start.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_finalization_parent_start.py new file mode 100644 index 000000000..be548c5db --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_finalization_parent_start.py @@ -0,0 +1,163 @@ +"""Parent materialization failures preserve child and root finalization.""" + +from collections import Counter +from datetime import UTC, datetime, timedelta +from typing import Any + +import pytest +from aws_durable_execution_sdk_python.plugin import ( + InvocationEndInfo, + InvocationStartInfo, + InvocationStatus, + OperationInfo, + OperationStatus, + OperationType, +) +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 opentelemetry.trace import Span, StatusCode + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) +from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( + operation_id_to_span_id, +) + + +def _error_leaves(error: BaseException) -> list[BaseException]: + if isinstance(error, BaseExceptionGroup): + return [leaf for nested in error.exceptions for leaf in _error_leaves(nested)] + return [error] + + +@pytest.mark.parametrize("view", [InvocationOtelPlugin, ExecutionOtelPlugin]) +@pytest.mark.parametrize("permanent", [False, True], ids=["transient", "permanent"]) +def test_parent_start_failure_finishes_descendants_before_reporting_errors( + view: type[InvocationOtelPlugin] | type[ExecutionOtelPlugin], + permanent: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + 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) + execution_arn = "execution/finalization-parent-start" + invocation: dict[str, Any] = dict( + request_id="parent-start-failure", + execution_arn=execution_arn, + execution_start_time=now - timedelta(seconds=10), + is_first_invocation=False, + ) + parent = OperationInfo( + operation_id="parent", + operation_type=OperationType.CONTEXT, + sub_type=None, + name="parent", + parent_id=None, + start_time=now - timedelta(seconds=4), + end_time=now - timedelta(seconds=1), + is_replayed=False, + status=OperationStatus.SUCCEEDED, + ) + child = OperationInfo( + operation_id="child", + operation_type=OperationType.CALLBACK, + sub_type=None, + name="child", + parent_id="parent", + start_time=now - timedelta(seconds=3), + end_time=now - timedelta(seconds=2), + is_replayed=False, + status=OperationStatus.SUCCEEDED, + ) + operations = {"parent": parent, "child": child} + before = otel_context.get_current() + error = RuntimeError("parent span could not start") + attempts: Counter[str] = Counter() + original_start = plugin._start_span + + def start_span(*args: Any, **kwargs: Any) -> Span: + operation_id = kwargs.get("operation_id") + if operation_id == "parent": + attempts[operation_id] += 1 + if permanent or attempts[operation_id] == 1: + raise error + return original_start(*args, **kwargs) + + try: + plugin.on_invocation_start( + InvocationStartInfo( + **invocation, operations=operations, updated_operations=operations + ) + ) + monkeypatch.setattr(plugin, "_start_span", start_span) + with pytest.raises(Exception) as caught: + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.SUCCEEDED) + ) + + assert _error_leaves(caught.value) == [error] * (2 if permanent else 1) + assert attempts["parent"] == 2 + assert otel_context.get_current() == before + assert not plugin._tracing_enabled + assert plugin._context_tokens == {} + assert plugin._finalization is None + assert not plugin._pending_completions + assert not plugin._pending_operation_ends + assert not plugin._ending_operation_ids + + spans = exporter.get_finished_spans() + counts = Counter(span.name for span in spans) + assert counts["child"] == counts["Invocation"] == counts["Workflow"] == 1 + assert counts["parent"] == (0 if permanent else 1) + assert ("parent" in plugin._ended_operation_ids) is not permanent + child_span = next(span for span in spans if span.name == "child") + invocation_span = next(span for span in spans if span.name == "Invocation") + assert child_span.parent is not None + assert child_span.context is not None + assert invocation_span.context is not None + assert child_span.context.trace_id == invocation_span.context.trace_id + assert child_span.status.status_code is StatusCode.OK + assert child_span.attributes is not None + assert child_span.attributes["durable.operation.status"] == "SUCCEEDED" + assert spans.index(child_span) < spans.index(invocation_span) + + if permanent: + # Known completed input still supplies the logical parent identity, + # without falsely claiming that its recording span was exported. + assert child_span.parent.span_id == operation_id_to_span_id( + execution_arn, "parent" + ) + else: + parent_span = next(span for span in spans if span.name == "parent") + assert parent_span.context is not None + assert child_span.parent.span_id == parent_span.context.span_id + assert parent_span.status.status_code is StatusCode.OK + assert parent_span.attributes is not None + assert parent_span.attributes["durable.operation.status"] == "SUCCEEDED" + assert parent_span.start_time is not None + assert parent_span.end_time is not None + assert child_span.start_time is not None + assert child_span.end_time is not None + assert ( + parent_span.start_time + <= child_span.start_time + < child_span.end_time + <= parent_span.end_time + ) + assert spans.index(child_span) < spans.index(parent_span) + assert spans.index(parent_span) < spans.index(invocation_span) + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_finalization_failures.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_finalization_failures.py new file mode 100644 index 000000000..6f1dbd7a3 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_finalization_failures.py @@ -0,0 +1,306 @@ +"""A failing finalizer cannot prevent unrelated spans from ending.""" + +from collections import Counter +from datetime import UTC, datetime, timedelta + +import pytest +from aws_durable_execution_sdk_python.plugin import ( + InvocationEndInfo, + InvocationStartInfo, + InvocationStatus, + OperationInfo, + OperationStatus, + OperationType, +) +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 opentelemetry.trace import StatusCode + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) +from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( + operation_id_to_span_id, +) + + +class _FailingProcessor(SimpleSpanProcessor): + def __init__(self, exporter, error, *, permanent=False): + super().__init__(exporter) + self.error = error + self.permanent = permanent + self.attempts = Counter() + self.failed_contexts = [] + self.failed_end_times = [] + + def on_end(self, span): + self.attempts[span.name] += 1 + if span.name == "a" and (self.permanent or self.attempts[span.name] == 1): + # Real SDK Span.end has already marked this span ended. Reusing + # end() on the same object cannot deliver it to this exporter. + self.failed_contexts.append(span.context) + self.failed_end_times.append(span.end_time) + raise self.error + super().on_end(span) + + +def _operation(name, now, *, parent_id=None): + return OperationInfo( + operation_id=name, + operation_type=OperationType.CALLBACK, + sub_type=None, + name=name, + parent_id=parent_id, + start_time=now - timedelta(seconds=2), + end_time=now - timedelta(seconds=1), + is_replayed=False, + status=OperationStatus.SUCCEEDED, + ) + + +def _error_leaves(error): + if isinstance(error, BaseExceptionGroup): + return [leaf for nested in error.exceptions for leaf in _error_leaves(nested)] + return [error] + + +@pytest.mark.parametrize("view", [InvocationOtelPlugin, ExecutionOtelPlugin]) +@pytest.mark.parametrize("permanent", [False, True], ids=["transient", "permanent"]) +@pytest.mark.parametrize("missing_parent", [False, True], ids=["root", "fallback"]) +def test_end_drain_reports_failure_after_other_spans_finish( + view, permanent, missing_parent +): + provider = TracerProvider() + exporter = InMemorySpanExporter() + error = RuntimeError("completion a could not be exported") + processor = _FailingProcessor(exporter, error, permanent=permanent) + provider.add_span_processor(processor) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + now = datetime.now(UTC) + invocation = dict( + request_id="finalization-failure", + execution_arn="execution/finalization-failure", + execution_start_time=now - timedelta(seconds=10), + is_first_invocation=False, + ) + parent_id = "missing-parent" if missing_parent else None + operations = { + name: _operation(name, now, parent_id=parent_id) for name in ("a", "b") + } + before = otel_context.get_current() + try: + plugin.on_invocation_start( + InvocationStartInfo( + **invocation, operations=operations, updated_operations=operations + ) + ) + # Neither callback is visited by the application. Both first reach the + # exporter from invocation End, including the missing-parent queue. + assert processor.attempts["a"] == processor.attempts["b"] == 0 + with pytest.raises(Exception) as caught: + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.SUCCEEDED) + ) + assert error in _error_leaves(caught.value) + if not permanent: + assert caught.value is error + assert otel_context.get_current() == before + assert not plugin._tracing_enabled + assert plugin._context_tokens == {} + assert plugin._pending_completions == plugin._pending_operation_ends == {} + assert not plugin._ending_operation_ids + + spans = exporter.get_finished_spans() + counts = Counter(span.name for span in spans) + assert counts["b"] == counts["Invocation"] == counts["Workflow"] == 1 + assert processor.attempts["a"] == 2 + assert processor.attempts["b"] == 1 + assert counts["a"] == (0 if permanent else 1) + invocation_span = next(span for span in spans if span.name == "Invocation") + workflow_span = next(span for span in spans if span.name == "Workflow") + expected_parent = ( + invocation_span.context.span_id + if view is InvocationOtelPlugin + else operation_id_to_span_id(invocation["execution_arn"], parent_id) + if parent_id is not None + else workflow_span.context.span_id + ) + for span in spans: + if span.name in ("a", "b"): + assert span.status.status_code is StatusCode.OK + assert span.attributes["durable.operation.status"] == "SUCCEEDED" + assert span.parent.span_id == expected_parent + assert span.context.trace_id == invocation_span.context.trace_id + assert span.start_time < span.end_time <= invocation_span.end_time + if not permanent: + recovered = next(span for span in spans if span.name == "a") + assert recovered.context == processor.failed_contexts[0] + assert recovered.end_time == processor.failed_end_times[0] + # A retained failure cannot create work after this invocation closes. + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.SUCCEEDED) + ) + assert processor.attempts["a"] == 2 + finally: + provider.shutdown() + + +@pytest.mark.parametrize("view", [InvocationOtelPlugin, ExecutionOtelPlugin]) +@pytest.mark.parametrize("error", [SystemExit("stop"), KeyboardInterrupt()]) +def test_control_exception_is_not_retried_and_propagates_after_finalization( + view, error +): + provider = TracerProvider() + exporter = InMemorySpanExporter() + processor = _FailingProcessor(exporter, error) + provider.add_span_processor(processor) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + now = datetime.now(UTC) + invocation = dict( + request_id="control-finalization", + execution_arn="execution/control-finalization", + execution_start_time=now - timedelta(seconds=10), + is_first_invocation=False, + ) + operations = {name: _operation(name, now) for name in ("a", "b")} + before = otel_context.get_current() + try: + plugin.on_invocation_start( + InvocationStartInfo( + **invocation, operations=operations, updated_operations=operations + ) + ) + with pytest.raises(type(error)) as caught: + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.SUCCEEDED) + ) + assert caught.value is error + assert otel_context.get_current() == before + assert not plugin._tracing_enabled + assert plugin._context_tokens == {} + assert not plugin._pending_completions + assert processor.attempts["a"] == 1 + counts = Counter(span.name for span in exporter.get_finished_spans()) + assert counts["a"] == 0 + assert counts["b"] == counts["Invocation"] == counts["Workflow"] == 1 + finally: + provider.shutdown() + + +@pytest.mark.parametrize("view", [InvocationOtelPlugin, ExecutionOtelPlugin]) +@pytest.mark.parametrize("control", [False, True], ids=["transient", "control"]) +def test_parent_registration_child_failure_keeps_recursive_drain_bounded(view, control): + provider = TracerProvider() + exporter = InMemorySpanExporter() + error = ( + SystemExit("stop child export") + if control + else RuntimeError("retry child export") + ) + parent_contexts = [] + + class Processor(_FailingProcessor): + def on_end(self, span): + if span.name == "a": + # The child first reaches the processor inside the parent's + # registration-driven drain, while that parent is recording. + parent_span = plugin._get_span("parent") + assert parent_span is not None and parent_span.is_recording() + parent_contexts.append(parent_span.get_span_context()) + super().on_end(span) + + processor = Processor(exporter, error) + provider.add_span_processor(processor) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + now = datetime.now(UTC) + invocation = dict( + request_id="nested-finalization", + execution_arn="execution/nested-finalization", + execution_start_time=now - timedelta(seconds=10), + is_first_invocation=False, + ) + parent = OperationInfo( + operation_id="parent", + operation_type=OperationType.CONTEXT, + sub_type=None, + name="parent", + parent_id=None, + start_time=now - timedelta(seconds=4), + end_time=now - timedelta(seconds=1), + is_replayed=False, + status=OperationStatus.SUCCEEDED, + ) + # Parent-first input makes materializing its recording span recursively + # drain both children before the outer pending-completion loop reaches them. + operations = { + "parent": parent, + "a": _operation("a", now, parent_id="parent"), + "b": _operation("b", now, parent_id="parent"), + } + before = otel_context.get_current() + try: + plugin.on_invocation_start( + InvocationStartInfo( + **invocation, operations=operations, updated_operations=operations + ) + ) + with pytest.raises(type(error)) as caught: + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.SUCCEEDED) + ) + assert caught.value is error + assert processor.attempts["a"] == (1 if control else 2) + assert processor.attempts["b"] == processor.attempts["parent"] == 1 + assert otel_context.get_current() == before + assert not plugin._tracing_enabled + assert plugin._context_tokens == {} + assert plugin._finalization is None + assert not plugin._pending_completions + assert not plugin._pending_operation_ends + assert not plugin._ending_operation_ids + + spans = exporter.get_finished_spans() + counts = Counter(span.name for span in spans) + assert counts["b"] == counts["parent"] == 1 + assert counts["Invocation"] == counts["Workflow"] == 1 + assert counts["a"] == (0 if control else 1) + parent_span = next(span for span in spans if span.name == "parent") + invocation_span = next(span for span in spans if span.name == "Invocation") + assert all(context == parent_span.context for context in parent_contexts) + assert parent_span.status.status_code is StatusCode.OK + assert spans.index(parent_span) < spans.index(invocation_span) + for child in (span for span in spans if span.name in ("a", "b")): + assert child.parent.span_id == parent_span.context.span_id + assert child.context.trace_id == parent_span.context.trace_id + assert child.status.status_code is StatusCode.OK + assert ( + parent_span.start_time + <= child.start_time + < child.end_time + <= parent_span.end_time + ) + assert spans.index(child) < spans.index(parent_span) + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py index ea68446e6..8228f048a 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py @@ -623,10 +623,14 @@ def test_finished_invocation_does_not_correlate_later_records(): class _SpanProcessorFailingOnEnd(SpanProcessor): """Raises when a span ends, standing in for customer span-processor code.""" + def __init__(self) -> None: + self.ended_names: list[str] = [] + def on_start(self, span, parent_context=None) -> None: pass def on_end(self, span) -> None: + self.ended_names.append(span.name) # The fallback root is exported during Start. Fail the later shutdown # path that this regression is intended to exercise. if span.name != "DurableExecutionRoot": @@ -648,7 +652,8 @@ def test_invocation_is_released_when_span_shutdown_fails(plugin_class): records to its finished spans. """ provider = TracerProvider() - provider.add_span_processor(_SpanProcessorFailingOnEnd()) + processor = _SpanProcessorFailingOnEnd() + provider.add_span_processor(processor) plugin = plugin_class( OtelPluginConfig( tracer_provider=provider, @@ -658,8 +663,14 @@ def test_invocation_is_released_when_span_shutdown_fails(plugin_class): ) plugin.on_invocation_start(_invocation_start_info()) - with pytest.raises(RuntimeError, match="span processor on_end failed"): + with pytest.raises(ExceptionGroup) as caught: plugin.on_invocation_end(_invocation_end_info()) + assert processor.ended_names == ["DurableExecutionRoot", "Invocation", "Workflow"] + assert len(caught.value.exceptions) == 2 + assert all( + isinstance(error, RuntimeError) and str(error) == "span processor on_end failed" + for error in caught.value.exceptions + ) assert plugin not in log_filter_module._open_invocations # This thread carried the claim, so it is the thread a later record would be From 3bf2bfa50d61928e8b8c845ad00054687ffd9502 Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Thu, 22 Oct 2026 16:20:17 -0700 Subject: [PATCH 3/3] fix(testing): match replay state to the service --- .github/scripts/install_otel_test_wheels.py | 5 +- .../otel_lifecycle_compatibility_test.py | 256 +++++++++++++++ .github/workflows/ci.yml | 5 +- .github/workflows/cloud-tests.yml | 5 +- .github/workflows/conformance-tests.yml | 5 +- .github/workflows/js-examples.yml | 5 +- .../opentelemetry-conformance-tests.yml | 5 +- CONTRIBUTING.md | 2 +- .../README.md | 4 +- .../test_cached_child_input_delivery_int.py | 295 ++++++++++++++++++ .../e2e/test_callback_input_delivery_int.py | 207 ++++++++++++ .../tests/e2e/test_late_child_parent_int.py | 245 +++++++++++++++ .../e2e/test_parent_creation_order_int.py | 262 ++++++++++++++++ .../tests/e2e/test_terminal_status_int.py | 181 +++++++++++ .../README.md | 25 ++ .../docs/architecture.md | 16 +- .../checkpoint/processor.py | 5 +- .../checkpoint/processors/base.py | 8 +- .../execution.py | 58 ++-- .../executor.py | 10 +- .../invoker.py | 10 +- .../model.py | 35 ++- .../operation_state.py | 64 ++++ .../e2e/callback_updated_operations_test.py | 284 +++++++++++++++++ .../e2e/invocation_input_delivery_test.py | 154 +++++++++ ...replay_children_service_parity_int_test.py | 256 +++++++++++++++ .../tests/event_factory_test.py | 61 ++++ .../tests/execution_test.py | 114 +++++++ .../tests/executor_checkpoint_test.py | 44 +++ .../tests/executor_test.py | 4 + .../tests/invocation_input_delivery_test.py | 98 ++++++ .../tests/model_test.py | 114 +++++++ .../replay_children_service_parity_test.py | 239 ++++++++++++++ .../tests/runner_test.py | 76 +++++ .../e2e/plugin_operation_type_int_test.py | 12 +- 35 files changed, 3122 insertions(+), 47 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_cached_child_input_delivery_int.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_callback_input_delivery_int.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_late_child_parent_int.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_parent_creation_order_int.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_terminal_status_int.py create mode 100644 packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/operation_state.py create mode 100644 packages/aws-durable-execution-sdk-python-testing/tests/e2e/callback_updated_operations_test.py create mode 100644 packages/aws-durable-execution-sdk-python-testing/tests/e2e/invocation_input_delivery_test.py create mode 100644 packages/aws-durable-execution-sdk-python-testing/tests/e2e/replay_children_service_parity_int_test.py create mode 100644 packages/aws-durable-execution-sdk-python-testing/tests/invocation_input_delivery_test.py create mode 100644 packages/aws-durable-execution-sdk-python-testing/tests/replay_children_service_parity_test.py diff --git a/.github/scripts/install_otel_test_wheels.py b/.github/scripts/install_otel_test_wheels.py index 938996aca..d4f136867 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] + # Both lanes use the built local runner, including its real callback-update + # producer. The legacy lane keeps the actual published OTel plugin unchanged. + packages = [CORE, TESTING] if args.legacy_plugin else [CORE, OTEL, TESTING] wheels = [built_wheel(package) for package in packages] subprocess.run( [ diff --git a/.github/tests/otel_lifecycle_compatibility_test.py b/.github/tests/otel_lifecycle_compatibility_test.py index e315464ba..ae40189d4 100644 --- a/.github/tests/otel_lifecycle_compatibility_test.py +++ b/.github/tests/otel_lifecycle_compatibility_test.py @@ -18,6 +18,7 @@ OperationStatus, OperationType, ) +from aws_durable_execution_sdk_python.plugin import InvocationStartInfo from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin from aws_durable_execution_sdk_python_otel.plugin_factory import ( ExecutionOtelPluginFactory, @@ -171,3 +172,258 @@ def handler(event, durable_context): provider.shutdown() contextvars.Context().run(run) + + +@pytest.mark.parametrize( + "view", [InvocationOtelPluginFactory, ExecutionOtelPluginFactory] +) +def test_installed_pair_preserves_nested_callback_terminal_on_resume(view): + """The installed Core 3 / OTel 2 pair retains callback completion on resume.""" + from aws_durable_execution_sdk_python_testing.runner import ( + DurableFunctionTestRunner, + ) + import aws_durable_execution_sdk_python_testing as installed_testing + + import aws_durable_execution_sdk_python_otel as installed_otel + + assert "site-packages" in Path(installed_testing.__file__).resolve().parts + assert "site-packages" in Path(core_execution.__file__).resolve().parts + assert "site-packages" in Path(installed_otel.__file__).resolve().parts + assert version("aws-durable-execution-sdk-python-otel") == "2.0.0" + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + enrich_logger=False, + context_extractor=lambda _: None, + ) + ) + submitted = threading.Event() + suspended = threading.Event() + callback_ids = [] + updated_ids = [] + invocation_statuses = [] + side_effects = [] + + class Observer(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + updated_ids.append(list(info.updated_operations)) + + def on_invocation_end(self, info): + invocation_statuses.append(info.status.value) + if info.status.value == "PENDING": + suspended.set() + + class ObserverFactory: + def create_plugin( + self, info: InvocationStartInfo + ) -> DurableInstrumentationPlugin: + return Observer() + + def submit(callback_id, _context): + side_effects.append("submit") + callback_ids.append(callback_id) + submitted.set() + + def saved(_context): + side_effects.append("saved") + return "saved" + + def body(_event, context): + context.step(saved, name="saved-before") + return context.wait_for_callback(submit, name="await-external") + + handler = durable_execution(body, plugins=[plugin, ObserverFactory()]) + try: + with DurableFunctionTestRunner( + handler=handler, skip_time=False, poll_interval=0.01, execution_timeout=20 + ) as runner: + arn = runner.run_async(input="{}") + assert submitted.wait(8) and suspended.wait(8) + runner.send_callback_success(callback_ids[0], result=b"done") + result = runner.wait_for_result(arn, timeout=10) + assert result.status.value == "SUCCEEDED" + assert result.result == '"done"' + assert invocation_statuses == ["PENDING", "SUCCEEDED"] + assert side_effects == ["saved", "submit"] + assert updated_ids[0] == [] and len(updated_ids[1]) == 1 + spans = exporter.get_finished_spans() + callback_spans = [ + span for span in spans if span.name == "await-external create callback id" + ] + terminal = [ + span + for span in callback_spans + if span.attributes["durable.operation.status"] == "SUCCEEDED" + ] + assert len(terminal) == 1 + assert terminal[0].attributes["durable.operation.id"] == updated_ids[1][0] + parent_spans = [span for span in spans if span.name == "await-external"] + assert terminal[0].parent in [span.context for span in parent_spans] + assert all( + span.context.trace_id == terminal[0].context.trace_id + for span in callback_spans + ) + finally: + provider.shutdown() + + +@pytest.mark.parametrize( + "view", [InvocationOtelPluginFactory, ExecutionOtelPluginFactory] +) +def test_installed_pair_waits_for_completed_replaychildren_parent(view): + """A real large child result retains and replays its callback's parent.""" + import json + import time + + from aws_durable_execution_sdk_python.config import CallbackConfig + from aws_durable_execution_sdk_python.serdes import JsonSerDes + from aws_durable_execution_sdk_python_testing.runner import ( + DurableFunctionTestRunner, + ) + import aws_durable_execution_sdk_python_otel as installed_otel + import aws_durable_execution_sdk_python_testing as installed_testing + + for module in (core_execution, installed_otel, installed_testing): + assert "site-packages" in Path(module.__file__).resolve().parts + assert version("aws-durable-execution-sdk-python-otel") == "2.0.0" + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + enrich_logger=False, + context_extractor=lambda _: None, + ) + ) + starts, ends, input_parents = [], [], [] + child_calls, saved_calls, marker_calls = [], [], [] + exports_at_marker = [] + + class Observer(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + starts.append(info) + + def on_invocation_end(self, info): + ends.append(info.status.value) + + class ObserverFactory: + def create_plugin( + self, info: InvocationStartInfo + ) -> DurableInstrumentationPlugin: + return Observer() + + def callback_spans(): + return [ + span + for span in exporter.get_finished_spans() + if span.name == "b" + and span.attributes.get("durable.operation.status") == "SUCCEEDED" + ] + + def saved(_context): + saved_calls.append(1) + return "saved" + + def child(context): + child_calls.append(1) + callback = context.create_callback( + name="b", config=CallbackConfig(serdes=JsonSerDes()) + ) + assert context.step(saved, name="saved-child") == "saved" + return {"id": callback.callback_id, "large": "x" * (300 * 1024)} + + def marker(_context): + marker_calls.append(1) + exports_at_marker.extend(callback_spans()) + return "marked" + + def body(_event, context): + a = context.create_callback( + name="a", config=CallbackConfig(serdes=JsonSerDes()) + ) + b_id = context.run_in_child_context(child, name="large-parent")["id"] + result = a.result() + assert context.step(marker, name="after-a") == "marked" + return {"a": result, "b_id": b_id} + + wrapped = durable_execution(body, plugins=[ObserverFactory(), plugin]) + + def entrypoint(event, context): + input_parents.append( + [ + (op.operation_id, op.status.value, op.context_details.replay_children) + for op in event.initial_execution_state.operations + if op.name == "large-parent" and op.context_details is not None + ] + ) + 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 ( + len(ends) >= count + and sum(e.event_type == "InvocationCompleted" for e in events) >= count + ): + return + time.sleep(0.01) + raise AssertionError("Invocation did not finish") + + try: + with DurableFunctionTestRunner( + handler=entrypoint, + 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) + completed(runner, arn, 1) + runner.send_callback_success(b_id, result=b'"B"') + completed(runner, arn, 2) + middle_exports = callback_spans() + runner.send_callback_success(a_id, result=b'"A"') + result = runner.wait_for_result(arn, timeout=10) + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == {"a": "A", "b_id": b_id} + assert ends == ["PENDING", "PENDING", "SUCCEEDED"] + assert child_calls == [1, 1, 1] and saved_calls == [1] and marker_calls == [1] + b_info = next( + info for info in starts[1].updated_operations.values() if info.name == "b" + ) + assert input_parents[1] == [(b_info.parent_id, "SUCCEEDED", True)] + terminal = callback_spans() + assert len(middle_exports) == len(terminal) == 1 + assert exports_at_marker == middle_exports == terminal + assert terminal[0].attributes["durable.operation.id"] == b_info.operation_id + parents = [ + span + for span in exporter.get_finished_spans() + if span.attributes.get("durable.operation.id") == b_info.parent_id + and span.attributes.get("durable.operation.status") == "SUCCEEDED" + ] + assert terminal[0].parent in [span.context for span in parents] + assert all( + span.context.trace_id == terminal[0].context.trace_id for span in parents + ) + if view is InvocationOtelPluginFactory: + invocations = [ + span + for span in exporter.get_finished_spans() + if span.name == "Invocation" + ] + assert len(parents) == len(invocations) == 3 + assert terminal[0].parent == parents[1].context + assert parents[1].parent == invocations[1].context + else: + assert len(parents) == 1 + finally: + provider.shutdown() diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index dfc01a16e..c9ee85ce9 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -10,7 +10,10 @@ on: branches: [ main ] pull_request: - branches: [ main ] + branches: + - main + - test/otel-conformance-feasibility + - fix/python-testing-replay-state-20261021 jobs: lint-commits: diff --git a/.github/workflows/cloud-tests.yml b/.github/workflows/cloud-tests.yml index 91e9f0a27..3cb2f2d63 100644 --- a/.github/workflows/cloud-tests.yml +++ b/.github/workflows/cloud-tests.yml @@ -2,7 +2,10 @@ name: Cloud tests on: pull_request: - branches: ["main"] + branches: + - main + - test/otel-conformance-feasibility + - fix/python-testing-replay-state-20261021 workflow_dispatch: inputs: lambda_endpoint: diff --git a/.github/workflows/conformance-tests.yml b/.github/workflows/conformance-tests.yml index b3d70c205..9b201a14a 100644 --- a/.github/workflows/conformance-tests.yml +++ b/.github/workflows/conformance-tests.yml @@ -8,7 +8,10 @@ on: push: branches: ["main"] pull_request: - branches: ["main"] + branches: + - main + - test/otel-conformance-feasibility + - fix/python-testing-replay-state-20261021 paths: - "packages/aws-durable-execution-sdk-python/**" - "packages/aws-durable-execution-sdk-python-conformance-tests/**" diff --git a/.github/workflows/js-examples.yml b/.github/workflows/js-examples.yml index 1b911c7bb..3439d3be9 100644 --- a/.github/workflows/js-examples.yml +++ b/.github/workflows/js-examples.yml @@ -20,7 +20,10 @@ on: - ".github/workflows/js-examples.yml" - "pyproject.toml" pull_request: - branches: ["main"] + branches: + - main + - test/otel-conformance-feasibility + - fix/python-testing-replay-state-20261021 paths: - "packages/aws-durable-execution-sdk-python-testing/**" - ".github/scripts/js_examples/**" diff --git a/.github/workflows/opentelemetry-conformance-tests.yml b/.github/workflows/opentelemetry-conformance-tests.yml index aed97f6e8..097fed102 100644 --- a/.github/workflows/opentelemetry-conformance-tests.yml +++ b/.github/workflows/opentelemetry-conformance-tests.yml @@ -8,7 +8,10 @@ name: OpenTelemetry Conformance Tests on: pull_request: - branches: ["main"] + branches: + - main + - test/otel-conformance-feasibility + - fix/python-testing-replay-state-20261021 paths: - "packages/aws-durable-execution-sdk-python/**" - "packages/aws-durable-execution-sdk-python-otel/**" diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index ce5faaf30..2f90f9d3c 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -72,7 +72,7 @@ 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 diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index fd99e5cb4..7078bd39a 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -489,7 +489,9 @@ durable checkpoint observations still do not close or re-export that segment. The service omits descendants of completed contexts unless their ancestors retain children through `ReplayChildren`. An `UpdatedOperationIds` entry alone does not provide missing operation metadata, and the plugin does not synthesize an -operation span from that missing record. +operation 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 descendants delivered by later checkpoint responses. Both views retain actual 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..8185da3a7 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_cached_child_input_delivery_int.py @@ -0,0 +1,295 @@ +"""Visible terminal input updates export 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 InvocationStartInfo +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 ( + ExecutionOtelPluginFactory, + InvocationOtelPluginFactory, + OtelPluginConfig, +) + + +@pytest.mark.parametrize( + ("view", "filesystem", "failed_delivery"), + [ + (ExecutionOtelPluginFactory, False, False), + (InvocationOtelPluginFactory, False, False), + (ExecutionOtelPluginFactory, True, False), + (InvocationOtelPluginFactory, True, False), + (ExecutionOtelPluginFactory, 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)) + + class ObserveFactory: + def create_plugin( + self, info: InvocationStartInfo + ) -> DurableInstrumentationPlugin: + return Observe() + + 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=[ObserveFactory(), 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" + ] + b_inputs = [ + (invocation, info) + for invocation, start in enumerate(starts, start=1) + for info in start.updated_operations.values() + 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_input_deliveries": [ + (n, i.operation_id, i.parent_id, i.is_replayed) + for n, i in b_inputs + ], + "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" + ) + # Input-only completions arrive in the invocation snapshot; the core + # does not synthesize an End hook for an unvisited callback result. + expected_delivery = ( + 2, + b_info.operation_id, + b_info.parent_id, + "SUCCEEDED", + False, + ) + assert [ + ( + invocation, + info.operation_id, + info.parent_id, + info.status.value, + info.is_replayed, + ) + for invocation, info in b_inputs + ] == [expected_delivery] + # Any natural replay/checkpoint End notifications may repeat only in + # this invocation. Neither delivery path may report B on invocation 3. + assert { + ( + invocation, + info.operation_id, + info.parent_id, + info.status.value, + info.is_replayed, + ) + for invocation, info in b_ends + } <= {expected_delivery} + assert len(terminal_spans) == 1 + assert terminal_spans[0].parent is not None + if view is ExecutionOtelPluginFactory: + 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..17a943209 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_callback_input_delivery_int.py @@ -0,0 +1,207 @@ +"""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 InvocationStartInfo +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 ( + ExecutionOtelPluginFactory, + InvocationOtelPluginFactory, + OtelPluginConfig, +) + + +@pytest.mark.parametrize( + "view,nested,page_bytes,fresh_executor,concurrent_completion", + [ + (ExecutionOtelPluginFactory, False, None, False, False), + (InvocationOtelPluginFactory, False, None, False, False), + (ExecutionOtelPluginFactory, True, None, False, False), + (InvocationOtelPluginFactory, True, None, False, False), + (InvocationOtelPluginFactory, True, 1, False, False), + (ExecutionOtelPluginFactory, False, None, True, False), + (InvocationOtelPluginFactory, 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) + + class ObserveFactory: + def create_plugin( + self, info: InvocationStartInfo + ) -> DurableInstrumentationPlugin: + return Observe() + + observer = ObserveFactory() + + 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] + # Replay reads and a later checkpoint may both report the same terminal + # operation. It must not be delivered again in a later invocation. + assert set(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 InvocationOtelPluginFactory: + 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_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..25268edfe --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_late_child_parent_int.py @@ -0,0 +1,245 @@ +"""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 InvocationStartInfo +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 ( + ExecutionOtelPluginFactory, + InvocationOtelPluginFactory, + OtelPluginConfig, +) + + +@pytest.mark.parametrize( + "view", [ExecutionOtelPluginFactory, InvocationOtelPluginFactory] +) +@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) + + class ObserveFactory: + def create_plugin( + self, info: InvocationStartInfo + ) -> DurableInstrumentationPlugin: + return Observe() + + 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=[ObserveFactory(), 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 InvocationOtelPluginFactory: + 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..45be9dd87 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_parent_creation_order_int.py @@ -0,0 +1,262 @@ +"""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, + OperationChangeInfo, + OperationEndInfo, + OperationStatus, + 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, + InvocationOtelPluginFactory, + 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)) + instances: list[InvocationOtelPlugin] = [] + + class RecordingFactory(InvocationOtelPluginFactory): + def create_plugin(self, info: InvocationStartInfo) -> InvocationOtelPlugin: + instance = super().create_plugin(info) + instances.append(instance) + return instance + + plugin = RecordingFactory( + OtelPluginConfig( + tracer_provider=provider, + enrich_logger=False, + context_extractor=lambda _: None, + ) + ) + second_start = threading.Event() + both_completed = threading.Event() + completion_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: + assert not any( + operation.name == "target-1" + for operation in info.updated_operations.values() + ) + 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 completion_observed.wait(10) + hook_order.append("parent-entered") + parent_entered.set() + + class GateFactory: + def create_plugin( + self, info: InvocationStartInfo + ) -> DurableInstrumentationPlugin: + return Gate() + + class ObserveCompletion(DurableInstrumentationPlugin): + def on_operation_change(self, info: OperationChangeInfo) -> None: + if generations[0] != 2: + return + for operation in info.updated_operations.values(): + if ( + operation.name == "target-1" + and operation.operation_type is OperationType.CALLBACK + and operation.status is OperationStatus.SUCCEEDED + ): + # OTel has retained this checkpoint snapshot before the + # sibling registers its replay parent. The end hook waits + # for the sibling to read its result, so cannot release it. + assert not parent_entered.is_set() + assert ( + len( + [ + span + for span in exporter.get_finished_spans() + if span.name == "target-1" + ] + ) + == 1 + ) + hook_order.append("early-completion-observed") + completion_observed.set() + + def on_operation_end(self, info: OperationEndInfo) -> None: + if ( + generations[0] == 2 + and info.name == "target-1" + and info.operation_type is OperationType.CALLBACK + ): + assert parent_entered.is_set() + hook_order.append("terminal-hook-observed") + + class ObserveCompletionFactory: + def create_plugin( + self, info: InvocationStartInfo + ) -> DurableInstrumentationPlugin: + return ObserveCompletion() + + 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=[GateFactory(), plugin, ObserveCompletionFactory()] + ) + 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 completion_observed.is_set() and parent_entered.is_set() + assert ( + hook_order.index("early-completion-observed") + < hook_order.index("parent-entered") + < hook_order.index("terminal-hook-observed") + ) + 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 len(instances) == 4 + assert len({id(instance) for instance in instances}) == 4 + assert all(instance._context_tokens == {} for instance in instances) + finally: + both_completed.set() + completion_observed.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..f7368d22a --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_terminal_status_int.py @@ -0,0 +1,181 @@ +"""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 InvocationStartInfo +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 ( + ExecutionOtelPluginFactory, + InvocationOtelPluginFactory, + OtelPluginConfig, +) + + +@pytest.mark.parametrize( + "view", [ExecutionOtelPluginFactory, InvocationOtelPluginFactory] +) +@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() + + class ObserverFactory: + def create_plugin( + self, info: InvocationStartInfo + ) -> DurableInstrumentationPlugin: + return Observer() + + @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, ObserverFactory()]) + 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) + # Core exposes natural replay/checkpoint observations. OTel owns export + # deduplication, while the runner must not redeliver in a later invocation. + assert target_ends + expected_status = "SUCCEEDED" if outcome == "success" else "FAILED" + for invocation_number, info in target_ends: + assert invocation_number == (1 if completion == "initial" else 2) + assert info.status.value == expected_status + assert info.error == error + spans = [ + span for span in exporter.get_finished_spans() if span.name == "target" + ] + expected_statuses = [expected_status] + if view is InvocationOtelPluginFactory 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-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 0a33e4ec1..4780b5a97 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 5190b9b69..9da824e1a 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, @@ -479,15 +482,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 = [] @@ -692,6 +691,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( @@ -708,8 +708,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( @@ -718,6 +721,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( @@ -744,6 +748,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( @@ -897,10 +902,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 @@ -920,11 +924,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 ------------------------------------------------- @@ -953,14 +966,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 4c3a0a71f..1092291db 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 @@ -1551,7 +1551,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 @@ -1618,6 +1619,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 4992212f4..aec058fe5 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 @@ -27,6 +27,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, @@ -511,10 +514,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 5c82e4d75..df41a8441 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..aded098fa --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/callback_updated_operations_test.py @@ -0,0 +1,284 @@ +"""Callback completions are observed before resumed code without later redelivery.""" + +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 CallbackObserverFactory: + def __init__(self) -> None: + self.invocations: list[InvocationStartInfo] = [] + self.ends: list[tuple[int, OperationEndInfo]] = [] + self.order: list[str] = [] + + def create_plugin(self, info: InvocationStartInfo) -> DurableInstrumentationPlugin: + return CallbackObserver(self) + + +class CallbackObserver(DurableInstrumentationPlugin): + def __init__(self, records: CallbackObserverFactory) -> None: + self.invocations = records.invocations + self.ends = records.ends + self.order = records.order + self.invocation_number = 0 + + def on_invocation_start(self, info: InvocationStartInfo) -> None: + self.invocations.append(info) + self.invocation_number = len(self.invocations) + + def on_operation_end(self, info: OperationEndInfo) -> None: + if info.name == "target": + self.ends.append((self.invocation_number, 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 = CallbackObserverFactory() + 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 + # Multiple natural hooks in the delivery invocation are allowed, but an + # acknowledged callback must never be delivered in later invocations. + assert {invocation for invocation, _ in observer.ends} == {2} + terminal = observer.ends[0][1] + assert ( + terminal.status.value + == {"success": "SUCCEEDED", "failure": "FAILED", "timeout": "TIMED_OUT"}[ + outcome + ] + ) + assert all( + end.operation_id == terminal.operation_id + and end.status is terminal.status + and end.is_replayed is False + and end.result == terminal.result + and end.error == terminal.error + for _, end in observer.ends + ) + 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.count("observed-start") == 1 + assert observer.order.index("target-end") < observer.order.index("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 = CallbackObserverFactory() + 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 {invocation for invocation, _ in observer.ends} == {1} + terminal = observer.ends[0][1] + target_id = terminal.operation_id + assert all( + end.operation_id == target_id + and end.status is terminal.status + and end.is_replayed is False + and end.result == terminal.result + and end.error == terminal.error + for _, end in observer.ends + ) + 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..6eaacf5bf --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/invocation_input_delivery_test.py @@ -0,0 +1,154 @@ +"""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 InvocationStartInfo +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)) + + class ObserveFactory: + def create_plugin( + self, info: InvocationStartInfo + ) -> DurableInstrumentationPlugin: + return Observe() + + @durable_step + def marker(context, value): + marker_calls.append(1) + return value + + @durable_execution(plugins=[ObserveFactory()]) + 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", + ] + # Replay and checkpoint processing can each report the same terminal state. + # Acknowledgment prevents either callback from being redelivered to a later + # invocation; deduplicating exports within an invocation belongs to plugins. + assert set(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 08edbac70..3f2af3677 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 @@ -20,6 +20,7 @@ InvocationStatus, ) from aws_durable_execution_sdk_python.lambda_service import ( + CallbackDetails, ErrorObject, Operation as SvcOperation, OperationAction, @@ -133,6 +134,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 89a5618dd..392ad864b 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 @@ -736,6 +736,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 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 a7d2576ca..cb57f7a10 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/tests/e2e/plugin_operation_type_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_operation_type_int_test.py index ffabc11bf..bc4b1d65e 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 @@ -123,10 +123,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,