feat(profiling): link task samples to origin task (#19082)

[PROF-15417](https://datadoghq.atlassian.net/jira/software/c/projects/PROF/boards/11?selectedIssue=PROF-15417)
local
## Description

When an asyncio task offloads blocking work to a `ThreadPoolExecutor` (`await loop.run_in_executor(None, fn)`), that work runs on a pool worker thread that has no asyncio task of its own. Today those samples can't be tied back to the task that submitted and is awaiting them.

This adds two pprof labels on the executor worker thread's samples:

- origin task id -- `id()` of the submitting asyncio task
- origin task name -- that task's name

origin task id is deliberately the same value the profiler already emits as task id for the task on its event-loop thread, so the offloaded work correlates back to the awaiting task with no new join key.

The `futures` integration already wraps `ThreadPoolExecutor.submit` and runs the work through an intermediate `_wrap_execution` on the worker thread. We capture the current asyncio task at submit time, and at execution time register a `thread_id` -> `{task id, name}` entry in a native per-thread map for the duration of the job, clearing it in a `finally`. The stack renderer stamps the labels by looking that map up by sampled thread id. The worker thread isnt predicted ahead of time, rather, the link is made on whichever worker the pool assigns and torn down when the job completes, so pool reuse is handled correctly.

## Testing

<!-- Describe your testing strategy or note what tests are included -->

## Risks

<!-- Note any risks associated with this change, or "None" if no risks -->

## Additional Notes

<!-- Any other information that would be helpful for reviewers -->


[PROF-15417]: https://datadoghq.atlassian.net/browse/PROF-15417?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ

Co-authored-by: gyuheon.oh <gyuheon.oh@datadoghq.com>

Author: gyuheon0h

ddtrace/contrib/internal/futures/threading.py

@@ -1,3 +1,5 @@
+import sys
+from types import ModuleType
 from typing import Optional
 
 import ddtrace
@@ -6,6 +8,57 @@
 from ddtrace.trace import Context
 
 
+_PROFILING_STACK_MODULE = "ddtrace.internal.datadog.profiling.stack"
+
+
+def _profiling_stack() -> Optional[ModuleType]:
+    """
+    Return the profiling stack module iff OriginTaskLinks is enabled.
+
+    We deliberately don't import it. The futures patch stays on for APM even when
+    profiling is off; settings.profiling may still import the stack module to
+    probe availability, so is_available alone is not enough. OriginTaskLinks is
+    enabled after stack sampler start (Profiler.start() / DD_PROFILING_ENABLED)
+    and disabled on stop.
+    """
+    stack = sys.modules.get(_PROFILING_STACK_MODULE)
+    if (
+        stack is not None
+        and getattr(stack, "is_available", False)
+        and getattr(stack, "is_origin_task_linking_enabled", lambda: False)()
+    ):
+        return stack
+    return None
+
+
+def _current_origin_task() -> tuple[Optional[int], Optional[str]]:
+    """
+    Identity of the asyncio task submitting this work, as (task_id, task_name).
+
+    `task_id` is `id(task)`, which matches the profiler's "task id" label (the
+    task object's address), so executor worker-thread samples can be correlated
+    back to the awaiting task's samples on the event-loop thread
+    """
+    if _profiling_stack() is None:
+        return None, None
+    asyncio = sys.modules.get("asyncio")
+
+    if asyncio is None:
+        # Not an asyncio application; there is no task to attribute this work to.
+        return None, None
+
+    try:
+        task = asyncio.current_task()
+    except RuntimeError:
+        # No running event loop on the submitting thread.
+        return None, None
+
+    if task is None:
+        return None, None
+
+    return id(task), task.get_name()
+
+
 def _wrap_submit(func, args, kwargs):
     """
     Wrap `Executor` method used to submit a work executed in another
@@ -18,6 +71,10 @@ def _wrap_submit(func, args, kwargs):
         "threading.submit", ()
     ).llmobs_ctx.value
 
+    # Capture the submitting asyncio task so the profiler can link the
+    # offloaded work back to it
+    origin_task_id, origin_task_name = _current_origin_task()
+
     # Shallow copy for cross-thread handoff. current_trace_context() returns the
     # active Context (or the current span's .context), not a detached copy.
     if current_ctx is not None and current_ctx.trace_id is not None and current_ctx.span_id is not None:
@@ -37,21 +94,34 @@ def _wrap_submit(func, args, kwargs):
         fn_args = args[1:]
     else:
         fn, fn_args = args[1], args[2:]
-    return func(self, _wrap_execution, (current_ctx, llmobs_ctx), fn, fn_args, kwargs)
+    return func(self, _wrap_execution, (current_ctx, llmobs_ctx, origin_task_id, origin_task_name), fn, fn_args, kwargs)
 
 
-def _wrap_execution(ctx: tuple[Optional[Context], Optional[Context]], fn, args, kwargs):
+def _wrap_execution(
+    ctx: tuple[Optional[Context], Optional[Context], Optional[int], Optional[str]],
+    fn,
+    args,
+    kwargs,
+):
     """
     Intermediate target function that is executed in a new thread;
     it receives the original function with arguments and keyword
     arguments, including our tracing `Context`. The current context
     provider sets the Active context in a thread local storage
     variable because it's outside the asynchronous loop.
     """
-    current_ctx, llmobs_ctx = ctx
+    current_ctx, llmobs_ctx, origin_task_id, origin_task_name = ctx
     if llmobs_ctx is not None:
         core.dispatch("threading.execution", (llmobs_ctx,))
-    if current_ctx is not None:
-        with ddtrace.tracer._activate_context(current_ctx):
-            return fn(*args, **kwargs)
-    return fn(*args, **kwargs)
+
+    stack = _profiling_stack() if origin_task_id is not None else None
+    if stack is not None:
+        stack.link_origin_task(origin_task_id, origin_task_name)
+    try:
+        if current_ctx is not None:
+            with ddtrace.tracer._activate_context(current_ctx):
+                return fn(*args, **kwargs)
+        return fn(*args, **kwargs)
+    finally:
+        if stack is not None:
+            stack.unlink_origin_task()

ddtrace/internal/datadog/profiling/dd_wrapper/include/libdatadog_helpers.hpp

@@ -54,6 +54,8 @@ intern_string(std::string_view s);
     X(thread_name, "thread name")                                                                                      \
     X(task_id, "task id")                                                                                              \
     X(task_name, "task name")                                                                                          \
+    X(origin_task_id, "origin task id")                                                                                \
+    X(origin_task_name, "origin task name")                                                                            \
     X(span_id, "span id")                                                                                              \
     X(local_root_span_id, "local root span id")                                                                        \
     X(trace_type, "trace type")                                                                                        \

ddtrace/internal/datadog/profiling/dd_wrapper/include/sample.hpp

@@ -129,6 +129,8 @@ class Sample
     bool push_threadinfo(int64_t thread_id, int64_t thread_native_id, std::string_view thread_name);
     bool push_task_id(uint64_t task_id);
     bool push_task_name(std::string_view task_name);
+    bool push_origin_task_id(uint64_t origin_task_id);
+    bool push_origin_task_name(std::string_view origin_task_name);
     bool push_span_id(uint64_t span_id);
     bool push_local_root_span_id(uint64_t local_root_span_id);
     bool push_trace_type(std::string_view trace_type);

ddtrace/internal/datadog/profiling/dd_wrapper/src/sample.cpp

@@ -641,6 +641,35 @@ Datadog::Sample::push_task_name(std::string_view task_name)
     return true;
 }
 
+bool
+Datadog::Sample::push_origin_task_id(uint64_t origin_task_id)
+{
+    static bool already_warned = false; // cppcheck-suppress threadsafety-threadsafety
+    const int64_t recoded_id =
+      reinterpret_cast<int64_t&>(origin_task_id); // NOLINT (cppcoreguidelines-pro-type-reinterpret-cast)
+    if (!push_label(ExportLabelKey::origin_task_id, recoded_id)) {
+        if (!already_warned) {
+            already_warned = true;
+            std::cerr << "bad push" << std::endl;
+        }
+        return false;
+    }
+    return true;
+}
+bool
+Datadog::Sample::push_origin_task_name(std::string_view origin_task_name)
+{
+    static bool already_warned = false; // cppcheck-suppress threadsafety-threadsafety
+    if (!push_label(ExportLabelKey::origin_task_name, origin_task_name)) {
+        if (!already_warned) {
+            already_warned = true;
+            std::cerr << "bad push" << std::endl;
+        }
+        return false;
+    }
+    return true;
+}
+
 bool
 Datadog::Sample::push_span_id(uint64_t span_id)
 {

ddtrace/internal/datadog/profiling/stack/CMakeLists.txt

@@ -69,6 +69,7 @@ add_library(
     src/echion/tasks.cc
     src/echion/threads.cc
     src/echion/vm.cc
+    src/origin_task_links.cpp
     src/sampler.cpp
     src/stack.cpp
     src/stack_renderer.cpp

ddtrace/internal/datadog/profiling/stack/__init__.py

@@ -41,5 +41,17 @@ def link_span(span: typing.Optional[typing.Union[context.Context, ddspan.Span]])
             _propagated_root.span_type = span_type
             _stack.link_span(span.span_id, local_root_span_id, span_type)
 
+    def link_origin_task(task_id: int, task_name: str) -> None:
+        """
+        Record, for the current thread, the asyncio task that submitted the work now running on it.
+        """
+        _stack.link_origin_task(task_id, task_name)
+
+    def unlink_origin_task() -> None:
+        """
+        Clear the originating asyncio task for the current thread.
+        """
+        _stack.unlink_origin_task()
+
 except Exception as e:
     failure_msg = str(e)

ddtrace/internal/datadog/profiling/stack/__init__.pyi

@@ -10,8 +10,13 @@ from ddtrace._trace import span as ddspan
 # Core stack v2 functions
 def start(min_interval: float = ...) -> bool: ...
 def stop() -> None: ...
+def is_origin_task_linking_enabled() -> bool: ...
 def _native_call_registry_size() -> int: ...
 
+# executor worker thread <-> originating asyncio task association
+def link_origin_task(task_id: int, task_name: str) -> None: ...
+def unlink_origin_task() -> None: ...
+
 # Sampling configuration
 def set_adaptive_sampling(do_adaptive_sampling: bool = False) -> None: ...
 def set_target_overhead(target_overhead: float) -> None: ...

ddtrace/internal/datadog/profiling/stack/_stack.pyi

@@ -10,6 +10,7 @@ from typing import Union
 # Core stack v2 functions
 def start(min_interval: float = ...) -> bool: ...
 def stop() -> None: ...
+def is_origin_task_linking_enabled() -> bool: ...
 
 # Sampling configuration
 def set_adaptive_sampling(do_adaptive_sampling: bool = False) -> None: ...
@@ -30,6 +31,13 @@ def link_span(
     span_type: Optional[str] = None,
 ) -> None: ...
 
+# executor worker thread <-> originating asyncio task association
+def link_origin_task(
+    task_id: int,
+    task_name: Optional[str] = None,
+) -> None: ...
+def unlink_origin_task() -> None: ...
+
 # Thread management
 def register_thread(python_thread_id: int, native_id: int, name: str) -> None: ...
 def unregister_thread(python_thread_id: int) -> None: ...

ddtrace/internal/datadog/profiling/stack/include/origin_task_links.hpp

@@ -0,0 +1,68 @@
+#pragma once
+
+#include <atomic>
+#include <mutex>
+#include <optional>
+#include <stdint.h>
+#include <string>
+#include <unordered_map>
+
+namespace Datadog {
+
+// Identity of the asyncio task that submitted work to a ThreadPoolExecutor.
+// task_id is the task object's address (id(task) in Python)
+struct OriginTask
+{
+    uint64_t task_id;
+    std::string task_name;
+
+    OriginTask(uint64_t _task_id, std::string _task_name)
+      : task_id(_task_id)
+      , task_name(std::move(_task_name))
+    {
+    }
+
+    // for testing
+    bool operator==(const OriginTask& other) const { return task_id == other.task_id && task_name == other.task_name; }
+};
+
+// Per-(worker)-thread registry of the asyncio task that offloaded the work
+// currently running on that thread. Populated by the futures integration for
+// the duration of an executor job and accessed by stack renderer.
+//
+// Lifecycle matches the stack sampler: enable() after Sampler::start() succeeds,
+// disable_and_reset() before Sampler::stop(). While disabled, link/unlink/get
+// are no-ops so the default-on futures patch is cheap when profiling is off.
+class OriginTaskLinks
+{
+  public:
+    static OriginTaskLinks& get_instance()
+    {
+        static OriginTaskLinks instance;
+        return instance;
+    }
+
+    OriginTaskLinks(OriginTaskLinks const&) = delete;
+    OriginTaskLinks& operator=(OriginTaskLinks const&) = delete;
+
+    void enable();
+    void disable_and_reset();
+    bool is_enabled() const { return enabled_.load(std::memory_order_acquire); }
+
+    void link_origin_task(uint64_t thread_id, uint64_t task_id, std::string task_name);
+    const std::optional<OriginTask> get_origin_task(uint64_t thread_id);
+    void unlink_origin_task(uint64_t thread_id);
+
+    static void postfork_child();
+
+  private:
+    std::atomic<bool> enabled_{ false };
+    std::mutex mtx;
+    std::unordered_map<uint64_t, OriginTask> thread_id_to_origin_task;
+
+    // Private Constructor/Destructor
+    OriginTaskLinks() = default;
+    ~OriginTaskLinks() = default;
+};
+
+}

ddtrace/internal/datadog/profiling/stack/include/sampler.hpp

@@ -160,8 +160,9 @@ class Sampler
     // Restart the sampling thread in the parent after fork
     void postfork_parent();
 
-    // Restart the sampler after fork if it was running
-    void restart_after_fork();
+    // Restart the sampler after fork if it was running.
+    // Returns true if start() was invoked and succeeded.
+    bool restart_after_fork();
 };
 
 } // namespace Datadog