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