fix(profiling): correctly unwind on-CPU Tasks (#15562)

## What does this PR do?

More details and investigation results available in
https://datadoghq.atlassian.net/browse/PROF-13137

Related PRs
* Dependency: #15552
* Dependent: #15579
* Echion PR: P403n1x87/echion#198

This PR fixes the way we sample and capture stacks for on-CPU asyncio
Tasks. This is a first step towards fixing _all_ weird (sub-)stacks we
are seeing (both in Echion and dd-trace-py), for example this one...

```
background_math_function
Task-background_wait
background_wait_function
sleep
```

... where in practice, `background_math_function` was executed as part
of a Task that was completely unrelated to `Task-background_wait` and
should never have appeared in the same stack.

This by itself doesn't fix everything (see **The fix(es, actually)**),
but, as I said, it's a first step.

See [this
document](https://docs.google.com/document/d/10HmKhY6tM5j7ojdk7CG8anWsN7qP7IyeZlj3-zAXudM/edit?pli=1&tab=t.9pzrl0uli4g)
for more explanations around how this works / is supposed to work, the
fix is explained at the bottom here.

## The fix(es, actually)

### Flawed existing logic...

The fix logic is the following:
* To correctly interleave Task Stacks and Task names, we need to
determine the "synchronous Python top-of-stack" that goes on top of the
on-CPU Task (if there is an on-CPU Task).
* This includes the Python stack that eventually led to doing something
like `asyncio.run`, as well as the internal `asyncio` Runner machinery.
* We will have to push this on top of every Task Stack we capture (note
that the on-CPU Stack will already have it by definition).
* To do that, we need to determine which Task is "the on CPU Task" (if
any) and process it first, so that we know how "high" the "upper Python
stack" is and we can re-use it for subsequent Task unwindings.
* To ensure the first Task we unwind is the on-CPU one, we first sort
the Leaf Tasks array by on-CPU-ness
* Note I'm saying _the_ on CPU Task because, as always in Python, there
can only be one thing running at a time.
* (Well, that's the case at least in theory – in practice on-CPU-ness
may change as we sample and so we could end up with more than one on-CPU
Task, but that is fine because it won't change how deep the "upper
stack" is)
* Once we've done this, we know how high the "top of Stack is" because
it is the height of the "pure Python" Stack we get by unwinding the
first (top-most) Frame of the on-CPU Task. We need to remember that
"height" to push that many items from the Python Thread Stack on top of
each Task's Stack.
* One final thing to keep in mind is that we "stack" Stacks (hehe) such
that if `Task-1` is awaiting `Task-2`, the Stack for `Task-2` goes under
that for `Task-1`. We only ever want to add the "pure Python
top-of-stack" once per Stack, so we need to do that _once per leaf Task_
after we've recursively unwound awaiting Tasks for the leaf Task.

### ... and a race condition?

On top of that – because it's not enough – I think I discovered a rare
but very real race condition (which actually was the root cause of that
weird, rare but common-in-CI bug with the stack I showed on top). The
race condition happens when the pure Python Stack we capture involves
some asynchronous Frames, but the associated Task/Coroutine completes or
yields before we get to unwinding it. In that case, we assume in Stack
interleaving that the Task is running, by doing so we include some
Frames from the Python Stack that actually _are not_ in the Stack
generated by unwinding the Task, and because we share the same Python
Stack for all Tasks (we only use it's "upper part", which is the same
for every Task) some Frames from the
previously-running-but-now-paused-or-completed Coroutine are leaked to
other Tasks.

Working around this is not simple; also honestly we should have a
discussion around whether we want to properly work around it (which is
going to be somewhat difficult, imperfect and probably somewhat costly)
or just give up on sampling when we detect something is off.
In the interest of keeping things somewhat simple, this will be handled
in a separate PR. But I described the bug here anyway, for clarity
(temporary PR: KowalskiThomas/echion#43)

### Caveats

I realised after the fact that an assumption we implicitly make (that
_the on-CPU Task – if it exists – is a Leaf Task_) is actually not
generally true. In the case of certain `asyncio` utils, the Parent Task
will actually from time to time execute code/be on-CPU to manage the
Child Tasks it is awaiting on. This is an edge case, and something that
happens rarely, but saying it can never happen is false (and it shows in
Profiles).
We will also need to figure something out for that – it probably isn't
_that_ hard but will require some refactors (probably split the
unwinding in two steps: (1) only determine upper Python Stack (2) unwind
Tasks strictly speaking, starting from Leaf Tasks).


## Testing

The PR adds new tests around on-CPU Tasks and what their Stacks look
like

The PR also updates a lot of existing tests to output their
`DataSummary` as JSON. This is immensely useful to understand what is
actually going on when we don't see the expected Stacks (or see the
unexpected ones) and I do think it'd be worth keeping.

After these changes, the tests that still fail locally and in the CI
(flakily) are unrelated to on-CPU Tasks which, I guess, means _mission
failed successfully_?

Author: KowalskiThomas

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

@@ -182,27 +182,26 @@ class TaskInfo
     PyObject* origin = NULL;
     PyObject* loop = NULL;
 
-    GenInfo::Ptr coro = nullptr;
-
     StringTable::Key name;
+    bool is_on_cpu = false;
+    GenInfo::Ptr coro = nullptr;
 
     // Information to reconstruct the async stack as best as we can
     TaskInfo::Ptr waiter = nullptr;
-    bool is_on_cpu = false;
 
     [[nodiscard]] static Result<TaskInfo::Ptr> create(TaskObj*);
     TaskInfo(PyObject* origin, PyObject* loop, GenInfo::Ptr coro, StringTable::Key name, TaskInfo::Ptr waiter)
       : origin(origin)
       , loop(loop)
-      , coro(std::move(coro))
       , name(name)
+      , is_on_cpu(coro && coro->is_running)
+      , coro(std::move(coro))
       , waiter(std::move(waiter))
-      , is_on_cpu(this->coro && this->coro->is_running)
     {
     }
 
     [[nodiscard]] static Result<TaskInfo::Ptr> current(PyObject*);
-    inline size_t unwind(FrameStack&);
+    inline size_t unwind(FrameStack&, size_t& upper_python_stack_size);
 };
 
 inline std::unordered_map<PyObject*, PyObject*> task_link_map;
@@ -344,7 +343,7 @@ inline std::vector<std::unique_ptr<StackInfo>> current_tasks;
 // ----------------------------------------------------------------------------
 
 inline size_t
-TaskInfo::unwind(FrameStack& stack)
+TaskInfo::unwind(FrameStack& stack, size_t& upper_python_stack_size)
 {
     // TODO: Check for running task.
     std::stack<PyObject*> coro_frames;
@@ -355,14 +354,37 @@ TaskInfo::unwind(FrameStack& stack)
             coro_frames.push(py_coro->frame);
     }
 
-    int count = 0;
+    // Total number of frames added to the Stack
+    size_t count = 0;
 
     // Unwind the coro frames
     while (!coro_frames.empty()) {
         PyObject* frame = coro_frames.top();
         coro_frames.pop();
 
-        count += unwind_frame(frame, stack);
+        auto new_frames = unwind_frame(frame, stack);
+
+        // If we failed to unwind the Frame, stop unwinding the coroutine chain; otherwise we could
+        // end up with Stacks with missing Frames between two coroutines Frames.
+        if (new_frames == 0) {
+            break;
+        }
+
+        // If this is the first Frame being unwound (we have not added any Frames to the Stack yet),
+        // use the number of Frames added to the Stack to determine the size of the upper Python stack.
+        if (count == 0) {
+            // 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;
+
+            // 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.
+            for (size_t i = 0; i < upper_python_stack_size; i++) {
+                stack.pop_back();
+            }
+        }
+
+        count += new_frames;
     }
 
     return count;

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

@@ -264,31 +264,45 @@ ThreadInfo::unwind_tasks()
         }
     }
 
+    // Make sure the on CPU task is first
+    for (size_t i = 0; i < leaf_tasks.size(); i++) {
+        if (leaf_tasks[i].get().is_on_cpu) {
+            if (i > 0) {
+                std::swap(leaf_tasks[i], leaf_tasks[0]);
+            }
+            break;
+        }
+    }
+
+    // The size of the "pure Python" stack (before asyncio Frames), computed later by TaskInfo::unwind
+    size_t upper_python_stack_size = 0;
+    // Unused variable, will be used later by TaskInfo::unwind
+    size_t unused;
+
+    bool on_cpu_task_seen = false;
     for (auto& leaf_task : leaf_tasks) {
+        on_cpu_task_seen = on_cpu_task_seen || leaf_task.get().is_on_cpu;
+
         auto stack_info = std::make_unique<StackInfo>(leaf_task.get().name, leaf_task.get().is_on_cpu);
         auto& stack = stack_info->stack;
+
         for (auto current_task = leaf_task;;) {
             auto& task = current_task.get();
 
-            size_t stack_size = task.unwind(stack);
-
+            // 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);
             if (task.is_on_cpu) {
-                // Undo the stack unwinding
-                // TODO[perf]: not super-efficient :(
-                for (size_t i = 0; i < stack_size; i++)
-                    stack.pop_back();
-
-                // Instead we get part of the thread stack
-                FrameStack temp_stack;
-                size_t nframes = (python_stack.size() > stack_size) ? python_stack.size() - stack_size : 0;
-                for (size_t i = 0; i < nframes; i++) {
-                    auto python_frame = python_stack.front();
-                    temp_stack.push_front(python_frame);
-                    python_stack.pop_front();
-                }
-                while (!temp_stack.empty()) {
-                    stack.push_front(temp_stack.front());
-                    temp_stack.pop_front();
+                // 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
+                // The number of Frames to push is the total number of Frames in the Python stack, from which we
+                // subtract the number of Frames in the "upper Python stack" (asyncio machinery + sync entrypoint)
+                // This gives us [outermost coroutine, ... , innermost coroutine, outermost sync function, ... ,
+                // innermost sync function]
+                size_t frames_to_push =
+                  (python_stack.size() > task_stack_size) ? python_stack.size() - task_stack_size : 0;
+                for (size_t i = 0; i < frames_to_push; i++) {
+                    const auto& python_frame = python_stack[frames_to_push - i - 1];
+                    stack.push_front(python_frame);
                 }
             }
 
@@ -317,8 +331,21 @@ ThreadInfo::unwind_tasks()
         }
 
         // Finish off with the remaining thread stack
-        for (auto p = python_stack.begin(); p != python_stack.end(); p++)
-            stack.push_back(*p);
+        // If we have seen an on-CPU Task, then upper_python_stack_size will be set and will include the sync entry
+        // point and the asyncio machinery Frames. Otherwise, we are in `select` (idle) and we should push all the
+        // Frames.
+
+        // There could be a race condition where relevant partial Python Thread Stack ends up being different from the
+        // one we saw in TaskInfo::unwind. This is extremely unlikely, I believe, but failing to account for it would
+        // cause an underflow, so let's be conservative.
+        size_t start_index = 0;
+        if (on_cpu_task_seen && python_stack.size() >= upper_python_stack_size) {
+            start_index = python_stack.size() - upper_python_stack_size;
+        }
+        for (size_t i = start_index; i < python_stack.size(); i++) {
+            const auto& python_frame = python_stack[i];
+            stack.push_back(python_frame);
+        }
 
         current_tasks.push_back(std::move(stack_info));
     }

releasenotes/notes/profiling-fix-on-cpu-tasks-unwinding-0ac6efcad8150f1c.yaml

@@ -0,0 +1,3 @@
+fixes:
+  - |
+    profiling: This fix makes stack sampling more accurate for on-CPU asyncio Tasks.