refactor(profiling): visit task identities during discovery

Author: taegyunkim

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

@@ -121,17 +121,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<std::vector<TaskObj*>> get_all_task_addresses(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_task_addresses_from_thread_linked_list(std::vector<TaskObj*>& tasks);
-    [[nodiscard]] Result<void> get_task_addresses_from_interpreter_linked_list(PyThreadState* tstate,
-                                                                               std::vector<TaskObj*>& tasks);
-    [[nodiscard]] Result<void> get_task_addresses_from_linked_list(uintptr_t head_addr, std::vector<TaskObj*>& 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

@@ -413,30 +413,30 @@ ThreadInfo::unwind_tasks(EchionSampler& echion, PyThreadState* tstate, microseco
 // ----------------------------------------------------------------------------
 #if PY_VERSION_HEX >= 0x030e0000
 Result<void>
-ThreadInfo::get_task_addresses_from_thread_linked_list(std::vector<TaskObj*>& tasks)
+ThreadInfo::visit_thread_task_addresses(const TaskAddressCallback& callback)
 {
     if (this->tstate_addr == 0 || this->asyncio_loop == 0) {
         return ErrorKind::TaskInfoError;
     }
 
     constexpr size_t asyncio_tasks_head_offset = offsetof(_PyThreadStateImpl, asyncio_tasks_head);
-    return get_task_addresses_from_linked_list(this->tstate_addr + asyncio_tasks_head_offset, tasks);
+    return visit_linked_task_addresses(this->tstate_addr + asyncio_tasks_head_offset, callback);
 }
 
 Result<void>
-ThreadInfo::get_task_addresses_from_interpreter_linked_list(PyThreadState* tstate, std::vector<TaskObj*>& 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);
     const uintptr_t head_addr = reinterpret_cast<uintptr_t>(tstate->interp) + asyncio_tasks_head_offset;
-    return get_task_addresses_from_linked_list(head_addr, tasks);
+    return visit_linked_task_addresses(head_addr, callback);
 }
 
 Result<void>
-ThreadInfo::get_task_addresses_from_linked_list(uintptr_t head_addr, std::vector<TaskObj*>& tasks)
+ThreadInfo::visit_linked_task_addresses(uintptr_t head_addr, const TaskAddressCallback& callback)
 {
     if (head_addr == 0 || this->asyncio_loop == 0) {
         return ErrorKind::TaskInfoError;
@@ -462,7 +462,7 @@ ThreadInfo::get_task_addresses_from_linked_list(uintptr_t head_addr, std::vector
 
         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);
-        tasks.push_back(reinterpret_cast<TaskObj*>(task_addr));
+        callback(reinterpret_cast<TaskObj*>(task_addr));
 
         if (copy_type(reinterpret_cast<void*>(next_node_addr), current_node)) {
             return ErrorKind::TaskInfoError;
@@ -472,28 +472,25 @@ ThreadInfo::get_task_addresses_from_linked_list(uintptr_t head_addr, std::vector
     return Result<void>::ok();
 }
 
-Result<std::vector<TaskObj*>>
-ThreadInfo::get_all_task_addresses(EchionSampler& echion, PyThreadState* tstate)
+Result<void>
+ThreadInfo::for_each_task_address(EchionSampler& echion, PyThreadState* tstate, const TaskAddressCallback& callback)
 {
-    std::vector<TaskObj*> tasks;
     if (this->asyncio_loop == 0) {
-        return tasks;
+        return Result<void>::ok();
     }
 
     // Native tasks can appear in both lists. Consumers already tolerate duplicate task objects.
-    if (this->tstate_addr != 0) {
-        (void)get_task_addresses_from_thread_linked_list(tasks);
-    }
-    if (tstate != nullptr) {
-        (void)get_task_addresses_from_interpreter_linked_list(tstate, tasks);
+    if (tstate != nullptr && this->tstate_addr != 0) {
+        (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) {
-                    tasks.push_back(reinterpret_cast<TaskObj*>(task));
+                    callback(reinterpret_cast<TaskObj*>(task));
                 }
             }
         }
@@ -509,19 +506,18 @@ ThreadInfo::get_all_task_addresses(EchionSampler& echion, PyThreadState* tstate)
             return ErrorKind::TaskInfoError;
         }
         for (auto task : *maybe_tasks) {
-            tasks.push_back(reinterpret_cast<TaskObj*>(task));
+            callback(reinterpret_cast<TaskObj*>(task));
         }
     }
 
-    return tasks;
+    return Result<void>::ok();
 }
 #else
-Result<std::vector<TaskObj*>>
-ThreadInfo::get_all_task_addresses(EchionSampler& echion, PyThreadState*)
+Result<void>
+ThreadInfo::for_each_task_address(EchionSampler& echion, PyThreadState*, const TaskAddressCallback& callback)
 {
-    std::vector<TaskObj*> tasks;
     if (this->asyncio_loop == 0) {
-        return tasks;
+        return Result<void>::ok();
     }
 
     auto maybe_set = MirrorSet::create(echion.asyncio_scheduled_tasks());
@@ -535,7 +531,7 @@ ThreadInfo::get_all_task_addresses(EchionSampler& echion, PyThreadState*)
     for (auto weakref_addr : *maybe_scheduled) {
         PyWeakReference weakref;
         if (!copy_type(weakref_addr, weakref)) {
-            tasks.push_back(reinterpret_cast<TaskObj*>(weakref.wr_object));
+            callback(reinterpret_cast<TaskObj*>(weakref.wr_object));
         }
     }
 
@@ -549,28 +545,26 @@ ThreadInfo::get_all_task_addresses(EchionSampler& echion, PyThreadState*)
             return ErrorKind::TaskInfoError;
         }
         for (auto task : *maybe_eager) {
-            tasks.push_back(reinterpret_cast<TaskObj*>(task));
+            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)
 {
-    auto maybe_addresses = get_all_task_addresses(echion, tstate);
-    if (!maybe_addresses) {
-        return maybe_addresses.error();
-    }
-
     std::vector<TaskInfo::Ptr> tasks;
-    for (TaskObj* task_addr : *maybe_addresses) {
+    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;
 }
@@ -1169,17 +1163,16 @@ ThreadInfo::sample_cpu_timer(EchionSampler& echion,
             return;
         }
 
-        auto maybe_addresses = get_all_task_addresses(echion, tstate);
-        if (maybe_addresses) {
-            std::vector<TaskIdentity> tasks;
-            tasks.reserve(maybe_addresses->size());
-            for (TaskObj* address : *maybe_addresses) {
-                TaskObj task;
-                if (!copy_type(address, task) && reinterpret_cast<uintptr_t>(task.task_loop) == asyncio_loop) {
-                    tasks.push_back({ address, task.task_coro, task.task_fut_waiter });
-                }
+        // Snapshot each lightweight identity as its address is visited. Retaining bare addresses for a later pass
+        // would widen the race with task completion and address reuse.
+        std::vector<TaskIdentity> tasks;
+        auto visit_result = for_each_task_address(echion, tstate, [&](TaskObj* address) {
+            TaskObj task;
+            if (!copy_type(address, task) && reinterpret_cast<uintptr_t>(task.task_loop) == asyncio_loop) {
+                tasks.push_back({ address, task.task_coro, task.task_fut_waiter });
             }
-
+        });
+        if (visit_result) {
             if (TaskObj* address = find_captured_task(tasks, raw)) {
                 auto maybe_task = TaskInfo::create(echion, address);
                 if (maybe_task) {