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.