How To Create Custom Async Return Types In Unity

Introduction

In the previous post Await Anything In Unity With AwaitableCompletionSource And Custom Awaiters I explained how to use the GetAwaiter method to enable the await keyword syntax, allowing code to be written more cleanly instead of relying on events and if statements.

While that post covered how to make various objects compatible with the await keyword, it didn’t show how to create types that can serve as return types in async methods.

C# provides the Task and Task<TResult> classes, which the compiler supports for use in async methods. It also provides a mechanism to create custom types that can act as async return types. Unity utilizes this mechanism with its Awaitable and Awaitable<T0> classes, allowing task-like behavior that integrates smoothly with Unity’s game loop.

This post takes a more theoretical approach since there aren’t many common use cases for creating custom async-returnable classes (though I’ll provide an example). However, understanding this concept may be valuable, whether for a scenario I haven’t considered or to gain insight into the internals of Unity’s Awaitable implementation.

The Advantage of Async/Await

The real advantage of the async/await pattern isn’t merely in writing asynchronous code, this was possible before async/await and even before Task. Nor is it solely about making asynchronous code readable in a synchronous style; although async/await improves readability, similar behavior could be achieved with Task’s ContinueWith method.

The main benefit of async/await lies in the ability to create an asynchronous operation, start it, and decide on its continuation later, in a different part of the program. Unlike the ContinueWith method, async/await allows an operation to be passed around as a Task and awaited at any convenient point. After the await expression, the code that follows serves as the continuation, but this continuation does not need to be defined at the task’s creation.

Async/await accomplishes this by generating a state machine. The first state represents the asynchronous operation, and once completed, the next state is everything after the await expression. If a method has multiple await expressions, the state machine will have more states; one for each await and one for everything after the last await.

The compiler can create a state machine for custom types using a builder type. A builder type is a class or struct marked with the AsyncMethodBuilder attribute, which tells the compiler how to construct the type that will serve as the async method’s return type.

The Builder Type

The builder type has certain methods, that are used by the compiler to generate the code for the state machine of an async method. Here are the methods that are being expected from the compiler for a hypothetical MyAwaiterBuilder<T> type that is responsible for building a MyAwaiter type instead of Unity’s Awaitable type:

class MyAwaiterBuilder<T>
{
    public static MyAwaiterBuilder<T> Create();

    public void Start<TStateMachine>(ref TStateMachine stateMachine)
        where TStateMachine : IAsyncStateMachine;

    public void SetStateMachine(IAsyncStateMachine stateMachine);
    public void SetException(Exception exception);
    public void SetResult(T result);

    public void AwaitOnCompleted<TAwaiter, TStateMachine>(
        ref TAwaiter awaiter, ref TStateMachine stateMachine)
        where TAwaiter : INotifyCompletion
        where TStateMachine : IAsyncStateMachine;
    public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(
        ref TAwaiter awaiter, ref TStateMachine stateMachine)
        where TAwaiter : ICriticalNotifyCompletion
        where TStateMachine : IAsyncStateMachine;

    public MyAwaiter<T> Task { get; }
}

The Create method exists to create an instance of the MyAwaiterBuilder type.

The Start method, is used to connect our builder type with the state machine generated by the compiler. On this method, we call the MoveNext() method of the state machine for the first time to advance the state machine. After that, each subsequent call of the MoveNext() method, advances the state machine to the next stage.

If the state machine is implemented as a struct, then the SetStateMachine is used because it accepts a boxed instance of the state machine that the builder can cache instead of continuously boxing the struct.

As the state machine advances stages, through the MoveNext() method, will eventually either complete successfully or throw an exception. This is where the SetResult and SetException methods come in. (For those wondering about cancellation, cancellation is a special exception)

When the state machine reaches a part of code that has an await expression, then the GetAwaiter method of that type will be called, as we saw in the previous post Await Anything In Unity With AwaitableCompletionSource And Custom Awaiters. The rest are the responsibility of the Awaiter; if its IsCompleted property is false, the state machine executes the AwaitUnsafeOnCompleted or the AwaitOnCompleted methods, depending on if the awaitable type implements the ICriticalNotifyCompletion or INotifyCompletion interfaces.

Both of those methods, are responsible for advancing the state machine by calling its MoveNext method.

All of this may seem complex for creating a custom async return type, especially considering that our implementation also needs to work within Unity’s game loop. Fortunately, we don’t have to handle all of this on our own. Instead, we can leverage Unity’s implementation of the AsyncMethodBuilder type.

This type also uses a state machine class, which is a nested type within AsyncMethodBuilder and is designed to be poolable for performance optimization.

Leveraging The AwaitableTaskMethodBuilder

Instead of implementing all those methods ourselves, we’ll create our own builder type, where each method simply calls the corresponding method in AsyncMethodBuilder.

Let’s assume we want to create extension methods for the Task class that complete at specific points in Unity’s game loop. For example, imagine we want a Task to execute and then wait until the end of the current frame in the game loop before continuing.

Our extension method, may look something like this:

public static async AwaitableTask<T> OnEndOfFrame<T>(this Task<T> task, CancellationToken cancellationToken = default)
{
    var result = await task;
    await UnityEngine.Awaitable.EndOfFrameAsync(cancellationToken);
    return result;
}

The AwaitableTask<T> is the type we want to create.

Would it work to return an Awaitable<T>? Technically, yes, the code wouldn’t throw an error. However, this would be poor programming practice.

Unity’s Awaitable respects Unity’s game loop, and as mentioned, one of the benefits of async/await is the flexibility to pass types around and await them as needed without determining their continuation upon creation. If this method were to return an Awaitable, it would alter the expected behavior of the type.

With a variable of type Awaitable (let’s call it MyAwaitable), there would no longer be a clear expectation of its behavior. Specifically, it would be unclear whether MyAwaitable would stop executing when the game stops in the editor. Unity’s Awaitable is expected to halt any asynchronous code execution when the game stops in the editor, but this may not be the case if its behavior has been modified. For example:

async void Start()
{
    Task<int> myTask = TestExample();
    MyAwaitable = myTask.OnEndOfFrame();
          
    var result = await MyAwaitable;

    Debug.Log(result);
}

private async Task<int> TestExample()
{
    await Task.Delay(10000);
    Debug.Log("Task Completed");
    return 42;
}

If we stop the game before ten seconds have passed, the task continues running, and eventually, the Task Completed message will appear in the console, even if the game has already stopped. Worse still, the number 42 will be logged to the console at a seemingly random time; when the game is restarted or at some point while using the editor.

If MyAwaitable is of type Awaitable<int> and is passed to us by a library, framework, or even a colleague’s code, we may be unpleasantly surprised, as we would expect it to behave like an Awaitable and not a Task. On the other hand, we don’t expect a Task type to interact with Unity’s game loop. We know that it won’t stop executing code asynchronously when the game stops in the editor, but we also don’t expect it to recognize specific moments in Unity’s game loop, such as the end of the current frame.

Both of these characteristics could be true of a new type we implement, which we’ll call AwaitableTask<T>. This type, along with AwaitableTask, can be the return type of an async method that executes a Task asynchronously, and like any other Task, we are responsible for stopping it. However, it also has knowledge of Unity’s game loop stages.

Implementation Of AwaitableTask And AwaitableTaskMethodBuilder

Here, I will provide the implementation for AwaitableTask<T> and AwaitableTaskMethodBuilder<T>. The non-generic versions use essentially the same code but without the generic parameter.

For the AwaitableTask<T> we can do:

[AsyncMethodBuilder(typeof(AwaitableTaskMethodBuilder<>))]
public readonly struct AwaitableTask<T>
{
    private readonly Awaitable<T> _awaitable;

    public Awaitable<T>.Awaiter GetAwaiter() => AsAwaitable().GetAwaiter();

    internal AwaitableTask(Awaitable<T> awaitable) => _awaitable = awaitable;

    private Awaitable<T> AsAwaitable() => _awaitable ?? throw new InvalidOperationException();
}

Essentially, we are using Unity’s Awaitable internally. This approach not only respects Unity’s game loop but also takes advantage of any performance optimizations it offers.

As we saw in the previous post, the only required method is GetAwaiter. Beyond that, we decorate our type with the AsyncMethodBuilder attribute, which specifies the type that will serve as the builder for the async state machine.

Here’s the implementation of the AwaitableTaskMethodBuilder<T> type:

public struct AwaitableTaskMethodBuilder<T>
{
    private static UnityEngine.Awaitable.AwaitableAsyncMethodBuilder<T> s_methodBuilder;

    static AwaitableTaskMethodBuilder() => s_methodBuilder = new UnityEngine.Awaitable.AwaitableAsyncMethodBuilder<T>();

    public static AwaitableTaskMethodBuilder<T> Create() => default;

    public void SetResult(T result) => s_methodBuilder.SetResult(result);

    public void SetException(Exception exception) => s_methodBuilder.SetException(exception);

    public AwaitableTask<T> Task => new(s_methodBuilder.Task);

    public void AwaitOnCompleted<TAwaiter, TStateMachine>(
        ref TAwaiter awaiter, ref TStateMachine stateMachine)
        where TAwaiter : INotifyCompletion
        where TStateMachine : IAsyncStateMachine =>
        s_methodBuilder.AwaitOnCompleted(ref awaiter, ref stateMachine);

    public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(
        ref TAwaiter awaiter, ref TStateMachine stateMachine)
        where TAwaiter : ICriticalNotifyCompletion
        where TStateMachine : IAsyncStateMachine =>
        s_methodBuilder.AwaitUnsafeOnCompleted(ref awaiter, ref stateMachine);

    public void SetStateMachine(IAsyncStateMachine stateMachine) => s_methodBuilder.SetStateMachine(stateMachine);

    public void Start<TStateMachine>(ref TStateMachine stateMachine)
        where TStateMachine : IAsyncStateMachine =>
        s_methodBuilder.Start(ref stateMachine);
}

Here, we are following the same approach as before by using Unity’s Awaitable.AwaitableAsyncMethodBuilder<T> implementations for the required methods. This allows us to leverage performance advantages, such as Unity’s state machine pooling, while also respecting Unity’s game loop.

Conclusion

Now we can declare MyAwaitable as public AwaitableTask<int> MyAwaitable. This type still encapsulates a Task and behaves like a Task, while also leveraging the stages of Unity’s game loop. Instead of using Awaitable, which could confuse anyone interacting with it, we’ve created our own type that can be returned from async methods.

Writing the rest of the extension methods for Task is now trivial, for example:

public static async AwaitableTask<T> OnFixedUpdate<T>(this Task<T> task, CancellationToken cancellationToken = default)
{
    var result = await task;
    await UnityEngine.Awaitable.FixedUpdateAsync(cancellationToken);
    return result;
}

We still need to cancel the Task when necessary, but now we can obtain the result at a specific stage of the game loop and continue executing code afterward.

As always, thank you for reading, and if you have any questions or comments you can use the comments section, or contact me directly via the contact form or by email. Also, if you don’t want to miss any of the new blog posts, you can subscribe to my newsletter or the RSS feed.