fix(otel): scope deterministic IDs to plugin tracers

Author: zhongkechen

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

@@ -6,23 +6,26 @@
 import hashlib
 import os
 import re
+from collections.abc import Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass
 from datetime import UTC, datetime
 from typing import TYPE_CHECKING
 
 from opentelemetry.sdk.trace import IdGenerator, RandomIdGenerator
 
 
 if TYPE_CHECKING:
-    from opentelemetry.sdk.trace import TracerProvider
+    from opentelemetry.sdk.trace import Tracer as SdkTracer
 
 
 HASHED_ID_PATTERN = re.compile(r"^[0-9a-f]{16}$")
 
-# Scoping the pending span ID to the execution context ensures concurrent
-# operations cannot consume each other's deterministic span ID.
-_next_span_id: contextvars.ContextVar[int | None] = contextvars.ContextVar(
-    "next_span_id", default=None
-)
+
+@dataclass(frozen=True)
+class _IdOverride:
+    trace_id: int | None
+    span_id: int | None
 
 
 def _parse_xray_root_trace_id(trace_header: str | None) -> str | None:
@@ -112,13 +115,12 @@ def derive_workflow_span_id(durable_execution_arn: str) -> int:
 
 
 class DeterministicIdGenerator(RandomIdGenerator):
-    """An ID generator that produces deterministic span IDs when a pending
-    operation ID is set, and falls back to the provided generator otherwise.
+    """An ID generator with invocation-scoped deterministic ID overrides.
 
-    Trace IDs are deterministic when an execution ARN is set, ensuring all
-    invocations of the same durable execution share a single trace. When no
-    deterministic ID is available, generation is delegated to the fallback
-    generator (the tracer provider's original ID generator by default).
+    Deterministic IDs are active only inside :meth:`use_ids`. All other
+    generation is delegated to the fallback generator. The override is stored
+    in a context variable so concurrent threads and async tasks cannot consume
+    or overwrite each other's IDs.
 
     Trace IDs embed a real timestamp so they satisfy the X-Ray format
     requirement (first 8 hex chars = Unix epoch seconds).
@@ -129,54 +131,56 @@ class DeterministicIdGenerator(RandomIdGenerator):
     """
 
     def __init__(self, fallback_id_generator: IdGenerator | None = None) -> None:
-        self._execution_trace_id: int | None = None
         self._fallback_id_generator = fallback_id_generator or RandomIdGenerator()
+        self._id_override: contextvars.ContextVar[_IdOverride | None] = (
+            contextvars.ContextVar("durable_execution_id_override", default=None)
+        )
 
     @classmethod
-    def install_on_provider(cls, provider: TracerProvider) -> DeterministicIdGenerator:
-        """Return the provider's deterministic generator, installing one if needed.
+    def install_on_tracer(cls, tracer: SdkTracer) -> DeterministicIdGenerator:
+        """Return the tracer's deterministic generator, installing one if needed.
 
-        OpenTelemetry tracers capture the provider's ID generator when they are
-        created. Reusing an installed generator ensures multiple plugin instances
-        with the same instrumentation scope configure the generator referenced by
-        the provider's cached tracer.
+        Installing on the plugin's tracer keeps unrelated instrumentation scopes
+        on the provider's original generator. Reusing an installed generator also
+        supports SDK versions that cache tracers by instrumentation scope.
         """
-        current_generator = provider.id_generator
+        current_generator = tracer.id_generator
         if isinstance(current_generator, cls):
             return current_generator
 
         generator = cls(fallback_id_generator=current_generator)
-        provider.id_generator = generator
+        tracer.id_generator = generator
         return generator
 
-    def set_next_span_id(self, span_id: int | None) -> None:
-        """Set the operation ID to use for the next span's ID.
-
-        After one span is created, it resets to random.
-        """
-        _next_span_id.set(span_id)
-
-    def set_trace_id(
-        self, execution_arn: str, start_timestamp: datetime | None
-    ) -> None:
-        """Compute and cache the deterministic trace ID for this execution.
-
-        Args:
-            execution_arn: The durable execution ARN (used for the hash portion).
-            start_timestamp: start time of invocation
-        """
-        self._execution_trace_id = _to_otel_trace_id(execution_arn, start_timestamp)
+    @contextmanager
+    def use_ids(self, *, trace_id: int | None, span_id: int | None) -> Iterator[None]:
+        """Temporarily override IDs generated in the current execution context."""
+        token = self._id_override.set(_IdOverride(trace_id, span_id))
+        try:
+            yield
+        finally:
+            self._id_override.reset(token)
 
     def generate_trace_id(self) -> int:
         """Generate a 128-bit trace ID."""
-        return (
-            self._execution_trace_id or self._fallback_id_generator.generate_trace_id()
-        )
+        override = self._id_override.get()
+        if override is not None and override.trace_id is not None:
+            return override.trace_id
+        return self._fallback_id_generator.generate_trace_id()
 
     def generate_span_id(self) -> int:
         """Generate a 64-bit span ID."""
-        span_id = _next_span_id.get()
-        # Consume once: the deterministic ID applies only to the next span
-        # created in this context; subsequent spans fall back to random.
-        _next_span_id.set(None)
-        return span_id or self._fallback_id_generator.generate_span_id()
+        override = self._id_override.get()
+        if override is not None and override.span_id is not None:
+            return override.span_id
+        return self._fallback_id_generator.generate_span_id()
+
+    def is_trace_id_random(self) -> bool:
+        """Report whether the current trace ID is randomly generated."""
+        override = self._id_override.get()
+        if override is not None and override.trace_id is not None:
+            return False
+        fallback_method = getattr(
+            self._fallback_id_generator, "is_trace_id_random", None
+        )
+        return bool(fallback_method()) if fallback_method is not None else False

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

@@ -57,6 +57,7 @@
 )
 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,
 )
@@ -104,40 +105,39 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None:
         )
         self._workflow_span_name = self._config.workflow_span_name
 
-        self._id_generator = DeterministicIdGenerator()
-        result = create_tracer_provider(
-            self._config,
-            id_generator=self._id_generator,
-        )
+        result = create_tracer_provider(self._config)
         self._provider = result.tracer_provider
         # GLOBAL (ADOT) mode parents the Invocation span to the ambient Lambda
         # invocation span instead of the Workflow span (see
         # _start_invocation_span).
         self._provider_source = result.source
 
-        # Deterministic stitching requires an SDK provider exposing id_generator.
-        from opentelemetry.sdk.trace import TracerProvider as SdkTracerProvider
+        self._tracer: Tracer = self._provider.get_tracer(self._config.instrument_name)
+
+        # Deterministic stitching is scoped to this instrumentation tracer so
+        # unrelated tracers on the same provider keep their original generator.
+        from opentelemetry.sdk.trace import Tracer as SdkTracer
 
-        if isinstance(self._provider, SdkTracerProvider):
-            self._id_generator = DeterministicIdGenerator.install_on_provider(
-                self._provider
+        self._id_generator = DeterministicIdGenerator()
+        if isinstance(self._tracer, SdkTracer):
+            self._id_generator = DeterministicIdGenerator.install_on_tracer(
+                self._tracer
             )
         else:
             logger.warning(
-                "ExecutionOtelPlugin expected an SDK TracerProvider but got %s; "
+                "ExecutionOtelPlugin expected an SDK Tracer but got %s; "
                 "spans will not use deterministic IDs.",
-                type(self._provider).__name__,
+                type(self._tracer).__name__,
             )
 
-        self._tracer: Tracer = self._provider.get_tracer(self._config.instrument_name)
-
         try:
             register_standalone_instrumentations(self._config, result)
         except Exception:
             logger.exception("Failed to register standalone instrumentations")
 
         # Per-invocation state.
         self._execution_arn = ""
+        self._execution_trace_id: int | None = None
         self._extracted_context: Context | None = None
         self._workflow_span: Span | None = None
         self._invocation_span: Span | None = None
@@ -205,8 +205,10 @@ def _resolve_parent(self, parent_id: str | None) -> Span | None:
     def on_invocation_start(self, info: InvocationStartInfo) -> None:
         logger.debug("Durable invocation started: %s", info)
         self._execution_arn = info.execution_arn or ""
+        self._execution_trace_id = _to_otel_trace_id(
+            self._execution_arn, info.execution_start_time
+        )
         self._extracted_context = self._context_extractor(info)
-        self._id_generator.set_trace_id(self._execution_arn, info.execution_start_time)
 
         self._start_workflow_span(info)
         # Create the Invocation span in both modes. In default-provider mode it
@@ -224,23 +226,23 @@ def _start_workflow_span(self, info: InvocationStartInfo) -> None:
         if not self._execution_arn:
             logger.warning("No execution ARN; skipping Workflow span creation")
             return
-        self._id_generator.set_next_span_id(
-            derive_workflow_span_id(self._execution_arn)
-        )
         start_time = _to_otel_timestamp(
             info.execution_start_time
         ) or _to_otel_timestamp(datetime.datetime.now(datetime.UTC))
         # Empty context => root span with no parent.
-        self._workflow_span = self._tracer.start_span(
-            name=self._workflow_span_name,
-            kind=SpanKind.INTERNAL,
-            attributes={"durable.execution.arn": self._execution_arn},
-            start_time=start_time,
-            context=Context(),
-        )
+        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(
+                name=self._workflow_span_name,
+                kind=SpanKind.INTERNAL,
+                attributes={"durable.execution.arn": self._execution_arn},
+                start_time=start_time,
+                context=Context(),
+            )
 
     def _start_invocation_span(self, info: InvocationStartInfo) -> None:
-        self._id_generator.set_next_span_id(None)
         attributes: dict[str, Any]
         if self._provider_source is ProviderSource.GLOBAL:
             # Default-provider mode: parent the Invocation span to the ambient
@@ -264,12 +266,15 @@ def _start_invocation_span(self, info: InvocationStartInfo) -> None:
                 "durable.execution.arn": self._execution_arn,
                 "durable.invocation.first": info.is_first_invocation,
             }
-        self._invocation_span = self._tracer.start_span(
-            name="Invocation",
-            kind=SpanKind.INTERNAL,
-            attributes=attributes,
-            context=parent_ctx,
-        )
+        with self._id_generator.use_ids(
+            trace_id=self._execution_trace_id, span_id=None
+        ):
+            self._invocation_span = self._tracer.start_span(
+                name="Invocation",
+                kind=SpanKind.INTERNAL,
+                attributes=attributes,
+                context=parent_ctx,
+            )
         self._set_span(_INVOCATION_KEY, self._invocation_span)
 
     def on_invocation_end(self, info: InvocationEndInfo) -> None:
@@ -330,6 +335,7 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None:
 
     def _reset_state(self) -> None:
         self._execution_arn = ""
+        self._execution_trace_id = None
         self._extracted_context = None
         self._workflow_span = None
         self._invocation_span = None
@@ -399,27 +405,26 @@ def _start_span(
         key = span_key if span_key is not None else operation_id
         with self._lock:
             links = self._build_invocation_links()
-            if deterministic:
-                # Operation spans always use the deterministic logical-operation
-                # span ID so a suspended-then-completed operation exports a
-                # single span (on completion) with a stable ID across invocations.
-                self._id_generator.set_next_span_id(
-                    operation_id_to_span_id(self._execution_arn, operation_id)
-                )
-            else:
-                self._id_generator.set_next_span_id(None)
+            span_id = (
+                operation_id_to_span_id(self._execution_arn, operation_id)
+                if deterministic
+                else None
+            )
 
             if parent is None:
                 parent_ctx = self._extracted_context or Context()
             else:
                 parent_ctx = trace.set_span_in_context(parent, self._extracted_context)
-            span = self._tracer.start_span(
-                name=name,
-                attributes=self._operation_attributes(info),
-                start_time=_to_otel_timestamp(start_time),
-                context=parent_ctx,
-                links=links,
-            )
+            with self._id_generator.use_ids(
+                trace_id=self._execution_trace_id, span_id=span_id
+            ):
+                span = self._tracer.start_span(
+                    name=name,
+                    attributes=self._operation_attributes(info),
+                    start_time=_to_otel_timestamp(start_time),
+                    context=parent_ctx,
+                    links=links,
+                )
             self._operation_spans[key] = span
         return span