From efc48e4ce5c384cab30f523080765e62e26fd15d Mon Sep 17 00:00:00 2001 From: Ayushi Ahjolia Date: Tue, 1 Sep 2026 14:29:57 -0700 Subject: [PATCH 1/8] fix(otel): end recording spans on non-terminal invocations --- .../execution_plugin.py | 209 ++++++++---- .../invocation_plugin.py | 91 ++++-- .../e2e/test_invocation_wait_resume_int.py | 5 +- .../tests/test_execution_plugin.py | 302 ++++++++++++++++-- .../tests/test_invocation_plugin.py | 104 +++++- 5 files changed, 605 insertions(+), 106 deletions(-) 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() From 734cf293b0df95ff9f220a06ee787387cd0d33dd Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 3 Sep 2026 21:33:20 +0000 Subject: [PATCH 2/8] fix(otel): support vendor parent span processors --- .../durable_parent_span.py | 143 ++++++++++++++++++ .../execution_plugin.py | 75 +++++++-- .../invocation_plugin.py | 13 +- .../tests/test_execution_plugin.py | 69 ++++++++- .../test_execution_plugin_integration.py | 60 +++++++- 5 files changed, 333 insertions(+), 27 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/durable_parent_span.py diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/durable_parent_span.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/durable_parent_span.py new file mode 100644 index 00000000..a561c6fb --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/durable_parent_span.py @@ -0,0 +1,143 @@ +"""OpenTelemetry parent span used while a durable span is not recording.""" + +from __future__ import annotations + +import datetime +from collections.abc import Mapping +from typing import Any + +from opentelemetry.trace import ( + Span, + SpanContext, + SpanKind, + Status, + StatusCode, +) + + +class DurableParentSpan(Span): + """A non-recording parent compatible with SDK span processors. + + ``NonRecordingSpan`` is the standard OpenTelemetry representation for a + span context without a recording span. Some vendor span processors, + including the ADOT processor used by Lambda, nevertheless inspect SDK-only + fields such as ``kind`` and ``attributes`` on every parent. This class + carries the same immutable context while exposing those read-only fields + with safe defaults. + + The span also remembers the earliest and latest timestamps observed for + descendants. Durable operation spans are materialized only when the + operation completes, so this lets the eventual recording span enclose + attempts that started before that materialization point. + """ + + def __init__( + self, + span_context: SpanContext, + *, + start_time: datetime.datetime | None = None, + ) -> None: + self._span_context = span_context + self.kind = SpanKind.INTERNAL + self.attributes: Mapping[str, Any] = {} + self.name = "" + self.status = Status(StatusCode.UNSET) + self._earliest_start_time = start_time + self._latest_end_time: datetime.datetime | None = None + + def get_span_context(self) -> SpanContext: + return self._span_context + + def is_recording(self) -> bool: + return False + + def end(self, end_time: int | None = None) -> None: + return + + def set_attributes(self, attributes: Mapping[str, Any]) -> None: + return + + def set_attribute(self, key: str, value: Any) -> None: + return + + def add_event( + self, + name: str, + attributes: Mapping[str, Any] | None = None, + timestamp: int | None = None, + ) -> None: + return + + def update_name(self, name: str) -> None: + return + + def set_status( + self, + status: Status | StatusCode, + description: str | None = None, + ) -> None: + return + + def record_exception( + self, + exception: BaseException, + attributes: Mapping[str, Any] | None = None, + timestamp: int | None = None, + escaped: bool = False, + ) -> None: + return + + def note_start_time(self, timestamp: datetime.datetime | None) -> None: + """Include a descendant or operation start timestamp.""" + if timestamp is None: + return + if self._earliest_start_time is None or timestamp < self._earliest_start_time: + self._earliest_start_time = timestamp + + def note_end_time(self, timestamp: datetime.datetime | None) -> None: + """Include a descendant or operation end timestamp.""" + if timestamp is None: + return + if self._latest_end_time is None or timestamp > self._latest_end_time: + self._latest_end_time = timestamp + + def normalized_start_time( + self, timestamp: datetime.datetime | None + ) -> datetime.datetime | None: + """Return a start that encloses all observed descendants.""" + if timestamp is None: + return self._earliest_start_time + if self._earliest_start_time is None: + return timestamp + return min(timestamp, self._earliest_start_time) + + def normalized_end_time( + self, + timestamp: datetime.datetime | None, + *, + start_time: datetime.datetime | None = None, + ) -> datetime.datetime | None: + """Return an end that encloses all observed descendants.""" + if timestamp is None: + normalized = self._latest_end_time + elif self._latest_end_time is None: + normalized = timestamp + else: + normalized = max(timestamp, self._latest_end_time) + if ( + start_time is not None + and normalized is not None + and normalized <= start_time + ): + return start_time + datetime.timedelta(microseconds=1) + return normalized + + +def ensure_end_after_start( + start_time: datetime.datetime | None, + end_time: datetime.datetime | None, +) -> datetime.datetime | None: + """Prevent a completed span from ending at or before its start.""" + if start_time is not None and end_time is not None and end_time <= start_time: + return start_time + datetime.timedelta(microseconds=1) + return end_time 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 12699663..9724c6e4 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 @@ -53,7 +53,6 @@ from opentelemetry.sdk.trace.sampling import Sampler from opentelemetry.trace import ( Link, - NonRecordingSpan, Span, SpanContext, SpanKind, @@ -80,6 +79,10 @@ resolve_sampling_result, store_sampling_intent, ) +from aws_durable_execution_sdk_python_otel.durable_parent_span import ( + DurableParentSpan, + ensure_end_after_start, +) from aws_durable_execution_sdk_python_otel.execution_trace_context import ( ExecutionTraceContext, canonical_trace_id, @@ -326,15 +329,39 @@ def _operation_span_context(self, operation_id: str) -> SpanContext | None: trace_state=self._resolved_trace_state(), ) - def _register_operation_placeholder(self, operation_id: str) -> Span | None: + def _register_operation_placeholder( + self, + operation_id: str, + start_time: datetime.datetime | None = None, + ) -> DurableParentSpan | 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) + placeholder = DurableParentSpan(span_context, start_time=start_time) self._set_span(operation_id, placeholder) return placeholder + def _note_parent_start( + self, + parent_id: str | None, + timestamp: datetime.datetime | None, + ) -> None: + """Include a child start time in a deferred parent placeholder.""" + parent = self._resolve_parent(parent_id) + if isinstance(parent, DurableParentSpan): + parent.note_start_time(timestamp) + + def _note_parent_end( + self, + parent_id: str | None, + timestamp: datetime.datetime | None, + ) -> None: + """Include a child end time in a deferred parent placeholder.""" + parent = self._resolve_parent(parent_id) + if isinstance(parent, DurableParentSpan): + parent.note_end_time(timestamp) + def _invocation_parent_context(self) -> Context: """Return same-trace ambient context, else execution ancestor context.""" execution_trace_context = self._execution_trace_context @@ -351,7 +378,7 @@ def _invocation_parent_context(self) -> Context: trace.set_span_in_context(ambient_span, Context()) ) - ancestor = NonRecordingSpan(execution_trace_context.execution_ancestor) + ancestor = DurableParentSpan(execution_trace_context.execution_ancestor) return self._with_sampling(trace.set_span_in_context(ancestor, Context())) def _with_sampling(self, parent_context: Context) -> Context: @@ -459,7 +486,10 @@ def _start_workflow_span(self, info: InvocationStartInfo) -> None: trace_flags=self._execution_trace_context.trace_flags, trace_state=self._resolved_trace_state(), ) - self._workflow_span = NonRecordingSpan(workflow_span_context) + self._workflow_span = DurableParentSpan( + workflow_span_context, + start_time=self._execution_start_time, + ) def _export_workflow_span(self, info: InvocationEndInfo) -> None: """Create and end the recording Workflow span once, on a terminal status. @@ -473,7 +503,7 @@ def _export_workflow_span(self, info: InvocationEndInfo) -> None: return parent_context = self._with_sampling( trace.set_span_in_context( - NonRecordingSpan(self._execution_trace_context.execution_ancestor), + DurableParentSpan(self._execution_trace_context.execution_ancestor), Context(), ) ) @@ -605,7 +635,8 @@ def on_operation_start(self, info: OperationStartInfo) -> None: 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) + self._register_operation_placeholder(info.operation_id, info.start_time) + self._note_parent_start(info.parent_id, info.start_time) def on_operation_end(self, info: OperationEndInfo) -> None: logger.debug("Durable operation ended: %s", info) @@ -618,14 +649,24 @@ def on_operation_end(self, info: OperationEndInfo) -> None: 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. + placeholder = self._get_span(info.operation_id) self._pop_span(info.operation_id) parent = self._resolve_parent(info.parent_id) + start_time = info.start_time + end_time = info.end_time + if isinstance(placeholder, DurableParentSpan): + start_time = placeholder.normalized_start_time(start_time) + end_time = placeholder.normalized_end_time( + end_time, + start_time=start_time, + ) + end_time = ensure_end_after_start(start_time, end_time) 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, + start_time=start_time, ) if info.error: @@ -636,9 +677,7 @@ def on_operation_end(self, info: OperationEndInfo) -> None: else: span.set_status(StatusCode.OK) - end_time = info.end_time - if end_time is not None and end_time == info.start_time: - end_time += datetime.timedelta(microseconds=1) + 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)) @@ -712,12 +751,17 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None: span_key=key, deterministic=False, ) + self._note_parent_start(info.operation_id, info.start_time) else: # CONTEXT # 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) + span = self._register_operation_placeholder( + info.operation_id, + info.start_time, + ) + self._note_parent_start(info.parent_id, info.start_time) if span is not None: self._attach_context( key, trace.set_span_in_context(span, otel_context.get_current()) @@ -737,6 +781,9 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: raise RuntimeError( "on_user_function_end without matching on_user_function_start" ) + end_time = ensure_end_after_start(info.start_time, info.end_time) + self._note_parent_end(info.operation_id, end_time) + self._note_parent_end(info.parent_id, end_time) if ( info.operation_type is OperationType.STEP and info.outcome is not UserFunctionOutcome.INCOMPLETE @@ -756,9 +803,7 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: else: span.set_status(StatusCode.OK) - end_time = info.end_time - if end_time is not None and end_time == info.start_time: - end_time += datetime.timedelta(microseconds=1) + end_time = ensure_end_after_start(info.start_time, info.end_time) popped = self._pop_span(key) if popped is not None: popped.end(end_time=_to_otel_timestamp(end_time)) 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 3bf07c95..8c1da2dd 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 @@ -25,7 +25,6 @@ from opentelemetry.sdk.trace.sampling import Sampler from opentelemetry.trace import ( Link, - NonRecordingSpan, Span, SpanContext, SpanKind, @@ -52,6 +51,9 @@ resolve_sampling_result, store_sampling_intent, ) +from aws_durable_execution_sdk_python_otel.durable_parent_span import ( + DurableParentSpan, +) from aws_durable_execution_sdk_python_otel.execution_trace_context import ( ExecutionTraceContext, canonical_trace_id, @@ -341,7 +343,7 @@ def _invocation_parent_context(self) -> Context: trace.set_span_in_context(ambient_span, Context()) ) - ancestor = NonRecordingSpan(execution_trace_context.execution_ancestor) + ancestor = DurableParentSpan(execution_trace_context.execution_ancestor) return self._with_sampling(trace.set_span_in_context(ancestor, Context())) def _with_sampling(self, parent_context: Context) -> Context: @@ -589,7 +591,10 @@ def _start_workflow_span(self, info: InvocationStartInfo) -> None: trace_flags=self._execution_trace_context.trace_flags, trace_state=self._resolved_trace_state(), ) - self._workflow_span = NonRecordingSpan(workflow_span_context) + self._workflow_span = DurableParentSpan( + workflow_span_context, + start_time=self._execution_start_time, + ) def _export_workflow_span(self, info: InvocationEndInfo) -> None: """Create and end the recording Workflow span once, on a terminal status. @@ -603,7 +608,7 @@ def _export_workflow_span(self, info: InvocationEndInfo) -> None: return parent_context = self._with_sampling( trace.set_span_in_context( - NonRecordingSpan(self._execution_trace_context.execution_ancestor), + DurableParentSpan(self._execution_trace_context.execution_ancestor), Context(), ) ) 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 5afcd3de..8c5a0c82 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 @@ -4,7 +4,7 @@ import threading from concurrent.futures import ThreadPoolExecutor -from datetime import UTC, datetime +from datetime import UTC, datetime, timedelta from types import SimpleNamespace import opentelemetry.context as otel_context @@ -33,6 +33,7 @@ from opentelemetry.trace import ( NonRecordingSpan, SpanContext, + SpanKind, TraceFlags, TraceState, ) @@ -44,6 +45,9 @@ operation_id_to_span_id, ) from aws_durable_execution_sdk_python_otel.execution_plugin import ExecutionOtelPlugin +from aws_durable_execution_sdk_python_otel.durable_parent_span import ( + DurableParentSpan, +) from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig @@ -62,9 +66,9 @@ def _assert_otel_context_balanced(): """ before = otel_context.get_current() yield - assert otel_context.get_current() == before, ( - "test leaked OTel context state: an attach() was not detached" - ) + assert ( + otel_context.get_current() == before + ), "test leaked OTel context state: an attach() was not detached" def _create_plugin( @@ -304,6 +308,63 @@ def test_workflow_span_not_exported_on_non_terminal_status(): assert "Workflow" not in names +def test_workflow_placeholder_exposes_vendor_parent_span_fields(): + plugin, _ = _create_plugin() + + plugin.on_invocation_start(_invocation_start_info()) + + placeholder = plugin._workflow_span + assert isinstance(placeholder, DurableParentSpan) + assert not placeholder.is_recording() + assert placeholder.kind is SpanKind.INTERNAL + assert placeholder.attributes.get("durable.execution.arn") is None + + plugin.on_invocation_end(_invocation_end_info(status=InvocationStatus.PENDING)) + + +def test_deferred_operation_encloses_attempt_timestamps(): + """A materialized operation must contain attempts emitted before its end.""" + plugin, exporter = _create_plugin() + operation_start_time = START_TIME + timedelta(seconds=1) + + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_operation_start( + OperationStartInfo( + operation_id="step-1", + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name="step-1", + parent_id=None, + start_time=operation_start_time, + is_replayed=False, + status=OperationStatus.STARTED, + ) + ) + plugin.on_user_function_start(_step_start_info("step-1")) + plugin.on_user_function_end(_step_end_info("step-1")) + plugin.on_operation_end( + OperationEndInfo( + operation_id="step-1", + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name="step-1", + parent_id=None, + start_time=operation_start_time, + is_replayed=False, + status=OperationStatus.SUCCEEDED, + end_time=END_TIME, + error=None, + ) + ) + plugin.on_invocation_end(_invocation_end_info()) + + spans = {span.name: span for span in exporter.get_finished_spans()} + operation = spans["step-1"] + attempt = spans["step-1 attempt 1"] + assert operation.start_time <= attempt.start_time + assert operation.end_time >= attempt.end_time + + def test_operation_parented_under_workflow_and_linked_to_invocation(): plugin, exporter = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py index 90493613..cb03ca0d 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py @@ -33,11 +33,12 @@ ) from opentelemetry import trace from opentelemetry.context import Context -from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor, TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.trace import ( ProxyTracerProvider, + SpanKind, TracerProvider as ApiTracerProvider, ) @@ -68,9 +69,9 @@ def _assert_otel_context_balanced(): """Assert each test leaves the OTel thread-local context as it found it.""" before = otel_context.get_current() yield - assert otel_context.get_current() == before, ( - "test leaked OTel context state: an attach() was not detached" - ) + assert ( + otel_context.get_current() == before + ), "test leaked OTel context state: an attach() was not detached" def _provider() -> tuple[TracerProvider, InMemorySpanExporter]: @@ -80,6 +81,31 @@ def _provider() -> tuple[TracerProvider, InMemorySpanExporter]: return provider, exporter +class _AdotParentInspectionProcessor(SpanProcessor): + """Exercise the SDK-only parent fields accessed by the ADOT processor.""" + + def on_start(self, span, parent_context=None) -> None: + parent = trace.get_current_span(parent_context) + if not parent.get_span_context().is_valid: + return + if isinstance(parent, ReadableSpan): + _ = parent.attributes + else: + _ = parent.kind + _ = parent.attributes.get("aws.trace.id") + if parent.kind is SpanKind.SERVER: + _ = parent.kind + + def on_end(self, span: ReadableSpan) -> None: + return + + def shutdown(self) -> None: + return + + def force_flush(self, timeout_millis: int = 30000) -> bool: + return True + + def _invocation_start() -> InvocationStartInfo: return InvocationStartInfo( request_id="request-1", @@ -242,6 +268,32 @@ def test_global_proxy_binds_sdk_provider_before_first_invocation( } +def test_parent_placeholder_supports_adot_style_parent_inspection() -> None: + """A vendor processor can inspect a deferred parent without an exception.""" + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(_AdotParentInspectionProcessor()) + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = ExecutionOtelPlugin( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + + plugin.on_invocation_start(_invocation_start()) + _run_step_lifecycle(plugin) + plugin.on_invocation_end(_invocation_end()) + + assert {span.name for span in exporter.get_finished_spans()} == { + "Invocation", + "Workflow", + OP_NAME, + f"{OP_NAME} attempt 1", + } + + def test_global_proxy_disables_entire_invocation_until_sdk_provider_is_ready( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture ) -> None: From 178e30c5d97c4ee91c7c1ceedcbe4872a53645e0 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 4 Sep 2026 14:17:20 +0000 Subject: [PATCH 3/8] fix(otel): address PR 696 review findings --- .../durable_parent_span.py | 39 +++--- .../execution_plugin.py | 21 ++++ .../tests/test_execution_plugin.py | 116 ++++++++++++++++-- .../test_execution_plugin_integration.py | 16 +-- 4 files changed, 163 insertions(+), 29 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/durable_parent_span.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/durable_parent_span.py index a561c6fb..46bcbdbc 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/durable_parent_span.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/durable_parent_span.py @@ -3,6 +3,7 @@ from __future__ import annotations import datetime +import threading from collections.abc import Mapping from typing import Any @@ -44,6 +45,7 @@ def __init__( self.status = Status(StatusCode.UNSET) self._earliest_start_time = start_time self._latest_end_time: datetime.datetime | None = None + self._timestamp_lock = threading.Lock() def get_span_context(self) -> SpanContext: return self._span_context @@ -91,25 +93,31 @@ def note_start_time(self, timestamp: datetime.datetime | None) -> None: """Include a descendant or operation start timestamp.""" if timestamp is None: return - if self._earliest_start_time is None or timestamp < self._earliest_start_time: - self._earliest_start_time = timestamp + with self._timestamp_lock: + if ( + self._earliest_start_time is None + or timestamp < self._earliest_start_time + ): + self._earliest_start_time = timestamp def note_end_time(self, timestamp: datetime.datetime | None) -> None: """Include a descendant or operation end timestamp.""" if timestamp is None: return - if self._latest_end_time is None or timestamp > self._latest_end_time: - self._latest_end_time = timestamp + with self._timestamp_lock: + if self._latest_end_time is None or timestamp > self._latest_end_time: + self._latest_end_time = timestamp def normalized_start_time( self, timestamp: datetime.datetime | None ) -> datetime.datetime | None: """Return a start that encloses all observed descendants.""" - if timestamp is None: - return self._earliest_start_time - if self._earliest_start_time is None: - return timestamp - return min(timestamp, self._earliest_start_time) + with self._timestamp_lock: + if timestamp is None: + return self._earliest_start_time + if self._earliest_start_time is None: + return timestamp + return min(timestamp, self._earliest_start_time) def normalized_end_time( self, @@ -118,12 +126,13 @@ def normalized_end_time( start_time: datetime.datetime | None = None, ) -> datetime.datetime | None: """Return an end that encloses all observed descendants.""" - if timestamp is None: - normalized = self._latest_end_time - elif self._latest_end_time is None: - normalized = timestamp - else: - normalized = max(timestamp, self._latest_end_time) + with self._timestamp_lock: + if timestamp is None: + normalized = self._latest_end_time + elif self._latest_end_time is None: + normalized = timestamp + else: + normalized = max(timestamp, self._latest_end_time) if ( start_time is not None and normalized is not None 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 9724c6e4..2d19ad04 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 @@ -338,6 +338,10 @@ def _register_operation_placeholder( span_context = self._operation_span_context(operation_id) if span_context is None: return None + existing = self._get_span(operation_id) + if isinstance(existing, DurableParentSpan): + existing.note_start_time(start_time) + return existing placeholder = DurableParentSpan(span_context, start_time=start_time) self._set_span(operation_id, placeholder) return placeholder @@ -548,6 +552,9 @@ def _end_open_recording_spans(self) -> None: continue popped = self._pop_span(key) if popped is not None: + popped.set_attribute( + "durable.span.truncated_at_invocation_boundary", True + ) popped.end() def _start_invocation_span(self, info: InvocationStartInfo) -> None: @@ -642,6 +649,11 @@ def on_operation_end(self, info: OperationEndInfo) -> None: 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: @@ -660,6 +672,11 @@ def on_operation_end(self, info: OperationEndInfo) -> None: end_time, start_time=start_time, ) + if start_time is None and end_time is not None: + # Checkpointless child contexts report no durable start timestamp. + # Use their callback end as the lower bound rather than allowing + # the tracer to choose a later wall-clock start. + start_time = end_time end_time = ensure_end_after_start(start_time, end_time) span = self._start_span( operation_id=info.operation_id, @@ -700,6 +717,10 @@ def _start_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 = ( operation_id_to_span_id(self._execution_arn, operation_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 8c5a0c82..a2489830 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 @@ -66,9 +66,9 @@ def _assert_otel_context_balanced(): """ before = otel_context.get_current() yield - assert ( - otel_context.get_current() == before - ), "test leaked OTel context state: an attach() was not detached" + assert otel_context.get_current() == before, ( + "test leaked OTel context state: an attach() was not detached" + ) def _create_plugin( @@ -365,6 +365,25 @@ def test_deferred_operation_encloses_attempt_timestamps(): assert operation.end_time >= attempt.end_time +def test_deferred_parent_timestamps_are_thread_safe(): + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + + parent = plugin._workflow_span + assert isinstance(parent, DurableParentSpan) + + start_times = [START_TIME + timedelta(seconds=offset) for offset in (3, 1, 2, 4)] + end_times = [END_TIME + timedelta(seconds=offset) for offset in (2, 4, 1, 3)] + with ThreadPoolExecutor(max_workers=4) as executor: + list(executor.map(parent.note_start_time, start_times)) + list(executor.map(parent.note_end_time, end_times)) + + assert parent.normalized_start_time(None) == START_TIME + assert parent.normalized_end_time(END_TIME) == END_TIME + timedelta(seconds=4) + + plugin.on_invocation_end(_invocation_end_info(status=InvocationStatus.PENDING)) + + def test_operation_parented_under_workflow_and_linked_to_invocation(): plugin, exporter = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) @@ -748,6 +767,30 @@ def test_suspend_then_resume_operation_exports_one_deterministic_span(): EXECUTION_ARN, operation_id ) + # ReplayChildren/virtual child completion callbacks are replay-only and + # must not re-export the terminal deterministic span in a later invocation. + plugin.on_invocation_start(_invocation_start_info()) + 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=True, + status=OperationStatus.SUCCEEDED, + end_time=END_TIME, + error=None, + ) + ) + plugin.on_invocation_end(_invocation_end_info(status=InvocationStatus.PENDING)) + + operation_spans = [ + s for s in exporter.get_finished_spans() if s.name == "long-wait" + ] + assert len(operation_spans) == 1 + def test_suspended_child_context_exports_one_span_on_replay(): """A child context that suspends then replays exports a single span.""" @@ -792,6 +835,49 @@ def test_suspended_child_context_exports_one_span_on_replay(): ) +def test_checkpointless_context_end_uses_a_non_negative_duration(): + plugin, exporter = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + + plugin.on_operation_end( + OperationEndInfo( + operation_id="virtual-context", + operation_type=OperationType.CONTEXT, + sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT, + name="virtual-context", + parent_id=None, + start_time=None, + is_replayed=False, + status=OperationStatus.SUCCEEDED, + end_time=END_TIME, + error=None, + ) + ) + plugin.on_invocation_end(_invocation_end_info()) + + span = next( + span for span in exporter.get_finished_spans() if span.name == "virtual-context" + ) + assert span.start_time == int(END_TIME.timestamp() * 1_000_000_000) + assert span.end_time > span.start_time + + +def test_incomplete_attempt_is_marked_when_invocation_ends(): + plugin, exporter = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_user_function_start(_step_start_info("step-suspends")) + plugin.on_user_function_end(_step_incomplete_info("step-suspends")) + + plugin.on_invocation_end(_invocation_end_info(status=InvocationStatus.PENDING)) + + attempt = next( + span + for span in exporter.get_finished_spans() + if span.name == "step-suspends attempt 1" + ) + assert attempt.attributes["durable.span.truncated_at_invocation_boundary"] is True + + def test_duplicate_operation_end_exports_span_once(): """A repeated on_operation_end for one operation exports a single span.""" plugin, exporter = _create_plugin() @@ -1049,7 +1135,9 @@ def _context_incomplete_info( def _context_start_info( - operation_id: str, parent_id: str | None = None + operation_id: str, + parent_id: str | None = None, + start_time: datetime = START_TIME, ) -> UserFunctionStartInfo: return UserFunctionStartInfo( operation_id=operation_id, @@ -1057,7 +1145,7 @@ def _context_start_info( sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT, name=operation_id, parent_id=parent_id, - start_time=START_TIME, + start_time=start_time, is_replayed=False, status=OperationStatus.STARTED, is_replay_children=False, @@ -1303,8 +1391,15 @@ def test_reentered_child_context_does_not_leave_abandoned_span_current(): assert suspended_span is not None # Timed in-process resume re-enters the same operation. - plugin.on_user_function_start(_context_start_info(context_id)) + plugin.on_user_function_start( + _context_start_info( + context_id, + start_time=START_TIME + timedelta(seconds=1), + ) + ) assert len([key for key in plugin._context_tokens if key == context_id]) == 1 + assert plugin._get_span(context_id) is suspended_span + assert suspended_span.normalized_start_time(None) == START_TIME plugin.on_user_function_end(_context_end_info(context_id)) @@ -1325,7 +1420,12 @@ def test_reentered_step_attempt_releases_the_previous_scope(): before_context = otel_context.get_current() plugin.on_user_function_start(_step_start_info("step-1")) + first_attempt = plugin._get_span("step-1:attempt:1") + assert first_attempt is not None + plugin.on_user_function_start(_step_start_info("step-1")) + assert not first_attempt.is_recording() + plugin.on_user_function_end(_step_end_info("step-1")) assert otel_context.get_current() == before_context @@ -1423,7 +1523,9 @@ def test_nested_suspension_unwinds_scopes_in_reverse_order(): ) resumed_inner = plugin._get_span("ctx-inner") assert resumed_outer is not None - assert resumed_outer is not suspended_outer + # Re-entry reuses the deferred placeholder so timestamps from the + # suspended run remain available when the context eventually completes. + assert resumed_outer is suspended_outer assert trace.get_current_span() is resumed_inner plugin.on_user_function_end(_context_end_info("ctx-inner", parent_id="ctx-outer")) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py index cb03ca0d..df3e4b65 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py @@ -69,9 +69,9 @@ def _assert_otel_context_balanced(): """Assert each test leaves the OTel thread-local context as it found it.""" before = otel_context.get_current() yield - assert ( - otel_context.get_current() == before - ), "test leaked OTel context state: an attach() was not detached" + assert otel_context.get_current() == before, ( + "test leaked OTel context state: an attach() was not detached" + ) def _provider() -> tuple[TracerProvider, InMemorySpanExporter]: @@ -91,10 +91,12 @@ def on_start(self, span, parent_context=None) -> None: if isinstance(parent, ReadableSpan): _ = parent.attributes else: - _ = parent.kind - _ = parent.attributes.get("aws.trace.id") - if parent.kind is SpanKind.SERVER: - _ = parent.kind + parent_kind = getattr(parent, "kind", None) + parent_attributes = getattr(parent, "attributes", {}) + _ = parent_kind + _ = parent_attributes.get("aws.trace.id") + if getattr(parent, "kind", None) is SpanKind.SERVER: + _ = getattr(parent, "kind", None) def on_end(self, span: ReadableSpan) -> None: return From d2b20cb219dcdaf60ad9ec1cc9c27a9ed4599a83 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 4 Sep 2026 14:56:37 +0000 Subject: [PATCH 4/8] fix(otel): handle virtual context replay spans --- .../execution_plugin.py | 76 ++++++++++-- .../invocation_plugin.py | 24 +++- .../tests/test_execution_plugin.py | 61 +++++++++ .../tests/test_invocation_plugin.py | 116 ++++++++++++++++++ 4 files changed, 264 insertions(+), 13 deletions(-) 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 2d19ad04..35d401b0 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 @@ -10,9 +10,11 @@ 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. 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. +checkpoint-backed 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. Checkpointless virtual contexts use fresh per-invocation segment IDs +linked to their deterministic logical operation context. This is the Python adaptation of the JS ``ExecutionOtelPlugin`` from aws-durable-execution-sdk-js#729. Because the Python plugin interface differs @@ -147,6 +149,10 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: self._workflow_span: Span | None = None self._invocation_span: Span | None = None self._operation_spans: dict[str, Span] = {} + # CONTEXT operations that emitted a durable START hook this invocation. + # 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() @@ -316,14 +322,23 @@ def _resolved_trace_state(self) -> TraceState: 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.""" + def _operation_span_context( + self, + operation_id: str, + *, + span_id: int | None = None, + ) -> SpanContext | None: + """Return a SpanContext for an operation or invocation-local segment.""" 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), + span_id=( + span_id + if span_id is not None + else 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(), @@ -333,9 +348,14 @@ def _register_operation_placeholder( self, operation_id: str, start_time: datetime.datetime | None = None, + *, + deterministic: bool = True, ) -> DurableParentSpan | None: """Register a non-recording placeholder holding the operation context.""" - span_context = self._operation_span_context(operation_id) + span_context = self._operation_span_context( + operation_id, + span_id=None if deterministic else self._id_generator.generate_span_id(), + ) if span_context is None: return None existing = self._get_span(operation_id) @@ -628,6 +648,7 @@ def _reset_state(self) -> None: self._invocation_span = None with self._lock: self._operation_spans = {} + self._checkpointed_context_ids = set() self._ended_operation_ids = set() self._tracing_enabled = False @@ -639,7 +660,9 @@ def on_operation_start(self, info: OperationStartInfo) -> None: if not self._tracing_enabled: return if info.operation_type is OperationType.CONTEXT: - return # tracked via on_user_function_start + with self._lock: + self._checkpointed_context_ids.add(info.operation_id) + return # span 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, info.start_time) @@ -678,12 +701,30 @@ def on_operation_end(self, info: OperationEndInfo) -> None: # the tracer to choose a later wall-clock start. start_time = end_time end_time = ensure_end_after_start(start_time, end_time) + placeholder_span_id = ( + placeholder.get_span_context().span_id + if isinstance(placeholder, DurableParentSpan) + else None + ) + logical_span_id = operation_id_to_span_id( + self._execution_arn, info.operation_id + ) + checkpointless_context = ( + info.operation_type is OperationType.CONTEXT and info.start_time is None + ) + span_id_override = placeholder_span_id + if span_id_override is None and checkpointless_context: + span_id_override = self._id_generator.generate_span_id() span = self._start_span( operation_id=info.operation_id, name=info.name or info.operation_id, info=info, parent=parent, start_time=start_time, + span_id_override=span_id_override, + link_logical_operation=( + span_id_override is not None and span_id_override != logical_span_id + ), ) if info.error: @@ -709,6 +750,8 @@ def _start_span( start_time: datetime.datetime | None, span_key: str | None = None, deterministic: bool = True, + span_id_override: int | None = None, + link_logical_operation: bool = False, ) -> Span: """Start a recording span for an operation/attempt and register it. @@ -723,10 +766,18 @@ def _start_span( existing.end() links = self._build_invocation_links() span_id = ( - operation_id_to_span_id(self._execution_arn, operation_id) - if deterministic - else None + 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 link_logical_operation: + logical_context = self._operation_span_context(operation_id) + if logical_context is not None and logical_context.is_valid: + links = [*links, Link(context=logical_context)] if parent is None: parent_ctx = self._with_sampling(Context()) @@ -778,9 +829,12 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None: # 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. + with self._lock: + checkpointed = info.operation_id in self._checkpointed_context_ids span = self._register_operation_placeholder( info.operation_id, info.start_time, + deterministic=checkpointed or info.is_replay_children, ) self._note_parent_start(info.parent_id, info.start_time) if span is not None: diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py index 8c1da2dd..4c4356f9 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 @@ -148,6 +148,9 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: self._span_time_floor_ns: int | None = None # Maps operation ID (None for root) to the active span. self._operation_spans: dict[str | None, Span] = {} + # Replay state supplied by CONTEXT operation START hooks. Missing + # entries identify checkpointless contexts such as FLAT branches. + self._context_operation_replays: dict[str, bool] = {} # 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 @@ -649,6 +652,11 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: operation_ids = list(reversed(self._operation_spans)) for operation_id in operation_ids: if operation_id: + span = self._get_span(operation_id) + if span is not None and span.is_recording(): + span.set_attribute( + "durable.span.truncated_at_invocation_boundary", True + ) self._end_span(operation_id) invocation_span = self._get_span(None) @@ -697,6 +705,7 @@ def _reset_state(self) -> None: self._span_time_floor_ns = None with self._operation_spans_lock: self._operation_spans = {} + self._context_operation_replays = {} self._tracing_enabled = False def on_operation_start(self, info: OperationStartInfo) -> None: @@ -705,7 +714,10 @@ def on_operation_start(self, info: OperationStartInfo) -> None: if not self._tracing_enabled: return if info.operation_type is OperationType.CONTEXT: - # Context operations are tracked using on_user_function_start. + # The user-function hook owns the span, but this durable START hook + # distinguishes checkpoint-backed contexts from virtual branches. + with self._operation_spans_lock: + self._context_operation_replays[info.operation_id] = info.is_replayed return parent_span = self._resolve_parent_span(info.parent_id) attributes = self._extract_attributes(info) @@ -794,13 +806,21 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None: if info.operation_type is OperationType.CONTEXT else info.start_time ) + if info.operation_type is OperationType.CONTEXT: + with self._operation_spans_lock: + context_replay = self._context_operation_replays.get(info.operation_id) + existed = ( + context_replay is None or context_replay or info.is_replay_children + ) + else: + existed = False span = self._start_span( operation_id=info.operation_id, name=span_name, attributes=attributes, start_time=span_start_time, parent_span=parent_span, - existed=info.attempt != 1 and info.operation_type is not OperationType.STEP, + existed=existed, span_key=span_key, deterministic_span_id=info.operation_type is not OperationType.STEP, ) 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 a2489830..9d96e586 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 @@ -799,6 +799,18 @@ def test_suspended_child_context_exports_one_span_on_replay(): # Invocation 1: the child context starts and suspends (no end hook). plugin.on_invocation_start(_invocation_start_info()) + plugin.on_operation_start( + OperationStartInfo( + 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.STARTED, + ) + ) plugin.on_user_function_start(_context_start_info(context_id)) placeholder = plugin._get_span(context_id) assert placeholder is not None @@ -810,6 +822,18 @@ def test_suspended_child_context_exports_one_span_on_replay(): # Invocation 2: the context replays and completes. plugin.on_invocation_start(_invocation_start_info()) + plugin.on_operation_start( + OperationStartInfo( + 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=True, + status=OperationStatus.STARTED, + ) + ) plugin.on_user_function_start(_context_start_info(context_id)) plugin.on_user_function_end(_context_end_info(context_id)) plugin.on_operation_end( @@ -862,6 +886,43 @@ def test_checkpointless_context_end_uses_a_non_negative_duration(): assert span.end_time > span.start_time +def test_virtual_context_replay_uses_unique_linked_segments(): + plugin, exporter = _create_plugin() + context_id = "flat-branch" + logical_span_id = operation_id_to_span_id(EXECUTION_ARN, context_id) + + for _ in range(2): + plugin.on_invocation_start(_invocation_start_info()) + # Virtual contexts have no durable START hook. + 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.PARALLEL, + name=context_id, + parent_id=None, + start_time=None, + is_replayed=False, + status=OperationStatus.SUCCEEDED, + end_time=END_TIME, + error=None, + ) + ) + plugin.on_invocation_end(_invocation_end_info(status=InvocationStatus.PENDING)) + + contexts = [ + span for span in exporter.get_finished_spans() if span.name == context_id + ] + assert len(contexts) == 2 + assert len({span.context.span_id for span in contexts}) == 2 + assert all( + logical_span_id in {link.context.span_id for link in span.links} + for span in contexts + ) + + def test_incomplete_attempt_is_marked_when_invocation_ends(): plugin, exporter = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) 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 cf53db27..ea5aa564 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 @@ -1703,6 +1703,122 @@ def test_replayed_context_span_links_previous_logical_operation(): } +def test_checkpointed_context_first_span_uses_deterministic_id(): + plugin, exporter = _create_plugin() + operation_id = "child-context" + span_name = f"step-{operation_id}" + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_operation_start( + OperationStartInfo( + operation_id=operation_id, + operation_type=OperationType.CONTEXT, + sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT, + name=span_name, + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.STARTED, + ) + ) + plugin.on_user_function_start( + _user_function_start_info( + operation_id, + operation_type=OperationType.CONTEXT, + ) + ) + plugin.on_user_function_end( + _user_function_end_info( + operation_id, + operation_type=OperationType.CONTEXT, + ) + ) + plugin.on_operation_end( + OperationEndInfo( + operation_id=operation_id, + operation_type=OperationType.CONTEXT, + sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT, + name=span_name, + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.SUCCEEDED, + end_time=END_TIME, + error=None, + ) + ) + + span = next( + span for span in exporter.get_finished_spans() if span.name == span_name + ) + assert span.context.span_id == operation_id_to_span_id(EXECUTION_ARN, operation_id) + plugin.on_invocation_end(_invocation_end_info()) + + +def test_virtual_context_replay_uses_unique_linked_segments(): + plugin, exporter = _create_plugin() + operation_id = "flat-branch" + span_name = f"step-{operation_id}" + logical_span_id = operation_id_to_span_id(EXECUTION_ARN, operation_id) + + for _ in range(2): + plugin.on_invocation_start(_invocation_start_info()) + # Virtual contexts have no durable START hook. + plugin.on_user_function_start( + _user_function_start_info( + operation_id, + operation_type=OperationType.CONTEXT, + ) + ) + plugin.on_user_function_end( + _user_function_end_info( + operation_id, + operation_type=OperationType.CONTEXT, + ) + ) + plugin.on_operation_end( + OperationEndInfo( + operation_id=operation_id, + operation_type=OperationType.CONTEXT, + sub_type=OperationSubType.PARALLEL, + name=span_name, + parent_id=None, + start_time=None, + is_replayed=False, + status=OperationStatus.SUCCEEDED, + end_time=END_TIME, + error=None, + ) + ) + plugin.on_invocation_end(_invocation_end_info(status=InvocationStatus.PENDING)) + + contexts = [ + span for span in exporter.get_finished_spans() if span.name == span_name + ] + assert len(contexts) == 2 + assert len({span.context.span_id for span in contexts}) == 2 + assert all( + logical_span_id in {link.context.span_id for link in span.links} + for span in contexts + ) + + +def test_incomplete_attempt_is_marked_when_invocation_ends(): + plugin, exporter = _create_plugin() + operation_id = "step-suspends" + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_user_function_start(_user_function_start_info(operation_id)) + plugin.on_user_function_end(_user_function_incomplete_info(operation_id)) + + plugin.on_invocation_end(_invocation_end_info(status=InvocationStatus.PENDING)) + + attempt = next( + span + for span in exporter.get_finished_spans() + if span.name == f"step-{operation_id} attempt 1" + ) + assert attempt.attributes["durable.span.truncated_at_invocation_boundary"] is True + + def test_workflow_span_name_is_configurable(): """The Workflow span name can be overridden via constructor kwarg.""" exporter = InMemorySpanExporter() From f51711173d28372801e5093f581e3baf542610dc Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 4 Sep 2026 15:43:03 +0000 Subject: [PATCH 5/8] fix(otel): implement current span interface --- .../durable_parent_span.py | 7 +++++++ .../tests/test_execution_plugin.py | 16 ++++++++++++++++ 2 files changed, 23 insertions(+) diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/durable_parent_span.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/durable_parent_span.py index 46bcbdbc..57ed4978 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/durable_parent_span.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/durable_parent_span.py @@ -70,6 +70,13 @@ def add_event( ) -> None: return + def add_link( + self, + context: SpanContext, + attributes: Mapping[str, Any] | None = None, + ) -> None: + return + def update_name(self, name: str) -> None: return 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 9d96e586..a9f6b48e 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 @@ -322,6 +322,22 @@ def test_workflow_placeholder_exposes_vendor_parent_span_fields(): plugin.on_invocation_end(_invocation_end_info(status=InvocationStatus.PENDING)) +def test_durable_parent_span_implements_current_span_interface(): + span_context = SpanContext( + trace_id=1, + span_id=1, + is_remote=False, + trace_flags=TraceFlags(TraceFlags.SAMPLED), + trace_state=TraceState(), + ) + + placeholder = DurableParentSpan(span_context) + placeholder.add_link(span_context) + + assert placeholder.get_span_context() == span_context + assert not placeholder.is_recording() + + def test_deferred_operation_encloses_attempt_timestamps(): """A materialized operation must contain attempts emitted before its end.""" plugin, exporter = _create_plugin() From 0090a019f73af8fed0417775342d5a410553fe87 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 4 Sep 2026 16:12:53 +0000 Subject: [PATCH 6/8] fix(otel): reuse reentered context spans --- .../invocation_plugin.py | 31 +++++++++++----- .../tests/test_invocation_plugin.py | 37 ++++++++++++++++++- 2 files changed, 56 insertions(+), 12 deletions(-) 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 4c4356f9..6424bbd4 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 @@ -814,16 +814,27 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None: ) else: existed = False - span = self._start_span( - operation_id=info.operation_id, - name=span_name, - attributes=attributes, - start_time=span_start_time, - parent_span=parent_span, - existed=existed, - span_key=span_key, - deterministic_span_id=info.operation_type is not OperationType.STEP, - ) + existing_span = self._get_span(span_key) + if ( + info.operation_type is OperationType.CONTEXT + and existing_span is not None + and existing_span.is_recording() + ): + # A timed in-process resume can re-enter a CONTEXT before its + # previous user-function scope reports an end. Continue the same + # invocation segment instead of overwriting and abandoning it. + span = existing_span + else: + span = self._start_span( + operation_id=info.operation_id, + name=span_name, + attributes=attributes, + start_time=span_start_time, + parent_span=parent_span, + existed=existed, + span_key=span_key, + deterministic_span_id=info.operation_type is not OperationType.STEP, + ) self._attach_context( span_key, trace.set_span_in_context(span, context.get_current()) ) 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 ea5aa564..a1304258 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 @@ -1995,6 +1995,19 @@ def test_reentered_child_context_does_not_leave_abandoned_span_current(): before_context = otel_context.get_current() context_id = "ctx-1" + plugin.on_operation_start( + OperationStartInfo( + 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.STARTED, + ) + ) + # First run: the child context suspends, so no end hook fires. plugin.on_user_function_start( _user_function_start_info(context_id, operation_type=OperationType.CONTEXT) @@ -2007,6 +2020,7 @@ def test_reentered_child_context_does_not_leave_abandoned_span_current(): _user_function_start_info(context_id, operation_type=OperationType.CONTEXT) ) assert len([key for key in plugin._context_tokens if key == context_id]) == 1 + assert plugin._get_span(context_id) is suspended_span plugin.on_user_function_end( _user_function_end_info(context_id, operation_type=OperationType.CONTEXT) @@ -2019,6 +2033,22 @@ def test_reentered_child_context_does_not_leave_abandoned_span_current(): != suspended_span.get_span_context().span_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, + ) + ) + assert not suspended_span.is_recording() + plugin.on_invocation_end(_invocation_end_info()) @@ -2110,7 +2140,9 @@ def test_nested_suspension_unwinds_scopes_in_reverse_order(): "ctx-inner", parent_id="ctx-outer", operation_type=OperationType.CONTEXT ) ) + suspended_inner = plugin._get_span("ctx-inner") assert suspended_outer is not None + assert suspended_inner is not None # Both contexts suspend: the inner one unwinds first. plugin.on_user_function_end( @@ -2142,7 +2174,8 @@ def test_nested_suspension_unwinds_scopes_in_reverse_order(): ) resumed_inner = plugin._get_span("ctx-inner") assert resumed_outer is not None - assert resumed_outer is not suspended_outer + assert resumed_outer is suspended_outer + assert resumed_inner is suspended_inner assert trace.get_current_span() is resumed_inner plugin.on_user_function_end( @@ -2151,7 +2184,7 @@ def test_nested_suspension_unwinds_scopes_in_reverse_order(): ) ) - # The resumed outer scope is restored, not the one from the suspended run. + # The reused outer span's scope is restored. assert trace.get_current_span() is resumed_outer plugin.on_user_function_end( From e0ecd23c3dc019f9ee647fdce16303b9435349d5 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 4 Sep 2026 16:57:24 +0000 Subject: [PATCH 7/8] fix(otel): align virtual context links --- .../execution_plugin.py | 14 +------------- .../invocation_plugin.py | 14 +++++++++++++- .../tests/test_execution_plugin.py | 8 ++++++-- .../tests/test_invocation_plugin.py | 4 ++-- 4 files changed, 22 insertions(+), 18 deletions(-) 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 35d401b0..1a46cafc 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 @@ -14,7 +14,7 @@ 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. Checkpointless virtual contexts use fresh per-invocation segment IDs -linked to their deterministic logical operation context. +and retain their durable operation ID as an attribute for correlation. This is the Python adaptation of the JS ``ExecutionOtelPlugin`` from aws-durable-execution-sdk-js#729. Because the Python plugin interface differs @@ -706,9 +706,6 @@ def on_operation_end(self, info: OperationEndInfo) -> None: if isinstance(placeholder, DurableParentSpan) else None ) - logical_span_id = operation_id_to_span_id( - self._execution_arn, info.operation_id - ) checkpointless_context = ( info.operation_type is OperationType.CONTEXT and info.start_time is None ) @@ -722,9 +719,6 @@ def on_operation_end(self, info: OperationEndInfo) -> None: parent=parent, start_time=start_time, span_id_override=span_id_override, - link_logical_operation=( - span_id_override is not None and span_id_override != logical_span_id - ), ) if info.error: @@ -751,7 +745,6 @@ def _start_span( span_key: str | None = None, deterministic: bool = True, span_id_override: int | None = None, - link_logical_operation: bool = False, ) -> Span: """Start a recording span for an operation/attempt and register it. @@ -774,11 +767,6 @@ def _start_span( else None ) ) - if link_logical_operation: - logical_context = self._operation_span_context(operation_id) - if logical_context is not None and logical_context.is_valid: - links = [*links, Link(context=logical_context)] - if parent is None: parent_ctx = self._with_sampling(Context()) else: 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 6424bbd4..e0f46daa 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 @@ -401,6 +401,7 @@ def _start_span( existed: bool = False, span_key: str | None = None, deterministic_span_id: bool = True, + link_logical_operation: bool | None = None, ) -> Span: """Start and store a span for an invocation or durable operation. @@ -418,6 +419,8 @@ def _start_span( deterministic_span_id: Whether to use the deterministic operation span ID. Attempt spans set this to ``False`` so they can be separate children of the logical operation span. + link_logical_operation: Whether to link a fresh segment to the + deterministic logical operation context. Defaults to ``existed``. Returns: The started OpenTelemetry span. @@ -441,7 +444,10 @@ def _start_span( if operation_id else None ) - if existed and operation_id is not None: + 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)] @@ -809,11 +815,16 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None: if info.operation_type is OperationType.CONTEXT: with self._operation_spans_lock: context_replay = self._context_operation_replays.get(info.operation_id) + checkpointless_context = ( + context_replay is None and not info.is_replay_children + ) existed = ( context_replay is None or context_replay or info.is_replay_children ) + link_logical_operation = existed and not checkpointless_context else: existed = False + link_logical_operation = None existing_span = self._get_span(span_key) if ( info.operation_type is OperationType.CONTEXT @@ -834,6 +845,7 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None: existed=existed, span_key=span_key, deterministic_span_id=info.operation_type is not OperationType.STEP, + link_logical_operation=link_logical_operation, ) self._attach_context( span_key, trace.set_span_in_context(span, context.get_current()) 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 a9f6b48e..5c3cc3ec 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 @@ -905,7 +905,6 @@ def test_checkpointless_context_end_uses_a_non_negative_duration(): def test_virtual_context_replay_uses_unique_linked_segments(): plugin, exporter = _create_plugin() context_id = "flat-branch" - logical_span_id = operation_id_to_span_id(EXECUTION_ARN, context_id) for _ in range(2): plugin.on_invocation_start(_invocation_start_info()) @@ -931,10 +930,15 @@ def test_virtual_context_replay_uses_unique_linked_segments(): contexts = [ span for span in exporter.get_finished_spans() if span.name == context_id ] + invocation_ids = { + span.context.span_id + for span in exporter.get_finished_spans() + if span.name == "Invocation" + } assert len(contexts) == 2 assert len({span.context.span_id for span in contexts}) == 2 assert all( - logical_span_id in {link.context.span_id for link in span.links} + len(span.links) == 1 and span.links[0].context.span_id in invocation_ids for span in contexts ) 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 a1304258..bf9ac028 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 @@ -1758,7 +1758,7 @@ def test_virtual_context_replay_uses_unique_linked_segments(): plugin, exporter = _create_plugin() operation_id = "flat-branch" span_name = f"step-{operation_id}" - logical_span_id = operation_id_to_span_id(EXECUTION_ARN, operation_id) + workflow_span_id = derive_workflow_span_id(EXECUTION_ARN) for _ in range(2): plugin.on_invocation_start(_invocation_start_info()) @@ -1797,7 +1797,7 @@ def test_virtual_context_replay_uses_unique_linked_segments(): assert len(contexts) == 2 assert len({span.context.span_id for span in contexts}) == 2 assert all( - logical_span_id in {link.context.span_id for link in span.links} + len(span.links) == 1 and span.links[0].context.span_id == workflow_span_id for span in contexts ) From ab539a717c3c09a1cd9228f7caf47b1e523c138d Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 4 Sep 2026 18:29:01 +0000 Subject: [PATCH 8/8] fix(otel): mark only incomplete attempts truncated --- .../invocation_plugin.py | 18 +++++++++++++++++- .../tests/test_invocation_plugin.py | 4 ++++ 2 files changed, 21 insertions(+), 1 deletion(-) 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 e0f46daa..012f17f2 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 @@ -151,6 +151,9 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: # Replay state supplied by CONTEXT operation START hooks. Missing # entries identify checkpointless contexts such as FLAT branches. self._context_operation_replays: dict[str, bool] = {} + # STEP attempt spans whose user function explicitly ended INCOMPLETE. + # Only these spans are marked truncated during invocation cleanup. + self._incomplete_attempt_span_keys: 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 @@ -656,10 +659,15 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: # order to keep every child contained within its parent. with self._operation_spans_lock: operation_ids = list(reversed(self._operation_spans)) + incomplete_attempt_span_keys = set(self._incomplete_attempt_span_keys) for operation_id in operation_ids: if operation_id: span = self._get_span(operation_id) - if span is not None and span.is_recording(): + if ( + operation_id in incomplete_attempt_span_keys + and span is not None + and span.is_recording() + ): span.set_attribute( "durable.span.truncated_at_invocation_boundary", True ) @@ -712,6 +720,7 @@ def _reset_state(self) -> None: with self._operation_spans_lock: self._operation_spans = {} self._context_operation_replays = {} + self._incomplete_attempt_span_keys = set() self._tracing_enabled = False def on_operation_start(self, info: OperationStartInfo) -> None: @@ -874,6 +883,13 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: "on_user_function_end called without matching on_user_function_start" ) + if info.operation_type is OperationType.STEP: + with self._operation_spans_lock: + if info.outcome is UserFunctionOutcome.INCOMPLETE: + self._incomplete_attempt_span_keys.add(span_key) + else: + self._incomplete_attempt_span_keys.discard(span_key) + if ( info.operation_type is OperationType.STEP and info.outcome is not UserFunctionOutcome.INCOMPLETE 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 bf9ac028..f108ef72 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 @@ -1594,6 +1594,10 @@ def test_open_operation_reference_is_non_recording_after_non_terminal(status): plugin.on_invocation_end(_invocation_end_info(status)) assert not operation_reference.is_recording() + assert ( + "durable.span.truncated_at_invocation_boundary" + not in operation_reference.attributes + ) def test_operation_span_links_to_workflow_span():