diff --git a/ddtrace/internal/datadog/profiling/stack_v2/echion/echion/long.h b/ddtrace/internal/datadog/profiling/stack_v2/echion/echion/long.h index 0f66203dfd1..512fd81fb3e 100644 --- a/ddtrace/internal/datadog/profiling/stack_v2/echion/echion/long.h +++ b/ddtrace/internal/datadog/profiling/stack_v2/echion/echion/long.h @@ -5,6 +5,7 @@ #include #if PY_VERSION_HEX >= 0x030c0000 +#define Py_BUILD_CORE #include // 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 diff --git a/ddtrace/internal/datadog/profiling/stack_v2/echion/echion/tasks.h b/ddtrace/internal/datadog/profiling/stack_v2/echion/echion/tasks.h index 87d86c470e7..848af75aef2 100644 --- a/ddtrace/internal/datadog/profiling/stack_v2/echion/echion/tasks.h +++ b/ddtrace/internal/datadog/profiling/stack_v2/echion/echion/tasks.h @@ -4,6 +4,8 @@ #pragma once +#include + #define PY_SSIZE_T_CLEAN #include #include @@ -111,8 +113,8 @@ GenInfo::create(PyObject* gen_addr) // Type-pun the PyGenObject to a PyAsyncGenASend. *gen_addr was actually never a PyGenObject to begin with, // but we do not care as the only thing we will use from it is the ags_gen field. PyAsyncGenASend* asend = reinterpret_cast(&gen); - PyAsyncGenObject* gen = asend->ags_gen; - auto asend_yf = reinterpret_cast(gen); + PyAsyncGenObject* gen_ptr = asend->ags_gen; + auto asend_yf = reinterpret_cast(gen_ptr); auto result = GenInfo::create(asend_yf); recursion_depth--; return result; @@ -176,6 +178,7 @@ class TaskInfo // Information to reconstruct the async stack as best as we can TaskInfo::Ptr waiter = nullptr; + std::optional is_on_cpu_ = std::nullopt; [[nodiscard]] static Result create(TaskObj*); TaskInfo(PyObject* origin, PyObject* loop, GenInfo::Ptr coro, StringTable::Key name, TaskInfo::Ptr waiter) @@ -189,6 +192,27 @@ class TaskInfo [[nodiscard]] static Result 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_; + } + + auto* last_coro = static_cast(nullptr); + auto* current_coro = this->coro.get(); + + // Check if the innermost coroutine is running. + // We don't need to test all the other coroutines as they are awaiting, by definition. + while (current_coro) { + last_coro = current_coro; + current_coro = current_coro->await.get(); + } + + is_on_cpu_ = last_coro && last_coro->is_running; + return *is_on_cpu_; + } }; inline std::unordered_map task_link_map; diff --git a/ddtrace/internal/datadog/profiling/stack_v2/echion/echion/threads.h b/ddtrace/internal/datadog/profiling/stack_v2/echion/echion/threads.h index 4bffbcd0dc5..7eed87f849a 100644 --- a/ddtrace/internal/datadog/profiling/stack_v2/echion/echion/threads.h +++ b/ddtrace/internal/datadog/profiling/stack_v2/echion/echion/threads.h @@ -252,8 +252,17 @@ 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(); + on_cpu_task_seen = on_cpu; + } + auto stack_info = std::make_unique(leaf_task.get().name, on_cpu); auto& stack = stack_info->stack; for (auto current_task = leaf_task;;) { diff --git a/releasenotes/notes/profiling-fix-detection-of-on-cpu-tasks-a2eb47163c3df850.yaml b/releasenotes/notes/profiling-fix-detection-of-on-cpu-tasks-a2eb47163c3df850.yaml new file mode 100644 index 00000000000..c7e066d8893 --- /dev/null +++ b/releasenotes/notes/profiling-fix-detection-of-on-cpu-tasks-a2eb47163c3df850.yaml @@ -0,0 +1,5 @@ +fixes: + - | + profiling: This fix improves the detection of on-CPU asyncio Tasks. Previously, the Profiler would only + consider a Task as running if its coroutine was running. The Profiler now recursively checks if any coroutine in the + await chain of the Task's coroutine is running.