Await, SynchronizationContext, and Console Apps: Part 3

In Part 1 and Part 2 of this short series, I demonstrated how you can build a SynchronizationContext and use it run an async method such that all of the continuations in that method will run on serialized on the current thread.  This can be helpful when executing async methods in a console app, or in a unit test framework that doesn’t directly support async methods.  However, the support I showed thus far targets async methods that return Task… what about async methods that return void?

C# and Visual Basic support two flavors of async methods: ones that return tasks (either Task or Task<T>) and ones that return void.  The former use the returned Task to represent the completion of the async method. In the case of an “async void” method, however, there is no returned Task to represent the method’s processing.  Instead, “async void” methods interact with the current SynchronizationContext to alert the context to the async method’s execution status.  Before entering the body of the async method, if there is a current SynchronizationContext, it is retrieved and its OperationStarted method is called.  And after the async method has completed, that same context has its OperationCompleted method called.  Further, if an exception goes unhandled in the body of the async void method, the throwing of that exception is Post to the SynchronizationContext, so that the exception escapes back to the context for it to handle as it pleases.

All of this means that if we want our AsyncPump to be able to handle “async void” methods in addition to “async Task” methods, we need to augment the type slightly.  First, we need to augment our SingleThreadSynchronizationContext to react appropriate to calls to OperationStarted and OperationCompleted.  These methods need to maintain a count of how many outstanding operations there are, such that when the count reaches 0, we call Complete, just as before we called Complete when the async method’s Task completed.  We do this by adding three members to the custom context:

private int m_operationCount = 0;

public override void OperationStarted()
{
    Interlocked.Increment(ref m_operationCount);
}

public override void OperationCompleted()
{
    if (Interlocked.Decrement(ref m_operationCount) == 0)
        Complete();
}

Then we need to add a new AsyncPump.Run overload that works with Action (for “async void” methods) instead of with Func<Task> (for “async Task” methods). As a reminder, here’s the existing Run method from our AsyncPump class:

public static void Run(Func<Task> asyncMethod)
{
    var prevCtx = SynchronizationContext.Current;
   
try
    {
        var syncCtx = new SingleThreadSynchronizationContext(false);
        SynchronizationContext.SetSynchronizationContext(syncCtx);

        var t = asyncMethod();
        t.ContinueWith(delegate { syncCtx.Complete(); }, TaskScheduler.Default);

        syncCtx.RunOnCurrentThread();
        t.GetAwaiter().GetResult();
    }
    finally
    {
        SynchronizationContext.SetSynchronizationContext(prevCtx);
    }
}

Most of this will remain the same for our new Action-based variant.  In fact, for the new one, we primarily just need to delete all the code having to do with the returned task, since there isn’t one, and all notion of completion is being handled by the OperationStarted and OperationCompleted methods we added to the context.  We do surround the asyncMethod invocation with calls to OperationStarted and OperationCompleted, just in case the asyncMethod is actually just a void method and not an “async void” method, in which case we need to make sure the operation count is greater than 0 for the duration of the invocation in order to avoid races that could result if the delegate invoked other async void methods.

public static void Run(Action asyncMethod)
{
    var prevCtx = SynchronizationContext.Current;
    try
    {
        var syncCtx = new SingleThreadSynchronizationContext(true);
        SynchronizationContext.SetSynchronizationContext(syncCtx);

        syncCtx.OperationStarted();
        asyncMethod();
        syncCtx.OperationCompleted();

        syncCtx.RunOnCurrentThread();
    }
    finally
    {
        SynchronizationContext.SetSynchronizationContext(prevCtx);
    }
}

That’s it (note that I’ve added a parameter to SingleThreadSynchronizationContext’s constructor, which allows me to specify whether operation count tracking should be performed: we want it for this new Run method, but not for the previously described ones). We’re now able to use our AsyncPump to run “async void” methods synchronously with all continuations executed on the current thread, e.g.

static void Main()
{
    AsyncPump.Run((Action)FooAsync);
}

static async void FooAsync()
{
    Foo1Async();
    await Foo2Async();
    Foo3Async();
}

static async void Foo1Async()
{
    await Task.Delay(1000);
    Console.WriteLine(1);
}

static async Task Foo2Async()
{
    await Task.Delay(1000);
    Console.WriteLine(2);
}

static async void Foo3Async()
{
    await Task.Delay(1000);
    Console.WriteLine(3);
}

Happy async’ing.

AsyncPump.cs