AsyncStackTrace class that can walk continuations

[AArnott]

Background and Motivation

The StackTrace class can be instantiated at any time to get the CLR to talk the thread's callstack and capture the result in an object that can be stored or shared. Similar to stacks captured by thrown exceptions, this can be very useful--useful enough for the API in the first place, so I'll forebear explaining all the reasons why it's good.

But its usefulness has substantially degraded in the async world, where the literal callstack is often not helpful. With stack traces in exceptions it's better (although very hard to read as captured by #4996) for the final rethrown exception because as the exception is thrown and rethrown down the async continuation chain, the logical "return stack" is reassembled so the stack trace contains the whole picture.

The VS debugger has gotten better. Its Callstack tool window actually does some work to walk not only the callstack but also the async continuation stack/tree to suggest who the async callers are. The Parallel Stacks tool window has a new Tasks view that is also pretty good at representing this information.

But when just using new StackTrace() none of this is available. So when I try to use this (e.g. under special debug/instrumented mode for an application in order to establish ownership of an object) it can be useless.

What I propose is an AsyncStackTrace class that understands common continuations/awaiter patterns just like the VS debugger does to reconstruct the async return stack.

See also: #14434

Proposed API

I'm not sure yet what would make sense to expose. We might want to allow custom awaiters to register themselves as helpers so that async stacks can be stitched together by this class as well as the built-in recognized ones. Hey, if we do this well, maybe the debugger can even use this to help with stitching to do a more complete job.

namespace System.Diagnostics
{
+    public class AsyncStackTrace {
+        public AsyncStackTrace() { }
+        public static void RegisterCustomAwaiter<TAwaiter>(Func<TAwaiter, Action> awaiterToContinuation);
+        /// <summary>Gets the async and sync frames in the call/return stack.</summary>
+        public ReadOnlyMemory<StackFrame> Frames { get; }
+    }
}

Usage Examples

Consumption of the API should be straightforward:

class Foo
{
    AsyncStackTrace owner;

    public Foo()
    {
        this.owner = new AsyncStackTrace();
    }
}

Later, when debugging and finding a Foo instance somewhere, it's trivial to see what stack (including logical async frames) created the object by inspecting the AsyncStackTrace field.

Custom awaiters should define a static constructor to register themselves. This ensures that the right owner registers each custom awaiter and that it happens exactly once:

public struct MyCustomAwaiter : INotifyCompletion
{
    private Action continuation;
    static MyCustomAwaiter()
    {
        AsyncStackTrace.RegisterCustomAwaiter<MyCustomAwaiter>(self => self.continuation);
    }

    public bool IsCompleted => false;
    public void OnCompleted(Action continuation) => this.continuation = continuation;
    public void GetResult() { }
}

Alternative Designs

Awaiters may have OnCompleted(Action) called multiple times, so we may want to offer an API that acknowledges that. And maybe the AsyncStackTrace class should represent a tree of continuation frames instead of assuming a stack representation as StackTrace does.

Risks

TBD after some initial discussion