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 0352ebfa..12699663 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 @@ -9,7 +9,10 @@ ancestor: a propagated backend parent when present, otherwise a deterministic synthetic root. The Workflow span is exported exactly once, when the execution reaches a terminal status. Operations are parented under the Workflow span (or -their parent operation) and *linked* to the current Invocation span. +their parent operation) and *linked* to the current Invocation span. Each +operation is likewise exported exactly once, on its deterministic span ID, when +it reaches a terminal status; while it spans invocations it is held as a +non-recording placeholder so no recording span is abandoned. This is the Python adaptation of the JS ``ExecutionOtelPlugin`` from aws-durable-execution-sdk-js#729. Because the Python plugin interface differs @@ -56,6 +59,7 @@ SpanKind, StatusCode, Tracer, + TraceState, ) from aws_durable_execution_sdk_python_otel.context_extractors import ( @@ -133,12 +137,16 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: # Per-invocation state. self._execution_arn = "" self._execution_trace_id: int | None = None + self._execution_start_time: datetime.datetime | None = None self._extracted_context: ExtractedContext | None = None self._execution_trace_context: ExecutionTraceContext | None = None self._sampling_intent: DurableSamplingIntent | None = None self._workflow_span: Span | None = None self._invocation_span: Span | None = None self._operation_spans: dict[str, Span] = {} + # 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 @@ -292,6 +300,41 @@ def _resolve_parent(self, parent_id: str | None) -> Span | None: return existing return self._workflow_span + def _resolved_trace_state(self) -> TraceState: + """Return the resolved sampling trace state, else the ancestor state. + + The sampling result preserves a same-trace ambient ``tracestate`` that + the empty ancestor state would drop. + """ + intent = self._sampling_intent + if intent is not None and intent.result.trace_state is not None: + return intent.result.trace_state + if self._execution_trace_context is not None: + return self._execution_trace_context.execution_ancestor.trace_state + return TraceState() + + def _operation_span_context(self, operation_id: str) -> SpanContext | None: + """Return the deterministic SpanContext for a logical operation.""" + execution_trace_context = self._execution_trace_context + if execution_trace_context is None: + return None + return SpanContext( + trace_id=execution_trace_context.trace_id, + span_id=operation_id_to_span_id(self._execution_arn, operation_id), + is_remote=False, + trace_flags=execution_trace_context.trace_flags, + trace_state=self._resolved_trace_state(), + ) + + def _register_operation_placeholder(self, operation_id: str) -> Span | None: + """Register a non-recording placeholder holding the operation context.""" + span_context = self._operation_span_context(operation_id) + if span_context is None: + return None + placeholder = NonRecordingSpan(span_context) + self._set_span(operation_id, placeholder) + return placeholder + def _invocation_parent_context(self) -> Context: """Return same-trace ambient context, else execution ancestor context.""" execution_trace_context = self._execution_trace_context @@ -346,6 +389,7 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: ) self._tracing_enabled = False return + self._execution_start_time = info.execution_start_time self._extracted_context = _ensure_extracted_context( self._context_extractor(info) ) @@ -393,11 +437,40 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: ) def _start_workflow_span(self, info: InvocationStartInfo) -> None: + """Install a non-recording placeholder for the execution-scoped Workflow span. + + The Workflow span spans the whole durable execution and is exported once, + on the terminal invocation. During every invocation the plugin only needs + its deterministic SpanContext -- to parent operation spans, to keep the + Workflow current so auto-instrumented spans join the execution trace, and + for log correlation. A non-recording placeholder fills that role so a + non-terminal invocation never abandons a recording span. The recording + span is created and ended once by :meth:`_export_workflow_span`. + """ if not self._execution_arn: logger.warning("No execution ARN; skipping Workflow span creation") return if self._execution_trace_context is None: return + workflow_span_context = SpanContext( + trace_id=self._execution_trace_context.trace_id, + span_id=derive_workflow_span_id(self._execution_arn), + is_remote=False, + trace_flags=self._execution_trace_context.trace_flags, + trace_state=self._resolved_trace_state(), + ) + self._workflow_span = NonRecordingSpan(workflow_span_context) + + def _export_workflow_span(self, info: InvocationEndInfo) -> None: + """Create and end the recording Workflow span once, on a terminal status. + + Uses the same deterministic span ID as the placeholder and the shared + execution ancestor as its parent, so the exported Workflow span stays on + the execution trace and correlates with every operation span across all + invocations. Anchored at the execution start time. + """ + if not self._execution_arn or self._execution_trace_context is None: + return parent_context = self._with_sampling( trace.set_span_in_context( NonRecordingSpan(self._execution_trace_context.execution_ancestor), @@ -408,13 +481,44 @@ def _start_workflow_span(self, info: InvocationStartInfo) -> None: trace_id=None, span_id=derive_workflow_span_id(self._execution_arn), ): - self._workflow_span = self._tracer.start_span( + workflow_span = self._tracer.start_span( name=self._workflow_span_name, kind=SpanKind.INTERNAL, - attributes={"durable.execution.arn": self._execution_arn}, - start_time=_to_otel_timestamp(info.execution_start_time), + attributes={ + "durable.execution.arn": self._execution_arn, + "durable.execution.status": ( + info.status.value if info.status else "" + ), + }, + start_time=_to_otel_timestamp(self._execution_start_time), context=parent_context, ) + if info.status is InvocationStatus.FAILED: + workflow_span.set_status( + StatusCode.ERROR, info.error.message if info.error else "" + ) + elif info.status is InvocationStatus.SUCCEEDED: + workflow_span.set_status(StatusCode.OK) + workflow_span.end() + + def _end_open_recording_spans(self) -> None: + """End recording user-function spans left open by a suspended operation. + + Operation placeholders are non-recording and export their span from + on_operation_end, so they are skipped. Reverse order keeps each child + contained within its parent; the invocation span is ended by the caller. + """ + with self._lock: + keys = list(reversed(self._operation_spans)) + for key in keys: + if key == _INVOCATION_KEY: + continue + span = self._get_span(key) + if span is None or not span.is_recording(): + continue + popped = self._pop_span(key) + if popped is not None: + popped.end() def _start_invocation_span(self, info: InvocationStartInfo) -> None: self._invocation_span = self._tracer.start_span( @@ -434,12 +538,6 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: self._reset_state() return - # Operation spans still open here belong to operations that suspended - # (e.g. PENDING/RETRYING) rather than completed this invocation. They are - # ended only by on_operation_end; drop the references without ending them - # so they are not exported as if completed. _reset_state - # clears the span map below. - # 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 @@ -462,23 +560,16 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: ) self._invocation_span.end() - # The Workflow span (execution view) is exported only on a terminal - # status; otherwise its reference is dropped without ending it. Its span - # status reflects the execution outcome: SUCCEEDED -> OK, FAILED -> ERROR - # (RETRY/PENDING are non-terminal and never reach here -> UNSET). - if self._workflow_span is not None: - if info.status in _TERMINAL_INVOCATION_STATUSES: - self._workflow_span.set_attribute( - "durable.execution.status", - info.status.value if info.status else "", - ) - if info.status is InvocationStatus.FAILED: - self._workflow_span.set_status( - StatusCode.ERROR, info.error.message if info.error else "" - ) - elif info.status is InvocationStatus.SUCCEEDED: - self._workflow_span.set_status(StatusCode.OK) - self._workflow_span.end() + # End recording user-function spans left open by a suspended operation. + self._end_open_recording_spans() + + # The Workflow span (execution view) is a non-recording placeholder + # during the invocation, so only a terminal status materializes and ends + # the recording span. Its span status reflects the execution outcome: + # SUCCEEDED -> OK, FAILED -> ERROR (RETRY/PENDING are non-terminal and + # leave the Workflow span unexported until a later terminal invocation). + if info.status in _TERMINAL_INVOCATION_STATUSES: + self._export_workflow_span(info) self._reset_state() @@ -495,10 +586,12 @@ def _reset_state(self) -> None: self._extracted_context = None self._execution_trace_context = None self._sampling_intent = None + self._execution_start_time = None self._workflow_span = None self._invocation_span = None with self._lock: self._operation_spans = {} + self._ended_operation_ids = set() self._tracing_enabled = False # ------------------------------------------------------------------ @@ -510,8 +603,24 @@ def on_operation_start(self, info: OperationStartInfo) -> None: return if info.operation_type is OperationType.CONTEXT: return # tracked via on_user_function_start + # Hold a non-recording placeholder while the operation is open; its + # recording span is exported once on terminal on_operation_end. + self._register_operation_placeholder(info.operation_id) + + def on_operation_end(self, info: OperationEndInfo) -> None: + logger.debug("Durable operation ended: %s", info) + if not self._tracing_enabled: + return + # Export the span only on the first end for this operation. + with self._lock: + if info.operation_id in self._ended_operation_ids: + return + self._ended_operation_ids.add(info.operation_id) + # An open operation is held as a non-recording placeholder; drop it and + # create the single recording span for the operation now. + self._pop_span(info.operation_id) parent = self._resolve_parent(info.parent_id) - self._start_span( + span = self._start_span( operation_id=info.operation_id, name=info.name or info.operation_id, info=info, @@ -519,25 +628,6 @@ def on_operation_start(self, info: OperationStartInfo) -> None: start_time=info.start_time, ) - def on_operation_end(self, info: OperationEndInfo) -> None: - logger.debug("Durable operation ended: %s", info) - if not self._tracing_enabled: - return - span = self._get_span(info.operation_id) - if span is None: - # Cross-invocation stitching: operation started in a prior - # invocation. Create + immediately end a linked span. - parent = self._resolve_parent(info.parent_id) - span = self._start_span( - operation_id=info.operation_id, - name=info.name or info.operation_id, - info=info, - parent=parent, - start_time=info.start_time, - ) - else: - span.set_attributes(self._operation_attributes(info)) - if info.error: span.set_status(StatusCode.ERROR, info.error.message or "") span.record_exception( @@ -564,7 +654,11 @@ def _start_span( span_key: str | None = None, deterministic: bool = True, ) -> Span: - """Start a span for an operation/attempt and register it.""" + """Start a recording span for an operation/attempt and register it. + + Operation spans use the deterministic operation span ID; attempt spans + 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: links = self._build_invocation_links() @@ -603,6 +697,7 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None: "on_user_function_start only supports CONTEXT and STEP operations" ) key = self._user_function_key(info) + span: Span | None if info.operation_type is OperationType.STEP: parent = self._get_span(info.operation_id) or self._resolve_parent( info.parent_id @@ -618,17 +713,15 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None: deterministic=False, ) else: # CONTEXT - parent = self._resolve_parent(info.parent_id) - span = self._start_span( - operation_id=info.operation_id, - name=info.name or info.operation_id, - info=info, - parent=parent, - start_time=info.start_time, + # A child context can suspend before completing, so hold a + # non-recording placeholder while it runs; on_operation_end + # materializes its single recording span. This keeps a suspended + # context from being exported early and re-exported on replay. + span = self._register_operation_placeholder(info.operation_id) + if span is not None: + self._attach_context( + key, trace.set_span_in_context(span, otel_context.get_current()) ) - self._attach_context( - key, trace.set_span_in_context(span, otel_context.get_current()) - ) def on_user_function_end(self, info: UserFunctionEndInfo) -> None: logger.debug("Durable user function ended: %s", info) 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 f9e390c4..3bf07c95 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 @@ -31,6 +31,7 @@ SpanKind, StatusCode, Tracer, + TraceState, ) from aws_durable_execution_sdk_python_otel.context_extractors import ( @@ -137,6 +138,7 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: # per invocation status: self._execution_arn = "" self._execution_trace_id: int | None = None + self._execution_start_time: datetime.datetime | None = None self._extracted_context: ExtractedContext | None = None self._execution_trace_context: ExecutionTraceContext | None = None self._sampling_intent: DurableSamplingIntent | None = None @@ -358,6 +360,19 @@ def _next_ordered_timestamp( self._span_time_floor_ns = candidate return candidate + def _resolved_trace_state(self) -> TraceState: + """Return the resolved sampling trace state, else the ancestor state. + + The sampling result preserves a same-trace ambient ``tracestate`` that + the empty ancestor state would drop. + """ + intent = self._sampling_intent + if intent is not None and intent.result.trace_state is not None: + return intent.result.trace_state + if self._execution_trace_context is not None: + return self._execution_trace_context.execution_ancestor.trace_state + return TraceState() + def _operation_link_context(self, operation_id: str) -> SpanContext | None: """Return the deterministic logical operation context for links.""" execution_trace_context = self._execution_trace_context @@ -368,7 +383,7 @@ def _operation_link_context(self, operation_id: str) -> SpanContext | None: span_id=operation_id_to_span_id(self._execution_arn, operation_id), is_remote=False, trace_flags=execution_trace_context.trace_flags, - trace_state=execution_trace_context.execution_ancestor.trace_state, + trace_state=self._resolved_trace_state(), ) def _start_span( @@ -511,6 +526,7 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: ) self._tracing_enabled = False return + self._execution_start_time = info.execution_start_time self._extracted_context = _ensure_extracted_context( self._context_extractor(info) ) @@ -550,20 +566,41 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: ) def _start_workflow_span(self, info: InvocationStartInfo) -> None: - """Create the deterministic, execution-scoped Workflow span. + """Install a non-recording placeholder for the execution-scoped Workflow span. The Workflow span is keyed to a deterministic span ID derived from the execution ARN, so every invocation of the same durable execution - contributes to one Workflow span. It is parented to the shared execution - ancestor and exported once, on a terminal invocation. Operation and - attempt spans link to it while remaining parented to the invocation - span. + contributes to one Workflow span. During each invocation the plugin only + needs its deterministic SpanContext so operation and attempt spans can + link to it while remaining parented to the invocation span; a + non-recording placeholder fills that role so a non-terminal invocation + never abandons a recording span. The recording span is created and ended + once, on a terminal status, by :meth:`_export_workflow_span`. """ if not self._execution_arn: logger.warning("No execution ARN; skipping Workflow span creation") return if self._execution_trace_context is None: return + workflow_span_context = SpanContext( + trace_id=self._execution_trace_context.trace_id, + span_id=derive_workflow_span_id(self._execution_arn), + is_remote=False, + trace_flags=self._execution_trace_context.trace_flags, + trace_state=self._resolved_trace_state(), + ) + self._workflow_span = NonRecordingSpan(workflow_span_context) + + def _export_workflow_span(self, info: InvocationEndInfo) -> None: + """Create and end the recording Workflow span once, on a terminal status. + + Uses the same deterministic span ID as the placeholder and the shared + execution ancestor as its parent, so the exported Workflow span stays on + the execution trace and correlates with every operation span across all + invocations. Anchored at the execution start time. + """ + if not self._execution_arn or self._execution_trace_context is None: + return parent_context = self._with_sampling( trace.set_span_in_context( NonRecordingSpan(self._execution_trace_context.execution_ancestor), @@ -574,13 +611,25 @@ def _start_workflow_span(self, info: InvocationStartInfo) -> None: trace_id=None, span_id=derive_workflow_span_id(self._execution_arn), ): - self._workflow_span = self._tracer.start_span( + workflow_span = self._tracer.start_span( name=self._workflow_span_name, kind=SpanKind.INTERNAL, - attributes={"durable.execution.arn": self._execution_arn}, - start_time=_to_otel_timestamp(info.execution_start_time), + attributes={ + "durable.execution.arn": self._execution_arn, + "durable.execution.status": ( + info.status.value if info.status else "" + ), + }, + start_time=_to_otel_timestamp(self._execution_start_time), context=parent_context, ) + if info.status is InvocationStatus.FAILED: + workflow_span.set_status( + StatusCode.ERROR, info.error.message if info.error else "" + ) + elif info.status is InvocationStatus.SUCCEEDED: + workflow_span.set_status(StatusCode.OK) + workflow_span.end() def on_invocation_end(self, info: InvocationEndInfo) -> None: """Called at the end of each invocation. Ends the invocation span and flushes.""" @@ -617,23 +666,12 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: # end the invocation span self._end_span(None) - # The Workflow span (execution view) is exported only on a terminal - # status; on non-terminal statuses its reference is dropped without - # ending it (so it is not exported yet). SUCCEEDED -> OK, FAILED -> ERROR; - # RETRY/PENDING are non-terminal and leave it unexported. - if self._workflow_span is not None: - if info.status in _TERMINAL_INVOCATION_STATUSES: - self._workflow_span.set_attribute( - "durable.execution.status", - info.status.value if info.status else "", - ) - if info.status is InvocationStatus.FAILED: - self._workflow_span.set_status( - StatusCode.ERROR, info.error.message if info.error else "" - ) - elif info.status is InvocationStatus.SUCCEEDED: - self._workflow_span.set_status(StatusCode.OK) - self._workflow_span.end() + # The Workflow span (execution view) is a non-recording placeholder + # during the invocation, so only a terminal status materializes and ends + # the recording span. SUCCEEDED -> OK, FAILED -> ERROR; RETRY/PENDING are + # non-terminal and leave it unexported until a later terminal invocation. + if info.status in _TERMINAL_INVOCATION_STATUSES: + self._export_workflow_span(info) self._reset_state() @@ -649,6 +687,7 @@ def _reset_state(self) -> None: self._extracted_context = None self._execution_trace_context = None self._sampling_intent = None + self._execution_start_time = None self._workflow_span = None self._span_time_floor_ns = None with self._operation_spans_lock: diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py index 07513662..63563493 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py @@ -226,7 +226,10 @@ def handler_impl(_event: Any, context: DurableContext) -> str: after_resume = next(span for span in spans if span.name == "otel-after-resume") assert len(invocations) >= 2 - assert len(waits) >= 2 + if plugin_type is InvocationOtelPlugin: + assert len(waits) >= 2 # one segment per invocation + else: + assert len(waits) == 1 # one span per operation assert workflow.context.span_id == derive_workflow_span_id(EXECUTION_ARN) assert workflow.parent is not None assert workflow.parent.span_id == XRAY_PARENT_SPAN_ID diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py index 6d8ad688..5afcd3de 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py @@ -26,11 +26,19 @@ UserFunctionStartInfo, ) from opentelemetry import baggage, trace +from opentelemetry.context import Context from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import ( + NonRecordingSpan, + SpanContext, + TraceFlags, + TraceState, +) from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( + _to_otel_trace_id, derive_execution_root_span_id, derive_workflow_span_id, operation_id_to_span_id, @@ -282,15 +290,16 @@ def test_explicit_mode_invocation_span_ignores_different_trace_ambient_span(): assert invocation.parent.span_id == workflow.parent.span_id -def test_workflow_span_dropped_on_non_terminal_status(): +def test_workflow_span_not_exported_on_non_terminal_status(): plugin, exporter = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) plugin.on_invocation_end(_invocation_end_info(status=InvocationStatus.PENDING)) names = [s.name for s in exporter.get_finished_spans()] - # Invocation span is always ended/exported; the Workflow span is dropped - # (not ended) on a non-terminal status, so it must not be exported. + # Invocation span is always ended/exported. The Workflow span is a + # non-recording placeholder during the invocation and is only materialized + # (created + ended) on a terminal status, so it is not exported here. assert "Invocation" in names assert "Workflow" not in names @@ -343,6 +352,7 @@ def test_operation_parented_under_workflow_and_linked_to_invocation(): def test_cross_invocation_operation_end_uses_deterministic_span_id(): + """An operation completing in a later invocation exports one deterministic span.""" plugin, exporter = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) @@ -364,8 +374,7 @@ def test_cross_invocation_operation_end_uses_deterministic_span_id(): plugin.on_invocation_end(_invocation_end_info()) matching = [s for s in exporter.get_finished_spans() if s.name == "earlier-step"] - # Exported exactly once, using the deterministic logical-operation span ID - # (no separate continuation span). + # Exported once, using the deterministic operation span ID. assert len(matching) == 1 assert matching[0].context.span_id == operation_id_to_span_id( EXECUTION_ARN, "step-earlier" @@ -418,12 +427,11 @@ def test_context_span_waits_for_terminal_operation_status( attempt=1, ) ) + # The open child context is held as a non-recording placeholder until + # on_operation_end materializes its recording span. active_span = plugin._get_span(operation_id) assert active_span is not None - assert ( - active_span.attributes["durable.operation.status"] - == OperationStatus.STARTED.value - ) + assert not active_span.is_recording() plugin.on_user_function_end( UserFunctionEndInfo( @@ -443,11 +451,9 @@ def test_context_span_waits_for_terminal_operation_status( ) ) + # Still a placeholder, still not exported, until the terminal operation end. assert plugin._get_span(operation_id) is active_span - assert ( - active_span.attributes["durable.operation.status"] - == OperationStatus.STARTED.value - ) + assert not active_span.is_recording() assert not exporter.get_finished_spans() plugin.on_operation_end( @@ -585,12 +591,8 @@ def test_default_mode_invocation_span_ignores_different_trace_ambient_span(monke assert invocation.context.trace_id != ambient.get_span_context().trace_id -def test_open_operation_span_not_exported_at_invocation_end(): - """A suspended operation (started, not ended) must not be exported. - - on_invocation_end drops the reference without ending it; the - span is ended only when on_operation_end fires in a later invocation. - """ +def test_suspended_operation_held_as_non_recording_placeholder(): + """A suspended operation is a non-recording placeholder, not exported.""" plugin, exporter = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) @@ -606,14 +608,274 @@ def test_open_operation_span_not_exported_at_invocation_end(): status=OperationStatus.STARTED, ) ) + # The registered span is a non-recording placeholder on the deterministic ID. + placeholder = plugin._get_span("wait-1") + assert placeholder is not None + assert not placeholder.is_recording() + assert placeholder.get_span_context().span_id == operation_id_to_span_id( + EXECUTION_ARN, "wait-1" + ) + # No on_operation_end: the operation suspended. plugin.on_invocation_end(_invocation_end_info(status=InvocationStatus.PENDING)) + # Nothing is exported for the suspended operation. exported = {s.name for s in exporter.get_finished_spans()} - # The open operation span is NOT exported (never ended). assert "wait-for-signal" not in exported +def test_suspend_then_resume_operation_exports_one_deterministic_span(): + """An operation spanning invocations exports one deterministic span.""" + plugin, exporter = _create_plugin() + operation_id = "wait-across-invocations" + + # Invocation N: operation starts and suspends (non-terminal invocation end). + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_operation_start( + OperationStartInfo( + operation_id=operation_id, + operation_type=OperationType.WAIT, + sub_type=OperationSubType.WAIT, + name="long-wait", + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.STARTED, + ) + ) + plugin.on_invocation_end(_invocation_end_info(status=InvocationStatus.PENDING)) + + # Nothing exported for the operation at the non-terminal boundary. + assert not [s for s in exporter.get_finished_spans() if s.name == "long-wait"] + + # Invocation N+1: the still-open operation is replayed, then completes. + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_operation_start( + OperationStartInfo( + operation_id=operation_id, + operation_type=OperationType.WAIT, + sub_type=OperationSubType.WAIT, + name="long-wait", + parent_id=None, + start_time=START_TIME, + is_replayed=True, + status=OperationStatus.STARTED, + ) + ) + plugin.on_operation_end( + OperationEndInfo( + operation_id=operation_id, + operation_type=OperationType.WAIT, + sub_type=OperationSubType.WAIT, + name="long-wait", + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.SUCCEEDED, + end_time=END_TIME, + error=None, + ) + ) + plugin.on_invocation_end(_invocation_end_info(status=InvocationStatus.SUCCEEDED)) + + operation_spans = [ + s for s in exporter.get_finished_spans() if s.name == "long-wait" + ] + # Exported exactly once, using the deterministic operation span ID. + assert len(operation_spans) == 1 + assert operation_spans[0].context.span_id == operation_id_to_span_id( + EXECUTION_ARN, operation_id + ) + + +def test_suspended_child_context_exports_one_span_on_replay(): + """A child context that suspends then replays exports a single span.""" + plugin, exporter = _create_plugin() + context_id = "ctx-1" + + # Invocation 1: the child context starts and suspends (no end hook). + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_user_function_start(_context_start_info(context_id)) + placeholder = plugin._get_span(context_id) + assert placeholder is not None + assert not placeholder.is_recording() + plugin.on_invocation_end(_invocation_end_info(status=InvocationStatus.PENDING)) + + # Nothing exported for the suspended context. + assert not [s for s in exporter.get_finished_spans() if s.name == context_id] + + # Invocation 2: the context replays and completes. + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_user_function_start(_context_start_info(context_id)) + plugin.on_user_function_end(_context_end_info(context_id)) + plugin.on_operation_end( + OperationEndInfo( + operation_id=context_id, + operation_type=OperationType.CONTEXT, + sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT, + name=context_id, + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.SUCCEEDED, + end_time=END_TIME, + error=None, + ) + ) + plugin.on_invocation_end(_invocation_end_info()) + + contexts = [s for s in exporter.get_finished_spans() if s.name == context_id] + assert len(contexts) == 1 + assert contexts[0].context.span_id == operation_id_to_span_id( + EXECUTION_ARN, context_id + ) + + +def test_duplicate_operation_end_exports_span_once(): + """A repeated on_operation_end for one operation exports a single span.""" + plugin, exporter = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + + end_info = OperationEndInfo( + operation_id="wait-1", + operation_type=OperationType.WAIT, + sub_type=OperationSubType.WAIT, + name="otel-long-wait", + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.SUCCEEDED, + end_time=END_TIME, + error=None, + ) + plugin.on_operation_end(end_info) + plugin.on_operation_end(end_info) + plugin.on_invocation_end(_invocation_end_info()) + + waits = [s for s in exporter.get_finished_spans() if s.name == "otel-long-wait"] + assert len(waits) == 1 + assert waits[0].context.span_id == operation_id_to_span_id(EXECUTION_ARN, "wait-1") + + +def test_pre_terminal_placeholder_preserves_same_trace_tracestate(): + """Placeholder and operation spans carry a same-trace ambient tracestate.""" + plugin, exporter = _create_plugin() + canonical = _to_otel_trace_id(EXECUTION_ARN, START_TIME) + trace_state = TraceState([("vendor", "opaque")]) + ambient_context = SpanContext( + trace_id=canonical, + span_id=int("1234567890abcdef", 16), + is_remote=False, + trace_flags=TraceFlags(TraceFlags.SAMPLED), + trace_state=trace_state, + ) + ambient = NonRecordingSpan(ambient_context) + token = otel_context.attach(trace.set_span_in_context(ambient, Context())) + try: + plugin.on_invocation_start(_invocation_start_info()) + assert plugin._workflow_span is not None + assert plugin._workflow_span.get_span_context().trace_state == trace_state + plugin.on_operation_end( + OperationEndInfo( + operation_id="wait-existing", + operation_type=OperationType.WAIT, + sub_type=OperationSubType.WAIT, + name="existing-wait", + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.SUCCEEDED, + end_time=END_TIME, + error=None, + ) + ) + plugin.on_invocation_end(_invocation_end_info(status=InvocationStatus.PENDING)) + finally: + otel_context.detach(token) + + span = next(s for s in exporter.get_finished_spans() if s.name == "existing-wait") + assert span.context.trace_state == trace_state + + +@pytest.mark.parametrize( + ("status", "expected_code"), + [ + (InvocationStatus.SUCCEEDED, trace.StatusCode.OK), + (InvocationStatus.FAILED, trace.StatusCode.ERROR), + ], +) +def test_workflow_span_exported_once_on_terminal(status, expected_code): + """A terminal invocation materializes and ends the Workflow span exactly once.""" + plugin, exporter = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_invocation_end(_invocation_end_info(status=status)) + + workflows = [s for s in exporter.get_finished_spans() if s.name == "Workflow"] + assert len(workflows) == 1 + workflow = workflows[0] + # Shared execution trace: the Workflow span is parented to the synthetic + # execution root, not a parentless root. + assert workflow.parent is not None + assert workflow.parent.span_id == derive_execution_root_span_id(EXECUTION_ARN) + assert workflow.kind is trace.SpanKind.INTERNAL + assert workflow.context.span_id == derive_workflow_span_id(EXECUTION_ARN) + assert workflow.attributes["durable.execution.status"] == status.value + assert workflow.status.status_code is expected_code + # Anchored to the execution start time. + assert workflow.start_time == int(START_TIME.timestamp() * 1_000_000_000) + + +@pytest.mark.parametrize( + "status", + [ + InvocationStatus.PENDING, + InvocationStatus.RETRY, + InvocationStatus.SUCCEEDED, + InvocationStatus.FAILED, + ], +) +def test_workflow_reference_is_non_recording_after_cleanup(status): + """The retained Workflow span reference is never a recording span. + + During the invocation it is a non-recording deterministic placeholder, so + invocation cleanup on any status leaves no recording span abandoned. + """ + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + workflow_reference = plugin._workflow_span + assert workflow_reference is not None + assert not workflow_reference.is_recording() + + plugin.on_invocation_end(_invocation_end_info(status=status)) + + assert not workflow_reference.is_recording() + + +@pytest.mark.parametrize("status", [InvocationStatus.PENDING, InvocationStatus.RETRY]) +def test_open_operation_reference_is_non_recording_after_non_terminal(status): + """A suspended operation's retained span reference is a non-recording placeholder.""" + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_operation_start( + OperationStartInfo( + operation_id="wait-1", + operation_type=OperationType.WAIT, + sub_type=OperationSubType.WAIT, + name="wait-for-signal", + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.STARTED, + ) + ) + operation_reference = plugin._get_span("wait-1") + assert operation_reference is not None + + plugin.on_invocation_end(_invocation_end_info(status=status)) + + assert not operation_reference.is_recording() + + @pytest.mark.parametrize( ("status", "expected_code"), [ diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py index 66a6300f..cf53db27 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py @@ -417,6 +417,53 @@ def test_invocation_span_parents_to_same_trace_ambient_span(): assert workflow.parent.span_id == derive_execution_root_span_id(EXECUTION_ARN) +def test_pre_terminal_placeholder_preserves_same_trace_tracestate(): + """The Workflow placeholder and operation links carry ambient tracestate.""" + plugin, exporter = _create_plugin() + canonical_trace_id = _to_otel_trace_id(EXECUTION_ARN, START_TIME) + trace_state = TraceState([("vendor", "opaque")]) + ambient_context = SpanContext( + trace_id=canonical_trace_id, + span_id=int("1234567890abcdef", 16), + is_remote=False, + trace_flags=TraceFlags(TraceFlags.SAMPLED), + trace_state=trace_state, + ) + ambient = NonRecordingSpan(ambient_context) + token = otel_context.attach(trace.set_span_in_context(ambient, Context())) + try: + plugin.on_invocation_start(_invocation_start_info()) + assert plugin._workflow_span is not None + assert plugin._workflow_span.get_span_context().trace_state == trace_state + # A cross-invocation completion links the deterministic operation context. + plugin.on_operation_end( + OperationEndInfo( + operation_id="wait-existing", + operation_type=OperationType.WAIT, + sub_type=OperationSubType.WAIT, + name="existing-wait", + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.SUCCEEDED, + end_time=END_TIME, + error=None, + ) + ) + plugin.on_invocation_end(_invocation_end_info(status=InvocationStatus.PENDING)) + finally: + otel_context.detach(token) + + span = next(s for s in exporter.get_finished_spans() if s.name == "existing-wait") + operation_link = next( + link + for link in span.links + if link.context.span_id + == operation_id_to_span_id(EXECUTION_ARN, "wait-existing") + ) + assert operation_link.context.trace_state == trace_state + + def test_extracted_remote_parent_is_execution_ancestor(): remote_trace_id = int("5759e988bd862e3fe1be46a994272793", 16) remote_parent_id = int("53995c3f42cd8ad8", 16) @@ -1484,7 +1531,11 @@ def test_workflow_span_exported_on_terminal(status, expected_code): @pytest.mark.parametrize("status", [InvocationStatus.PENDING, InvocationStatus.RETRY]) def test_workflow_span_not_exported_on_non_terminal(status): - """Non-terminal invocations do not export (end) the Workflow span.""" + """Non-terminal invocations do not materialize (export) the Workflow span. + + The Workflow span is a non-recording placeholder during the invocation, so a + non-terminal status leaves nothing to export and no recording span to abandon. + """ plugin, exporter = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) plugin.on_invocation_end(_invocation_end_info(status)) @@ -1494,6 +1545,57 @@ def test_workflow_span_not_exported_on_non_terminal(status): assert "Invocation" in names +@pytest.mark.parametrize( + "status", + [ + InvocationStatus.PENDING, + InvocationStatus.RETRY, + InvocationStatus.SUCCEEDED, + InvocationStatus.FAILED, + ], +) +def test_workflow_reference_is_non_recording_after_cleanup(status): + """The retained Workflow span reference is never a recording span. + + During the invocation it is a non-recording deterministic placeholder, so + invocation cleanup on any status leaves no recording span abandoned. + """ + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + workflow_reference = plugin._workflow_span + assert workflow_reference is not None + assert not workflow_reference.is_recording() + + plugin.on_invocation_end(_invocation_end_info(status)) + + assert not workflow_reference.is_recording() + + +@pytest.mark.parametrize("status", [InvocationStatus.PENDING, InvocationStatus.RETRY]) +def test_open_operation_reference_is_non_recording_after_non_terminal(status): + """A suspended operation's retained span reference is ended, not abandoned.""" + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_operation_start( + OperationStartInfo( + operation_id="wait-1", + operation_type=OperationType.WAIT, + sub_type=OperationSubType.WAIT, + name="wait-for-signal", + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.STARTED, + ) + ) + operation_reference = plugin._get_span("wait-1") + assert operation_reference is not None + + plugin.on_invocation_end(_invocation_end_info(status)) + + assert not operation_reference.is_recording() + + def test_operation_span_links_to_workflow_span(): """Operation spans link to the Workflow span while parented to invocation.""" plugin, exporter = _create_plugin()