fix(otel): end recording spans on non-terminal

Author: Ayushi Ahjolia

packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py

@@ -133,6 +133,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: ExtractedContext | None = None
         self._execution_trace_context: ExecutionTraceContext | None = None
         self._sampling_intent: DurableSamplingIntent | None = None
@@ -346,6 +347,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 +395,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._execution_trace_context.execution_ancestor.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 +439,41 @@ 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_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()
 
     def _start_invocation_span(self, info: InvocationStartInfo) -> None:
         self._invocation_span = self._tracer.start_span(
@@ -434,12 +493,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 +515,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()
 
@@ -495,6 +545,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._invocation_span = None
         with self._lock:

packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py

@@ -137,6 +137,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
@@ -511,6 +512,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 +552,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._execution_trace_context.execution_ancestor.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 +597,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 +652,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 +673,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: