Async method builders and where the allocations go

This post is part of the series 'Async/await in C#'. Be sure to check out the rest of the blog posts of the series!

• Why async exists: the cost of blocking a thread
• There is no thread: OS async I/O and the .NET thread pool
• The Task lifecycle and TaskCompletionSource
• The awaitable pattern: writing your own awaitable
• How the C# compiler rewrites an async method
• Writing your own Task from scratch
• Writing your own CancellationToken from scratch
• Async method builders and where the allocations go (this post)

The post on the compiler rewrite pulled an async method apart and found a generated struct with a MoveNext method. One field in that struct was left unexplained:

public AsyncTaskMethodBuilder<int> <>t__builder;

The builder is the other half of the machinery. The state machine knows where it is; the builder knows what to hand back to the caller and how to get running again. It is also the piece you can replace, and the piece that decides what your async method allocates.

#What the builder is asked to do

The generated code only ever calls five members on it:

// in the stub method
stateMachine.<>t__builder = AsyncTaskMethodBuilder<int>.Create();
stateMachine.<>t__builder.Start(ref stateMachine);
return stateMachine.<>t__builder.Task;

// in MoveNext
<>t__builder.AwaitUnsafeOnCompleted(ref awaiter, ref this);
<>t__builder.SetResult(result);
<>t__builder.SetException(exception);

That is the entire contract, and each member has one job:

The return type of your method selects the builder: Task gets AsyncTaskMethodBuilder, Task<T> gets AsyncTaskMethodBuilder<T>, ValueTask<T> gets AsyncValueTaskMethodBuilder<T>, and void gets AsyncVoidMethodBuilder.

#Where the allocation happens

Here is the part worth internalising. Start runs MoveNext synchronously, on the caller's stack, with the state machine still a struct sitting in the caller's frame. If the method runs to completion without ever suspending, that struct is simply discarded. Nothing was ever put on the heap for it.

The moment the method does suspend, the stack frame is about to disappear, so the state machine has to be preserved. That is what AwaitUnsafeOnCompleted does on its first call: it boxes the state machine onto the heap.

And it does something clever while it is there. A naive implementation would need three heap objects: a box for the state machine, a Task to return, and an Action delegate for the continuation. The BCL uses a single object, AsyncStateMachineBox<TStateMachine>, which is the Task<T>, holds the state machine, and is the work item queued to the thread pool. One allocation instead of three.

#The cost, measured

Two groups, because the two paths behave completely differently. First, methods that complete without ever suspending:

Task<int> costs 72 bytes even though the method never suspended, because the builder still has to produce a Task object to return. ValueTask<int> costs nothing at all: on the synchronous path it just wraps the value in a struct. That is the entire reason ValueTask exists.

Now the same three, with a method that always suspends:

The differences have almost vanished. Once you suspend, you pay for boxing the state machine, queueing a continuation and getting scheduled back, and that dwarfs everything else. ValueTask is slightly worse here, because it has to allocate a backing object anyway and the struct wrapper is pure overhead on top.

The practical rule that falls out of these two tables: ValueTask is a win when the method usually completes synchronously, and worth nothing when it usually does not.

#Replacing the builder

[AsyncMethodBuilder] lets you override the compiler's choice, on a single method or on a whole return type:

[AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))]
public static async ValueTask<int> SuspendsAsync()
{
    await Task.Yield();
    return 42;
}

PoolingAsyncValueTaskMethodBuilder<T> keeps a pool of state machine boxes and reuses them instead of allocating one per suspension. In the numbers above it saves a couple of bytes and a few percent, which is a fair summary of what it usually buys: something, but not much, and only on the suspending path.

It also comes with rules, and they are strict. The sample's tests check each one:

// It may be awaited exactly once.
var valueTask = PooledOperations.WaitForAsync(gate.Task);
Assert.Equal(42, await valueTask);
await Assert.ThrowsAsync<InvalidOperationException>(async () => await valueTask);

// It may not be blocked on. This does not wait, it throws.
Assert.Throws<InvalidOperationException>(() => valueTask.GetAwaiter().GetResult());

// AsTask is the escape hatch when you need to await more than once.
var task = PooledOperations.WaitForAsync(gate.Task).AsTask();
Assert.Equal(42, await task);
Assert.Equal(42, await task);

The reason is the pooling itself: the box is handed back to the pool as soon as its result has been read, and a version token guards against a second read. Touch it twice and you are reading state that now belongs to somebody else's operation, which is exactly the kind of bug you never want. The runtime throws rather than let that happen.

Given the size of the win and the size of the footgun, the honest advice is: measure first, and only reach for the pooling builder on a method you have proven is both hot and suspending.

#Why the cached tasks matter here

AsyncTaskMethodBuilder has one more trick worth knowing, because it explains a number from the task lifecycle post.

When SetResult is called on a builder that never suspended, it does not always allocate. For Task it uses Task.CompletedTask. For Task<bool> it uses cached instances of true and false. For Task<int> it uses a small cache of low values. Only outside those ranges does it construct a new Task<T>.

That is why an async Task method returning early is cheaper than an async Task<Guid> method returning early, and why the Task.FromResult benchmark there showed true costing nothing while 42 cost 72 bytes.

#What to take away

• The builder owns the returned task and drives MoveNext. The state machine only tracks position.
• Start runs the method synchronously on the caller's thread. Nothing reaches the heap until the method actually suspends.
• On the first suspension, the state machine is boxed into a single object that is simultaneously the box, the Task and the queued work item.
• A synchronously completing async Task<T> still costs 72 bytes. The same method returning ValueTask<T> costs nothing.
• Once a method suspends, the return type barely matters. Scheduling dominates.
• [AsyncMethodBuilder] can swap in PoolingAsyncValueTaskMethodBuilder, at the price of strict single-await rules that throw when broken.

The next post moves from what is allocated to what is carried: the ExecutionContext, and what actually flows across an await.

#Additional resources

• Code sample
• Benchmark

Do you have a question or a suggestion about this post? Contact me!