fix(futures): ensure child spans are correctly parented in Threadpools (#13168)

Introduces a private API for activating a trace `Context` in a
contextmanager.

## Checklist
- [x] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist
- [x] Reviewer has checked that all the criteria below are met 
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)

Author: mabdinur

ddtrace/_trace/context.py

@@ -52,6 +52,7 @@ class Context(object):
         "_span_links",
         "_baggage",
         "_is_remote",
+        "_reactivate",
         "__weakref__",
     ]
 
@@ -75,6 +76,7 @@ def __init__(
         self.trace_id: Optional[int] = trace_id
         self.span_id: Optional[int] = span_id
         self._is_remote: bool = is_remote
+        self._reactivate: bool = False
 
         if dd_origin is not None and _DD_ORIGIN_INVALID_CHARS_REGEX.search(dd_origin) is None:
             self._meta[_ORIGIN_KEY] = dd_origin

ddtrace/_trace/provider.py

@@ -14,7 +14,8 @@
 log = get_logger(__name__)
 
 
-_DD_CONTEXTVAR: contextvars.ContextVar[Optional[Union[Context, Span]]] = contextvars.ContextVar(
+ActiveTrace = Union[Span, Context]
+_DD_CONTEXTVAR: contextvars.ContextVar[Optional[ActiveTrace]] = contextvars.ContextVar(
     "datadog_contextvar", default=None
 )
 
@@ -37,11 +38,11 @@ def _has_active_context(self) -> bool:
         pass
 
     @abc.abstractmethod
-    def activate(self, ctx: Optional[Union[Context, Span]]) -> None:
+    def activate(self, ctx: Optional[ActiveTrace]) -> None:
         self._hooks.emit(self.activate, ctx)
 
     @abc.abstractmethod
-    def active(self) -> Optional[Union[Context, Span]]:
+    def active(self) -> Optional[ActiveTrace]:
         pass
 
     def _on_activate(
@@ -70,37 +71,14 @@ def _deregister_on_activate(
         self._hooks.deregister(self.activate, func)
         return func
 
-    def __call__(self, *args: Any, **kwargs: Any) -> Optional[Union[Context, Span]]:
+    def __call__(self, *args: Any, **kwargs: Any) -> Optional[ActiveTrace]:
         """Method available for backward-compatibility. It proxies the call to
         ``self.active()`` and must not do anything more.
         """
         return self.active()
 
 
-class DatadogContextMixin(object):
-    """Mixin that provides active span updating suitable for synchronous
-    and asynchronous executions.
-    """
-
-    def activate(self, ctx: Optional[Union[Context, Span]]) -> None:
-        raise NotImplementedError
-
-    def _update_active(self, span: Span) -> Optional[Span]:
-        """Updates the active span in an executor.
-
-        The active span is updated to be the span's parent if the span has
-        finished until an unfinished span is found.
-        """
-        if span.finished:
-            new_active: Optional[Span] = span
-            while new_active and new_active.finished:
-                new_active = new_active._parent
-            self.activate(new_active)
-            return new_active
-        return span
-
-
-class DefaultContextProvider(BaseContextProvider, DatadogContextMixin):
+class DefaultContextProvider(BaseContextProvider):
     """Context provider that retrieves contexts from a context variable.
 
     It is suitable for synchronous programming and for asynchronous executors
@@ -115,14 +93,31 @@ def _has_active_context(self) -> bool:
         ctx = _DD_CONTEXTVAR.get()
         return ctx is not None
 
-    def activate(self, ctx: Optional[Union[Span, Context]]) -> None:
+    def activate(self, ctx: Optional[ActiveTrace]) -> None:
         """Makes the given context active in the current execution."""
         _DD_CONTEXTVAR.set(ctx)
         super(DefaultContextProvider, self).activate(ctx)
 
-    def active(self) -> Optional[Union[Context, Span]]:
+    def active(self) -> Optional[ActiveTrace]:
         """Returns the active span or context for the current execution."""
         item = _DD_CONTEXTVAR.get()
         if isinstance(item, Span):
             return self._update_active(item)
         return item
+
+    def _update_active(self, span: Span) -> Optional[ActiveTrace]:
+        """Updates the active trace in an executor.
+
+        When a span finishes, the active span becomes its parent.
+        If no parent exists and the context is reactivatable, that context is restored.
+        """
+        if span.finished:
+            new_active: Optional[Span] = span
+            while new_active and new_active.finished:
+                if new_active._parent is None and new_active._parent_context and new_active._parent_context._reactivate:
+                    self.activate(new_active._parent_context)
+                    return new_active._parent_context
+                new_active = new_active._parent
+            self.activate(new_active)
+            return new_active
+        return span

ddtrace/_trace/span.py

@@ -129,6 +129,7 @@ class Span(object):
         "duration_ns",
         # Internal attributes
         "_context",
+        "_parent_context",
         "_local_root_value",
         "_parent",
         "_ignored_exceptions",
@@ -211,6 +212,7 @@ def __init__(
         self.parent_id: Optional[int] = parent_id
         self._on_finish_callbacks = [] if on_finish is None else on_finish
 
+        self._parent_context: Optional[Context] = context
         self._context = context.copy(self.trace_id, self.span_id) if context else None
 
         self._links: List[Union[SpanLink, _SpanPointer]] = []

ddtrace/_trace/tracer.py

@@ -1,3 +1,4 @@
+from contextlib import contextmanager
 import functools
 from inspect import iscoroutinefunction
 from itertools import chain
@@ -290,6 +291,17 @@ def _sample_before_fork(self) -> None:
         if span is not None and span.context.sampling_priority is None:
             self.sample(span)
 
+    @contextmanager
+    def _activate_context(self, context: Context):
+        prev_active = self.context_provider.active()
+        context._reactivate = True
+        self.context_provider.activate(context)
+        try:
+            yield
+        finally:
+            context._reactivate = False
+            self.context_provider.activate(prev_active)
+
     @property
     def _sampler(self):
         return self._span_aggregator.sampling_processor.sampler

ddtrace/contrib/internal/futures/threading.py

@@ -34,8 +34,9 @@ def _wrap_execution(ctx: Tuple[Optional[Context], Optional[Context]], fn, args,
     provider sets the Active context in a thread local storage
     variable because it's outside the asynchronous loop.
     """
-    if ctx[0] is not None:
-        ddtrace.tracer.context_provider.activate(ctx[0])
     if ctx[1] is not None:
         core.dispatch("threading.execution", (ctx[1],))
+    if ctx[0] is not None:
+        with ddtrace.tracer._activate_context(ctx[0]):
+            return fn(*args, **kwargs)
     return fn(*args, **kwargs)

releasenotes/notes/fix-futures-parenting-3a53b94ee05dff93.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    futures: Resolves an edge where trace context was not propagated to set to spans generated by the ThreadPoolExecutor.

tests/contrib/futures/test_propagation.py

@@ -91,6 +91,35 @@ def fn(value, key=None):
         assert spans[0].trace_id == root.trace_id
         assert spans[0].parent_id == root.span_id
 
+    def test_propagation_with_sub_spans(self):
+        @self.tracer.wrap("executor.thread")
+        def fn1(value):
+            return value
+
+        @self.tracer.wrap("executor.thread")
+        def fn2(value):
+            return value * 100
+
+        def main(value, key=None):
+            return fn1(value) + fn2(value)
+
+        with self.override_global_tracer():
+            with self.tracer.trace("main.thread"):
+                with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
+                    future = executor.submit(main, 42)
+            value = future.result()
+            # assert the right result
+            self.assertEqual(value, 4242)
+
+        # the trace must be completed
+        spans = self.get_spans()
+        assert len(spans) == 3
+        assert spans[0].name == "main.thread"
+        for span in spans[1:]:
+            assert span.name == "executor.thread"
+            assert span.trace_id == spans[0].trace_id
+            assert span.parent_id == spans[0].span_id
+
     def test_propagation_with_kwargs(self):
         # instrumentation must work if only kwargs are provided
 

tests/tracer/test_tracer.py

@@ -1904,3 +1904,64 @@ def test_multiple_tracer_instances():
     log.error.assert_called_once_with(
         "Initializing multiple Tracer instances is not supported. Use ``ddtrace.trace.tracer`` instead.",
     )
+
+
+def test_span_creation_with_context(tracer):
+    """Test that a span created with a Context parent correctly sets its parent context.
+
+    This test verifies that when a span is created with a Context object as its parent,
+    the span's _parent_context attribute is properly set to that Context.
+    """
+
+    context = Context(trace_id=123, span_id=321)
+    with tracer.start_span("s1", child_of=context) as span:
+        assert span._parent_context == context
+
+
+def test_activate_context_propagates_to_child_spans(tracer):
+    """Test that activating a context properly propagates to child spans.
+
+    This test verifies that when a context is activated using _activate_context:
+    1. Child spans created within the context inherit the trace_id and parent_id
+    2. Multiple child spans within the same context maintain consistent parentage
+    3. Spans created outside the context have different trace_id and parent_id
+    """
+    with tracer._activate_context(Context(trace_id=1, span_id=1)):
+        with tracer.trace("child1") as c1:
+            assert c1.trace_id == 1
+            assert c1.parent_id == 1
+
+        with tracer.trace("child1") as c2:
+            assert c2.trace_id == 1
+            assert c2.parent_id == 1
+
+    with tracer.trace("root") as root:
+        assert root.trace_id != 1
+        assert root.parent_id != 1
+
+
+def test_activate_context_nesting_and_restoration(tracer):
+    """Test that context activation properly handles nesting and restoration.
+
+    This test verifies that:
+    1. A context can be activated and its values are accessible
+    2. A nested context can be activated and its values override the outer context
+    3. When the nested context exits, the outer context is properly restored
+    4. When all contexts exit, the active context is None
+    """
+
+    with tracer._activate_context(Context(trace_id=1, span_id=1)):
+        active = tracer.context_provider.active()
+        assert active.trace_id == 1
+        assert active.span_id == 1
+
+        with tracer._activate_context(Context(trace_id=2, span_id=2)):
+            active = tracer.context_provider.active()
+            assert active.trace_id == 2
+            assert active.span_id == 2
+
+        active = tracer.context_provider.active()
+        assert active.trace_id == 1
+        assert active.span_id == 1
+
+    assert tracer.context_provider.active() is None