fix(otel): address PR 696 review findings

Author: Frank Chen

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

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)

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,15 +1135,17 @@ 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,
         operation_type=OperationType.CONTEXT,
         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"))

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