fix(profiling): workaround for on-CPU Task race condition

Author: KowalskiThomas

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

@@ -201,7 +201,7 @@ class TaskInfo
     }
 
     [[nodiscard]] static Result<TaskInfo::Ptr> current(PyObject*);
-    inline size_t unwind(FrameStack&, size_t& upper_python_stack_size);
+    inline Result<size_t> unwind(FrameStack&, size_t& upper_python_stack_size);
 };
 
 inline std::unordered_map<PyObject*, PyObject*> task_link_map;
@@ -342,7 +342,7 @@ inline std::vector<std::unique_ptr<StackInfo>> current_tasks;
 
 // ----------------------------------------------------------------------------
 
-inline size_t
+inline Result<size_t>
 TaskInfo::unwind(FrameStack& stack, size_t& upper_python_stack_size)
 {
     // TODO: Check for running task.
@@ -370,6 +370,13 @@ TaskInfo::unwind(FrameStack& stack, size_t& upper_python_stack_size)
             // The first Frame is the coroutine Frame, so the Python stack size is the number of Frames - 1
             upper_python_stack_size = new_frames - 1;
 
+            // If the Task is on CPU, then we should have at least the asyncio runtime Frames on top of
+            // the asynchronous Stack. If we do not have any Python Frames, then the execution state changed
+            // (race condition) and we cannot recover.
+            if (this->is_on_cpu && upper_python_stack_size == 0) {
+                return ErrorKind::TaskInfoError;
+            }
+
             // Remove the Python Frames from the Stack (they will be added back later)
             // We cannot push those Frames now because otherwise they would be added once per Task,
             // we only want to add them once per Leaf Task, and on top of all non-leaf Tasks.

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

@@ -205,6 +205,36 @@ ThreadInfo::unwind(PyThreadState* tstate)
 inline Result<void>
 ThreadInfo::unwind_tasks()
 {
+    // Check if the Python stack contains "_run".
+    // To avoid having to do string comparisons every time we unwind Tasks, we keep track
+    // of the cache key of the "_run" Frame.
+    static std::optional<Frame::Key> frame_cache_key;
+    bool expect_at_least_one_running_task = false;
+    if (!frame_cache_key) {
+        for (size_t i = 0; i < python_stack.size(); i++) {
+            const auto& frame = python_stack[i].get();
+            const auto& frame_name = string_table.lookup(frame.name)->get();
+            if (frame_name.size() >= 4 && frame_name.rfind("_run") == frame_name.size() - 4) {
+
+                // Although Frames are stored in an LRUCache, the cache key is ALWAYS the same
+                // even if the Frame gets evicted from the cache.
+                // This means we can keep the cache key and re-use it to determine
+                // whether we see the "_run" Frame in the Python stack.
+                frame_cache_key = frame.cache_key;
+                expect_at_least_one_running_task = true;
+                break;
+            }
+        }
+    } else {
+        for (size_t i = 0; i < python_stack.size(); i++) {
+            const auto& frame = python_stack[i].get();
+            if (frame.cache_key == *frame_cache_key) {
+                expect_at_least_one_running_task = true;
+                break;
+            }
+        }
+    }
+
     std::vector<TaskInfo::Ref> leaf_tasks;
     std::unordered_set<PyObject*> parent_tasks;
     std::unordered_map<PyObject*, TaskInfo::Ref> waitee_map; // Indexed by task origin
@@ -217,6 +247,22 @@ ThreadInfo::unwind_tasks()
     }
 
     auto all_tasks = std::move(*maybe_all_tasks);
+
+    bool saw_at_least_one_running_task = false;
+    std::string running_task_name;
+    for (const auto& task_ref : all_tasks) {
+        const auto& task = task_ref.get();
+        if (task->is_on_cpu) {
+            running_task_name = string_table.lookup(task->name)->get();
+            saw_at_least_one_running_task = true;
+            break;
+        }
+    }
+
+    if (saw_at_least_one_running_task != expect_at_least_one_running_task) {
+        return ErrorKind::TaskInfoError;
+    }
+
     {
         std::lock_guard<std::mutex> lock(task_link_map_lock);
 
@@ -284,8 +330,13 @@ ThreadInfo::unwind_tasks()
         for (auto current_task = leaf_task;;) {
             auto& task = current_task.get();
 
+            auto maybe_task_stack_size = task.unwind(stack, task.is_on_cpu ? upper_python_stack_size : unused);
+            if (!maybe_task_stack_size) {
+                return ErrorKind::TaskInfoError;
+            }
+
             // The task_stack_size includes both the coroutines frames and the "upper" Python synchronous frames
-            size_t task_stack_size = task.unwind(stack, task.is_on_cpu ? upper_python_stack_size : unused);
+            size_t task_stack_size = std::move(*maybe_task_stack_size);
             if (task.is_on_cpu) {
                 // Get the "bottom" part of the Python synchronous Stack, that is to say the
                 // synchronous functions and coroutines called by the Task's outermost coroutine

releasenotes/notes/profiling-on-cpu-task-race-condition-workaround-93a01ac9f8d4983c.yaml

@@ -0,0 +1,4 @@
+fixes:
+  - |
+    profiling: This fix works around a race condition leading to incorrect stacks sometimes being reported when
+    profiling on-CPU asyncio Tasks.