From 626c17565679e661d6d85d969c56ef91a80a56b2 Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Wed, 28 Oct 2026 11:54:41 -0700 Subject: [PATCH] fix: export visible replay completions Retain terminal input and checkpoint updates until their parent registers or invocation finalization exports them. Preserve actual replay parent contexts and close each reopened invocation segment independently. Share completion claims, pending state, and successful-delivery tracking between both OTel views. Store payloads once, index children by ID, and retain completed input parents as an ID set. Validation: 901 OTel tests, type checking, formatting, and installed-wheel OTel/lifecycle validation. Twelve public-handler regressions fail against the base implementation because the callback completion is missing. Fixes #854 --- .../README.md | 25 + .../_invocation_state.py | 137 +++++ .../execution_plugin.py | 187 +++++-- .../invocation_plugin.py | 242 +++++++-- .../tests/e2e/test_completion_snapshot_int.py | 366 +++++++++++++ .../e2e/test_context_segment_reentry_int.py | 162 ++++++ .../tests/e2e/test_parent_registration_int.py | 187 +++++++ .../tests/test_completion_replay_parent.py | 132 +++++ .../tests/test_completion_snapshots.py | 388 ++++++++++++++ .../tests/test_context_segment_reentry.py | 250 +++++++++ .../tests/test_terminal_input_parent_refs.py | 481 ++++++++++++++++++ .../README.md | 12 +- 12 files changed, 2470 insertions(+), 99 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_invocation_state.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_completion_snapshot_int.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_context_segment_reentry_int.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_parent_registration_int.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/test_completion_replay_parent.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/test_completion_snapshots.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/test_context_segment_reentry.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/test_terminal_input_parent_refs.py diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 8e9ed0816..5ce25e899 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -349,6 +349,31 @@ Invocation span uses that ambient span as its parent. Ambient spans on a different trace are ignored for durable parenting so Invocation remains on the execution trace. +### Completion snapshots and replay + +The plugins consume visible terminal updates from invocation input and checkpoint +change notifications as well as natural operation End hooks. A completed parent +can return its cached result without revisiting an updated callback; the plugin +exports that unvisited completion before the Invocation span closes. Records in +the full replay snapshot alone are not new completion notifications. + +Pending children use a parent segment when it registers during replay, and retain +that actual parent context if it closes before a later update arrives. Completed +input-context metadata is used only after replay has had a chance to register its +current segment. At invocation End, pending parents are exported before their +children. If a parent's metadata and actual segment are unavailable, Invocation +view falls back to Invocation; Execution view uses the derived logical parent +context. The fallback is logged and cannot reconstruct a missing parent span. +Operation parent links are expected to form the durable service's tree. + +Successful completion observations are deduplicated within one invocation. +Invocation-view ReplayChildren and virtual contexts may still open distinct +segments; each segment closes with its own natural outcome. Completion metadata +and retained parent references are released at invocation End, so concurrent +invocations and warm reuse do not share the delivery state. This is not an +exactly-once export guarantee across failed or redelivered invocations, nor an +acknowledgment from the telemetry backend. + ### Synthetic fallback anchor Each sampled fallback invocation creates and ends `DurableExecutionRoot` at 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..486466a8d --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_invocation_state.py @@ -0,0 +1,137 @@ +"""Visible completion metadata owned by one instrumentation invocation.""" + +from collections.abc import Mapping +from graphlib import TopologicalSorter + +from aws_durable_execution_sdk_python.plugin import ( + InvocationStartInfo, + OperationInfo, + OperationStatus, + OperationType, +) + + +_TERMINAL_STATUSES = frozenset( + { + OperationStatus.SUCCEEDED, + OperationStatus.FAILED, + OperationStatus.TIMED_OUT, + OperationStatus.STOPPED, + OperationStatus.CANCELLED, + } +) + + +class CompletionState: + """Pending deliveries for one invocation, guarded by its span-registry lock. + + Callers hold that lock for every access, including parent registration and + dequeue. Span rendering runs outside the lock, between claim and release. + Only complete records successful delivery; a failed attempt stays pending. + """ + + def __init__(self) -> None: + self.pending: dict[str, OperationInfo] = {} + # An ordered ID index; completion payloads live only in pending. + self.children_by_parent: dict[str, dict[str, None]] = {} + self.in_flight: set[str] = set() + self.delivered: set[str] = set() + + def remember(self, updates: Mapping[str, OperationInfo]) -> None: + """Keep explicit terminal updates that have not already been delivered.""" + for operation_id, operation in updates.items(): + if ( + operation.status in _TERMINAL_STATUSES + and not operation.is_replayed + and operation_id not in self.delivered + ): + self.pending[operation_id] = operation + self._index_child(operation) + + def _index_child(self, operation: OperationInfo) -> None: + if operation.parent_id is not None: + self.children_by_parent.setdefault(operation.parent_id, {})[ + operation.operation_id + ] = None + + def take_children(self, parent_id: str) -> list[OperationInfo]: + """Dequeue children in observation order, retaining undelivered payloads.""" + return [ + self.pending[operation_id] + for operation_id in self.children_by_parent.pop(parent_id, {}) + ] + + def claim( + self, operation: OperationInfo, *, parent_available: bool, repeat: bool = False + ) -> bool: + """Claim delivery, or defer it until its parent is available. + + A newly opened context segment may repeat a delivered operation, but + it must still respect another caller's in-flight claim. + """ + operation_id = operation.operation_id + if operation_id in self.in_flight or ( + operation_id in self.delivered and not repeat + ): + return False + self.pending[operation_id] = operation + if not parent_available: + self._index_child(operation) + return False + self.in_flight.add(operation_id) + return True + + def complete(self, operation: OperationInfo) -> None: + """Record success and remove the payload and any remaining index entry.""" + operation_id = operation.operation_id + self.delivered.add(operation_id) + self.pending.pop(operation_id, None) + if operation.parent_id is not None: + siblings = self.children_by_parent.get(operation.parent_id) + if siblings is not None: + siblings.pop(operation_id, None) + if not siblings: + del self.children_by_parent[operation.parent_id] + + def release(self, operation_id: str) -> None: + """Release a claim after either successful or failed rendering.""" + self.in_flight.discard(operation_id) + + def clear(self) -> None: + """Release all invocation-owned delivery metadata.""" + self.pending.clear() + self.children_by_parent.clear() + self.in_flight.clear() + self.delivered.clear() + + +def completed_input_parent_ids(info: InvocationStartInfo) -> set[str]: + """Visible completed contexts may parent input or later checkpoint updates.""" + return { + operation.operation_id + for operation in info.operations.values() + if operation.operation_type is OperationType.CONTEXT + and operation.status in _TERMINAL_STATUSES + } + + +def parent_first( + operations: Mapping[str, OperationInfo], +) -> list[OperationInfo]: + """Materialize pending ancestors before referencing them for descendants. + + Durable operations form a parent tree. Missing ancestors remain unresolved; + their view-specific fallback is chosen by the plugin, without inventing an + operation or treating its non-recording reference as an exported completion. + """ + order = TopologicalSorter( + { + operation_id: ( + (operation.parent_id,) + if operation.parent_id is not None and operation.parent_id in operations + else () + ) + for operation_id, operation in operations.items() + } + ) + return [operations[operation_id] for operation_id in order.static_order()] 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 356990ede..411ccddca 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 @@ -42,6 +42,8 @@ InvocationStatus, InvocationStartInfo, OperationEndInfo, + OperationChangeInfo, + OperationInfo, OperationStartInfo, OperationStatus, OperationType, @@ -64,6 +66,11 @@ TraceState, ) +from aws_durable_execution_sdk_python_otel._invocation_state import ( + CompletionState, + completed_input_parent_ids, + parent_first, +) from aws_durable_execution_sdk_python_otel.context_extractors import ( ContextExtractor, ExtractedContext, @@ -192,13 +199,15 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: # A context absent from this set is checkpointless (for example, a FLAT # map/parallel branch) and needs a fresh per-invocation span identity. self._checkpointed_context_ids: set[str] = set() - # Operations whose span was already exported this invocation, so a - # repeated on_operation_end does not export it twice. - self._ended_operation_ids: set[str] = set() # Tokens returned by context.attach(), keyed by the span registry key, # paired with the thread that attached them. Every attach the plugin # owns is released through _detach_context so the plugin never leaves a # scope on the context stack. + # Completion delivery is invocation-owned, separate from active context + # segments. Registration and dequeue use the same span-registry lock. + self._completions = CompletionState() + self._completed_input_parent_ids: set[str] = set() + self._completed_parent_spans: dict[str, Span] = {} self._context_tokens: dict[str, tuple[int, object]] = {} self._lock = threading.RLock() self._tracing_enabled = False @@ -235,6 +244,8 @@ 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._completions.take_children(key) + self._drain_registered_children(pending) def _get_span(self, key: str | None) -> Span | None: if key is None: @@ -596,6 +607,8 @@ 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) + self._completed_input_parent_ids = completed_input_parent_ids(info) + 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 @@ -739,6 +752,7 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: def _end_invocation_spans(self, info: InvocationEndInfo) -> None: """End this invocation's open spans and export the Workflow span.""" + self._drain_pending_completions() # 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 @@ -815,6 +829,10 @@ def _release_invocation_scope(self) -> None: self._detach_remaining_contexts() unbind_invocation(self) self._tracing_enabled = False + with self._lock: + self._completions.clear() + self._completed_input_parent_ids.clear() + self._completed_parent_spans.clear() # ------------------------------------------------------------------ # Operation lifecycle @@ -840,24 +858,103 @@ 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: + """Observe one natural completion, sharing delivery with snapshots.""" + self._render_operation_end(info) + + def on_operation_change(self, info: OperationChangeInfo) -> None: + """Retain checkpoint completions until their natural End or finalization.""" + if self._tracing_enabled: + self._remember_completions(info.updated_operations) + + def _remember_completions(self, updates: dict[str, OperationInfo]) -> None: + with self._lock: + self._completions.remember(updates) + + def _completion_parent_span(self, parent_id: str | None) -> Span | None: + """Prefer the active replay segment, then its retained completed context.""" + with self._lock: + return ( + self._operation_spans.get(parent_id) + or self._completed_parent_spans.get(parent_id) + if parent_id is not None + else self._workflow_span + ) + + def _drain_pending_completions(self) -> None: + """Export pending parents before children while Invocation is still open.""" + with self._lock: + pending = dict(self._completions.pending) + for operation in parent_first(pending): + parent = self._completion_parent_span(operation.parent_id) + if parent is None: + parent = self._fallback_parent(operation) + self._render_operation_end(operation, fallback_parent=parent) + + def _render_operation_end( + self, info: OperationInfo, *, fallback_parent: Span | None = None + ) -> None: + """Claim a completion outside customer callbacks; record success last.""" logger.debug("Durable operation ended: %s", info) if not self._tracing_enabled: return - # ReplayChildren and virtual child contexts intentionally re-execute - # without creating a new terminal checkpoint. Their end callback is - # 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: + parent_available = ( + info.parent_id is None + or fallback_parent is not None + or self._completion_parent_span(info.parent_id) is not None + ) + if not self._completions.claim( + info, + parent_available=parent_available, + ): return - self._ended_operation_ids.add(info.operation_id) + try: + self._finish_operation_end(info, fallback_parent) + with self._lock: + self._completions.complete(info) + finally: + with self._lock: + self._completions.release(info.operation_id) + + def _drain_registered_children(self, pending: list[OperationInfo]) -> None: + """A child exporter must not interrupt its parent's context attachment.""" + for operation in pending: + try: + self._render_operation_end(operation) + except Exception: + logger.exception( + "Completion export failed while registering parent %s", + operation.parent_id, + ) + + def _fallback_parent(self, operation: OperationInfo) -> Span: + parent_id = operation.parent_id + parent_context = ( + self._operation_span_context(parent_id) if parent_id is not None else None + ) + if parent_context is None: + raise ValueError( + "No execution trace context available for completion fallback" + ) + if parent_id not in self._completed_input_parent_ids: + logger.warning( + "Using fallback parent context for operation %s: parent %s " + "has no available span", + operation.operation_id, + parent_id, + ) + return DurableParentSpan(parent_context) + + def _finish_operation_end( + self, info: OperationInfo, fallback_parent: Span | None + ) -> None: + """Materialize the existing terminal span without a registry gap.""" # An open operation is held as a non-recording placeholder; drop it and # create the single recording span for the operation now. 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._completion_parent_span(info.parent_id) start_time = info.start_time end_time = info.end_time if isinstance(placeholder, DurableParentSpan): @@ -899,9 +996,14 @@ def on_operation_end(self, info: OperationEndInfo) -> None: 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)) + span.end(end_time=_to_otel_timestamp(end_time)) + with self._lock: + if self._operation_spans.get(info.operation_id) is span: + del self._operation_spans[info.operation_id] + if info.operation_type is OperationType.CONTEXT: + self._completed_parent_spans[info.operation_id] = DurableParentSpan( + span.get_span_context(), start_time=start_time + ) def _start_span( self, @@ -921,36 +1023,35 @@ def _start_span( pass ``deterministic=False`` for a fresh ID beneath the operation span. """ key = span_key if span_key is not None else operation_id - with self._lock: - existing = self._operation_spans.get(key) - if existing is not None and existing.is_recording(): - existing.set_attribute("durable.span.replaced_on_reentry", True) - existing.end() - links = self._build_invocation_links() - span_id = ( - span_id_override - if span_id_override is not None - else ( - operation_id_to_span_id(self._execution_arn, operation_id) - if deterministic - else None - ) + existing = self._get_span(key) + if existing is not None and existing.is_recording(): + existing.set_attribute("durable.span.replaced_on_reentry", True) + existing.end() + links = self._build_invocation_links() + span_id = ( + span_id_override + if span_id_override is not None + else ( + operation_id_to_span_id(self._execution_arn, operation_id) + if deterministic + else None ) - if parent is None: - parent_ctx = self._with_sampling(Context()) - else: - parent_ctx = self._with_sampling( - trace.set_span_in_context(parent, Context()) - ) - with self._id_generator.use_ids(trace_id=None, span_id=span_id): - span = self._tracer.start_span( - name=name, - attributes=self._operation_attributes(info), - start_time=_to_otel_timestamp(start_time), - context=parent_ctx, - links=links, - ) - self._operation_spans[key] = span + ) + if parent is None: + parent_ctx = self._with_sampling(Context()) + else: + parent_ctx = self._with_sampling( + trace.set_span_in_context(parent, Context()) + ) + with self._id_generator.use_ids(trace_id=None, span_id=span_id): + span = self._tracer.start_span( + name=name, + attributes=self._operation_attributes(info), + start_time=_to_otel_timestamp(start_time), + context=parent_ctx, + links=links, + ) + self._set_span(key, span) return span # ------------------------------------------------------------------ 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 4eab6c7c3..4d564baed 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 @@ -13,6 +13,8 @@ InvocationStatus, InvocationStartInfo, OperationEndInfo, + OperationChangeInfo, + OperationInfo, OperationStartInfo, OperationStatus, OperationType, @@ -34,6 +36,11 @@ TraceState, ) +from aws_durable_execution_sdk_python_otel._invocation_state import ( + CompletionState, + completed_input_parent_ids, + parent_first, +) from aws_durable_execution_sdk_python_otel.context_extractors import ( ContextExtractor, ExtractedContext, @@ -197,6 +204,11 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: # 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. + # Completion delivery is invocation-owned, separate from active context + # segments. Registration and dequeue use the same span-registry lock. + self._completions = CompletionState() + self._completed_input_parent_ids: set[str] = set() + self._completed_parent_spans: dict[str, Span] = {} self._context_tokens: dict[str, tuple[int, object]] = {} self._operation_spans_lock = threading.RLock() self._tracing_enabled = False @@ -229,9 +241,15 @@ def _bind_sdk_tracer(self) -> bool: return True def _set_span(self, operation_id: str | None, span: Span) -> None: - """Register the active span for an operation ID.""" + """Register a segment and dequeue its children in the same critical section.""" with self._operation_spans_lock: self._operation_spans[operation_id] = span + pending = ( + self._completions.take_children(operation_id) + if operation_id is not None + else [] + ) + self._drain_registered_children(pending) def _delete_span(self, operation_id: str | None) -> None: """Remove the active span for an operation ID if one is stored.""" @@ -510,54 +528,53 @@ def _start_span( parent_span, ) registry_key = span_key if span_key is not None else operation_id - with self._operation_spans_lock: - links: list[Link] - if not deterministic_span_id or existed: - links = [] - span_id = None - else: - links = [] - span_id = ( - operation_id_to_span_id(self._execution_arn, operation_id) - if operation_id - else None - ) - should_link_logical_operation = ( - existed if link_logical_operation is None else link_logical_operation + links: list[Link] + if not deterministic_span_id or existed: + links = [] + span_id = None + else: + links = [] + span_id = ( + operation_id_to_span_id(self._execution_arn, operation_id) + if operation_id + else None ) - if should_link_logical_operation and operation_id is not None: - operation_context = self._operation_link_context(operation_id) - if operation_context is not None and operation_context.is_valid: - links = [*links, Link(context=operation_context)] - # Operation and attempt spans link to the execution-scoped Workflow - # span (the invocation span itself, operation_id=None, does not). - if self._workflow_span is not None and operation_id is not None: - workflow_ctx = self._workflow_span.get_span_context() - if workflow_ctx and workflow_ctx.is_valid: - links = [*links, Link(context=workflow_ctx)] - if parent_span is None: - parent_context = self._invocation_parent_context() - else: - parent_context = self._with_sampling( - trace.set_span_in_context(parent_span, Context()) - ) - span_start_time = ( - self._next_ordered_timestamp(start_time) - if operation_id is not None - else _to_otel_timestamp(start_time) + should_link_logical_operation = ( + existed if link_logical_operation is None else link_logical_operation + ) + if should_link_logical_operation and operation_id is not None: + operation_context = self._operation_link_context(operation_id) + if operation_context is not None and operation_context.is_valid: + links = [*links, Link(context=operation_context)] + # Operation and attempt spans link to the execution-scoped Workflow + # span (the invocation span itself, operation_id=None, does not). + if self._workflow_span is not None and operation_id is not None: + workflow_ctx = self._workflow_span.get_span_context() + if workflow_ctx and workflow_ctx.is_valid: + links = [*links, Link(context=workflow_ctx)] + if parent_span is None: + parent_context = self._invocation_parent_context() + else: + parent_context = self._with_sampling( + trace.set_span_in_context(parent_span, Context()) ) - with self._id_generator.use_ids(trace_id=None, span_id=span_id): - span = self._tracer.start_span( - name=name, - kind=SpanKind.INTERNAL, - attributes=attributes, - start_time=span_start_time, - context=parent_context, - links=links, - ) - self._operation_spans[registry_key] = span - if operation_id is None: - self._span_time_floor_ns = span_start_time + span_start_time = ( + self._next_ordered_timestamp(start_time) + if operation_id is not None + else _to_otel_timestamp(start_time) + ) + with self._id_generator.use_ids(trace_id=None, span_id=span_id): + span = self._tracer.start_span( + name=name, + kind=SpanKind.INTERNAL, + attributes=attributes, + start_time=span_start_time, + context=parent_context, + links=links, + ) + self._set_span(registry_key, span) + if operation_id is None: + self._span_time_floor_ns = span_start_time logger.debug("Started OTel span: %s", span) return span @@ -700,6 +717,10 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: name="Invocation", attributes=self._extract_attributes(info), ) + # A completed input context may re-enter through ReplayChildren. Keep + # its logical reference separate until an actual segment can register. + self._completed_input_parent_ids = completed_input_parent_ids(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 @@ -814,6 +835,7 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: def _end_invocation_spans(self, info: InvocationEndInfo) -> None: """End this invocation's open spans and export the Workflow span.""" + self._drain_pending_completions() # Spans are registered parent-first, so close pending spans in reverse # order to keep every child contained within its parent. with self._operation_spans_lock: @@ -903,6 +925,10 @@ def _release_invocation_scope(self) -> None: self._detach_remaining_contexts() unbind_invocation(self) self._tracing_enabled = False + with self._operation_spans_lock: + self._completions.clear() + self._completed_input_parent_ids.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.""" @@ -935,22 +961,123 @@ def on_operation_start(self, info: OperationStartInfo) -> None: ) def on_operation_end(self, info: OperationEndInfo) -> None: - """Called when an operation reaches a terminal durable status. + """Observe one natural completion, sharing delivery with snapshots.""" + self._render_operation_end(info) - Non-user-function operations are started by ``on_operation_start``. If - an operation end is observed without a matching in-memory span, this - invocation is completing an operation that began earlier, so a short - continuation span with a fresh span ID is created and ended. - """ + def on_operation_change(self, info: OperationChangeInfo) -> None: + """Retain checkpoint completions until their natural End or finalization.""" + if self._tracing_enabled: + self._remember_completions(info.updated_operations) + + def _remember_completions(self, updates: dict[str, OperationInfo]) -> None: + with self._operation_spans_lock: + self._completions.remember(updates) + + def _completion_parent_span(self, parent_id: str | None) -> Span | None: + """Prefer the active replay segment, then its retained completed context.""" + with self._operation_spans_lock: + return self._operation_spans.get(parent_id) or ( + self._completed_parent_spans.get(parent_id) + if parent_id is not None + else None + ) + + def _drain_pending_completions(self) -> None: + """Export pending parents before children while Invocation is still open.""" + with self._operation_spans_lock: + pending = dict(self._completions.pending) + for operation in parent_first(pending): + parent = self._completion_parent_span(operation.parent_id) + if parent is None: + parent = self._fallback_parent(operation) + self._render_operation_end(operation, fallback_parent=parent) + + def _render_operation_end( + self, info: OperationInfo, *, fallback_parent: Span | None = None + ) -> None: + """Claim a completion outside customer callbacks; record success last.""" 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: + span = self._operation_spans.get(info.operation_id) + # A replayed/virtual context opens a new invocation segment without + # a new durable checkpoint. Only its natural End may close it; + # a duplicate durable checkpoint must leave the new segment open. + 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() + ) + parent_available = ( + span is not None + or info.parent_id is None + or fallback_parent is not None + or self._completion_parent_span(info.parent_id) is not None + ) + if not self._completions.claim( + info, + parent_available=parent_available, + repeat=ends_active_context, + ): + return + try: + self._finish_operation_end(info, span, fallback_parent) + with self._operation_spans_lock: + self._completions.complete(info) + finally: + with self._operation_spans_lock: + self._completions.release(info.operation_id) + + def _drain_registered_children(self, pending: list[OperationInfo]) -> None: + """A child exporter must not interrupt its parent's context attachment.""" + for operation in pending: + try: + self._render_operation_end(operation) + except Exception: + logger.exception( + "Completion export failed while registering parent %s", + operation.parent_id, + ) + + def _fallback_parent(self, operation: OperationInfo) -> Span: + parent_id = operation.parent_id + # Referencing completed input metadata is safe only after every pending + # ancestor was materialized and actual replay parents had a chance to start. + if parent_id in self._completed_input_parent_ids: + parent_context = self._operation_link_context(parent_id) + if parent_context is not None: + return DurableParentSpan(parent_context) + invocation_span = self._get_span(None) + if invocation_span is None: + raise ValueError("No invocation span available for completion fallback") + logger.warning( + "Using invocation fallback parent for operation %s: parent %s " + "has no available span", + operation.operation_id, + parent_id, + ) + return invocation_span + + def _finish_operation_end( + self, info: OperationInfo, span: Span | None, fallback_parent: Span | None + ) -> None: + """End the claimed segment using the existing terminal rendering path.""" 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, @@ -970,6 +1097,11 @@ def on_operation_end(self, info: OperationEndInfo) -> None: span.set_status(StatusCode.OK) self._end_span(info.operation_id, info.end_time) + 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 on_user_function_start(self, info: UserFunctionStartInfo) -> None: """Called when a context or step operation starts user code. 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..53b6fc8db --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_completion_snapshot_int.py @@ -0,0 +1,366 @@ +"""Unvisited callback updates survive the public handler's invocation lifecycle.""" + +from datetime import UTC, datetime, timedelta +from concurrent.futures import ThreadPoolExecutor +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.identifier import OperationIdNamespace +from aws_durable_execution_sdk_python.lambda_service import ( + CallbackDetails, + ContextDetails, + OperationSubType, + ErrorObject, + ExecutionDetails, + Operation, + OperationStatus, + OperationType, +) +from aws_durable_execution_sdk_python.plugin import InvocationStartInfo +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, + ExecutionOtelPlugin, + InvocationOtelPluginFactory, + InvocationOtelPlugin, + 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("parent_visibility", ["root", "missing", "completed"]) +def test_unvisited_input_completion_exports_before_invocation_closes( + factory_type: type[ExecutionOtelPluginFactory] | type[InvocationOtelPluginFactory], + status: OperationStatus, + parent_visibility: str, + 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 + ) + parent_id = OperationIdNamespace().create_id_for_step(1) + 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=( + parent_id + if parent_visibility == "completed" + else "missing-context" + if parent_visibility == "missing" + 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, + ), + ), + ] + if parent_visibility == "completed": + operations.append( + Operation( + operation_id=parent_id, + operation_type=OperationType.CONTEXT, + sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT, + name="cached parent", + status=OperationStatus.SUCCEEDED, + start_timestamp=start, + end_timestamp=start + timedelta(seconds=1), + context_details=ContextDetails(result='"winner"'), + ) + ) + invocation: dict[str, Any] = { + "DurableExecutionArn": EXECUTION_ARN, + "CheckpointToken": "resume-token", + "InitialExecutionState": { + "Operations": [operation.to_json_dict() for operation in operations], + "NextMarker": "", + }, + "UpdatedOperationIds": [CALLBACK_ID, "unknown-callback"], + } + 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 completed_parent(_context: DurableContext) -> str: + pytest.fail( + "A completed parent must return its saved result without repeating its body" + ) + + 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() + ) + ) + if parent_visibility == "completed": + return context.run_in_child_context(completed_parent, name="cached parent") + 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 parent_visibility == "completed": + expected_parent_id = operation_id_to_span_id(EXECUTION_ARN, parent_id) + elif factory_type is InvocationOtelPluginFactory: + expected_parent_id = invocation_span.context.span_id + elif parent_visibility == "missing": + 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(parent_visibility == "missing") + assert otel_context.get_current() == before + finally: + provider.shutdown() + + +@pytest.mark.parametrize( + "factory_type", [ExecutionOtelPluginFactory, InvocationOtelPluginFactory] +) +def test_overlapping_invocations_and_warm_reuse_own_their_pending_completions( + factory_type: type[ExecutionOtelPluginFactory] | type[InvocationOtelPluginFactory], + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + factory = factory_type( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + instances: dict[str, ExecutionOtelPlugin | InvocationOtelPlugin] = {} + original_create = factory.create_plugin + + def create(info: InvocationStartInfo) -> ExecutionOtelPlugin | InvocationOtelPlugin: + plugin = original_create(info) + assert info.request_id is not None + instances[info.request_id] = plugin + return plugin + + monkeypatch.setattr(factory, "create_plugin", create) + ready = {request: Event() for request in ("a", "b")} + release = {request: Event() for request in ("a", "b")} + + def body(event: dict[str, str], _context: DurableContext) -> str: + request = event["request"] + if request in ready: + ready[request].set() + assert release[request].wait(5) + return request + + handler = durable_execution(body, plugins=[factory]) + start = datetime.now(UTC) - timedelta(seconds=5) + + def invoke(request: str) -> dict[str, Any]: + operations = [ + Operation( + operation_id="execution-id-" + request, + operation_type=OperationType.EXECUTION, + status=OperationStatus.STARTED, + start_timestamp=start, + execution_details=ExecutionDetails( + input_payload='{"request":"' + request + '"}' + ), + ), + Operation( + operation_id=CALLBACK_ID, + operation_type=OperationType.CALLBACK, + name="unvisited callback", + status=OperationStatus.SUCCEEDED, + start_timestamp=start, + end_timestamp=start + timedelta(seconds=1), + callback_details=CallbackDetails( + callback_id="callback", result='"saved"' + ), + ), + ] + event = { + "DurableExecutionArn": EXECUTION_ARN + "-" + request, + "CheckpointToken": "resume-token", + "InitialExecutionState": { + "Operations": [operation.to_json_dict() for operation in operations], + }, + "UpdatedOperationIds": [CALLBACK_ID], + } + 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, + ) + return handler(event, context) + + service = Mock() + before = otel_context.get_current() + try: + with ( + patch( + "aws_durable_execution_sdk_python.execution.LambdaClient.initialize_client", + return_value=service, + ), + ThreadPoolExecutor(max_workers=2) as callers, + ): + first = callers.submit(invoke, "a") + second = callers.submit(invoke, "b") + try: + assert ready["a"].wait(5) and ready["b"].wait(5) + release["a"].set() + assert first.result(5) == {"Status": "SUCCEEDED", "Result": '"a"'} + assert not instances["a"]._completions.pending + assert CALLBACK_ID in instances["b"]._completions.pending + assert instances["b"]._tracing_enabled + release["b"].set() + assert second.result(5) == {"Status": "SUCCEEDED", "Result": '"b"'} + assert invoke("c") == {"Status": "SUCCEEDED", "Result": '"c"'} + finally: + release["a"].set() + release["b"].set() + callbacks = [ + span + for span in exporter.get_finished_spans() + if span.name == "unvisited callback" + ] + assert len(callbacks) == 3 + assert { + span.attributes["durable.execution.arn"] + for span in callbacks + if span.attributes + } == {EXECUTION_ARN + "-" + request for request in ("a", "b", "c")} + assert len({span.context.trace_id for span in callbacks if span.context}) == 3 + for plugin in instances.values(): + assert not plugin._completions.pending + assert not plugin._completions.children_by_parent + assert not plugin._completed_input_parent_ids + assert not plugin._completed_parent_spans + assert not plugin._completions.in_flight + assert not plugin._completions.delivered + assert not plugin._tracing_enabled + service.checkpoint.assert_not_called() + assert otel_context.get_current() == before + finally: + release["a"].set() + release["b"].set() + 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_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..2daf744fe --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_parent_registration_int.py @@ -0,0 +1,187 @@ +"""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 +): + 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 + ] + 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/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..f077c13f8 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_completion_snapshots.py @@ -0,0 +1,388 @@ +"""Input, replay and checkpoint observations share an invocation's export ledger.""" + +from concurrent.futures import ThreadPoolExecutor +from dataclasses import replace +from datetime import UTC, datetime, timedelta +from threading import Event +from typing import Any + +import pytest +from aws_durable_execution_sdk_python.plugin import ( + InvocationEndInfo, + InvocationStartInfo, + InvocationStatus, + OperationChangeInfo, + OperationEndInfo, + OperationInfo, + OperationStartInfo, + OperationStatus, + OperationType, + UserFunctionStartInfo, + UserFunctionEndInfo, +) +from opentelemetry import trace +from opentelemetry.context import Context +from opentelemetry.sdk.trace import Span as SdkSpan, SpanProcessor, 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, +) + + +TraceCase = tuple[ + ExecutionOtelPlugin | InvocationOtelPlugin, + InMemorySpanExporter, + dict[str, Any], +] + + +@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( + "status", + [ + OperationStatus.SUCCEEDED, + OperationStatus.FAILED, + OperationStatus.TIMED_OUT, + OperationStatus.STOPPED, + OperationStatus.CANCELLED, + ], +) +@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, status +): + plugin, exporter, invocation = trace_case + end = replace(completion(), status=status) + 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() + (completed_span,) = [span for span in spans if span.name == "b"] + assert completed_span.attributes["durable.operation.status"] == status.value + assert not [span for span in spans if span.name == "active"] + assert plugin._completions.pending == plugin._completions.children_by_parent == {} + + +def test_replay_only_metadata_is_not_a_new_completion(trace_case: TraceCase) -> None: + plugin, exporter, invocation = trace_case + completed = completion() + plugin.on_invocation_start( + InvocationStartInfo( + **invocation, + operations={"b": completed}, + updated_operations={"b": replace(completed, is_replayed=True)}, + ) + ) + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.PENDING) + ) + plugin.on_operation_change( + OperationChangeInfo( + execution_arn=invocation["execution_arn"], + operations={"b": completed}, + updated_operations={"b": completed}, + ) + ) + assert not [span for span in exporter.get_finished_spans() if span.name == "b"] + assert not plugin._completions.pending + + +@pytest.mark.parametrize("fail_child", [False, True]) +def test_parent_registration_runs_child_processors_outside_lock_and_attaches_context( + trace_case: TraceCase, fail_child: bool, caplog: pytest.LogCaptureFixture +) -> None: + plugin, exporter, invocation = trace_case + provider = plugin._provider + assert isinstance(provider, TracerProvider) + child = completion(parent_id="parent") + with ThreadPoolExecutor(max_workers=1) as workers: + + class CheckParent(SpanProcessor): + reject = fail_child + + def on_start( + self, span: SdkSpan, parent_context: Context | None = None + ) -> None: + if span.name == "b": + # A processor can consult the registry on another thread. + # Neither parent registration nor terminal rendering owns + # that lock while executing the callback. + assert workers.submit( + plugin._completion_parent_span, "parent" + ).result(3) + if self.reject: + raise RuntimeError("child processor failed") + + processor = CheckParent() + provider.add_span_processor(processor) + plugin.on_invocation_start( + InvocationStartInfo(**invocation, updated_operations={"b": child}) + ) + fields: dict[str, Any] = dict( + operation_id="parent", + operation_type=OperationType.CONTEXT, + sub_type=None, + name="parent", + parent_id=None, + start_time=datetime.now(UTC), + is_replayed=True, + status=OperationStatus.STARTED, + ) + plugin.on_operation_start(OperationStartInfo(**fields)) + start = UserFunctionStartInfo(**fields) + plugin.on_user_function_start(start) + parent = plugin._get_span("parent") + assert parent is not None + assert trace.get_current_span().get_span_context() == parent.get_span_context() + with provider.get_tracer("application").start_as_current_span("user"): + pass + (user,) = [ + span for span in exporter.get_finished_spans() if span.name == "user" + ] + assert user.parent == parent.get_span_context() + assert len( + [span for span in exporter.get_finished_spans() if span.name == "b"] + ) == int(not fail_child) + assert ( + any( + "child processor failed" in record.getMessage() or record.exc_info + for record in caplog.records + ) + is fail_child + ) + processor.reject = False + plugin.on_user_function_end(UserFunctionEndInfo.from_start_info(start, None)) + plugin.on_operation_end( + OperationEndInfo( + **{**fields, "status": OperationStatus.SUCCEEDED, "is_replayed": False}, + end_time=datetime.now(UTC), + ) + ) + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.PENDING) + ) + # A failed child delivery remains available for finalization, while a + # successful registration drain is not repeated. + assert ( + len([span for span in exporter.get_finished_spans() if span.name == "b"]) + == 1 + ) + + +def test_natural_child_end_preserves_other_pending_children( + trace_case: TraceCase, +) -> None: + plugin, exporter, invocation = trace_case + plugin.on_invocation_start(InvocationStartInfo(**invocation)) + fields: dict[str, Any] = dict( + operation_id="parent", + operation_type=OperationType.CONTEXT, + sub_type=None, + name="parent", + parent_id=None, + start_time=datetime.now(UTC), + is_replayed=False, + status=OperationStatus.STARTED, + ) + plugin.on_operation_start(OperationStartInfo(**fields)) + start = UserFunctionStartInfo(**fields) + plugin.on_user_function_start(start) + parent_span = plugin._get_span("parent") + assert parent_span is not None + parent_context = parent_span.get_span_context() + updates = { + name: completion(name, parent_id="parent") for name in ("first", "second") + } + plugin.on_operation_change( + OperationChangeInfo( + execution_arn=invocation["execution_arn"], + operations=updates, + updated_operations=updates, + ) + ) + plugin.on_operation_end(updates["first"]) + assert [ + span.name for span in exporter.get_finished_spans() if span.name in updates + ] == ["first"] + plugin.on_user_function_end(UserFunctionEndInfo.from_start_info(start, None)) + plugin.on_operation_end( + OperationEndInfo( + **{**fields, "status": OperationStatus.SUCCEEDED}, + end_time=datetime.now(UTC), + ) + ) + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.PENDING) + ) + children = [span for span in exporter.get_finished_spans() if span.name in updates] + assert [span.name for span in children] == ["first", "second"] + assert all(span.parent == parent_context for span in children) + + +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 + ) + + +@pytest.mark.parametrize("full_snapshot", [False, True]) +def test_missing_ancestor_fallback_preserves_available_descendant_hierarchy( + trace_case, caplog, full_snapshot +): + 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, + operations={"b": child, "parent": parent} if full_snapshot else {}, + 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._completions.pending == plugin._completions.children_by_parent == {} + + +def test_partial_start_releases_pending_metadata_at_end( + trace_case: TraceCase, monkeypatch: pytest.MonkeyPatch +) -> None: + plugin, _exporter, invocation = trace_case + + def fail_extraction(_info: InvocationStartInfo) -> None: + 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")) + with pytest.raises(ValueError, match="available for completion fallback"): + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.PENDING) + ) + assert not plugin._tracing_enabled + assert not plugin._completions.pending + assert not plugin._completions.children_by_parent 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..7ed48fadd --- /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._completions.in_flight + 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_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..c0723e2c3 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_terminal_input_parent_refs.py @@ -0,0 +1,481 @@ +"""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._completions.pending == plugin._completions.children_by_parent == {} + ) + 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._completed_input_parent_ids == set() + 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_parent_ids, + ) + + 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_parent_ids( + InvocationStartInfo( + **invocation, + operations={"parent": parent, "unfinished": unfinished, "step": step}, + updated_operations={}, + ) + ) == {"parent"} + missing_parent_child = _info(OperationInfo, now) + assert ( + completed_input_parent_ids( + InvocationStartInfo( + **invocation, + operations={"b": missing_parent_child}, + updated_operations={"b": missing_parent_child}, + ) + ) + == set() + ) + + +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 dc84924b9..0eff21ce4 100644 --- a/packages/aws-durable-execution-sdk-python/README.md +++ b/packages/aws-durable-execution-sdk-python/README.md @@ -57,12 +57,22 @@ durable execution spans as many invocations as it waits, retries or resumes, and the instance is dropped when each of those returns — so anything a plugin keeps in its attributes is gone by the next invocation of the same execution. State that has to survive that has two honest homes: rebuild it from the operation map the -invocation hooks carry (`InvocationStartInfo.operations` is a full snapshot, +invocation hooks carry (`InvocationStartInfo.operations` is a snapshot of the +visible operation map, which is how the bundled Telemetry exporter reports operations that completed in an earlier invocation), or put it on the factory, which outlives every invocation — keyed by execution ARN, and pruned by the owner, because the SDK will not tell the factory when an execution ends for good. +`InvocationStartInfo.updated_operations` identifies visible operations updated +since the handler last observed state. `UpdatedOperationIds` alone cannot supply +metadata for an operation absent from the snapshot. A cached parent result can +skip a child's natural operation hooks during replay, even when that child has +a newly delivered terminal update. Instrumentation that needs those completions +must also consume invocation-start and checkpoint-change snapshots. The bundled +OTel plugins retain such updates until their natural End or invocation +finalization; this does not add synthetic lifecycle hooks for other plugins. + Invocation operation-map snapshots and terminal `OperationEndInfo` own a copy of the error's mutable stack-trace list, including virtual child completions without a checkpoint. A plugin may modify that list without changing the checkpointed