Extending the async methods in C#

Sergey Tepliakov

The async series

• Dissecting the async methods in C#.
• Extending the async methods in C#.
• The performance characteristics of the async methods in C#.
• One user scenario to rule them all.

In the previous blog post we discussed how the C# compiler transforms asynchronous methods. In this post, we’ll focus on extensibility points the C# compiler provides for customizing the behavior of async methods.

There are 3 ways how you can control the async method’s machinery:

1. Provide your own async method builder in the System.Runtime.CompilerServices namespace.
2. Use custom task awaiters.
3. Define your own task-like types.

Custom types fromm System.Runtime.CompilerServices namespace

As we know from the previous post, the C# compiler transforms async methods into a generated state machine that relies on some predefined types. But the C# compiler does not expect that these well-known types come from a specific assembly. For instance, you can provide your own implementation of AsyncVoidMethodBuilder in your project and the C# compiler will “bind” async machinery to your custom type.

This is a good way to explore what the underlying transformations are and to see what’s happening at runtime:

namespace System.Runtime.CompilerServices
{
    // AsyncVoidMethodBuilder.cs in your project
    public class AsyncVoidMethodBuilder
    {
        public AsyncVoidMethodBuilder()
            => Console.WriteLine(".ctor");
 
        public static AsyncVoidMethodBuilder Create()
            => new AsyncVoidMethodBuilder();
 
        public void SetResult() => Console.WriteLine("SetResult");
 
        public void Start<TStateMachine>(ref TStateMachine stateMachine)
            where TStateMachine : IAsyncStateMachine
        {
            Console.WriteLine("Start");
            stateMachine.MoveNext();
        }
 
        // AwaitOnCompleted, AwaitUnsafeOnCompleted, SetException 
        // and SetStateMachine are empty
    }   
}

Now, every async method in your project will use the custom version of AsyncVoidMethodBuilder. We can test this with a simple async method:

[Test]
public void RunAsyncVoid()
{
    Console.WriteLine("Before VoidAsync");
    VoidAsync();
    Console.WriteLine("After VoidAsync");
 
    async void VoidAsync() { }
}

The output of this test is:

Before VoidAsync
.ctor
Start
SetResult
After VoidAsync

You can implement UnsafeAwaitOnComplete method to test the behavior of an async method with await clause that returns non-completed task as well. The full example can be found at github.

To change the behavior for async Task and async Task<T> methods you should provide your own version of AsyncTaskMethodBuilder and AsyncTaskMethodBuilder<T>

The full example with these types can be found at my github project called EduAsync (*) in AsyncTaskBuilder.cs and AsyncTaskMethodBuilderOfT.cs respectively.

(*) Thanks Jon Skeet for inspiration for this project. This is a really good way to learn async machinery deeper.

Custom awaiters

The previous example is “hacky” and not suitable for production. We can learn the async machinery that way, but you definitely don’t want to see such a code in your codebase. The C# language authors built-in proper extensibility points into the compiler that allows to “await” different types in async methods.

In order for a type to be “awaitable” (i.e. to be valid in the context of an await expression) the type should follow a special pattern:

• Compiler should be able to find an instance or an extension method called GetAwaiter. The return type of this method should follow certain requirements:
• The type should implement INotifyCompletion interface.
• The type should have bool IsCompleted {get;} property and T GetResult() method.

This means that we can easily make Lazy<T> awaitable:

public struct LazyAwaiter<T> : INotifyCompletion
{
    private readonly Lazy<T> _lazy;
 
    public LazyAwaiter(Lazy<T> lazy) => _lazy = lazy;
 
    public T GetResult() => _lazy.Value;
 
    public bool IsCompleted => true;
 
    public void OnCompleted(Action continuation) { }
}
 
public static class LazyAwaiterExtensions
{
    public static LazyAwaiter<T> GetAwaiter<T>(this Lazy<T> lazy)
    {
        return new LazyAwaiter<T>(lazy);
    }
}

public static async Task Foo()
{
    var lazy = new Lazy<int>(() => 42);
    var result = await lazy;
    Console.WriteLine(result);
}

The example could be looked too contrived but this extensibility point is actually very helpful and is used in the wild. For instance, Reactive Extensions for .NET provides a custom awaiter for awaiting IObservable<T> instances in async methods. The BCL itself has YieldAwaitable used by Task.Yieldand HopToThreadPoolAwaitable:

public struct HopToThreadPoolAwaitable : INotifyCompletion
{
    public HopToThreadPoolAwaitable GetAwaiter() => this;
    public bool IsCompleted => false;
 
    public void OnCompleted(Action continuation) => Task.Run(continuation);
    public void GetResult() { }
}

The following unit test demonstrates the last awaiter in action:

[Test]
public async Task Test()
{
    var testThreadId = Thread.CurrentThread.ManagedThreadId;
    await Sample();
 
    async Task Sample()
    {
        Assert.AreEqual(Thread.CurrentThread.ManagedThreadId, testThreadId);
 
        await default(HopToThreadPoolAwaitable);
        Assert.AreNotEqual(Thread.CurrentThread.ManagedThreadId, testThreadId);
    }
}

The first part of any “async” method (before the first await statement) runs synchronously. In most cases, this is fine and desirable for eager argument validation, but sometimes we would like to make sure that the method body would not block the caller’s thread. HopToThreadPoolAwaitable makes sure that the rest of the method is executed in the thread pool thread rather than in the caller’s thread.

Task-like types

Custom awaiters were available from the very first version of the compiler that supported async/await (i.e. from C# 5). This extensibility point is very useful but limited because all the async methods should’ve returned void, Task or Task<T>. Starting from C# 7.2 the compiler support task-like types.

Task-like type is a class or a struct with an associated builder type identified by AsyncMethodBuilderAttribute (**). To make the task-like type useful it should be awaitable in a way we describe in the previous section. Basically, task-like types combine the first two extensibility points described before by making the first way officially supported one.

(**) Today you have to define this attribute yourself. The example can be found at my github repo.

Here is a simple example of a custom task-like type defined as a struct:

public sealed class TaskLikeMethodBuilder
{
    public TaskLikeMethodBuilder()
        => Console.WriteLine(".ctor");
 
    public static TaskLikeMethodBuilder Create()
        => new TaskLikeMethodBuilder();
 
    public void SetResult() => Console.WriteLine("SetResult");
 
    public void Start<TStateMachine>(ref TStateMachine stateMachine)
        where TStateMachine : IAsyncStateMachine
    {
        Console.WriteLine("Start");
        stateMachine.MoveNext();
    }
 
    public TaskLike Task => default(TaskLike);
 
    // AwaitOnCompleted, AwaitUnsafeOnCompleted, SetException 
    // and SetStateMachine are empty

}
 
[System.Runtime.CompilerServices.AsyncMethodBuilder(typeof(TaskLikeMethodBuilder))]
public struct TaskLike
{
    public TaskLikeAwaiter GetAwaiter() => default(TaskLikeAwaiter);
}
 
public struct TaskLikeAwaiter : INotifyCompletion
{
    public void GetResult() { }
 
    public bool IsCompleted => true;
 
    public void OnCompleted(Action continuation) { }
}

And now we can define a method that returns TaskLike type and even use different task-like types in the method body:

public async TaskLike FooAsync()
{
    await Task.Yield();
    await default(TaskLike);
}

The main reason for having task-like types is an ability to reduce the overhead of async operations. Every async operation that returns Task<T>allocates at least one object in the managed heap – the task itself. This is perfectly fine for a vast majority of applications especially when they deal with coarse-grained async operations. But this is not the case for infrastructure-level code that could span thousands of small tasks per second. For such kind of scenarios reducing one allocation per call could reasonably increase performance.

Async pattern extensibility 101

• The C# compiler provides various ways for extending async methods.
• You can change the behavior for existing Task-based async methods by providing your own version of AsyncTaskMethodBuilder type.
• You can make a type “awaitable” by implementing “awaitable pattern”.
• Starting from C# 7 you can build your own task-like types.

Additional references

• Dissecting the async methods in C#
• EduAsync repo on github.
• Task-like types

Next time we’ll discuss the perf characteristics of async methods and will see how the newest task-like value type called System.ValueTask affects performance.

Sergey Tepliakov Senior Software Engineer, Tools for Software Engineers

Follow