Use SynchronizationContext only on Task RanToCompletion

Question

The code below explains the idea

    private async void button1_Click(object sender, EventArgs e)
    {
        string result;
        CancellationToken cancellationToken = new CancellationTokenSource(3000).Token;

        try
        {
            result = await GetDataAsync(cancellationToken)
                .ContextIfSuccess();    // Should use SynchronizationContext only if Task status is RanToCompletion
        }
        catch(OperationCanceledException)
        {
            /* Context is not required */
            return;
        }
        catch (Exception ex)
        {
            /* Context is not required otherwise it can slow down UI Thread a little bit */
            Log(ex.ToString());
            return;
        }

        /* UI Thread only */
        button1.Text = result;
    }

The question is "Is it possible to make the method like ContextIfSuccess() ?"

Answer 1

In order to have the method that you want you'll need to create a custom awaiter. It's mostly boilerplate, the key is simply that, when asked to add a continuation, you add one using the current sync context to run on successful completion, and one using the default scheduler when it doesn't run to completion.

public struct CaptureContextOnSuccessAwaiter : INotifyCompletion
{
    private Task task;

    public CaptureContextOnSuccessAwaiter(Task task)
    {
        this.task = task;
    }

    public CaptureContextOnSuccessAwaiter GetAwaiter() { return this; }

    public void OnCompleted(Action continuation)
    {
        if (SynchronizationContext.Current != null)
        {
            task.ContinueWith(t => continuation(),
                CancellationToken.None,
                TaskContinuationOptions.OnlyOnRanToCompletion,
                TaskScheduler.FromCurrentSynchronizationContext());
            task.ContinueWith(t => continuation(),
                CancellationToken.None,
                TaskContinuationOptions.NotOnRanToCompletion,
                TaskScheduler.Default);
        }
        else
        {
            task.ContinueWith(t => continuation(),
                CancellationToken.None,
                TaskContinuationOptions.None,
                TaskScheduler.Default);
        }
    }

    public void GetResult() { task.GetAwaiter().GetResult(); }
    public bool IsCompleted { get { return task.GetAwaiter().IsCompleted; } }
}

public struct CaptureContextOnSuccessAwaiter<T> : INotifyCompletion
{
    private Task<T> task;

    public CaptureContextOnSuccessAwaiter(Task<T> task)
    {
        this.task = task;
    }

    public CaptureContextOnSuccessAwaiter<T> GetAwaiter() { return this; }

    public void OnCompleted(Action continuation)
    {
        if (SynchronizationContext.Current != null)
        {
            task.ContinueWith(t => continuation(),
                CancellationToken.None,
                TaskContinuationOptions.OnlyOnRanToCompletion,
                TaskScheduler.FromCurrentSynchronizationContext());
            task.ContinueWith(t => continuation(),
                CancellationToken.None,
                TaskContinuationOptions.NotOnRanToCompletion,
                TaskScheduler.Default);
        }
        else
        {
            task.ContinueWith(t => continuation(),
                CancellationToken.None,
                TaskContinuationOptions.None,
                TaskScheduler.Default);
        }
    }

    public T GetResult() { return task.GetAwaiter().GetResult(); }
    public bool IsCompleted { get { return task.GetAwaiter().IsCompleted; } }
}

public static CaptureContextOnSuccessAwaiter ContextIfSuccess(this Task task)
{
    return new CaptureContextOnSuccessAwaiter(task);
}

public static CaptureContextOnSuccessAwaiter<T> ContextIfSuccess<T>(this Task<T> task)
{
    return new CaptureContextOnSuccessAwaiter<T>(task);
}

Answer 2

await without context (ConfigureAwait(false)). Then, if the desired condition is true you switch into the context.

So you will need to capture SynchronizationContext.Current and Post to that.

This is very similar to what TaskAwaiter does anyway. It resumed without context, then switched back into the context if the caller desires so.

You should be able to make this into a ContextIfSuccess method. Basically, clone the TaskAwaiter source code and decide whether to Post or not in the completion notification. This feature is already there I assume. The code must look at the ConfigureAwait(...) value and conditionally apply the context or not.