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:

private MyAwaitable awaitable;

private async void buttonStart_Click(object sender, EventArgs e)
{
    awaitable = new MyAwaitable(false);
    progressBar1.Visible = true;

    // A regular Task can marshal the execution back to the UI thread
    // Here ConfigureAwait is not available and I don't know how to control the flow
    var result = await awaitable;

    // As a result, here comes the usual "Cross-thread operation not valid" exception
    // A [Begin]Invoke could help but regular Tasks also can handle this situation
    progressBar1.Visible = false;
}

private void buttonStop_Click(object sender, EventArgs e) => awaitable.Finish();

Here is the MyAwaitable class:

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

And the problematic custom awaiter:

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()
    {
        var wait = new SpinWait();
        while (!awaitable.IsFinished)
            wait.SpinOnce();
        return new Random().Next();
    }

    public void OnCompleted(Action continuation)
    {            
        // continuation(); // This would block the UI thread

        // Task constructor + Start was suggested by the references I saw,
        // Results with Task.Run/Task.Factory.StartNew are similar.
        var task = new Task(continuation, TaskCreationOptions.LongRunning);

        // If executed from a WinForms app, we have a WinFormsSyncContext here,
        // which is promising, still, it does not solve the problem.
        if (capturedContext != null)
            capturedContext.Post(state => task.Start(), null);
        else
            task.Start();
    }
}

I suspect that my OnCompleted implementation is not quite correct.

I tried to dig into the ConfiguredTaskAwaiter returned by the Task.ConfigureAwait(bool).GetAwaiter() method and could see that the black magic happens in a SynchronizationContextAwaitTaskContinuation class but that is an internal one, along with lot of other internally used types. Is there a way to refactor my OnCompleted implementation to work as expected?

Update: Note to downvoters: I know I do improper things in OnCompleted, that's why I ask. If you have concerns about quality (or anything else) please leave a comment and help me to improve the question so I also can help you to highlight the problem better. Thanks.

Note 2: I know I could use a workaround with a TaskCompletionSource<TResult> and its regular Task<TResult> result but I would like to understand the background. This is the only motivation. Pure curiosity.

Update 2: Notable references I investigated:

How awaiter works:

• https://stackoverflow.com/a/20573825/5114784
• https://stackoverflow.com/a/12661431/5114784

Some implementations:

• https://marcinotorowski.com/2018/03/13/tap-await-anything-not-only-tasks/ - wrong, blocks caller thread until GetResult returns
• https://blogs.msdn.microsoft.com/pfxteam/2011/01/13/await-anything - just calls other awaiters
• https://weblogs.asp.net/dixin/understanding-c-sharp-async-await-2-awaitable-awaiter-pattern - similar to mine, has the same issue

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.