fix(profiling): correctly detect on-cpu tasks

Author: KowalskiThomas

ddtrace/internal/datadog/profiling/stack_v2/echion/echion/long.h

@@ -5,6 +5,7 @@
 
 #include <Python.h>
 #if PY_VERSION_HEX >= 0x030c0000
+#define Py_BUILD_CORE
 #include <internal/pycore_long.h>
 // Note: Even if use the right PYLONG_BITS_IN_DIGIT that is specified in the
 // Python we use to build echion, it can be different from the Python that is

ddtrace/internal/datadog/profiling/stack_v2/echion/echion/tasks.h

@@ -4,6 +4,8 @@
 
 #pragma once
 
+#include <optional>
+
 #define PY_SSIZE_T_CLEAN
 #include <Python.h>
 #include <weakrefobject.h>
@@ -176,6 +178,7 @@ class TaskInfo
 
     // Information to reconstruct the async stack as best as we can
     TaskInfo::Ptr waiter = nullptr;
+    std::optional<bool> is_on_cpu_ = std::nullopt;
 
     [[nodiscard]] static Result<TaskInfo::Ptr> create(TaskObj*);
     TaskInfo(PyObject* origin, PyObject* loop, GenInfo::Ptr coro, StringTable::Key name, TaskInfo::Ptr waiter)
@@ -189,6 +192,25 @@ class TaskInfo
 
     [[nodiscard]] static Result<TaskInfo::Ptr> current(PyObject*);
     inline size_t unwind(FrameStack&);
+
+    // Check if any coroutine in the chain is currently running (on CPU)
+    inline bool is_on_cpu()
+    {
+        if (is_on_cpu_.has_value()) {
+            return *is_on_cpu_;
+        }
+
+        for (auto coroutine = this->coro.get(); coroutine != nullptr; coroutine = coroutine->await.get())
+        {
+            if (coroutine->is_running) {
+                is_on_cpu_ = true;
+                return true;
+            }
+        }
+        
+        is_on_cpu_.emplace(false);
+        return false;
+    }
 };
 
 inline std::unordered_map<PyObject*, PyObject*> task_link_map;

ddtrace/internal/datadog/profiling/stack_v2/echion/echion/threads.h

@@ -252,8 +252,19 @@ ThreadInfo::unwind_tasks()
         }
     }
 
+    // Only one Task can be on CPU at a time.
+    // Since determining if a task is on CPU is somewhat costly, we
+    // stop checking if Tasks are on CPU after seeing the first one.
+    bool on_cpu_task_seen = false;
     for (auto& leaf_task : leaf_tasks) {
-        bool on_cpu = leaf_task.get().coro->is_running;
+        bool on_cpu = false;
+        if (!on_cpu_task_seen) { 
+            on_cpu = leaf_task.get().is_on_cpu();
+            if (on_cpu) {
+                on_cpu_task_seen = true;
+            }
+        }
+
         auto stack_info = std::make_unique<StackInfo>(leaf_task.get().name, on_cpu);
         auto& stack = stack_info->stack;
         for (auto current_task = leaf_task;;) {