refactor(profiling): separate task address discovery

Author: taegyunkim

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

@@ -114,20 +114,21 @@ class ThreadInfo
     };
 
   private:
+    using TaskAddressCallback = std::function<void(TaskObj*)>;
+
     void reset_cycle_state() noexcept;
     void render_unwound_stacks(EchionSampler&);
     [[nodiscard]] Result<void> unwind_tasks(EchionSampler&, PyThreadState*, microsecond_t wall_time_us);
     void unwind_greenlets(EchionSampler&, PyThreadState*, unsigned long, microsecond_t wall_time_us);
     [[nodiscard]] Result<std::vector<TaskInfo::Ptr>> get_all_tasks(EchionSampler&, PyThreadState* tstate);
+    [[nodiscard]] Result<void> for_each_task_address(EchionSampler&,
+                                                     PyThreadState* tstate,
+                                                     const TaskAddressCallback& callback);
 #if PY_VERSION_HEX >= 0x030e0000
-    [[nodiscard]] Result<void> get_tasks_from_thread_linked_list(EchionSampler& echion,
-                                                                 std::vector<TaskInfo::Ptr>& tasks);
-    [[nodiscard]] Result<void> get_tasks_from_interpreter_linked_list(EchionSampler& echion,
-                                                                      PyThreadState* tstate,
-                                                                      std::vector<TaskInfo::Ptr>& tasks);
-    [[nodiscard]] Result<void> get_tasks_from_linked_list(EchionSampler& echion,
-                                                          uintptr_t head_addr,
-                                                          std::vector<TaskInfo::Ptr>& tasks);
+    [[nodiscard]] Result<void> visit_thread_task_addresses(const TaskAddressCallback& callback);
+    [[nodiscard]] Result<void> visit_interpreter_task_addresses(PyThreadState* tstate,
+                                                                const TaskAddressCallback& callback);
+    [[nodiscard]] Result<void> visit_linked_task_addresses(uintptr_t head_addr, const TaskAddressCallback& callback);
 #endif
 };
 

ddtrace/internal/datadog/profiling/stack/src/echion/threads.cc

@@ -411,236 +411,162 @@ ThreadInfo::unwind_tasks(EchionSampler& echion, PyThreadState* tstate, microseco
 // ----------------------------------------------------------------------------
 #if PY_VERSION_HEX >= 0x030e0000
 Result<void>
-ThreadInfo::get_tasks_from_thread_linked_list(EchionSampler& echion, std::vector<TaskInfo::Ptr>& tasks)
+ThreadInfo::visit_thread_task_addresses(const TaskAddressCallback& callback)
 {
     if (this->tstate_addr == 0 || this->asyncio_loop == 0) {
         return ErrorKind::TaskInfoError;
     }
 
-    // Calculate thread state's asyncio_tasks_head remote address
-    // Note: Since 3.13+, every PyThreadState is actually allocated as a _PyThreadStateImpl.
-    // We use PyThreadState* everywhere and cast to _PyThreadStateImpl* only when we need
-    // to access asyncio_tasks_head (which is only available in Python 3.14+).
-    // Since tstate_addr is a remote address, we calculate the offset and add it to the address.
-    // get_tasks_from_linked_list will handle copying the head node from remote memory internally.
     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 visit_linked_task_addresses(this->tstate_addr + asyncio_tasks_head_offset, callback);
 }
 
 Result<void>
-ThreadInfo::get_tasks_from_interpreter_linked_list(EchionSampler& echion,
-                                                   PyThreadState* tstate,
-                                                   std::vector<TaskInfo::Ptr>& tasks)
+ThreadInfo::visit_interpreter_task_addresses(PyThreadState* tstate, const TaskAddressCallback& callback)
 {
     if (tstate == nullptr || tstate->interp == nullptr || this->asyncio_loop == 0) {
         return ErrorKind::TaskInfoError;
     }
 
     constexpr size_t asyncio_tasks_head_offset = offsetof(PyInterpreterState, asyncio_tasks_head);
-    uintptr_t head_addr = reinterpret_cast<uintptr_t>(tstate->interp) + asyncio_tasks_head_offset;
-
-    return get_tasks_from_linked_list(echion, head_addr, tasks);
+    const uintptr_t head_addr = reinterpret_cast<uintptr_t>(tstate->interp) + asyncio_tasks_head_offset;
+    return visit_linked_task_addresses(head_addr, callback);
 }
 
 Result<void>
-ThreadInfo::get_tasks_from_linked_list(EchionSampler& echion, uintptr_t head_addr, std::vector<TaskInfo::Ptr>& tasks)
+ThreadInfo::visit_linked_task_addresses(uintptr_t head_addr, const TaskAddressCallback& callback)
 {
     if (head_addr == 0 || this->asyncio_loop == 0) {
         return ErrorKind::TaskInfoError;
     }
 
-    // Copy head node struct from remote memory to local memory
-    struct llist_node head_node_local;
-    if (copy_type(reinterpret_cast<void*>(head_addr), head_node_local)) {
+    struct llist_node current_node;
+    if (copy_type(reinterpret_cast<void*>(head_addr), current_node)) {
         return ErrorKind::TaskInfoError;
     }
 
-    // Check if list is empty (head points to itself in circular list)
-    uintptr_t head_addr_uint = head_addr;
-    uintptr_t next_as_uint = reinterpret_cast<uintptr_t>(head_node_local.next);
-    uintptr_t prev_as_uint = reinterpret_cast<uintptr_t>(head_node_local.prev);
-    if (next_as_uint == head_addr_uint && prev_as_uint == head_addr_uint) {
+    const uintptr_t head_addr_uint = head_addr;
+    if (reinterpret_cast<uintptr_t>(current_node.next) == head_addr_uint &&
+        reinterpret_cast<uintptr_t>(current_node.prev) == head_addr_uint) {
         return Result<void>::ok();
     }
 
-    struct llist_node current_node = head_node_local; // Start with head node
-
-    // Copied from CPython's _remote_debugging_module.c: MAX_ITERATIONS
-    const size_t MAX_ITERATIONS = 1 << 16;
+    constexpr size_t max_iterations = 1 << 16;
     size_t iteration_count = 0;
-
-    // Iterate over linked-list. The linked list is circular, so we stop
-    // when we're back at head.
     while (reinterpret_cast<uintptr_t>(current_node.next) != head_addr_uint) {
-        // Safety: prevent infinite loops
-        if (++iteration_count > MAX_ITERATIONS) {
+        if (++iteration_count > max_iterations || current_node.next == nullptr) {
             return ErrorKind::TaskInfoError;
         }
 
-        if (current_node.next == nullptr) {
-            return ErrorKind::TaskInfoError; // nullptr pointer - invalid list
-        }
-
-        uintptr_t next_node_addr = reinterpret_cast<uintptr_t>(current_node.next);
-
-        // Calculate task_addr from current_node.next
-        size_t task_node_offset_val = offsetof(TaskObj, task_node);
-        uintptr_t task_addr_uint = next_node_addr - task_node_offset_val;
+        const uintptr_t next_node_addr = reinterpret_cast<uintptr_t>(current_node.next);
+        const uintptr_t task_addr = next_node_addr - offsetof(TaskObj, task_node);
+        callback(reinterpret_cast<TaskObj*>(task_addr));
 
-        // Create TaskInfo for the task
-        auto maybe_task_info = TaskInfo::create(echion, reinterpret_cast<TaskObj*>(task_addr_uint));
-        if (maybe_task_info) {
-            auto& task_info = *maybe_task_info;
-            if (task_info->loop == reinterpret_cast<PyObject*>(this->asyncio_loop)) {
-                tasks.push_back(std::move(task_info));
-            }
-        }
-
-        // Read next node from current_node.next into current_node
         if (copy_type(reinterpret_cast<void*>(next_node_addr), current_node)) {
-            return ErrorKind::TaskInfoError; // Failed to read next node
+            return ErrorKind::TaskInfoError;
         }
     }
 
     return Result<void>::ok();
 }
 
-Result<std::vector<TaskInfo::Ptr>>
-ThreadInfo::get_all_tasks(EchionSampler& echion, PyThreadState* tstate)
+Result<void>
+ThreadInfo::for_each_task_address(EchionSampler& echion, PyThreadState* tstate, const TaskAddressCallback& callback)
 {
-    std::vector<TaskInfo::Ptr> tasks;
-    if (this->asyncio_loop == 0)
-        return tasks;
-
-    // Python 3.14+: Native tasks are in linked-list per thread AND per interpreter
-    // CPython iterates over both:
-    // 1. Per-thread list: tstate->asyncio_tasks_head (active tasks)
-    // 2. Per-interpreter list: interp->asyncio_tasks_head (lingering tasks)
-    // 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 (this->asyncio_loop == 0) {
+        return Result<void>::ok();
+    }
+
+    // Native tasks can appear in both lists. Consumers already tolerate duplicate task objects.
     if (tstate != nullptr && this->tstate_addr != 0) {
-        (void)get_tasks_from_thread_linked_list(echion, tasks);
-
-        // Second, get tasks from interpreter's linked-list (lingering tasks)
-        (void)get_tasks_from_interpreter_linked_list(echion, tstate, tasks);
-    }
-
-    // Handle third-party tasks from Python _scheduled_tasks WeakSet
-    // In Python 3.14+, _scheduled_tasks is a Python-level weakref.WeakSet() that only contains
-    // tasks that don't inherit from asyncio.Task. Native asyncio.Task instances are stored
-    // in linked-lists (handled above) and are NOT added to _scheduled_tasks.
-    // This is typically empty in practice, but we handle it for completeness.
-    auto asyncio_scheduled_tasks = echion.asyncio_scheduled_tasks();
-    if (asyncio_scheduled_tasks != nullptr) {
-        if (auto maybe_scheduled_tasks_set = MirrorSet::create(asyncio_scheduled_tasks)) {
-            auto scheduled_tasks_set = std::move(*maybe_scheduled_tasks_set);
-            if (auto maybe_scheduled_tasks = scheduled_tasks_set.as_unordered_set()) {
-                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<TaskObj*>(task_addr));
-                    if (maybe_task_info &&
-                        (*maybe_task_info)->loop == reinterpret_cast<PyObject*>(this->asyncio_loop)) {
-                        tasks.push_back(std::move(*maybe_task_info));
-                    }
+        (void)visit_thread_task_addresses(callback);
+        (void)visit_interpreter_task_addresses(tstate, callback);
+    }
+
+    // Python 3.14 stores only third-party Task implementations in _scheduled_tasks.
+    if (auto scheduled = echion.asyncio_scheduled_tasks(); scheduled != nullptr) {
+        if (auto maybe_set = MirrorSet::create(scheduled)) {
+            if (auto maybe_tasks = maybe_set->as_unordered_set()) {
+                for (auto task : *maybe_tasks) {
+                    callback(reinterpret_cast<TaskObj*>(task));
                 }
             }
         }
     }
 
-    auto asyncio_eager_tasks = echion.asyncio_eager_tasks();
-    if (asyncio_eager_tasks != nullptr) {
-        auto maybe_eager_tasks_set = MirrorSet::create(asyncio_eager_tasks);
-        if (!maybe_eager_tasks_set) {
+    if (auto eager = echion.asyncio_eager_tasks(); eager != nullptr) {
+        auto maybe_set = MirrorSet::create(eager);
+        if (!maybe_set) {
             return ErrorKind::TaskInfoError;
         }
-
-        auto eager_tasks_set = std::move(*maybe_eager_tasks_set);
-
-        auto maybe_eager_tasks = eager_tasks_set.as_unordered_set();
-        if (!maybe_eager_tasks) {
+        auto maybe_tasks = maybe_set->as_unordered_set();
+        if (!maybe_tasks) {
             return ErrorKind::TaskInfoError;
         }
-
-        auto eager_tasks = std::move(*maybe_eager_tasks);
-        for (auto task_addr : eager_tasks) {
-            auto maybe_task_info = TaskInfo::create(echion, reinterpret_cast<TaskObj*>(task_addr));
-            if (maybe_task_info) {
-                if ((*maybe_task_info)->loop == reinterpret_cast<PyObject*>(this->asyncio_loop)) {
-                    tasks.push_back(std::move(*maybe_task_info));
-                }
-            }
+        for (auto task : *maybe_tasks) {
+            callback(reinterpret_cast<TaskObj*>(task));
         }
     }
 
-    return tasks;
+    return Result<void>::ok();
 }
 #else
-// Pre-Python 3.14: get_all_tasks uses WeakSet approach
-Result<std::vector<TaskInfo::Ptr>>
-ThreadInfo::get_all_tasks(EchionSampler& echion, PyThreadState*)
+Result<void>
+ThreadInfo::for_each_task_address(EchionSampler& echion, PyThreadState*, const TaskAddressCallback& callback)
 {
-    std::vector<TaskInfo::Ptr> tasks;
-    if (this->asyncio_loop == 0)
-        return tasks;
+    if (this->asyncio_loop == 0) {
+        return Result<void>::ok();
+    }
 
-    auto asyncio_scheduled_tasks = echion.asyncio_scheduled_tasks();
-    auto maybe_scheduled_tasks_set = MirrorSet::create(asyncio_scheduled_tasks);
-    if (!maybe_scheduled_tasks_set) {
+    auto maybe_set = MirrorSet::create(echion.asyncio_scheduled_tasks());
+    if (!maybe_set) {
         return ErrorKind::TaskInfoError;
     }
-
-    auto scheduled_tasks_set = std::move(*maybe_scheduled_tasks_set);
-    auto maybe_scheduled_tasks = scheduled_tasks_set.as_unordered_set();
-    if (!maybe_scheduled_tasks) {
+    auto maybe_scheduled = maybe_set->as_unordered_set();
+    if (!maybe_scheduled) {
         return ErrorKind::TaskInfoError;
     }
-
-    auto scheduled_tasks = std::move(*maybe_scheduled_tasks);
-    for (auto task_wr_addr : scheduled_tasks) {
-        PyWeakReference task_wr;
-        if (copy_type(task_wr_addr, task_wr))
-            continue;
-
-        auto maybe_task_info = TaskInfo::create(echion, reinterpret_cast<TaskObj*>(task_wr.wr_object));
-        if (maybe_task_info) {
-            if (reinterpret_cast<uintptr_t>((*maybe_task_info)->loop) == this->asyncio_loop) {
-                tasks.push_back(std::move(*maybe_task_info));
-            }
+    for (auto weakref_addr : *maybe_scheduled) {
+        PyWeakReference weakref;
+        if (!copy_type(weakref_addr, weakref)) {
+            callback(reinterpret_cast<TaskObj*>(weakref.wr_object));
         }
     }
 
-    auto asyncio_eager_tasks = echion.asyncio_eager_tasks();
-    if (asyncio_eager_tasks != nullptr) {
-        auto maybe_eager_tasks_set = MirrorSet::create(asyncio_eager_tasks);
-        if (!maybe_eager_tasks_set) {
+    if (auto eager = echion.asyncio_eager_tasks(); eager != nullptr) {
+        auto maybe_eager_set = MirrorSet::create(eager);
+        if (!maybe_eager_set) {
             return ErrorKind::TaskInfoError;
         }
-
-        auto eager_tasks_set = std::move(*maybe_eager_tasks_set);
-
-        auto maybe_eager_tasks = eager_tasks_set.as_unordered_set();
-        if (!maybe_eager_tasks) {
+        auto maybe_eager = maybe_eager_set->as_unordered_set();
+        if (!maybe_eager) {
             return ErrorKind::TaskInfoError;
         }
-
-        auto eager_tasks = std::move(*maybe_eager_tasks);
-        for (auto task_addr : eager_tasks) {
-            auto maybe_task_info = TaskInfo::create(echion, reinterpret_cast<TaskObj*>(task_addr));
-            if (maybe_task_info) {
-                if (reinterpret_cast<uintptr_t>((*maybe_task_info)->loop) == this->asyncio_loop) {
-                    tasks.push_back(std::move(*maybe_task_info));
-                }
-            }
+        for (auto task : *maybe_eager) {
+            callback(reinterpret_cast<TaskObj*>(task));
         }
     }
 
-    return tasks;
+    return Result<void>::ok();
 }
 #endif // PY_VERSION_HEX >= 0x030e0000
 
+Result<std::vector<TaskInfo::Ptr>>
+ThreadInfo::get_all_tasks(EchionSampler& echion, PyThreadState* tstate)
+{
+    std::vector<TaskInfo::Ptr> tasks;
+    auto result = for_each_task_address(echion, tstate, [&](TaskObj* task_addr) {
+        auto maybe_task = TaskInfo::create(echion, task_addr);
+        if (maybe_task && reinterpret_cast<uintptr_t>((*maybe_task)->loop) == this->asyncio_loop) {
+            tasks.push_back(std::move(*maybe_task));
+        }
+    });
+    if (!result) {
+        return result.error();
+    }
+    return tasks;
+}
+
 // ----------------------------------------------------------------------------
 void
 ThreadInfo::unwind_greenlets(EchionSampler& echion,