Exploring the async/await State Machine – Conceptual Implementation

Intro

By now, you should have a very strong understanding of the async/await State Machine and its’ process flow.

It’s time to implement it.

Note that the previous article is a prerequisite for this one if you don’t already have a good conceptual understanding of the async/await State Machine.

The implementation in this post is still at a little higher abstraction level than what the compiler produces.

The reason is that I want to focus on semantics first. The actual code has some trickeries that are mostly related to performance optimizations. I find those optimizations pretty intriguing, but they hurt readability.

Once you’re comfortable with this somewhat simplified version, I’ll go the extra mile and discuss the actual compiler-generated code as closely as possible.

Still, don’t assume the code here is just a toy example. It is very similar to the actual implementation. I will mention some of the differences here and discuss them in detail in the next article.

This article’s content is influenced by the async/await material in the book C# in Depth by Jon Skeet and the blog post Dissecting the async methods in C# by Sergey Tepliakov.

For additional in-depth materials on async/await in C#, you can check this Pluralsight course.

Starter Code

I’ll use the same starter code as in the previos post:

public static async Task<int> MyAsyncMethod(int firstDelay, int secondDelay)
{
    Console.WriteLine("Before first await.");
    
    await Task.Delay(firstDelay);
    
    Console.WriteLine("Before second await.");
    
    await Task.Delay(secondDelay);
    
    Console.WriteLine("Done.");

    return 42;
}

Again, I find this small method appropriate because it is simple enough to understand and still allows us to focus on the crucial points of the async/await machinery.

State Machine Workflow Diagram

I include the State Machine diagram once again for convenience.

Spend a minute to make sure you’re comfortable with the workflow. After that, the code itself will be very intuitive to absorb.

Implementation

Below is the full implementation of the State Machine from Figure 1.

You can literally inspect MoveNext() side by side with the workflow diagram. This will help you appreciate how identical they are semantically.

One exercise I would recommend here is to follow the MoveNext() body and acknowledge how it will run synchronously if all the awaiters are completed. If this doesn’t make sense, visit the article on the Awaitable Pattern.

Although this implementation should be relatively easy for you to comprehend, there are a few points I’d like to stress on.

That’s coming in the next sections.

Exception Handling

MoveNext() would almost(*) never throw an exception directly. Even if an error occurs during the code execution, it will be propagated via the resulting task.

(*) Except for “special” cases like StackOverflowException.

That’s why the MoveNext() body is wrapped in a try-catch blog. If an exception occurs – it’s set on the task builder.

public void MoveNext()
{
    try
    {
        // …
    }
    catch (Exception ex)
    {
        _taskBuilder.SetException(ex);
    }
}

State Machine Instance Fields

The instance fields of the State Machine allow it to preserve its’ context across continuations.

Let’s look at them once again:

private class MyAsyncMethodStateMachine
{
    private readonly TaskCompletionSource<int> _taskBuilder = new TaskCompletionSource<int>();
    
    private readonly int _firstDelay;
    private readonly int _secondDelay;
    
    private int _state;
    private TaskAwaiter _awaiter;

    // …
}

In the following sections, I will describe each instance field and its’ purpose.

Original Async Method Parameters

The compiler would create a separate field for every input parameter of the original async method. In our example, for MyAsyncMethod:

public static async Task<int> MyAsyncMethod(int firstDelay, int secondDelay)

We get a couple of fields for firstDelay and secondDelay:

private class MyAsyncMethodStateMachine
{
    // …

    private readonly int _firstDelay;
    private readonly int _secondDelay;

    // …
}

Awaiters

Remember that we need to get the awaiter for every async call via the GetAwaiter() method on the awaitable:

_awaiter = Task.Delay(_firstDelay).GetAwaiter();
if (!_awaiter.IsCompleted)
{
    _awaiter.OnCompleted(MoveNext);
    return;
}

We store the awaiter in an instance field in order to retrieve the result when the State Machine resumes:

if (_state == 1)
{
    _awaiter.GetResult();
    // …
}

In our example, for both of the Task.Delay async calls, GetAwaiter() returns the same awaiter type, which is TaskAwaiter. This is why the State Machine can re-use the same TaskAwaiter field for both async calls:

private class MyAsyncMethodStateMachine
{
    // …

    private TaskAwaiter _awaiter;

    // …
}

However, if one of the methods was returning some other awaiter type, like TaskAwaiter<T>, the compiler would have created a separate field for it.

For example, if in our async method, we had a call like this:

int res = await SomeAsyncMethodReturningInt();

we would get a TaskAwaiter<int> field in the State Machine class to store the awaiter for this call.

State Variable

Of course, we also need a field for the state:

private class MyAsyncMethodStateMachine
{
    // …
    
    private int _state;

    // …
}

There is nothing special about it. It just stores the current state of the State Machine as an integer value.

Task Builder

This is the part I simplified the most by directly using TaskCompletionSource<int> It allows me to produce and manage the resulting task very conveniently.

private class MyAsyncMethodStateMachine
{
    private readonly TaskCompletionSource<int> _taskBuilder = new TaskCompletionSource<int>();

    // …

    public Task<int> ResultTask => _taskBuilder.Task;
}

As you can inspect, the output task’s lifecycle is controlled in MoveNext(). If MoveNext() completes successfully, the task gets assigned with the final result. If MoveNext() fails – it’s used to propagate the exception.

Logically, the compiler-generated version of the State Machine does the same. However, it needs to do a lot more housekeeping – things like boxing the State Machine on the first pausing, capturing the current ExecutionContext, and so on.

For all of this, rather than TaskCompletionSource<int>, it uses a special class called AsyncTaskMethodBuilder<T> that is meant for compiler use only.

You will learn a lot more about AsyncTaskMethodBuilder<T> in the next part. From a high-level perspective, you can just assume it has similar responsibilities as TaskCompletionSource<int>.

Summary

In this article, I presented a manual implementation of the async/await State Machine.

Although a little simplified, it’s fully functional and pretty close to the real compiler-generated code.

Now that you understand how this works from a higher abstraction level, I can show you some of the most intriguing details of the actual implementation.

Stay tuned, and thanks for reading!

Resources

  1. C# in Depth: Fourth Edition
  2. Dissecting the async methods in C#

Site Footer

Subscribe To My Newsletter

Email address