diff --git a/ddtrace/internal/datadog/profiling/stack/echion/echion/threads.h b/ddtrace/internal/datadog/profiling/stack/echion/echion/threads.h index cc53cfb61f3..9acd221367c 100644 --- a/ddtrace/internal/datadog/profiling/stack/echion/echion/threads.h +++ b/ddtrace/internal/datadog/profiling/stack/echion/echion/threads.h @@ -114,20 +114,21 @@ class ThreadInfo }; private: + using TaskAddressCallback = std::function; + void reset_cycle_state() noexcept; void render_unwound_stacks(EchionSampler&); [[nodiscard]] Result unwind_tasks(EchionSampler&, PyThreadState*, microsecond_t wall_time_us); void unwind_greenlets(EchionSampler&, PyThreadState*, unsigned long, microsecond_t wall_time_us); [[nodiscard]] Result> get_all_tasks(EchionSampler&, PyThreadState* tstate); + [[nodiscard]] Result for_each_task_address(EchionSampler&, + PyThreadState* tstate, + const TaskAddressCallback& callback); #if PY_VERSION_HEX >= 0x030e0000 - [[nodiscard]] Result get_tasks_from_thread_linked_list(EchionSampler& echion, - std::vector& tasks); - [[nodiscard]] Result get_tasks_from_interpreter_linked_list(EchionSampler& echion, - PyThreadState* tstate, - std::vector& tasks); - [[nodiscard]] Result get_tasks_from_linked_list(EchionSampler& echion, - uintptr_t head_addr, - std::vector& tasks); + [[nodiscard]] Result get_tasks_from_thread_linked_list(const TaskAddressCallback& callback); + [[nodiscard]] Result get_tasks_from_interpreter_linked_list(PyThreadState* tstate, + const TaskAddressCallback& callback); + [[nodiscard]] Result get_tasks_from_linked_list(uintptr_t head_addr, const TaskAddressCallback& callback); #endif }; diff --git a/ddtrace/internal/datadog/profiling/stack/src/echion/threads.cc b/ddtrace/internal/datadog/profiling/stack/src/echion/threads.cc index 86d26986e45..ab283f6e858 100644 --- a/ddtrace/internal/datadog/profiling/stack/src/echion/threads.cc +++ b/ddtrace/internal/datadog/profiling/stack/src/echion/threads.cc @@ -411,7 +411,7 @@ ThreadInfo::unwind_tasks(EchionSampler& echion, PyThreadState* tstate, microseco // ---------------------------------------------------------------------------- #if PY_VERSION_HEX >= 0x030e0000 Result -ThreadInfo::get_tasks_from_thread_linked_list(EchionSampler& echion, std::vector& tasks) +ThreadInfo::get_tasks_from_thread_linked_list(const TaskAddressCallback& callback) { if (this->tstate_addr == 0 || this->asyncio_loop == 0) { return ErrorKind::TaskInfoError; @@ -426,13 +426,11 @@ ThreadInfo::get_tasks_from_thread_linked_list(EchionSampler& echion, std::vector constexpr size_t asyncio_tasks_head_offset = offsetof(_PyThreadStateImpl, asyncio_tasks_head); uintptr_t head_addr = this->tstate_addr + asyncio_tasks_head_offset; - return get_tasks_from_linked_list(echion, head_addr, tasks); + return get_tasks_from_linked_list(head_addr, callback); } Result -ThreadInfo::get_tasks_from_interpreter_linked_list(EchionSampler& echion, - PyThreadState* tstate, - std::vector& tasks) +ThreadInfo::get_tasks_from_interpreter_linked_list(PyThreadState* tstate, const TaskAddressCallback& callback) { if (tstate == nullptr || tstate->interp == nullptr || this->asyncio_loop == 0) { return ErrorKind::TaskInfoError; @@ -441,11 +439,11 @@ ThreadInfo::get_tasks_from_interpreter_linked_list(EchionSampler& echion, constexpr size_t asyncio_tasks_head_offset = offsetof(PyInterpreterState, asyncio_tasks_head); uintptr_t head_addr = reinterpret_cast(tstate->interp) + asyncio_tasks_head_offset; - return get_tasks_from_linked_list(echion, head_addr, tasks); + return get_tasks_from_linked_list(head_addr, callback); } Result -ThreadInfo::get_tasks_from_linked_list(EchionSampler& echion, uintptr_t head_addr, std::vector& tasks) +ThreadInfo::get_tasks_from_linked_list(uintptr_t head_addr, const TaskAddressCallback& callback) { if (head_addr == 0 || this->asyncio_loop == 0) { return ErrorKind::TaskInfoError; @@ -489,14 +487,7 @@ ThreadInfo::get_tasks_from_linked_list(EchionSampler& echion, uintptr_t head_add size_t task_node_offset_val = offsetof(TaskObj, task_node); uintptr_t task_addr_uint = next_node_addr - task_node_offset_val; - // Create TaskInfo for the task - auto maybe_task_info = TaskInfo::create(echion, reinterpret_cast(task_addr_uint)); - if (maybe_task_info) { - auto& task_info = *maybe_task_info; - if (task_info->loop == reinterpret_cast(this->asyncio_loop)) { - tasks.push_back(std::move(task_info)); - } - } + callback(reinterpret_cast(task_addr_uint)); // Read next node from current_node.next into current_node if (copy_type(reinterpret_cast(next_node_addr), current_node)) { @@ -507,12 +498,11 @@ ThreadInfo::get_tasks_from_linked_list(EchionSampler& echion, uintptr_t head_add return Result::ok(); } -Result> -ThreadInfo::get_all_tasks(EchionSampler& echion, PyThreadState* tstate) +Result +ThreadInfo::for_each_task_address(EchionSampler& echion, PyThreadState* tstate, const TaskAddressCallback& callback) { - std::vector tasks; if (this->asyncio_loop == 0) - return tasks; + return Result::ok(); // Python 3.14+: Native tasks are in linked-list per thread AND per interpreter // CPython iterates over both: @@ -521,10 +511,10 @@ ThreadInfo::get_all_tasks(EchionSampler& echion, PyThreadState* tstate) // First, get tasks from this thread's linked-list (if tstate_addr is set) // Note: We continue processing even if one source fails to maximize partial results if (tstate != nullptr && this->tstate_addr != 0) { - (void)get_tasks_from_thread_linked_list(echion, tasks); + (void)get_tasks_from_thread_linked_list(callback); // Second, get tasks from interpreter's linked-list (lingering tasks) - (void)get_tasks_from_interpreter_linked_list(echion, tstate, tasks); + (void)get_tasks_from_interpreter_linked_list(tstate, callback); } // Handle third-party tasks from Python _scheduled_tasks WeakSet @@ -540,11 +530,7 @@ ThreadInfo::get_all_tasks(EchionSampler& echion, PyThreadState* tstate) auto scheduled_tasks = std::move(*maybe_scheduled_tasks); for (auto task_addr : scheduled_tasks) { // In WeakSet.data (set), elements are the Task objects themselves - auto maybe_task_info = TaskInfo::create(echion, reinterpret_cast(task_addr)); - if (maybe_task_info && - (*maybe_task_info)->loop == reinterpret_cast(this->asyncio_loop)) { - tasks.push_back(std::move(*maybe_task_info)); - } + callback(reinterpret_cast(task_addr)); } } } @@ -566,25 +552,19 @@ ThreadInfo::get_all_tasks(EchionSampler& echion, PyThreadState* tstate) auto eager_tasks = std::move(*maybe_eager_tasks); for (auto task_addr : eager_tasks) { - auto maybe_task_info = TaskInfo::create(echion, reinterpret_cast(task_addr)); - if (maybe_task_info) { - if ((*maybe_task_info)->loop == reinterpret_cast(this->asyncio_loop)) { - tasks.push_back(std::move(*maybe_task_info)); - } - } + callback(reinterpret_cast(task_addr)); } } - return tasks; + return Result::ok(); } #else -// Pre-Python 3.14: get_all_tasks uses WeakSet approach -Result> -ThreadInfo::get_all_tasks(EchionSampler& echion, PyThreadState*) +// Pre-Python 3.14: asyncio tracks tasks through Python weak sets. +Result +ThreadInfo::for_each_task_address(EchionSampler& echion, PyThreadState*, const TaskAddressCallback& callback) { - std::vector tasks; if (this->asyncio_loop == 0) - return tasks; + return Result::ok(); auto asyncio_scheduled_tasks = echion.asyncio_scheduled_tasks(); auto maybe_scheduled_tasks_set = MirrorSet::create(asyncio_scheduled_tasks); @@ -604,12 +584,7 @@ ThreadInfo::get_all_tasks(EchionSampler& echion, PyThreadState*) if (copy_type(task_wr_addr, task_wr)) continue; - auto maybe_task_info = TaskInfo::create(echion, reinterpret_cast(task_wr.wr_object)); - if (maybe_task_info) { - if (reinterpret_cast((*maybe_task_info)->loop) == this->asyncio_loop) { - tasks.push_back(std::move(*maybe_task_info)); - } - } + callback(reinterpret_cast(task_wr.wr_object)); } auto asyncio_eager_tasks = echion.asyncio_eager_tasks(); @@ -628,19 +603,30 @@ ThreadInfo::get_all_tasks(EchionSampler& echion, PyThreadState*) auto eager_tasks = std::move(*maybe_eager_tasks); for (auto task_addr : eager_tasks) { - auto maybe_task_info = TaskInfo::create(echion, reinterpret_cast(task_addr)); - if (maybe_task_info) { - if (reinterpret_cast((*maybe_task_info)->loop) == this->asyncio_loop) { - tasks.push_back(std::move(*maybe_task_info)); - } - } + callback(reinterpret_cast(task_addr)); } } - return tasks; + return Result::ok(); } #endif // PY_VERSION_HEX >= 0x030e0000 +Result> +ThreadInfo::get_all_tasks(EchionSampler& echion, PyThreadState* tstate) +{ + std::vector tasks; + auto result = for_each_task_address(echion, tstate, [&](TaskObj* task_addr) { + auto maybe_task_info = TaskInfo::create(echion, task_addr); + if (maybe_task_info && reinterpret_cast((*maybe_task_info)->loop) == this->asyncio_loop) { + tasks.push_back(std::move(*maybe_task_info)); + } + }); + if (!result) { + return result.error(); + } + return tasks; +} + // ---------------------------------------------------------------------------- void ThreadInfo::unwind_greenlets(EchionSampler& echion,