fix(tracing): correct otel thread context for active contexts, sampling and switches

ddtrace/_trace/context.py

@@ -8,6 +8,7 @@
 from ddtrace.constants import _ORIGIN_KEY
 from ddtrace.constants import _SAMPLING_PRIORITY_KEY
 from ddtrace.constants import _USER_ID_KEY
+from ddtrace.internal import core
 from ddtrace.internal.compat import NumericType
 from ddtrace.internal.constants import MAX_UINT_64BITS as _MAX_UINT_64BITS
 from ddtrace.internal.constants import W3C_TRACEPARENT_KEY
@@ -31,6 +32,9 @@
 
 _DD_ORIGIN_INVALID_CHARS_REGEX = re.compile(r"[^\x20-\x7E]+")
 
+# For listeners that mirror the W3C trace-flags byte outside the tracer.
+SAMPLING_DECISION_EVENT = "ddtrace.trace.sampling_decision"
+
 log = get_logger(__name__)
 
 
@@ -136,8 +140,12 @@ def sampling_priority(self, value: Optional[NumericType]) -> None:
             if value is None:
                 if _SAMPLING_PRIORITY_KEY in self._metrics:
                     del self._metrics[_SAMPLING_PRIORITY_KEY]
-                return
-            self._metrics[_SAMPLING_PRIORITY_KEY] = value
+            else:
+                self._metrics[_SAMPLING_PRIORITY_KEY] = value
+        # Trace sampling usually decides at trace-chunk finish, long after anything
+        # mirroring the trace-flags byte published it. Dispatched outside the lock
+        # because listeners run arbitrary code; runs once per trace, so not hot.
+        core.dispatch(SAMPLING_DECISION_EVENT)
 
     @property
     def _traceparent(self) -> str:

ddtrace/internal/ci_visibility/context.py

@@ -44,3 +44,9 @@ def active(self) -> ContextTypeValue:
         if isinstance(item, Span):
             return self._update_active(item)
         return item
+
+    def _peek_active(self) -> ContextTypeValue:
+        # Required because this provider has its own storage: the inherited
+        # implementation reads the default contextvar. Not read-only, unlike the
+        # inherited one -- nothing needs that of this provider yet.
+        return self.active()

ddtrace/internal/native/_native.pyi

@@ -191,6 +191,18 @@ if sys.platform == "linux":
         :param trace_flags: W3C Trace Context trace-flags byte (bit 0 = sampled).
         """
         ...
+    def update_otel_thread_context_ids(
+        trace_id: int, span_id: int, trace_flags: int, local_root_span_id: int
+    ) -> None:
+        """
+        Update the OTel thread context from raw ids, for an active ``Context``.
+
+        :param trace_id: 128-bit trace id.
+        :param span_id: The span this execution is attributable to.
+        :param trace_flags: W3C Trace Context trace-flags byte (bit 0 = sampled).
+        :param local_root_span_id: Local root span id.
+        """
+        ...
     def detach_otel_thread_context() -> None:
         """Detach the OTel thread context from the current thread."""
         ...
@@ -1627,6 +1639,13 @@ class BaseContextProvider(abc.ABC):
     def activate(self, ctx: Optional[ActiveTrace]) -> None: ...
     @abc.abstractmethod
     def active(self) -> Optional[ActiveTrace]: ...
+    def _peek_active(self) -> Optional[ActiveTrace]:
+        """``active()`` without its repair: no contextvar write, no activate dispatch.
+
+        For observers that must not perturb the active trace, such as anything
+        running from a CPython context-switch watcher.
+        """
+        ...
     def __call__(self, *args: Any, **kwargs: Any) -> Optional[ActiveTrace]: ...
 
 class DefaultContextProvider(BaseContextProvider):

ddtrace/internal/opentelemetry/thread_context.py

@@ -4,10 +4,12 @@
 from typing import Protocol
 from typing import Union
 
+from ddtrace._trace.context import SAMPLING_DECISION_EVENT
 from ddtrace._trace.context import Context
 from ddtrace._trace.provider import BaseContextProvider
 from ddtrace._trace.span import Span
 from ddtrace.internal import core
+from ddtrace.internal.compat import NumericType
 from ddtrace.internal.settings._config import config
 
 
@@ -16,36 +18,53 @@ class TracerProtocol(Protocol):
     def context_provider(self) -> BaseContextProvider: ...
 
 
-_ContextActivationListener = Callable[[BaseContextProvider, Optional[Union[Context, Span]]], None]
-_ContextSwitchListener = Callable[[], None]
-_ThreadContextListeners = tuple[_ContextActivationListener, _ContextSwitchListener]
+_ActiveTrace = Optional[Union[Context, Span]]
+_ContextActivationListener = Callable[[BaseContextProvider, _ActiveTrace], None]
+_ResyncListener = Callable[[], None]
+_ThreadContextListeners = tuple[_ContextActivationListener, _ResyncListener]
+
+
+def _w3c_trace_flags(sampling_priority: Optional[NumericType]) -> int:
+    return 1 if sampling_priority is not None and sampling_priority > 0 else 0
 
 
 if sys.platform == "linux":
     from ddtrace.internal.native._native import detach_otel_thread_context
     from ddtrace.internal.native._native import update_otel_thread_context
+    from ddtrace.internal.native._native import update_otel_thread_context_ids
 
     def register_otel_thread_context_listener(tracer: TracerProtocol) -> Optional[_ThreadContextListeners]:
         if not config._otel_thread_context_enabled:
             return None
 
-        def _sync_otel_thread_context(ctx: Optional[Union[Context, Span]]) -> None:
+        def _sync_otel_thread_context(ctx: _ActiveTrace) -> None:
             if type(ctx) is Span:
-                sampling_priority = ctx._local_root.context.sampling_priority
-                trace_flags = 1 if sampling_priority is not None and sampling_priority > 0 else 0
-                update_otel_thread_context(ctx, ctx._local_root_value, trace_flags)
+                update_otel_thread_context(
+                    ctx, ctx._local_root_value, _w3c_trace_flags(ctx._local_root.context.sampling_priority)
+                )
+            elif type(ctx) is Context and ctx.trace_id is not None and ctx.span_id is not None:
+                # A Context is a span this execution runs inside but does not own -- a
+                # remote parent, or the submitter of work handed to this thread. The
+                # execution is still attributable to it. Its local root is not knowable
+                # from here, so the span stands in for it, as it does for a root span.
+                update_otel_thread_context_ids(
+                    ctx.trace_id, ctx.span_id, _w3c_trace_flags(ctx.sampling_priority), ctx.span_id
+                )
             else:
                 detach_otel_thread_context()
 
         def _sync_active_otel_thread_context() -> None:
-            _sync_otel_thread_context(tracer.context_provider.active())
+            _sync_otel_thread_context(tracer.context_provider._peek_active())
 
-        def _on_context_provider_activate(provider: BaseContextProvider, ctx: Optional[Union[Context, Span]]) -> None:
+        def _on_context_provider_activate(provider: BaseContextProvider, ctx: _ActiveTrace) -> None:
             if provider is tracer.context_provider:
                 _sync_otel_thread_context(ctx)
 
         core.on("ddtrace.context_provider.activate", _on_context_provider_activate)
         core.on("python.context.switch", _sync_active_otel_thread_context)
+        # Sampling usually decides at trace-chunk finish, after the record was published
+        # with trace_flags=0.
+        core.on(SAMPLING_DECISION_EVENT, _sync_active_otel_thread_context)
 
         if sys.implementation.name == "cpython" and sys.version_info >= (3, 14):
             from ddtrace.internal.native._native import register_context_watcher

ddtrace/llmobs/_context.py

@@ -58,3 +58,9 @@ def active(self) -> ContextTypeValue:
         if isinstance(item, Span):
             return self._update_active(item)
         return item
+
+    def _peek_active(self) -> ContextTypeValue:
+        # Required because this provider has its own storage: the inherited
+        # implementation reads the default contextvar. Not read-only, unlike the
+        # inherited one -- nothing needs that of this provider yet.
+        return self.active()

releasenotes/notes/fix-otel-thread-context-correctness-0fb65fa39ac88d25.yaml

@@ -0,0 +1,16 @@
+---
+fixes:
+  - |
+    tracing: Fixes an issue where the OpenTelemetry thread context reported no active
+    trace whenever a trace was propagated without a local span, so compatible profiling
+    and security tools could not correlate samples taken in executor worker threads, or
+    on a distributed request before its first span started.
+  - |
+    tracing: Fixes an issue where the OpenTelemetry thread context kept reporting a trace
+    as unsampled for the whole life of the trace, because trace sampling usually decides
+    after the record is published. The record is now updated when the sampling decision
+    is made.
+  - |
+    tracing: Fixes an issue where publishing the OpenTelemetry thread context on an
+    asynchronous context switch could change which span is active in that context, and
+    publish the record twice for a single switch.

src/native/context_provider.rs

@@ -77,6 +77,57 @@ fn dispatch_activate(
     event_hub::dispatch(py, ACTIVATE_EVENT, Some(args.into_any().unbind()), false)
 }
 
+enum Resolved<'py> {
+    Unchanged,
+    /// May hold Python `None`, when the whole ancestor chain is finished.
+    Ancestor(Bound<'py, PyAny>),
+    ReactivatableContext(Bound<'py, PyAny>),
+}
+
+/// Where the active trace *should* be, walking past finished ancestors of `span`.
+///
+/// Shared by `_update_active`, which applies the result, and `_peek_active`, which
+/// only reports it, so the two cannot drift apart. Borrows `span` so that
+/// `Resolved::Unchanged` costs the caller nothing -- `_update_active` runs on every
+/// span finish, where an extra incref shows up.
+#[inline]
+fn resolve_active<'py>(py: Python<'py>, span: &Bound<'py, PyAny>) -> PyResult<Resolved<'py>> {
+    let mut current = span.clone();
+    loop {
+        // PERF: read `duration`, `_parent`, and `_parent_context` straight off the
+        // native SpanData fields in one borrow -- avoids three Python attribute
+        // lookups per ancestor hop.
+        let (parent, parent_context) = {
+            let Ok(sd) = current.cast::<SpanData>() else {
+                break; // not a Span (e.g. None) -- stop walking parents
+            };
+            let sd = sd.borrow();
+            if sd.duration.is_none() {
+                break; // unfinished span -- stop walking parents
+            }
+            (
+                sd._parent.as_ref().map(|p| p.bind(py).clone()),
+                sd._parent_context.as_ref().map(|c| c.bind(py).clone()),
+            )
+        };
+        if parent.is_none() {
+            if let Some(parent_context) = parent_context {
+                // `_reactivate` lives on the pure-Python Context -- still a getattr.
+                if parent_context.getattr("_reactivate")?.is_truthy()? {
+                    return Ok(Resolved::ReactivatableContext(parent_context));
+                }
+            }
+        }
+        // Advance to the parent; `None` ends the walk on the next iteration.
+        current = parent.unwrap_or_else(|| py.None().into_bound(py));
+    }
+    if current.is(span) {
+        Ok(Resolved::Unchanged)
+    } else {
+        Ok(Resolved::Ancestor(current))
+    }
+}
+
 /// `Some(v)` unless `v` is Python `None`, matching the `Optional[...]` return
 /// convention used throughout this module. Consumes `v` to avoid an incref.
 #[inline]
@@ -169,6 +220,12 @@ impl BaseContextProvider {
         Err(PyNotImplementedError::new_err(()))
     }
 
+    /// See `DefaultContextProvider::_peek_active`. A provider with its own storage
+    /// keeps today's behavior until it opts into a read-only path.
+    fn _peek_active(slf: &Bound<'_, Self>) -> PyResult<Option<Py<PyAny>>> {
+        slf.call_method0("active").map(none_or_unbind)
+    }
+
     /// Method available for backward-compatibility. It proxies the call to
     /// ``self.active()`` and must not do anything more.
     #[pyo3(signature = (*_args, **_kwargs))]
@@ -251,41 +308,43 @@ impl DefaultContextProvider {
         py: Python<'py>,
         span: Bound<'py, PyAny>,
     ) -> PyResult<Option<Py<PyAny>>> {
-        let original = span.clone();
-        let mut current = span;
-        loop {
-            // PERF: read `duration`, `_parent`, and `_parent_context` straight
-            // off the native SpanData fields in one borrow -- avoids three
-            // Python attribute lookups per ancestor hop.
-            let (parent, parent_context) = {
-                let Ok(sd) = current.cast::<SpanData>() else {
-                    break; // not a Span (e.g. None) -- stop walking parents
-                };
-                let sd = sd.borrow();
-                if sd.duration.is_none() {
-                    break; // unfinished span -- stop walking parents
-                }
-                (
-                    sd._parent.as_ref().map(|p| p.bind(py).clone()),
-                    sd._parent_context.as_ref().map(|c| c.bind(py).clone()),
-                )
-            };
-            if parent.is_none() {
-                if let Some(parent_context) = parent_context {
-                    // `_reactivate` lives on the pure-Python Context -- still a getattr.
-                    if parent_context.getattr("_reactivate")?.is_truthy()? {
-                        call_activate(slf, py, Some(parent_context.clone()))?;
-                        return Ok(Some(parent_context.unbind()));
-                    }
-                }
+        match resolve_active(py, &span)? {
+            Resolved::Unchanged => Ok(none_or_unbind(span)),
+            Resolved::ReactivatableContext(parent_context) => {
+                call_activate(slf, py, Some(parent_context.clone()))?;
+                Ok(Some(parent_context.unbind()))
             }
-            // Advance to the parent; `None` ends the walk on the next iteration.
-            current = parent.unwrap_or_else(|| py.None().into_bound(py));
+            Resolved::Ancestor(current) => {
+                call_activate(slf, py, none_or_clone(&current))?;
+                Ok(none_or_unbind(current))
+            }
+        }
+    }
+
+    /// `active` without its repair: no `activate`, so no contextvar write and no
+    /// `ddtrace.context_provider.activate` dispatch.
+    ///
+    /// For observers that must not perturb the active trace. A CPython
+    /// context-switch watcher is one: `activate` there would write to whichever
+    /// context the switch just made current, and re-enter the event hub mid-switch.
+    ///
+    /// Reads this contextvar and resolves the way `_update_active` does, so **a subclass
+    /// that overrides `active` or `_update_active` must override this too** -- otherwise
+    /// it reports from storage it does not use. Subclasses that only add behaviour
+    /// elsewhere inherit it safely.
+    fn _peek_active<'py>(_slf: &Bound<'py, Self>, py: Python<'py>) -> PyResult<Option<Py<PyAny>>> {
+        let item = contextvar_get(py, contextvar(py)?)?;
+        if item.is_none() {
+            return Ok(None);
+        }
+        if item.cast::<SpanData>().is_err() {
+            return Ok(Some(item.unbind()));
         }
-        if !current.is(&original) {
-            call_activate(slf, py, none_or_clone(&current))?;
+        match resolve_active(py, &item)? {
+            Resolved::Unchanged => Ok(none_or_unbind(item)),
+            Resolved::ReactivatableContext(parent_context) => Ok(Some(parent_context.unbind())),
+            Resolved::Ancestor(current) => Ok(none_or_unbind(current)),
         }
-        Ok(none_or_unbind(current))
     }
 }
 

src/native/lib.rs

@@ -71,6 +71,9 @@ fn _native(m: &Bound<'_, PyModule>) -> PyResult<()> {
         m.add_wrapped(wrap_pyfunction!(
             otel_thread_ctx::update_otel_thread_context
         ))?;
+        m.add_wrapped(wrap_pyfunction!(
+            otel_thread_ctx::update_otel_thread_context_ids
+        ))?;
         m.add_wrapped(wrap_pyfunction!(
             otel_thread_ctx::detach_otel_thread_context
         ))?;

src/native/otel_thread_ctx.rs

@@ -23,6 +23,23 @@ pub fn update_otel_thread_context(
     );
 }
 
+/// Publish from raw ids, for an active `Context` -- there is no `SpanData` to read.
+#[pyfunction]
+pub fn update_otel_thread_context_ids(
+    trace_id: u128,
+    span_id: u64,
+    trace_flags: u8,
+    local_root_span_id: u64,
+) {
+    ThreadContext::update(
+        trace_id.to_be_bytes(),
+        span_id.to_be_bytes(),
+        trace_flags,
+        local_root_span_id.to_be_bytes(),
+        &[],
+    );
+}
+
 #[pyfunction]
 pub fn detach_otel_thread_context() {
     ThreadContext::detach();