feat(otel): parent durable spans to shared trace

Author: Ayushi Ahjolia

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

@@ -3,11 +3,14 @@
 from aws_durable_execution_sdk_python_otel.__about__ import __version__
 from aws_durable_execution_sdk_python_otel.context_extractors import (
     ContextExtractor,
+    ExtractedContext,
+    Sampling,
     w3c_client_context_extractor,
     xray_context_extractor,
 )
 from aws_durable_execution_sdk_python_otel.deterministic_id_generator import (
     DeterministicIdGenerator,
+    derive_execution_root_span_id,
     derive_workflow_span_id,
     operation_id_to_span_id,
 )
@@ -35,11 +38,14 @@
     "ContextExtractor",
     "DeterministicIdGenerator",
     "ExecutionOtelPlugin",
+    "ExtractedContext",
     "OtelPluginConfig",
     "InvocationOtelPlugin",
     "OtelContextLogFilter",
+    "Sampling",
     "ProviderResult",
     "create_tracer_provider",
+    "derive_execution_root_span_id",
     "derive_workflow_span_id",
     "install_log_filter",
     "operation_id_to_span_id",

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

@@ -1,4 +1,4 @@
-"""Context extractors for propagating trace context into durable executions."""
+"""Trace-context extractors for durable execution telemetry."""
 
 from __future__ import annotations
 
@@ -7,12 +7,8 @@
 from enum import Enum
 from typing import TYPE_CHECKING, Callable
 
-from opentelemetry import context as otel_context, propagate
-
 
 if TYPE_CHECKING:
-    from opentelemetry.context import Context
-
     from aws_durable_execution_sdk_python.plugin import InvocationStartInfo
 
 
@@ -54,28 +50,83 @@ def has_complete_remote_parent(self) -> bool:
         return self.has_valid_trace_id and self.has_valid_parent_span_id
 
 
-ContextExtractor = Callable[["InvocationStartInfo"], "Context"]
+ContextExtractor = Callable[["InvocationStartInfo"], ExtractedContext | None]
+
+
+def _ensure_extracted_context(extracted: object) -> ExtractedContext | None:
+    """Validate a context extractor result."""
+    if extracted is None or isinstance(extracted, ExtractedContext):
+        return extracted
+    msg = "context extractor must return ExtractedContext or None"
+    raise TypeError(msg)
+
+
+def _parse_xray_trace_id(root: str | None) -> int | None:
+    if root is None:
+        return None
+    parts = root.split("-")
+    if len(parts) != 3 or parts[0] != "1":
+        return None
+    trace_id_hex = f"{parts[1]}{parts[2]}"
+    if len(trace_id_hex) != 32:
+        return None
+    try:
+        trace_id = int(trace_id_hex, 16)
+    except ValueError:
+        return None
+    return trace_id if 0 < trace_id < 2**128 else None
 
 
-def xray_context_extractor(info: "InvocationStartInfo") -> "Context":
-    """Read the X-Ray trace header from the _X_AMZN_TRACE_ID environment variable.
+def _parse_span_id(span_id_hex: str | None) -> int | None:
+    if span_id_hex is None or len(span_id_hex) != 16:
+        return None
+    try:
+        span_id = int(span_id_hex, 16)
+    except ValueError:
+        return None
+    return span_id if 0 < span_id < 2**64 else None
 
-    The durable execution backend propagates the same Root trace ID to every
-    invocation, so all invocations share one traceId.
+
+def _parse_sampling(value: str | None) -> Sampling:
+    if value == "1":
+        return Sampling.SAMPLED
+    if value == "0":
+        return Sampling.NOT_SAMPLED
+    return Sampling.UNDECIDED
+
+
+def xray_context_extractor(info: "InvocationStartInfo") -> ExtractedContext | None:
+    """Read durable execution trace context from ``_X_AMZN_TRACE_ID``.
+
+    The Lambda durable execution backend propagates an X-Ray style header. A
+    valid ``Root`` anchors the execution trace; a valid ``Parent`` becomes the
+    remote execution ancestor; and ``Sampled`` is preserved as the backend's
+    explicit sampling decision.
     """
     trace_header = os.environ.get("_X_AMZN_TRACE_ID")
     if not trace_header:
-        return otel_context.get_current()
-    return propagate.extract(
-        carrier={"X-Amzn-Trace-Id": trace_header},
-        context=otel_context.get_current(),
+        return None
+
+    parts: dict[str, str] = {}
+    for segment in trace_header.split(";"):
+        key, separator, value = segment.partition("=")
+        if separator:
+            parts[key.strip()] = value.strip()
+
+    trace_id = _parse_xray_trace_id(parts.get("Root"))
+    parent_span_id = _parse_span_id(parts.get("Parent"))
+    sampling = _parse_sampling(parts.get("Sampled"))
+    if trace_id is None and parent_span_id is None and sampling is Sampling.UNDECIDED:
+        return None
+    return ExtractedContext(
+        trace_id=trace_id,
+        parent_span_id=parent_span_id,
+        sampling=sampling,
     )
 
 
-def w3c_client_context_extractor(info: "InvocationStartInfo") -> "Context":
-    """Read W3C traceparent from context.clientContext.custom.traceparent.
-
-    Requires the backend clientContext propagation to be enabled.
-    This extractor is a placeholder for when backend propagation is supported.
-    """
-    return otel_context.get_current()
+def w3c_client_context_extractor(
+    info: "InvocationStartInfo",
+) -> ExtractedContext | None:
+    """Placeholder for future W3C traceparent propagation support."""
+    return None

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

@@ -5,11 +5,11 @@
     Workflow -> Operation -> Attempt
 
 that stitches a single trace across every Lambda invocation of one durable
-execution. The Workflow span is the root (created in an empty context so it
-never has a parent) and 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. The Invocation
-span belongs to the ambient Lambda trace instead of the Workflow trace.
+execution. Workflow and Invocation spans parent onto the same execution
+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.
 
 This is the Python adaptation of the JS ``ExecutionOtelPlugin`` from
 aws-durable-execution-sdk-js#729. Because the Python plugin interface differs
@@ -47,8 +47,10 @@
 from opentelemetry import trace
 from opentelemetry.context import Context
 from opentelemetry.sdk.trace import Tracer as SdkTracer
+from opentelemetry.sdk.trace.sampling import Sampler
 from opentelemetry.trace import (
     Link,
+    NonRecordingSpan,
     Span,
     SpanContext,
     SpanKind,
@@ -58,14 +60,26 @@
 
 from aws_durable_execution_sdk_python_otel.context_extractors import (
     ContextExtractor,
+    ExtractedContext,
+    _ensure_extracted_context,
     xray_context_extractor,
 )
 from aws_durable_execution_sdk_python_otel.deterministic_id_generator import (
     DeterministicIdGenerator,
-    _to_otel_trace_id,
     derive_workflow_span_id,
     operation_id_to_span_id,
 )
+from aws_durable_execution_sdk_python_otel.durable_sampling import (
+    DurableSampler,
+    DurableSamplingIntent,
+    is_sampled,
+    resolve_sampling_result,
+    store_sampling_intent,
+)
+from aws_durable_execution_sdk_python_otel.execution_trace_context import (
+    ExecutionTraceContext,
+    canonical_trace_id,
+)
 from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig
 from aws_durable_execution_sdk_python_otel.log_filter import install_log_filter
 from aws_durable_execution_sdk_python_otel.provider import create_tracer_provider
@@ -113,12 +127,15 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None:
 
         self._tracer: Tracer = self._provider.get_tracer(self._config.instrument_name)
         self._id_generator = DeterministicIdGenerator()
+        self._sampling_delegate: Sampler | None = None
         self._bind_sdk_tracer()
 
         # Per-invocation state.
         self._execution_arn = ""
         self._execution_trace_id: int | None = None
-        self._extracted_context: Context | 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] = {}
@@ -135,6 +152,7 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None:
 
     def _bind_sdk_tracer(self) -> bool:
         """Bind to an SDK tracer, retrying a deferred global provider."""
+        self._sampling_delegate = None
         tracer = self._tracer
         if not isinstance(tracer, SdkTracer):
             if self._uses_global_provider:
@@ -147,6 +165,7 @@ def _bind_sdk_tracer(self) -> bool:
         # Deterministic stitching is scoped to this instrumentation tracer so
         # unrelated tracers on the same provider keep their original generator.
         self._id_generator = DeterministicIdGenerator.install_on_tracer(tracer)
+        self._sampling_delegate = DurableSampler.install_on_tracer(tracer).delegate
         return True
 
     # ------------------------------------------------------------------
@@ -274,14 +293,26 @@ def _resolve_parent(self, parent_id: str | None) -> Span | None:
         return self._workflow_span
 
     def _invocation_parent_context(self) -> Context:
-        """Return the active ambient context, then extracted upstream context."""
-        ambient_context = otel_context.get_current()
-        ambient_span_context = trace.get_current_span(
-            ambient_context
-        ).get_span_context()
-        if ambient_span_context.is_valid:
-            return ambient_context
-        return self._extracted_context or ambient_context
+        """Return same-trace ambient context, else execution ancestor context."""
+        execution_trace_context = self._execution_trace_context
+        if execution_trace_context is None:
+            return self._with_sampling(Context())
+
+        ambient_span = trace.get_current_span()
+        ambient_context = ambient_span.get_span_context()
+        if (
+            ambient_context.is_valid
+            and ambient_context.trace_id == execution_trace_context.trace_id
+        ):
+            return self._with_sampling(
+                trace.set_span_in_context(ambient_span, Context())
+            )
+
+        ancestor = NonRecordingSpan(execution_trace_context.execution_ancestor)
+        return self._with_sampling(trace.set_span_in_context(ancestor, Context()))
+
+    def _with_sampling(self, parent_context: Context) -> Context:
+        return store_sampling_intent(parent_context, self._sampling_intent)
 
     # ------------------------------------------------------------------
     # Invocation lifecycle
@@ -307,13 +338,46 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None:
             return
 
         self._execution_arn = info.execution_arn or ""
-        self._execution_trace_id = _to_otel_trace_id(
-            self._execution_arn, info.execution_start_time
+        if not self._execution_arn:
+            logger.warning(
+                "ExecutionOtelPlugin requires InvocationStartInfo.execution_arn "
+                "to derive a deterministic execution root; telemetry is disabled "
+                "for this invocation."
+            )
+            self._tracing_enabled = False
+            return
+        self._extracted_context = _ensure_extracted_context(
+            self._context_extractor(info)
+        )
+        self._execution_trace_id = canonical_trace_id(
+            extracted=self._extracted_context,
+            execution_arn=self._execution_arn,
+            execution_start_time=info.execution_start_time,
+        )
+        if self._sampling_delegate is None:
+            logger.warning(
+                "No sampler available; telemetry is disabled for this invocation."
+            )
+            self._tracing_enabled = False
+            return
+        sampling_result = resolve_sampling_result(
+            extracted=self._extracted_context,
+            ambient_span=trace.get_current_span(),
+            canonical_trace_id=self._execution_trace_id,
+            sampler=self._sampling_delegate,
+            span_name=self._workflow_span_name,
+            attributes={"durable.execution.arn": self._execution_arn},
+        )
+        self._sampling_intent = DurableSamplingIntent(sampling_result)
+        self._execution_trace_context = ExecutionTraceContext.resolve(
+            extracted=self._extracted_context,
+            canonical_trace_id=self._execution_trace_id,
+            execution_arn=self._execution_arn,
+            root_sampled=lambda: is_sampled(sampling_result),
         )
-        self._extracted_context = self._context_extractor(info)
 
         self._start_workflow_span(info)
-        # Keep the invocation in the ambient Lambda trace in both provider modes.
+        # Keep the invocation on the shared execution trace.
         self._start_invocation_span(info)
 
         # Make the Workflow span the active span so auto-instrumented spans
@@ -323,24 +387,33 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None:
         if self._workflow_span is not None:
             self._attach_context(
                 _INVOCATION_CONTEXT_KEY,
-                trace.set_span_in_context(self._workflow_span, self._extracted_context),
+                trace.set_span_in_context(
+                    self._workflow_span, otel_context.get_current()
+                ),
             )
 
     def _start_workflow_span(self, info: InvocationStartInfo) -> None:
         if not self._execution_arn:
             logger.warning("No execution ARN; skipping Workflow span creation")
             return
-        # Empty context => root span with no parent.
+        if self._execution_trace_context is None:
+            return
+        parent_context = self._with_sampling(
+            trace.set_span_in_context(
+                NonRecordingSpan(self._execution_trace_context.execution_ancestor),
+                Context(),
+            )
+        )
         with self._id_generator.use_ids(
-            trace_id=self._execution_trace_id,
+            trace_id=None,
             span_id=derive_workflow_span_id(self._execution_arn),
         ):
             self._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),
-                context=Context(),
+                context=parent_context,
             )
 
     def _start_invocation_span(self, info: InvocationStartInfo) -> None:
@@ -420,6 +493,8 @@ def _reset_state(self) -> None:
         self._execution_arn = ""
         self._execution_trace_id = None
         self._extracted_context = None
+        self._execution_trace_context = None
+        self._sampling_intent = None
         self._workflow_span = None
         self._invocation_span = None
         with self._lock:
@@ -500,12 +575,12 @@ def _start_span(
             )
 
             if parent is None:
-                parent_ctx = self._extracted_context or Context()
+                parent_ctx = self._with_sampling(Context())
             else:
-                parent_ctx = trace.set_span_in_context(parent, self._extracted_context)
-            with self._id_generator.use_ids(
-                trace_id=self._execution_trace_id, span_id=span_id
-            ):
+                parent_ctx = self._with_sampling(
+                    trace.set_span_in_context(parent, Context())
+                )
+            with self._id_generator.use_ids(trace_id=None, span_id=span_id):
                 span = self._tracer.start_span(
                     name=name,
                     attributes=self._operation_attributes(info),
@@ -552,7 +627,7 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None:
                 start_time=info.start_time,
             )
         self._attach_context(
-            key, trace.set_span_in_context(span, self._extracted_context)
+            key, trace.set_span_in_context(span, otel_context.get_current())
         )
 
     def on_user_function_end(self, info: UserFunctionEndInfo) -> None: