fix(profiling): fix issue with boundary asyncio frames

Author: KowalskiThomas

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

@@ -17,6 +17,16 @@
 // Forward declaration
 class Frame;
 
+// Identity of the frame that separates the asyncio machinery from the pure Python stack.
+// We memoize the interned name and filename rather than a Frame cache key: the cache key
+// mixes in the bytecode offset, so it only matches the boundary frame while it sits on the
+// very instruction it happened to be on when we first identified it.
+struct BoundaryFrame
+{
+    StringTable::Key name = 0;
+    StringTable::Key filename = 0;
+};
+
 class EchionSampler
 {
     // Thread Info map (Thread ID -> ThreadInfo)
@@ -39,8 +49,8 @@ class EchionSampler
     PyObject* asyncio_eager_tasks_ = nullptr;
 
     // Task unwinding state
-    std::optional<Frame::Key> asyncio_frame_cache_key_;
-    std::optional<Frame::Key> uvloop_frame_cache_key_;
+    std::optional<BoundaryFrame> asyncio_boundary_frame_;
+    std::optional<BoundaryFrame> uvloop_boundary_frame_;
     std::unordered_set<PyObject*> previous_task_objects_;
 
     // Sampling-thread scratch buffer. Only the single sampling thread
@@ -99,8 +109,8 @@ class EchionSampler
         asyncio_eager_tasks_ = (eager_tasks != Py_None) ? eager_tasks : nullptr;
     }
 
-    std::optional<Frame::Key>& asyncio_frame_cache_key() { return asyncio_frame_cache_key_; }
-    std::optional<Frame::Key>& uvloop_frame_cache_key() { return uvloop_frame_cache_key_; }
+    std::optional<BoundaryFrame>& asyncio_boundary_frame() { return asyncio_boundary_frame_; }
+    std::optional<BoundaryFrame>& uvloop_boundary_frame() { return uvloop_boundary_frame_; }
     std::unordered_set<PyObject*>& previous_task_objects() { return previous_task_objects_; }
 
     std::unordered_set<PyObject*>& seen_frames_scratch() { return seen_frames_scratch_; }
@@ -149,8 +159,8 @@ class EchionSampler
         new (&greenlet_thread_map_) std::unordered_map<uintptr_t, GreenletInfo::ID>();
         new (&previous_task_objects_) std::unordered_set<PyObject*>();
 
-        asyncio_frame_cache_key_.reset();
-        uvloop_frame_cache_key_.reset();
+        asyncio_boundary_frame_.reset();
+        uvloop_boundary_frame_.reset();
         asyncio_task_count_ = 0;
         rng_ = std::minstd_rand{ std::random_device{}() };
 

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

@@ -63,6 +63,12 @@ class ThreadInfo
     [[nodiscard]] Result<void> sample(EchionSampler&, PyThreadState*, microsecond_t);
     void unwind(EchionSampler&, PyThreadState*, microsecond_t wall_time_us);
 
+    // Number of frames in python_stack from the asyncio boundary frame (inclusive) up to the root,
+    // that is to say the asyncio machinery plus the synchronous entry point. Returns the size of the
+    // whole stack when the boundary frame is not there, which is the case for a thread that is not
+    // running an event loop.
+    [[nodiscard]] size_t find_upper_python_stack_size(EchionSampler&) const;
+
     // ------------------------------------------------------------------------
 #if defined PL_LINUX
     ThreadInfo(uintptr_t thread_id, unsigned long native_id, const char* name, clockid_t cpu_clock_id)

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

@@ -39,37 +39,46 @@ ThreadInfo::unwind(EchionSampler& echion, PyThreadState* tstate, microsecond_t w
 }
 
 // ----------------------------------------------------------------------------
-Result<void>
-ThreadInfo::unwind_tasks(EchionSampler& echion, PyThreadState* tstate, microsecond_t wall_time_us)
+size_t
+ThreadInfo::find_upper_python_stack_size(EchionSampler& echion) const
 {
-    // The size of the "pure Python" stack (before asyncio Frames).
     // Defaults to the full Python stack size (and updated if we find the boundary frame)
     size_t upper_python_stack_size = python_stack.size();
 
     // Check if the Python stack contains the asyncio boundary frame.
     // For regular asyncio, this is "Handle._run" from asyncio/events.py.
     // For uvloop, this is "Runner.run" from asyncio/runners.py (uvloop uses asyncio.Runner internally).
-    // To avoid having to do string comparisons every time we unwind Tasks, we keep track
-    // of the cache key of the boundary frame.
+    // To avoid having to do string comparisons every time we unwind Tasks, we memoize the interned
+    // name and filename of the boundary Frame the first time we identify it. Those are already
+    // StringTable keys, so recognizing the boundary afterwards costs two integer comparisons.
+    //
+    // Do NOT memoize Frame::cache_key here: it mixes in the bytecode offset, so it would only match
+    // the boundary Frame while it sits on the instruction it happened to be on when we first saw it.
+    // A single sample taken while, say, Runner.run is still in its prologue would then hide the
+    // boundary for the rest of the process, and every Task stack would be rendered with the
+    // coroutine frames at the leaf and the whole thread stack appended below them.
 
-    // Note: We use separate cache keys for asyncio and uvloop because switching between them
+    // Note: We memoize asyncio and uvloop separately because switching between them
     // (though unlikely at runtime) would cause incorrect boundary detection otherwise.
-    auto& asyncio_frame_cache_key = echion.asyncio_frame_cache_key();
-    auto& uvloop_frame_cache_key = echion.uvloop_frame_cache_key();
+    auto& asyncio_boundary_frame = echion.asyncio_boundary_frame();
+    auto& uvloop_boundary_frame = echion.uvloop_boundary_frame();
+
+    auto& boundary_frame = using_uvloop ? uvloop_boundary_frame : asyncio_boundary_frame;
+
+    for (size_t i = 0; i < python_stack.size(); i++) {
+        const auto& frame = python_stack[i];
 
-    auto& frame_cache_key = using_uvloop ? uvloop_frame_cache_key : asyncio_frame_cache_key;
+        bool is_boundary_frame = false;
 
-    if (!frame_cache_key) {
-        for (size_t i = 0; i < python_stack.size(); i++) {
-            const auto& frame = python_stack[i];
+        if (boundary_frame) {
+            is_boundary_frame = frame.name == boundary_frame->name && frame.filename == boundary_frame->filename;
+        } else {
             auto maybe_frame_name = echion.string_table().lookup(frame.name);
             if (!maybe_frame_name) {
                 continue;
             }
             const auto& frame_name = maybe_frame_name->get();
 
-            bool is_boundary_frame = false;
-
             if (using_uvloop) {
                 // For uvloop, the boundary frame depends on the Python version:
                 // - Python 3.11+: Runner.run from asyncio/runners.py (uvloop uses asyncio.Runner)
@@ -113,25 +122,26 @@ ThreadInfo::unwind_tasks(EchionSampler& echion, PyThreadState* tstate, microseco
             }
 
             if (is_boundary_frame) {
-                // 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 reuse it to determine
-                // whether we see the boundary Frame in the Python stack.
-                frame_cache_key = frame.cache_key;
-                upper_python_stack_size = python_stack.size() - i;
-                break;
+                boundary_frame = BoundaryFrame{ frame.name, frame.filename };
             }
         }
-    } else {
-        for (size_t i = 0; i < python_stack.size(); i++) {
-            const auto& frame = python_stack[i];
-            if (frame.cache_key == *frame_cache_key) {
-                upper_python_stack_size = python_stack.size() - i;
-                break;
-            }
+
+        if (is_boundary_frame) {
+            upper_python_stack_size = python_stack.size() - i;
+            break;
         }
     }
 
+    return upper_python_stack_size;
+}
+
+// ----------------------------------------------------------------------------
+Result<void>
+ThreadInfo::unwind_tasks(EchionSampler& echion, PyThreadState* tstate, microsecond_t wall_time_us)
+{
+    // The size of the "pure Python" stack (before asyncio Frames).
+    const size_t upper_python_stack_size = find_upper_python_stack_size(echion);
+
     std::vector<TaskInfo::Ref> leaf_tasks;
     std::unordered_set<PyObject*> parent_tasks;
     std::unordered_map<PyObject*, TaskInfo::Ref> waitee_map; // Indexed by task origin

releasenotes/notes/fix-profiling-asyncio-boundary-frame-8f3b41c07d2a95e6.yaml

@@ -0,0 +1,7 @@
+---
+fixes:
+  - |
+    profiling: Fixes an issue where asyncio task samples could be reported with the coroutine
+    frames at the top of the stack and the calling frames duplicated underneath them, making
+    the flame graph unusable for the affected process. This was most likely to happen with
+    ``uvloop``.