diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 6693855c3..fd99e5cb4 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -405,13 +405,108 @@ 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. + +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 +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..abeca2dcd --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_invocation_state.py @@ -0,0 +1,92 @@ +"""Invocation-owned completion metadata and bounded finalization.""" + +from collections.abc import Callable, Mapping + +from aws_durable_execution_sdk_python.plugin import ( + InvocationStartInfo, + OperationInfo, + OperationStatus, + OperationType, +) + + +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]: + """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 2cc841a35..7f383678d 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,18 @@ TraceState, ) +from aws_durable_execution_sdk_python_otel._invocation_state import ( + InvocationFinalization, + 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 +202,17 @@ 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._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, # 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 +248,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 +612,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 @@ -688,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: - continue - span = self._get_span(key) - if span is None or not span.is_recording(): + if key == _INVOCATION_KEY or key in self._pending_completions: 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( @@ -728,23 +755,77 @@ 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.""" - # 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. + """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: + 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) + + def _finish_invocation_span(self, info: InvocationEndInfo) -> None: if self._invocation_span is not None: self._invocation_span.set_attribute( "durable.invocation.status", @@ -760,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. @@ -782,7 +852,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 +882,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 +913,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 +946,65 @@ 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 - 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. + # 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 + 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) + 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] + 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 + ) -> 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 +1019,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 +1033,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 +1169,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 +1254,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 d6ab7f792..b15b46e16 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,18 @@ TraceState, ) +from aws_durable_execution_sdk_python_otel._invocation_state import ( + InvocationFinalization, + 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 +204,16 @@ 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._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] = {} + 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 +424,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 +579,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 +609,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 +761,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 @@ -799,36 +869,106 @@ 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.""" - # Spans are registered parent-first, so close pending spans in reverse - # order to keep every child contained within its parent. + """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: - 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 + 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: + 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 ) - self._end_span(operation_id) + 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) + + def _finish_invocation_span(self, info: InvocationEndInfo) -> None: invocation_span = self._get_span(None) if invocation_span: invocation_span.set_attribute( @@ -852,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. @@ -870,7 +1003,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 +1033,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 +1081,122 @@ 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 + 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) + 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 as error: + with self._operation_spans_lock: + self._failed_operation_ids.add(info.operation_id) + 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 + ) -> 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 +1209,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 +1289,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 +1392,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 +1408,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_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/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_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_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_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_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_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 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..3c843e0cb 100644 --- a/packages/aws-durable-execution-sdk-python/README.md +++ b/packages/aws-durable-execution-sdk-python/README.md @@ -167,6 +167,13 @@ 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: cached parent scopes can skip their descendant hooks, and +a terminal operation can be observed again in a later checkpoint response. +Plugins that require completion delivery for unvisited operations or deduplication +manage that from 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 33780eb23..ee71c49b3 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py @@ -9,7 +9,7 @@ from collections.abc import 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 ( @@ -107,13 +107,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, @@ -136,7 +137,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. """ @@ -438,9 +439,12 @@ def on_operation_start(self, info: OperationStartInfo) -> None: def on_operation_end(self, info: OperationEndInfo) -> None: """ - Called when an operation reaches a terminal status. Terminal operations - are not emitted again during replay. Child contexts without a terminal - checkpoint may emit this from the thread that runs the operation. + Called when ordinary execution or replay observes a terminal operation. + Cached parent scopes can skip descendant hooks, and repeated checkpoint + observations can emit the same operation again. Plugins that need + unvisited completions or deduplication must use the invocation snapshots + and lifecycle hooks. Child contexts without a terminal checkpoint may + emit this from their operation thread. Args: info: Information about the operation. @@ -1017,8 +1021,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) @@ -1205,26 +1207,7 @@ def on_operation_update( ) for operation in updated_operations: if self._is_terminal_status(operation.status): - self.execute_plugins( - 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, - ), - ) + self.execute_plugins(OperationEndInfo.from_operation(operation)) if ( operations is None 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..9bf8420ca 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -13,6 +13,9 @@ 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 +394,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 +2282,151 @@ 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_each_reach_plugins(self): + """Each terminal observation reaches hooks; plugins own export deduplication.""" + 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"), 2) + + 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 in ( + ServiceOperationType.STEP, + ServiceOperationType.CALLBACK, + ServiceOperationType.CHAINED_INVOKE, + ServiceOperationType.CONTEXT, + ): + with self.subTest(operation_type=operation_type): + 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, []) + # Replay and checkpoint observations convert independently. + executor.on_operation_update(operation) + 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"]] * 3) + 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)