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

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();
        });
    }
}