TAP: Await anything (not only Tasks)

The async/await pattern is simple yet effective way to work with tasks. Its ingeniousity is that the programmer is focused on implementation of the function at the first place, leaving the technical details to the compiler which does some behind-the-scenes magic to make sure that async methods can:

  • run asynchronously,
  • capture or disregard the context (via ConfigureAwait(bool))
  • use all benefits of TAP (Task-based Asynchronous Pattern)

A nice touch of the async/await syntax is that it is easy to convert any synchronous method into its asynchronous counterpart, without sacrificing code clarity or increasing its complexity.

Interestingly enough, the await keyword is not really limited to Tasks, and as a matter of fact it can be used for any object for which it makes sense to wait for. This blog post is devoted to defining a class which can be awaited with await keyword, and yet is by no mean a Task at all!

So what is the compiler doing and how to make sure it understands our custom awaiting desire? In order to demonstrate, let’s try the following code which uses a plain AutoResetEvent. One thing to notice is that – instead of WaitOne() in line 21 – we try to await it by adding the await keyword before.

class Program
{
    static void Main(string[] args)
    {
        using (var waitHandle = new AutoResetEvent(false))
        {
            var t1 = DoSomething1Async(waitHandle);
            var t2 = DoSomething2Async(waitHandle);
            Task.WaitAll(t1, t2);
        }
    }

    private static async Task DoSomething1Async(AutoResetEvent waitHandle)
    {
        await Task.Delay(13000).ConfigureAwait(false);
        waitHandle.Set();
    }

    private static async Task DoSomething2Async(AutoResetEvent waitHandle)
    {
        await waitHandle;
    }
}

The code above will not compile, because the compiler cannot transpile the await syntax in combination with AutoResetEvent into proper IL. In order to await anything, the object must have a method GetAwaiter() which returns an instance implementing the INotifyCompletion interface. Extensions methods are also fine for this purpose, and are the only choice if we do not own the source of the original class or if we simply do not want to bloat it with non-business critical methods. Let’s fix this small console app by creating a new utility class which implements the required interface, and writing an extension method which provides the required GetAwaiter():

public static class AutoResetEventExtensions
{
    public static AutoResetEventAwaiter GetAwaiter(this AutoResetEvent source)
    {
        return new AutoResetEventAwaiter(source);
    }
}

public class AutoResetEventAwaiter : INotifyCompletion
{
    private readonly AutoResetEvent autoResetEvent;

    public AutoResetEventAwaiter(AutoResetEvent autoResetEvent)
    {
        this.autoResetEvent = autoResetEvent;
    }

    public void OnCompleted(Action continuation)
    {
        continuation();
    }
}

This is still not enough – first of all how to inform the compiler what are actually waiting for? Before we solve this, let’s make the code compilable at all, because right now the compiler must show two errors:

  • AutoResetEvent does not contain a definition for 'IsCompleted'
  • AutoResetEvent does not contain a definition for 'GetResult()'

Both members must belong to our AutoResetEventAwaiter although they do not belong to the INotifyCompletion interface. Let’s fix this by implementing them:

public class AutoResetEventAwaiter : INotifyCompletion
{
    private readonly AutoResetEvent autoResetEvent;

    public AutoResetEventAwaiter(AutoResetEvent autoResetEvent)
    {
        this.autoResetEvent = autoResetEvent;
    }

    public void OnCompleted(Action continuation)
    {
        continuation();
    }

    public bool IsCompleted { get; private set; }

    public void GetResult()
    {
        this.autoResetEvent.WaitOne();
        this.IsCompleted = true;
    }
}

This looks better now, the code can be compiled. In this simple implementation we wrapped the WaitOne() inside the GetResult() method which also sets the value of IsCompleted to true, so that any subsequent awaiting is done instantly without waiting for a signal. Finally, the only method required by the INotifyCompletion interface simply executes a given action when the awaiting is over. With a bit of tracing, the execution path looks like that:

In this example I didn’t care about capturing of the execution context and instead I focused myself only on implementing syntactic sugar of a custom awaitable object.

Awaiting without custom awaiter

For many classes (like Process, ProcessStartInfo) there is even no need to write custom awaiter at all. Since the awaitable actions can be expressed as tasks, we can simply wrap their part using the TaskCompletionSource class and return the GetAwaiter() of its Task property. Below is a quick sample how to wrap the Process to be able to await it:

class Program
{
	static void Main(string[] args)
	{
		Task processTask = StartProcessAsync();
		processTask.Wait();
	}

	public static async Task<int> StartProcessAsync()
	{
		return await Process.Start("notepad.exe");
	}
}

public static class ProcessExtensions
{
	public static TaskAwaiter<int> GetAwaiter(this Process source)
	{
		var taskCompletionSource = new TaskCompletionSource<int>();
		source.EnableRaisingEvents = true;
		source.Exited += (sender, args) =>
		{
			taskCompletionSource.SetResult(source.ExitCode);
		};

		return taskCompletionSource.Task.GetAwaiter();
	}
}

As a bonus, by awaiting the process we get its exit code, thanks to the usage of generics.

Summary

In order to await any object:

  1. Make sure that it has a method GetAwaiter() which returns an object implementing the INotifyCompletion interface (either member method or via extension method for types which are not in your control)
  2. In that class implement the property IsCompleted (public getter) which returns a boolean value determining if the awaitable execution has finished
  3. In the same class implement the method IsCompleted() which returns void and which makes sure that the awaitable execution is finished. It should also set IsCompleted to true.

For many situations, you can simply rely on GetAwaiter() method belonging to the Task class, and simply wrap the awaitable method using the TaskCompletionSource class to produce robust awaitable code.

Leave a Reply

Your email address will not be published. Required fields are marked *