How to implement the OnCompleted method of a custom awaiter correctly?

Question

I have a custom awaitable type and the problem is that the continuation resumes on a different thread, which causes problems in UIs such as WinForms/WPF/MVC/etc:

Answer 1

The MSDN explanation for OnCompleted method is:

Schedules the continuation action that's invoked when the instance completes.

Hence neither of the implementations of the OnCompleted is "correct", because the awaiter shouldn't execute the passed delegate during that call if the awaitable is not already complete, but register it to be executed when the awaitable completes.

The only unclear is what the method should do if the awaitable is already complete at the time the method is called (although the compiler generated code does not call it in such case) - ignore the continuation delegate or execute. According to the Task implementation, it should be the later (execute).

Of course there are exceptions of the rule (hence the word "correct"). For instance, the YieldAwaiter specifically always returns IsCompleted == false to force calling it's OnCompleted method, which immediately schedules the passed delegate on the thread pool. But "normally" you won't do that.

Usually (as with the standard Task implementation) the awaitable will perform the operation, provide the result, the wait mechanism, and will maintain/execute the continuations. Their awaiters are usually structs holding the reference to the shared awaitable (along with continuation options when needed) and will delegate the GetResult and OnCompleted method calls to the shared awaitable, and specifically for OnCompleted passing the continuation delegate as well as options to the awaitable internal method responsible for registering/executing them. The "configurable" awaitables will simply hold the shared awaitable plus the options and simply pass them to the created awaiters.

Since in your example the waiting and result are provided by the awaiter, the awaitable can simply provide completion event:

public class MyAwaitable
{
    private volatile bool finished;
    public bool IsFinished => finished;
    public event Action Finished;
    public MyAwaitable(bool finished) => this.finished = finished;
    public void Finish()
    {
        if (finished) return;
        finished = true;
        Finished?.Invoke();
    }
    public MyAwaiter GetAwaiter() => new MyAwaiter(this);
}

and the awaiters would subscribe on it:

public class MyAwaiter : INotifyCompletion
{
    private readonly MyAwaitable awaitable;
    private int result;

    public MyAwaiter(MyAwaitable awaitable)
    {
        this.awaitable = awaitable;
        if (IsCompleted)
            SetResult();

    }
    public bool IsCompleted => awaitable.IsFinished;

    public int GetResult()
    {
        if (!IsCompleted)
        {
            var wait = new SpinWait();
            while (!IsCompleted)
                wait.SpinOnce();
        }
        return result;
    }

    public void OnCompleted(Action continuation)
    {
        if (IsCompleted)
            {
                continuation();
                return;
            }
        var capturedContext = SynchronizationContext.Current;
        awaitable.Finished += () =>
        {
            SetResult();
            if (capturedContext != null)
                capturedContext.Post(_ => continuation(), null);
            else
                continuation();
        };
    }

    private void SetResult()
    {
        result = new Random().Next();
    }
}

When OnCompleted is called, first we check if we are complete. If yes, we simply execute the passed delegate and return. Otherwise, we capture the synchronization context, subscribe on the awaitable completion event, and inside that event execute the action either via the captured synchronization context or directly.

Again, in the real life scenarios the awaitable should perform the real work, provide the result and maintain continuation actions, while awaiters should only register the continuation actions, eventually abstracting the continuation execution strategy - directly, via captured synchronization context, via thread pool etc.

Answer 2

Note: Originally I put this answer as a summary in the end of the question after @IvanStoev gave the correct answer (many thanks for the enlightenment). Now I extracted that part into a real answer.

---

So based on Ivan's answer here is a small summary containing the missing parts, which I believe should be in the documentations. The example below mimics also the ConfigureAwait behavior of Task.

1. The test app

A WinForms app (could be other single-threaded UI as well) with a ProgressBar and 3 Button controls: one button simply starts an async operation (and a progress bar), the others finish it either in the UI thread or in a foreign thread.

public partial class Form1 : Form
{
    public Form1()
    {
        InitializeComponent();
        progressBar.Style = ProgressBarStyle.Marquee;
        progressBar.Visible = false;
    }

    private MyAwaitable awaitable;

    private async void buttonStart_Click(object sender, EventArgs e)
    {
        awaitable = new MyAwaitable();
        progressBar.Visible = true;
        var result = await awaitable; //.ConfigureAwait(false); from foreign thread this throws an exception
        progressBar.Visible = false;
        MessageBox.Show(result.ToString());
    }

    private void buttonStopUIThread_Click(object sender, EventArgs e) =>
        awaitable.Finish(new Random().Next());

    private void buttonStopForeignThread_Click(object sender, EventArgs e) =>
        Task.Run(() => awaitable.Finish(new Random().Next()));
 }

2. The custom awaitable class

As opposed to the original example in the question, here the awaitable class itself contains the continuation, which is invoked once the execution is finished. So an awaiter can just request to schedule the continuation for later execution.

And please note that ConfigureAwait and GetAwaiter are basically the same - the latter can use the default configuration.

public class MyAwaitable
{
    private volatile bool completed;
    private volatile int result;
    private Action continuation;

    public bool IsCompleted => completed;

    public int Result => RunToCompletionAndGetResult();

    public MyAwaitable(int? result = null)
    {
        if (result.HasValue)
        {
            completed = true;
            this.result = result.Value;
        }
    }

    public void Finish(int result)
    {
        if (completed)
            return;
        completed = true;
        this.result = result;

        continuation?.Invoke();
    }

    public MyAwaiter GetAwaiter() => ConfigureAwait(true);

    public MyAwaiter ConfigureAwait(bool captureContext)
        => new MyAwaiter(this, captureContext);

    internal void ScheduleContinuation(Action action) => continuation += action;

    internal int RunToCompletionAndGetResult()
    {
        var wait = new SpinWait();
        while (!completed)
            wait.SpinOnce();
        return result;
    }
}

3. The awaiter

OnCompleted now does not execute the continuation (unlike the examples I investigated) but registers it for later by calling MyAwaitable.ScheduleContinuation.

Secondly, please note that now the awaiter also has a GetAwaiter method that just returns itself. This is needed for the await myAwaitable.ConfigureAwait(bool) usage.

public class MyAwaiter : INotifyCompletion
{
    private readonly MyAwaitable awaitable;
    private readonly bool captureContext;

    public MyAwaiter(MyAwaitable awaitable, bool captureContext)
    {
        this.awaitable = awaitable;
        this.captureContext = captureContext;
    }

    public MyAwaiter GetAwaiter() => this;

    public bool IsCompleted => awaitable.IsCompleted;

    public int GetResult() => awaitable.RunToCompletionAndGetResult();

    public void OnCompleted(Action continuation)
    {
        var capturedContext = SynchronizationContext.Current;
        awaitable.ScheduleContinuation(() =>
        {
            if (captureContext && capturedContext != null)
                capturedContext.Post(_ => continuation(), null);
            else
                continuation();
        });
    }
}

Answer 3

This proves that the continuation runs on the captured context:

public class MyAwaitable
{
    private volatile bool finished;
    public bool IsFinished => finished;
    public MyAwaitable(bool finished) => this.finished = finished;
    public void Finish() => finished = true;
    public MyAwaiter GetAwaiter() => new MyAwaiter(this);
}

public class MyAwaiter : INotifyCompletion
{
    private readonly MyAwaitable awaitable;
    private readonly SynchronizationContext capturedContext = SynchronizationContext.Current;

    public MyAwaiter(MyAwaitable awaitable) => this.awaitable = awaitable;
    public bool IsCompleted => awaitable.IsFinished;

    public int GetResult()
    {
        SpinWait.SpinUntil(() => awaitable.IsFinished);
        return new Random().Next();
    }

    public void OnCompleted(Action continuation)
    {
        if (capturedContext != null) capturedContext.Post(state => continuation(), null);
        else continuation();
    }
}

public class MySynchronizationContext : SynchronizationContext
{
    public override void Post(SendOrPostCallback d, object state)
    {
        Console.WriteLine("Posted to synchronization context");
        d(state);
    }
}

class Program
{
    static async Task Main()
    {
        SynchronizationContext.SetSynchronizationContext(new MySynchronizationContext());

        var awaitable = new MyAwaitable(false);

        var timer = new Timer(_ => awaitable.Finish(), null, 100, -1);

        var result = await awaitable;

        Console.WriteLine(result);
    }
}

Output:

Posted to synchronization context
124762545

But you are not posting the continuation to the synchronization context.

You're posting scheduling the execution of the continuation on another thread.

The scheduling runs on the synchronization context but continuation itself doesn't. Thus your problems.

You can read this to understand how it works.