fix: correctly unwind on-CPU tasks

[KowalskiThomas]

What does this PR do?

dd-trace-py PR:

• fix(profiling): correctly unwind on-CPU Tasks DataDog/dd-trace-py#15562

Note this PR depends on

• refactor: change GenInfo::is_running #202
• refactor: make StringTable::lookup const #209

Note the diff says 22 23 Files changed but it's mostly tests don't worry :)

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 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#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?