Await Anything In Unity With AwaitableCompletionSource And Custom Awaiters

Posted by : on

Category : Unity

Introduction

Some time ago I wrote a post about Different Ways to Determine When an Event Occurred in Code. That post was useful because we often want to perform operations after an event has taken place. In this post, I’ll show how to create more readable code that’s easier to modify, debug, and test after an event has occurred.

This approach doesn’t change how we detect when something happens, it still uses events under the hood. However, it hides these events behind a state machine created when we use C#’s await keyword. This allows us to write code using C#’s async/await paradigm.

The benefit of async/await in this context is that we can write all our code in one place, with a synchronous style, even though the code executes asynchronously. This makes it easier to understand by reading our code from top to bottom.

While events are valuable, they can lead to complicated code when used for complex tasks. Event-driven code often spreads across multiple places and requires shared variables for communication between parts, leading to unnecessary states. Sometimes, overly decoupled code can be as challenging to manage as monolithic code.

In the following sections, I’ll describe a simple problem: a spawner in a game that spawns objects and raises an event whenever an object is spawned. One or more classes can subscribe to this event and perform their own actions each time something is spawned. I’ll then show how we can use the await keyword to make this code simpler to read.

Traditionally, if we wanted a class to perform specific actions depending on the number of objects spawned, for example, one action for the first spawn, a different action for the second, and other actions for the next ten spawns—we’d need multiple methods that unsubscribe and subscribe new handlers or use variables in a method with a switch statement. With the await approach, this complexity is no longer necessary.

There are two ways to use the await approach: by using the GetAwaiter method of an existing type, such as Unity’s Awaitable type, or by creating our own custom type. I’ll provide code examples for both approaches, and then demonstrate how we can apply the same method to existing types with events in Unity’s API.

If you want to refresh your knowledge about how Unity’s Awaitable works, you can read my previous blog post: Asynchronous Code In Unity Using Awaitable and AwaitableCompletionSource.

Finally, I’ll share a small code snippet that adds FromResult functionality from the Task class to Unity’s Awaitable class. To start, let’s look at the traditional solution with a simple Spawner class and a Circle class that responds whenever an object is spawned.

Creating A Spawner That Uses Events

The following is the simplest spawner I could think of. It takes a prefab, and whenever the space key is pressed, it spawns a game object from this prefab and triggers an event that includes the spawned game object as a parameter:

public class Spawner : MonoBehaviour
{
    public event Action<GameObject> GameObjectSpawned;
    
    [SerializeField] private GameObject objectToInstantiate;

    private void Update()
    {
        if (Keyboard.current.spaceKey.wasPressedThisFrame)
        {
            OnGameObjectSpawned(Spawn());
        }
    }

    private GameObject Spawn() => Instantiate(objectToInstantiate);

    private void OnGameObjectSpawned(GameObject go) => GameObjectSpawned?.Invoke(go);
}

This class offers the event that other classes can subscribe to, for example:

public class Circle : MonoBehaviour
{
    [SerializeField] private Spawner spawner;

    private void OnEnable() => spawner.GameObjectSpawned += ShowLog;

    private void Update() => Debug.Log("Doing stuff...");
    
    private void OnDisable() => spawner.GameObjectSpawned -= ShowLog;

    private void ShowLog(GameObject go) => Debug.Log($"{go.name} Spawned message from {gameObject.name}");
}

This setup is straightforward: our Circle class runs its Update method, and whenever an object is spawned, a message is logged. For simple cases like this, it works perfectly. However, problems arise when we want to add more complex logic, such as performing different actions based on the number of spawned objects, or worse, when multiple classes need to execute various methods on different spawned objects.

The code required for these scenarios can become difficult to follow, as we’d need to track code scattered across different parts of our program. Let’s simplify this by using await to write code that can be read in a linear, more comprehensible way.

Change The Spawner To Have A GetAwaiter Method That Uses AwaitableCompletionSource

First, let’s take a look at how our code in the Circle class will look after modifying the Spawner class:

public class CircleWithAwait : MonoBehaviour
{
    [SerializeField] private AwaitableSpawner spawner;

    private async void Start() => await WaitSpawner();

    private void Update() => Debug.Log("Doing stuff...");

    private async Awaitable WaitSpawner()
    {
       Debug.Log("Waiting...");
       var foo = await spawner;
       Debug.Log($"{foo.name} Spawned message from {gameObject.name}");
       var foo2 = await spawner;
       Debug.Log($"{foo2.name} Spawned again! Message from {gameObject.name}");
       var foo3 = await spawner;
       Debug.Log($"{foo3.name} Spawned for a third time! Different code executed after each spawn in contrast with events!");
       
       while (true)
       {
          var foo4 = await spawner;
          Debug.Log($"{foo4.name} Spawned!;");
       }
    }
}

We still have our Update method as before, but this time, instead of subscribing and unsubscribing in the OnEnable/OnDisable event methods, we start an asynchronous method in the Start event method. This method, called WaitSpawner, demonstrates the true power of writing asynchronous code in a synchronous style. When the first object is spawned, it logs a specific message, another message when the second object spawns, a different one for the third, and a unique message for each subsequent spawn.

This method avoids introducing unnecessary state management. We don’t need a counter to track the number of spawned objects, nor do we have to subscribe and unsubscribe from different methods each time the event is invoked to execute different code depending on the spawn count.

Most importantly, this method is easy to understand, just by reading it from top to bottom, we can follow what’s happening.

Writing code, like this, is feasible, by adding a GetAwaiter method ito our Spawner class. The await keyword only requires the type to have a GetAwaiter method that returns a type implementing the INotifyCompletion interface.

We’ll start by returning the Awaiter class from Unity’s existing Awaitable type, and in the next section, I’ll show how to create a custom type. Let’s look at how to implement the Spawner class now:

public class AwaitableSpawner : MonoBehaviour
{
   [SerializeField] private GameObject objectToInstantiate;

   private Action<GameObject> gameObjectSpawned;
   
   private void Update()
   {
      if (Keyboard.current.spaceKey.wasPressedThisFrame)
      {
         gameObjectSpawned?.Invoke(Spawn());
      }
   }

   private GameObject Spawn() => Instantiate(objectToInstantiate);
   
   public Awaitable<GameObject>.Awaiter GetAwaiter()
   {
      var acs = new AwaitableCompletionSource<GameObject>();
      
      gameObjectSpawned += SetResult;
      
      return acs.Awaitable.GetAwaiter();

      void SetResult(GameObject go)
      {
         gameObjectSpawned -= SetResult;
         acs.TrySetResult(go);
      }
   }
}

There are a few changes from the previous implementation. Now, we no longer need a public event since the delegate is only invoked within the class. We’ve also implemented the GetAwaiter method, which returns an Awaitable<GameObject>.Awaiter type to allow us to return the spawned game object each time. To enable this, we create an AwaitableCompletionSource<GameObject>(), which we’ll call acs.

We then subscribe to the delegate and return the awaiter of a new Awaitable.

The Awaitable property of AwaitableCompletionSource<GameObject>() is implemented by Unity to always return a new Awaitable (sometimes pooled). We simply return the GetAwaiter of this new Awaitable, so we don’t have to implement a custom type; we’re just using Unity’s built-in type.

The SetResult local method is what we subscribe to the delegate. When an object is spawned, this method will execute. First, it unsubscribes itself from the delegate, as each time await calls GetAwaiter, we’ll resubscribe to the delegate. This prevents unused methods from accumulating in the delegate.

Next, we call TrySetResult(go) on AwaitableCompletionSource<GameObject>. This step is only necessary if we want to return a result. If our return type for GetAwaiter were simply Awaitable.Awaiter, this step wouldn’t be required. Although this adds a few extra lines of boilerplate code (about six), it enables us to write code that executes after each spawn in a more readable way.

Changing The GetAwaiter To Return A Custom Awaiter

While using Awaitable.Awaiter as a return type is convenient, we can also create our own custom type that can be “awaited”. This approach not only provides greater control but also helps us understand exactly what happens in each await call.

All we need to do is create a type that implements the INotifyCompletion interface and return it from our GetAwaiter method. In the following implementation of the Spawner class, I’ve created this custom type as a nested class within Spawner, though it could also be a separate, non-nested class.

After explaining the implementation, I’ll also cover a Unity-specific issue that arises here. For now, take a look at the updated Spawner implementation, which remains mostly the same as before, except for the GetAwaiter method and the new SpawnerAwaiter class:

public class CustomAwaitableSpawner : MonoBehaviour
{
    [SerializeField] private GameObject objectToInstantiate;

    private Action<GameObject> gameObjectSpawned;
   
    private void Update()
    {
        if (Keyboard.current.spaceKey.wasPressedThisFrame)
        {
            gameObjectSpawned?.Invoke(Spawn());
        }
    }

    private GameObject Spawn() => Instantiate(objectToInstantiate);
    
    public SpawnerAwaiter GetAwaiter() => new(this);

    public class SpawnerAwaiter : INotifyCompletion
    {
        public bool IsCompleted { get; private set; }

        private readonly CustomAwaitableSpawner _spawner;
        private GameObject _go;
    
        public SpawnerAwaiter(CustomAwaitableSpawner spawner) => _spawner = spawner;
    
        public void OnCompleted(Action continuation)
        {
            _spawner.gameObjectSpawned += Spawned;
            return;

            void Spawned(GameObject go)
            {
                _go = go;
                _spawner.gameObjectSpawned -= Spawned;
                IsCompleted = true;
                continuation();
            }
        }

        public GameObject GetResult() => _go;
    }
}

The GetAwaiter method takes this as a parameter and returns a new instance of SpawnerAwaiter.

SpawnerAwaiter is our custom awaiter implementation that conforms to the INotifyCompletion interface. Here’s how it works:

The async keyword generates a state machine. For each await expression, there are two states: the first is the await expression itself, and the second is the code that follows it. If the boolean property IsCompleted is false, the state machine waits for the OnCompleted method to be called and then moves to the next state, where it sets the result using the SetResult method. Each await adds another state to the state machine. Let’s see how this theory applies to our SpawnerAwaiter implementation.

The public bool IsCompleted { get; private set; } property is initially set to false to ensure that the OnCompleted method will execute. We also have two private fields: private readonly CustomAwaitableSpawner _spawner, which holds the instance of CustomAwaitableSpawner, and private GameObject _go, which stores the result we want to return. The _spawner field is assigned in the constructor.

The public GameObject GetResult() => _go method simply returns the game object that was spawned, so the only remaining task is implementing the OnCompleted method.

The OnCompleted method is straightforward: it registers the Spawned method with the spawner’s delegate. Now, whenever an object is spawned, Spawned will execute. It assigns the spawned game object to the _go field for the GetResult method to return, unsubscribes itself from the delegate, sets IsCompleted to true, and then calls the continuation action.

The continuation parameter, populated automatically by the async state machine, contains all code after the await expression. After the await expression finishes, OnCompleted runs and, among other things, triggers continuation, allowing the code following await to proceed.

As we can see, implementing a custom awaiter isn’t overly complicated once we understand how the async state machine works.

Note About Monobehaviours And INotifyCompletion

You may read in some blog posts that the INotifyCompletion interface can be implemented by the same class that has the GetAwaiter method. In those cases the GetAwaiter method would return the class, but in Unity there is a problem. Our classes are Monobehaviours and this means that their lifetime is controlled by the game engine and not by us.

When I tried to do the above, I was getting very strange results. Probably because Monobehaviours also have a C++ representation that exists independently. In any case, everything works fine if the implementor of the INotifyCompletion is not a Monobehaviour, so keep this in mind.

Using GetAwaiter Extension Methods With Unity Types

The use of await is not limited to our own types. We can easily create an extension GetAwaiter method for any Unity type that provides an event to subscribe to, allowing us to take advantage of the improved syntax. For example, let’s consider the Button class from the Unity UI namespace. Typically, we would use it like this:

public class ButtonClickEvent : MonoBehaviour
{
    [SerializeField] private Button button;
    [SerializeField] private TextMeshProUGUI  text;

    private void OnEnable() => button.onClick.AddListener(ButtonClicked);

    private void OnDisable() => button.onClick.RemoveAllListeners();

    private void ButtonClicked() => text.text = "Button Clicked!";
}

For more complex behavior, we would typically need to increase the complexity of our code. However, by implementing a GetAwaiter extension method, we can simplify our code to look like this:

public class ButtonAwait : MonoBehaviour
{
    [SerializeField] private Button button;
    [SerializeField] private TextMeshProUGUI  text;

    private async void Start() => await ButtonClicked();

    private async Awaitable ButtonClicked()
    {
        int noOfTimes;
        
        await button;
        Debug.Log("button Clicked!");
        await button;
        Debug.Log("Button Clicked a second time!");
        
        do
        {
            noOfTimes = await button;
            Debug.Log($"Button has been awaited {noOfTimes} times!");
            
        } while (noOfTimes < 10);
        
        Debug.Log("Button has been awaited at least 10 times! It cannot be clicked anymore!");
        button.enabled = false;
    }
}

This code is easier to read, understand, and debug because it doesn’t introduce new fields and is contained within a single method, eliminating the need to navigate through different methods that subscribe to the button’s events.

To write code that invokes the onClick button event like this, we first need to create a Button extension method:

public static class ButtonExtensions
{
   public static ButtonAwaiter GetAwaiter(this Button button) => new(button);
}

Then we have to implement the ButtonAwaiter class:

public class ButtonAwaiter : INotifyCompletion
{
   private readonly Button _button;
   private static int _NoOfTimes;

   public ButtonAwaiter(Button button) => _button = button;

   public bool IsCompleted { get; private set; }

   public void OnCompleted(Action continuation)
   {
      _button.onClick.AddListener(ButtonClicked);
      return;

      void ButtonClicked()
      {
         _button.onClick.RemoveListener(ButtonClicked);
         _NoOfTimes++;
         IsCompleted = true;
         continuation();
      }
   }

   public int GetResult() => _NoOfTimes;
}

I don’t believe much explanation is necessary, as everything is similar to the previous implementation of CustomAwaitableSpawner. The only difference here is that I decided the result will be an int representing how many times the button has been awaited. In your implementations, you may choose not to return a result at all, or you might want to return a different value, such as the coordinates of the pointer when the button was pressed, or anything else you find necessary for your code.

Bonus: Implementing A FromResult Method For Unity’s Awaitable

Unity’s Awaitable does not have a FromResult method like Task, which creates a new Awaitable from a given result. However, we can create our own utility method to achieve this, allowing us to write code like this:

private async void Start()
{
    Awaitable<int> foo = AwaitableUtilities.FromResult(42);
    int bar = await foo;
    Debug.Log(bar);
}

This can be easily accomplished by using the AwaitableCompletionSource class:

public static class AwaitableUtilities
{
    public static Awaitable<T> FromResult<T>(T result)
    {
        var acs = new AwaitableCompletionSource<T>();
        Awaitable<T> awaitable = acs.Awaitable;
        acs.SetResult(result);
        return awaitable;
    }
}

The Awaitable in Unity offers much more than what is documented. If you have the time, I recommend checking out Unity’s GitHub repository, where the source code is available for viewing: Awaitable.cs.

Conclusion

If you noticed, in my examples, I wrote code like this: private async void Start() => await ButtonClicked() instead of this: private void Start() => ButtonClicked(), where ButtonClicked would have been implemented as private async void ButtonClicked. I chose this approach because it allows me to catch any exceptions that may occur in my async methods. So, for your own sanity, please avoid using async void methods.

That concludes our discussion on creating custom awaiters so that we can await any class that offers events to subscribe to. For simple cases, events are fine, but if you find yourself needing more complex logic based on how many times an event has been triggered, or if you realize you’re spending too much time navigating through different parts of your code due to excessive decoupling from events, consider implementing your own awaiters and using those types with the async/await workflow as described above. You may discover that your code’s readability improves significantly by writing just a few more lines of boilerplate code.

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.


Follow me: