Optimize async state machines

[timcassell]

Motivation

Consider ValueTask<TResult> that stores the TResult _result in the struct, so that synchronous value retrieval can occur without allocating on the heap. ValueTaskAwaiter<TResult> has a ValueTask<TResult> field, so it can do the same thing. AsyncValueTaskMethodBuilder<TResult> also stores a TResult _result field to be able to construct a completed ValueTask<TResult> without allocating. Currently, async state machines store both the entire async method builder, and the entire awaiter, whether it completes synchronously or asynchronously. If the ValueTask<TResult> is actually pending, that's wasted space being stored in the state machine. When the async continuation occurs, the result is retrieved from the backing heap object, not from the result stored in the builder or the awaiter.

The proposed optimization improvements come in 2 parts; awaiters and async method builders.

Optimized Awaiters

In order to prevent the result from being stored in the state machine which will never be used, the compiler may recognize a new pattern on awaiters:

public struct MyValueTaskAwaiter<T>
{
    public MyValueTaskAsyncResultGetter GetAsyncResultGetter();
    public static T GetResult(MyValueTaskAsyncResultGetter asyncResultGetter);
}

When awaiter.IsCompleted returns false, the compiler may call the new awaiter.GetAsyncResultGetter() and store the result in the state machine instead of the awaiter. When the state machine is continued, it calls the new static function MyValueTaskAwaiter<T>.GetResult(asyncResultGetter), passing in the stored getter.

Drawbacks

None

Optimized Async Method Builders

We can prevent the async method builder from being stored in the state machine altogether, and pass it into MoveNext instead.

New interface:

namespace System.Runtime.CompilerServices;

public interface IAsyncStateMachine<TBuilder>
{
    void MoveNext(TBuilder asyncMethodBuilder);
}

Notice how the new interface no longer contains a SetStateMachine method that IAsyncStateMachine has. That's because it's no longer necessary. Modern async method builders store the generic state machine directly without boxing.

New builder shape:

public struct NewAsyncValueTaskMethodBuilder<T>
{
    public static void CreateAndStart<TStateMachine>(ref TStateMachine stateMachine) where TStateMachine : IAsyncStateMachine<NewAsyncValueTaskMethodBuilder<T>>;
    public ValueTask<T> GetCompletedTask(T result);
    public ValueTask<T> GetTaskFromException(Exception exception);
    public ValueTask<T> GetPendingTask<TStateMachine>(ref TStateMachine stateMachine) where TStateMachine : IAsyncStateMachine<NewAsyncValueTaskMethodBuilder<T>>;;
    public void AwaitOnCompleted<TAwaiter>(ref TAwaiter awaiter) where TAwaiter : INotifyCompletion;
    public void AwaitUnsafeOnCompleted<TAwaiter>(ref TAwaiter awaiter) where TAwaiter : ICriticalNotifyCompletion;
    public void SetException(Exception exception);
    public void SetResult(T result);
}

The compiler generates 2 state machines for this builder instead of 1. The first state machine is passed to CreateAndStart, and it controls whether to call GetCompletedTask, GetTaskFromException, or GetPendingTask + Await(Unsafe)OnCompleted. The second state machine, which is passed to GetPendingTask, is almost exactly the same as existing state machines, and it controls the async state, passing awaiters to Await(Unsafe)OnCompleted and calling SetException or SetResult when it's complete.

Drawbacks

• Existing async method builders need to continue to work for older compilers, and the AsyncMethodBuilderAttribute does not allow multiple.

One option to get around this is a new attribute to tell the compiler to use the new async method builder type, which older compilers will ignore. Or possibly the existing attribute could be extended to accept 2 types, as long as older compilers can still read it as if it only has 1 type.

A second option is the compiler can infer the new builder type if the old builder type contains a CreateAndStart function:

public struct AsyncValueTaskMethodBuilder<T>
{
    public static void CreateAndStart<TStateMachine>(ref TStateMachine stateMachine)
        where TStateMachine : IAsyncStateMachine<NewAsyncValueTaskMethodBuilder<T>>;
}

• It requires a new interface, which means a change to the BCL. The compiler could recognize an interface with the same shape in the same namespace for polyfilling old runtimes.

An alternative solution would be to just store the new builder type in the state machine to avoid the new interface, but it's less optimal.

Benchmarks

Getting async results for object and int:

Method	Pending	Mean	Ratio	Code Size	Allocated	Alloc Ratio
Current	False	135.4 ns	1.00	573 B	-	NA
OptimizedAwait	False	129.3 ns	0.96	557 B	-	NA
OptimizedBuilder	False	118.3 ns	0.87	468 B	-	NA
OptimizedBoth	False	119.6 ns	0.88	468 B	-	NA
Current	True	813.1 ns	1.00	1,018 B	576 B	1.00
OptimizedAwait	True	749.1 ns	0.92	1,002 B	536 B	0.93
OptimizedBuilder	True	776.0 ns	0.95	913 B	536 B	0.93
OptimizedBoth	True	745.9 ns	0.92	913 B	496 B	0.86

Getting async results for Matrix3x2 and Matrix4x4:

Method	Pending	Mean	Ratio	Code Size	Allocated	Alloc Ratio
Current	False	210.9 ns	1.00	4,091 B	-	NA
OptimizedAwait	False	218.0 ns	1.03	4,476 B	-	NA
OptimizedBuilder	False	150.8 ns	0.72	2,848 B	-	NA
OptimizedBoth	False	141.2 ns	0.67	2,991 B	-	NA
Current	True	849.6 ns	1.00	8,511 B	808 B	1.00
OptimizedAwait	True	747.3 ns	0.88	8,821 B	688 B	0.85
OptimizedBuilder	True	765.0 ns	0.90	6,768 B	696 B	0.86
OptimizedBoth	True	667.1 ns	0.79	6,881 B	576 B	0.71

You can clearly see that the larger the type we are retrieving, the larger the gains are.

Benchmark code

Additional Notes

I am aware that the runtime is working on adding runtime-async feature. That feature is only supporting Task, Task<T>, ValueTask, and ValueTask<T>. Custom task-like types will not gain from that feature (at least not for now, don't know if there will be future plans to expand it). Additionally, the new runtime-async feature will interop with existing awaiters, so the awaiter improvements here will still be an improvement for the new feature, even if the async method builders won't be touched.