Async/Await equivalent to .ContinueWith with CancellationToken and TaskScheduler.FromCurrentSynchronizationContext() scheduler

Question

This is a follow-up to this question.

Question: What would be a succinct way to express the following using async/await instead of .ContinueWith()?:

var task = Task.Run(() => LongRunningAndMightThrow());

m_cts = new CancellationTokenSource();
CancellationToken ct = m_cts.Token;

var uiTaskScheduler = TaskScheduler.FromCurrentSynchronizationContext();
Task updateUITask = task.ContinueWith(t => UpdateUI(t), ct, TaskContinuationOptions.None, uiTaskScheduler);

I'm mainly interested in the case of a UI SynchronizationContext (e.g. for Winforms)

Note with this that the behavior has all the following desired behaviors:

1. When the CancellationToken is cancelled, the updateUITask ends up cancelled as soon as possible (i.e. the LongRunningAndMightThrow work may still be going on for quite some time).

2. The ct CancellationToken gets checked for cancellation on the UI thread prior to running the UpdateUI lambda (see this answer).

3. The updateUITask will end up cancelled in some cases where the task completed or faulted (since the ct CancellationToken is checked on the UI thread before executing the UpdateUI lambda.

4. There is no break in flow between the check of the CancellationToken on the UI thread and the running of the UpdateUI lambda. That is, if the CancellationTokenSource is only ever cancelled on the UI thread, then there is no race condition between the checking of the CancellationToken and the running of the UpdateUI lambda--nothing could trigger the CancellationToken in between those two events because the UI thread is not relinquished in between those two events.

Discussion:

• One of my main goals in moving this to async/await is to get the UpdateUI work out of a lambda (for ease of readability/debuggability).

• #1 above can be addressed by Stephen Toub's WithCancellation task extension method. (which you can feel free to use in the answers).

• The other requirements seemed difficult to encapsulate into a helper method without passing UpdateUI as a lambda since I cannot have a break (i.e. an await) between the checking of the CancellationToken and the executing of UpdateUI (because I assume I cannot rely on the implementation detail that await uses ExecuteSynchronously as mentioned here. This is where it seems that having the mythical Task extension method .ConfigureAwait(CancellationToken) that Stephen talks about would be very useful.

• I've posted the best answer I have right now, but I'm hoping that someone will come up with something better.

Sample Winforms Application demonstrating the usage:

public partial class Form1 : Form
{
    CancellationTokenSource m_cts = new CancellationTokenSource();

    private void Form1_Load(object sender, EventArgs e)
    {
        cancelBtn.Enabled = false;
    }

    private void cancelBtn_Click(object sender, EventArgs e)
    {
        m_cts.Cancel();
        cancelBtn.Enabled = false;
        doWorkBtn.Enabled = true;
    }

    private Task DoWorkAsync()
    {
        cancelBtn.Enabled = true;
        doWorkBtn.Enabled = false;

        var task = Task.Run(() => LongRunningAndMightThrow());

        m_cts = new CancellationTokenSource();
        CancellationToken ct = m_cts.Token;
        var uiTaskScheduler = TaskScheduler.FromCurrentSynchronizationContext();
        Task updateUITask = task.ContinueWith(t => UpdateUI(t), ct, TaskContinuationOptions.None, uiTaskScheduler);

        return updateUITask;
    }

    private async void doWorkBtn_Click(object sender, EventArgs e)
    {
        try
        {
            await DoWorkAsync();
            MessageBox.Show("Completed");
        }
        catch (OperationCanceledException)
        {
            MessageBox.Show("Cancelled");
        }
        catch
        {
            MessageBox.Show("Faulted");
        }
    }

    private void UpdateUI(Task<bool> t)
    {
        // We *only* get here when the cancel button was *not* clicked.
        cancelBtn.Enabled = false;
        doWorkBtn.Enabled = true;

        // Update the UI based on the results of the task (completed/failed)
        // ...
    }

    private bool LongRunningAndMightThrow()
    {
        // Might throw, might complete
        // ...
        return true;
    }
}

Stephen Toub's WithCancellation extension method:

public static async Task<T> WithCancellation<T>(this Task<T> task, CancellationToken cancellationToken) 
{ 
    var tcs = new TaskCompletionSource<bool>(); 
    using(cancellationToken.Register(s => ((TaskCompletionSource<bool>)s).TrySetResult(true), tcs)) 
    if (task != await Task.WhenAny(task, tcs.Task)) 
        throw new OperationCanceledException(cancellationToken); 
    return await task; 
}

Related Links:

• Equivalent of ContinueWith(delegate, CancellationToken) with await continuation
• http://blogs.msdn.com/b/pfxteam/archive/2012/10/05/how-do-i-cancel-non-cancelable-async-operations.aspx
• https://stackoverflow.com/a/15673072/495262
• https://stackoverflow.com/a/17560746/495262

Answer 1

Writing a WithCancellation method can be done much simpler, in just one line of code:

public static Task WithCancellation(this Task task,
    CancellationToken token)
{
    return task.ContinueWith(t => t.GetAwaiter().GetResult(), token);
}
public static Task<T> WithCancellation<T>(this Task<T> task,
    CancellationToken token)
{
    return task.ContinueWith(t => t.GetAwaiter().GetResult(), token);
}

As for the operation you want to do, just using await instead of ContinueWith is just as easy as it sounds; you replace the ContinueWith with an await. Most of the little pieces can be cleaned up a lot though.

m_cts.Cancel();
m_cts = new CancellationTokenSource();
var result = await Task.Run(() => LongRunningAndMightThrow())
    .WithCancellation(m_cts.Token);
UpdateUI(result);

The changes are not huge, but they're there. You [probably] want to cancel the previous operation when starting a new one. If that requirement doesn't exist, remove the corresponding line. The cancellation logic is all already handled by WithCancellation, there is no need to throw explicitly if cancellation is requested, as that will already happen. There's no real need to store the task, or the cancellation token, as local variables. UpdateUI shouldn't accept a Task<bool>, it should just accept a boolean. The value should be unwrapped from the task before callingUpdateUI.

Answer 2

The following should be equivalent:

var task = Task.Run(() => LongRunningAndMightThrow());

m_cts = new CancellationTokenSource();
CancellationToken ct = m_cts.Token;

try
{
    await task.WithCancellation(ct);
}
finally
{
    ct.ThrowIfCancellationRequested();
    UpdateUI(task);
}

Notice that the try/finally is required for the case where the LongRunningAndMightThrow method faults, but by the time we return to the UI thread the CancellationToken has been triggered. Without it the returned outer Task would be faulted where in the original ContinueWith case, it would have been cancelled.