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 3f7c10fb..8ccc3719 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 @@ -49,11 +49,13 @@ from opentelemetry.sdk.trace import Tracer as SdkTracer from opentelemetry.trace import ( Link, + NonRecordingSpan, Span, SpanContext, SpanKind, StatusCode, Tracer, + TraceFlags, ) from aws_durable_execution_sdk_python_otel.context_extractors import ( @@ -118,6 +120,7 @@ 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: Context | None = None self._workflow_span: Span | None = None self._invocation_span: Span | None = None @@ -166,6 +169,22 @@ def _pop_span(self, key: str) -> Span | None: with self._lock: return self._operation_spans.pop(key, None) + def _end_open_operation_spans(self) -> None: + """End every operation span still open at invocation end (except invocation). + + Spans are registered parent-first, so ending them in reverse keeps each + child contained within its parent. The invocation span is ended + separately by the caller. + """ + with self._lock: + keys = list(reversed(self._operation_spans)) + for key in keys: + if key == _INVOCATION_KEY: + continue + popped = self._pop_span(key) + if popped is not None: + popped.end() + @staticmethod def _attempt_key( info: UserFunctionStartInfo | UserFunctionEndInfo, @@ -310,6 +329,7 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: self._execution_trace_id = _to_otel_trace_id( self._execution_arn, info.execution_start_time ) + self._execution_start_time = info.execution_start_time self._extracted_context = self._context_extractor(info) self._start_workflow_span(info) @@ -327,21 +347,60 @@ 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. Using a non-recording span for that role means 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 + workflow_span_context = SpanContext( + trace_id=self._execution_trace_id or 0, + span_id=derive_workflow_span_id(self._execution_arn), + is_remote=False, + trace_flags=TraceFlags(TraceFlags.SAMPLED), + ) + 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 trace and span IDs as the placeholder so the + exported root correlates with every operation span across all + invocations, and anchors the span at the execution start time. + """ + if not self._execution_arn: + return # Empty context => root span with no parent. with self._id_generator.use_ids( trace_id=self._execution_trace_id, 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=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 _start_invocation_span(self, info: InvocationStartInfo) -> None: self._invocation_span = self._tracer.start_span( @@ -361,12 +420,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 @@ -389,23 +442,20 @@ 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() + # Operation spans still open here belong to operations that suspended + # (e.g. PENDING/RETRYING) rather than completed this invocation. End them + # so no recording span is abandoned; the authoritative span for an + # operation that resumes in a later invocation is created and ended by + # on_operation_end at that time. + self._end_open_operation_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() @@ -419,6 +469,7 @@ def _reset_state(self) -> None: self._detach_remaining_contexts() self._execution_arn = "" self._execution_trace_id = None + self._execution_start_time = None self._extracted_context = None self._workflow_span = None self._invocation_span = 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 3b95a5d5..bb9854c0 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 @@ -24,11 +24,13 @@ from opentelemetry.sdk.trace import Tracer as SdkTracer from opentelemetry.trace import ( Link, + NonRecordingSpan, Span, SpanContext, SpanKind, StatusCode, Tracer, + TraceFlags, ) from aws_durable_execution_sdk_python_otel.context_extractors import ( @@ -119,6 +121,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: Context | None = None self._workflow_span: Span | None = None # Maps operation ID (None for root) to the active span. @@ -432,6 +435,7 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: self._execution_trace_id = _to_otel_trace_id( self._execution_arn, info.execution_start_time ) + self._execution_start_time = info.execution_start_time self._extracted_context = self._context_extractor(info) self._start_workflow_span(info) @@ -443,32 +447,61 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: ) def _start_workflow_span(self, info: InvocationStartInfo) -> None: - """Create the deterministic, execution-scoped Workflow root span. + """Install a non-recording placeholder for the execution-scoped Workflow span. The Workflow span is a parentless root 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 exported once, on a - terminal invocation. Operation and attempt spans link to it while - remaining parented to the invocation span. It is created unconditionally - -- InvocationOtelPlugin has no default/owned tracer-provider distinction, - so the span is emitted whether the provider is the ambient (ADOT/global) - one or an explicitly supplied one. + terminal invocation, by :meth:`_export_workflow_span`. During each + invocation the plugin only needs the 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. """ if not self._execution_arn: logger.warning("No execution ARN; skipping Workflow span creation") return + workflow_span_context = SpanContext( + trace_id=self._execution_trace_id or 0, + span_id=derive_workflow_span_id(self._execution_arn), + is_remote=False, + trace_flags=TraceFlags(TraceFlags.SAMPLED), + ) + 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 trace and span IDs as the placeholder so the + exported root correlates with every operation span across all + invocations, and anchors the span at the execution start time. + """ + if not self._execution_arn: + return # Empty context => root span with no parent. with self._id_generator.use_ids( trace_id=self._execution_trace_id, 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=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.""" @@ -505,23 +538,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() @@ -534,6 +556,7 @@ def _reset_state(self) -> None: self._detach_remaining_contexts() self._execution_arn = "" self._execution_trace_id = None + self._execution_start_time = None self._extracted_context = None self._workflow_span = None with self._operation_spans_lock: 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 dc4a699c..47f174bb 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 @@ -214,19 +214,96 @@ def test_explicit_mode_invocation_span_parented_to_ambient_span(): assert workflow.context.trace_id != ambient.get_span_context().trace_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 +@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] + assert workflow.parent is None + 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 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=status)) + + assert not operation_reference.is_recording() + + def test_operation_parented_under_workflow_and_linked_to_invocation(): plugin, exporter = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) @@ -515,11 +592,12 @@ def test_default_mode_invocation_span_parented_to_ambient_span(monkeypatch): 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. +def test_open_operation_span_ended_at_invocation_end(): + """A suspended operation span is ended at invocation end, not abandoned. - on_invocation_end drops the reference without ending it; the - span is ended only when on_operation_end fires in a later invocation. + on_invocation_end ends the still-open operation span so no recording span + leaks across a non-terminal boundary; the authoritative span for an + operation that resumes later is created when on_operation_end fires. """ plugin, exporter = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) @@ -536,12 +614,17 @@ def test_open_operation_span_not_exported_at_invocation_end(): status=OperationStatus.STARTED, ) ) + open_span = plugin._get_span("wait-1") + assert open_span is not None + assert open_span.is_recording() + # No on_operation_end: the operation suspended. plugin.on_invocation_end(_invocation_end_info(status=InvocationStatus.PENDING)) + # The open operation span is ended (exported), not abandoned recording. + assert not open_span.is_recording() 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 + assert "wait-for-signal" in exported @pytest.mark.parametrize( 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 0f3bb62d..504660d0 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 @@ -1151,7 +1151,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)) @@ -1161,6 +1165,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()