DispatcherTaskWaiter

1. The problem

The concept of the [DispatcherWaiter] presented in my previous post Convert BeginInvoke to async/await solves the problem of switching back to GUI thread after await some task in a almost perfect way：

1
2
3
4
5
6
7
8
9
10
11
private async Task DoSomethingAsync()
{
  ////can be any thread
  await SomeTask.ConfigureAwait(false);

  ////can be any thread
  await waiter.WaitAsync();

  ////The thread associated with the Dispatcher of the waiter, aka GUI thread
  DoSomethingOnGUIThread();
}

But await waiter.WaitAsync() still seems clumsy, in this article we’ll push further to remove it.

To briefly recap the usages of ConfigureAwait(bool continueOnCapturedContext), it provides a hook to change the behavior of how the await behaves via this custom awaiter.

• ConfigureAwait(true) - it’s the default behavior and can be omitted. The state machine code generated by the compiler would automatically marshal the continuation back to the original context captured, aka if await SomeTask is called on GUI thread, after SomeTask is completed, execution would continue on GUI thread.

• ConfigureAwait(false) - needs to be specified explicitly. It’s recommended to always doing so in library codes to improve performance and avoid potential deadlocks. Basically the captured context would be ignored when the generated code sees this and it would continue execution on the same thread the task is complete.

While ConfigureAwait(bool) is enough to serve most of the cases, however, it’s tricky in the case of await several tasks, and use their results on the GUI thread.

In the below example, if DoSomethingAsync is called on GUI thread, aka Thread 1, then Thread 2 and Thread 3 will also be GUI thread, eventually DoSomethingOnGUIThread will be executed on GUI thread too, as expected:

1
2
3
4
5
6
7
8
9
10
private async Task DoSomethingAsyncV1()
{
  ////Thread 1 => GUI Thread
  var result1 = await BackgroundTask1; ////Equivalent to ConfigureAwait(true)
  ////Thread 2 => GUI Thread
  var result2 = await BackgroundTask2; ////Equivalent to ConfigureAwait(true)

  ////Thread 3 => GUI Thread
  DoSomethingOnGUIThread(result1, result2);
}

Since there’re two background tasks, it seems unnecessary to spin off the GUI thread and marshal back to GUI thread two times, we might want to save the first one by adding ConfigureAwait(false), but it might not work because Thread 2 and Thread 3 can be end up called on background thread, depends on how fast the BackgroundTask1 is complete.

To fix it, we need to refactor the code a little bit to extract the execution of BackgroundTask1 and BackgroundTask2 to a separate method GetResultAsync(), and ConfigureAwait(false) all the way inside it, as shown below. It has more codes to write but works, in fact it’s the correct way to do the job:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
private async Task DoSomethingAsyncV2()
{
  ////Thread 1 => GUI Thread
  var tuple = await GetResultAsync(); ////Equivalent to ConfigureAwait(true)

  ////Thread 2 => GUI Thread.
  DoSomethingOnGUIThread(tuple.Item1, tuple.Item2);
}

private async Task<Tuple<int, int>> GetResultAsync()
{
  ////Thread 1 => GUI Thread
  var result1 = await BackgroundTask1.ConfigureAwait(false);
  ////Thread 2 => might be background Thread
  var result2 = await BackgroundTask2.ConfigureAwait(false);

  ////Thread 3 => might be background Thread
  return Tuple.Create(result1, result2);
}

Although it’s the correct way but developers might not recognize it and stick to the original way of adding ConfigureAwait(false) to await BackgroundTask1 and get exceptions. By using [DispatcherWaiter] it saves the refactoring work of the DoSomethingAsyncV2 above and keeps the simplicity similar to DoSomethingAsyncV1:

1
2
3
4
5
6
7
8
9
10
11
12
13
private async Task DoSomethingAsyncV3()
{
  ////Thread 1 => GUI Thread
  var result1 = await BackgroundTask1.ConfigureAwait(false);
  ////Thread 2 => might be background Thread
  var result2 = await BackgroundTask2.ConfigureAwait(false);

  ////Thread 3 => might be background Thread
  await waiter.WaitAsync();

  ////Thread 4 => GUI Thread
  DoSomethingOnGUIThread(result1, result2);
}

So far it’s already pretty elegant to use [DispatcherWaiter] to simplify things, but what if we can use ConfigureAwait(false) all the way on all of the task calls and only add something like ConfigureAwait(DispatcherWaiter) to the last one to automatically marshal back to GUI thread (or whatever thread we want it be)?

1
2
3
4
5
6
7
8
9
10
private async Task DoSomethingAsync()
{
  ////Thread 1 => GUI Thread
  var result1 = await BackgroundTask1.ConfigureAwait(false);
  ////Thread 2 => might be background Thread
  var result2 = await BackgroundTask2.ConfigureAwait(DispatcherWaiter); ////Note the parameter is DispatcherWaiter

  ////Thread 4 => GUI Thread
  DoSomethingOnGUIThread(result1, result2);
}

2. Analysis

Given that [DispatcherWaiter] gets its behavior from Dispatcher, to simplify discussion we only look into implement ConfigureAwait(Dispatcher).

.net gives us two ways to implement custom awaitable behavior:

1) Match the pattern recognized by the compiler. See below from await anything.

The languages support awaiting any instance that exposes the right method (either instance method or extension method): GetAwaiter. A GetAwaiter needs to implement the INotifyCompletion interface (and optionally the ICriticalNotifyCompletion interface) and return a type that itself exposes three members:

1
2
3
bool IsCompleted { get; }
void OnCompleted(Action continuation);
TResult GetResult(); // TResult can also be void

Open source code of the ConfigureAwait(bool) method we can see that it creates an instance of ConfiguredTaskAwaitable, which has a GetAwaiter method that returns an instance of ConfiguredTaskAwaiter, the latter implements the ICriticalNotifyCompletion and INotifyCompletion interface, then compiler generates corresponding state machine codes to get the awaitable behavior.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
public ConfiguredTaskAwaitable ConfigureAwait(bool continueOnCapturedContext)
{
  return new ConfiguredTaskAwaitable(this, continueOnCapturedContext);
}

public struct ConfiguredTaskAwaitable
{
  internal ConfiguredTaskAwaitable(Task task, bool continueOnCapturedContext)
  {
    this.m_configuredTaskAwaiter = new ConfiguredTaskAwaitable.ConfiguredTaskAwaiter(task, continueOnCapturedContext);
  }

  public ConfiguredTaskAwaitable.ConfiguredTaskAwaiter GetAwaiter()
  {
    return this.m_configuredTaskAwaiter;
  }

  public struct ConfiguredTaskAwaiter : ICriticalNotifyCompletion, INotifyCompletion
  {
  }
}

2) Use AsyncMethodBuilder. AsyncMethodBuilder is used in advanced scenarios to implement some custom type, for example ValueTask, that is awaitable just like Task. Here’s some details from Async Task Types in C#:

A task type is a class or struct with an associated builder type identified with System.Runtime.CompilerServices.AsyncMethodBuilderAttribute. The task type may be non-generic, for async methods that do not return a value, or generic, for methods that return a value.

This article demonstrates one example of using AsyncMethodBuilder to achieve something similar to our goal. However, as mentioned by Stephen Toub in this dotnet/csharplang proposal, currently there’s no way to pass context to AsyncMethodBuilder because Attribute knows static information only.

1
2
3
4
5
6
7
8
9
10
11
public struct AsyncCoolTypeMethodBuilder
{
  public static AsyncCoolTypeMethodBuilder Create();
  ......
}

[AsyncMethodBuilder(typeof(AsyncCoolTypeMethodBuilder))]
public struct CoolType { ...... }

//// will implicitly use AsyncCoolTypeMethodBuilder.Create()
public async CoolType SomeMethodAsync() { ...... }

AsyncMethodBuilder can still be used for applications that only has one GUI thread - the custom Task Type can directly access Application.Current.Dispatcher. But for more advanced and complicated application that has multiple GUI threads running, Dispatcher of the current GUI thread has to be passed as context. So although AsyncMethodBuilder aim for advanced usages but it cannot help in our case.

To summaries, #1 is the only possible way to implement what we want under the current .net versions.

3. Implementations

Coding is fairly straightforward. Please not that the only difference between Task andTask<T> in terms of behavior is whether the task has return value or not, so we extract the common logic out to DispatchTaskAwaiterHelper to reuse. Non-generic version of DispatchTaskAwaiter and generic version of DispatchTaskAwaiter<T> just call void GetResult and GetResult<T> respectively.

One thing worth paying attention is the that even though ICriticalNotifyCompletion inherits from INotifyCompletion, codes generated by the default AsyncMethodBuilder shipped with .net would only call UnsafeOnCompleted(Action) for Tasks, so there’s no need to implement OnCompleted(Action).

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
public readonly struct DispatchTaskAwaiter : ICriticalNotifyCompletion
{
  private readonly DispatchTaskAwaiterHelper _helper;

  public DispatchTaskAwaiter(Task task, Dispatcher dispatcher)
  {
    _helper = new DispatchTaskAwaiterHelper(task, dispatcher);
  }
  
  public DispatchTaskAwaiter GetAwaiter()
  {
    return this;
  }
  
  public bool IsCompleted
  {
    get { return _helper.IsCompleted; }
  }
    
  public void OnCompleted(Action continuation)
  {
    ////This is not called, check out https://devblogs.microsoft.com/pfxteam/whats-new-for-parallelism-in-net-4-5-beta/ for more details
    throw new NotImplementedException();
  }

  public void UnsafeOnCompleted(Action continuation)
  {
    _helper.UnsafeOnCompleted(continuation);
  }

  public void GetResult()
  {
    _helper.GetResult();
  }
}

DispatchTaskAwaiterHelper does the real job:

1. Since a Task can be awaited multiple times, so the generated codes would access the awaiter’s IsCompleted property first to determine whether it needs to call UnsafeOnCompleted or not. If IsCompleted==true it directly jump to call GetResult. In the IsCompleted all we need to do is to check if the original task is completed and if we’re on currently the GUI Thread.

2. If IsCompleted==false then UnsafeOnCompleted would be called with a continuation delegate passed in as callback. In our case once the Task is complete, before calling continuation, we call Dispatcher.BeginInvoke to marshal back to the GUI Thread.

3. Exception handling. The codes need to make sure exceptions thrown by the original Task can be captured by the try...catch block in caller codes. One easy to think of approach is to throw in the UnsafeOnCompleted when the original Task has exceptions. But in fact because the exception handling logic is inside of the continuation, exceptions thrown in UnsafeOnCompleted won’t be captured and will become visible only to the event listener of TaskScheduler.UnobservedTaskException when GC happens. GetResult method is the only place can be used to handle the exception, exceptions thrown by VerifyException() would be captured by the caller on GUI Thread, which matches the same behavior of ConfigureAwait(true).

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
public readonly struct DispatchTaskAwaiterHelper
{
  private readonly Task _task;
  private readonly Dispatcher _dispatcher;
  
  public DispatchTaskAwaiterHelper(Task task, Dispatcher dispatcher)
  {
    _task = task;
    _dispatcher = dispatcher;
  }
  
  public bool IsCompleted
  {
    get { return _task.IsCompleted && _dispatcher.CheckAccess(); }
  }
    
  public void UnsafeOnCompleted(Action continuation)
  {
    var tmp = this;
    _task.ContinueWith(t => tmp._dispatcher.BeginInvoke(continuation));
  }
  
  /// <summary>
  /// Called by non-generic DispatchTaskAwaiter
  /// </summary>
  public void GetResult()
  {
    VerifyException();
  }
  
  /// <summary>
  /// Called by generic DispatchTaskAwaiter<T>
  /// </summary>
  public T GetResult<T>()
  {
    VerifyException();
    
    return ((Task<T>) _task).Result;
  }
  
  private void VerifyException()
  {
    if (!_task.IsCompleted)
    {
      throw new InvalidOperationException("Task is unexpectedly not completed");
    }
  
    if (_task.IsCanceled)
    {
      throw new TaskCanceledException();
    }

    if (_task.Exception != null)
    {
      throw _task.Exception;
    }
  }
}

4. Usages

To use DispatchTaskAwaiter we can use below extension methods for Task and Task<T>:

1
2
3
4
5
6
7
8
9
10
11
12
public static class TaskExt
{
  public static DispatchTaskAwaiter ConfigureAwait(this Task task, Dispatcher dispatcher)
  {
    return new DispatchTaskAwaiter(task, dispatcher);
  }
  
  public static DispatchTaskAwaiter<T> ConfigureAwait<T>(this Task<T> task, Dispatcher dispatcher)
  {
    return new DispatchTaskAwaiter<T>(task, dispatcher);
  }
}

And use them as simple as this:

1
2
3
4
5
6
7
8
9
10
11
12
private async Task DoSomethingAsync()
{
  ////Thread 1 => GUI Thread
  var result1 = await BackgroundTask1.ConfigureAwait(false);
  
  ////The _dispatcher is the reference to the GUI thread dispatcher saved when the view model is created
  ////Thread 2 => might be background Thread
  var result2 = await BackgroundTask2.ConfigureAwait(_dispatcher);

  ////Thread 3 => GUI Thread
  DoSomethingOnGUIThread(result1, result2);
}

It’s fairly easy to extend the theory to ConfigureAwait(DispatcherWaiter), that can be covered in some other posts.

Please visit github for the full implementations.

References

• await anything
• Convert BeginInvoke to async/await
• Async Task Types in C#
• Proposal: Allow [AsyncMethodBuilder(…)] on methods