Equivalent of ContinueWith(delegate, CancellationToken) with await continuation

Question

I have that situation:

private Task LongRunningTask = /* Something */;

private void DoSomethingMore(Task previousTask) { }

public Task IndependentlyCancelableSuccessorTask(CancellationToken cancellationToken)
{
    return LongRunningTask.ContinueWith(DoSomethingMore, cancellationToken);
}

In particular, the behavior that interests me here is detailed in MSDN's page about Continuation Tasks in the following terms:

A continuation goes into the Canceled state in these scenarios:

• [...]
• When the continuation was passed a System.Threading.CancellationToken as an argument and the IsCancellationRequested property of the token is true before the continuation runs. In such a case, the continuation does not start and it transitions to the Canceled state.

The code above works. However, I am in the process of converting as many as possible of my continuations to using the await keyword.

Is there an equivalent using await that would allow the continuation to be canceled before the awaited task completes?

Answer 1

The following should do it, albeit it looks a bit awkward:

private Task LongRunningTask = /* Something */;

private void DoSomethingMore() { }

public async Task IndependentlyCancelableSuccessorTask(
    CancellationToken cancellationToken)
{
    cancellationToken.ThrowIfCancellationRequested();

    var tcs = new TaskCompletionSource<bool>();
    using (cancellationToken.Register(() => tcs.TrySetCanceled()))
        await Task.WhenAny(LongRunningTask, tcs.Task);

    cancellationToken.ThrowIfCancellationRequested();
    DoSomethingMore();
}

[UPDATE] Following svick's suggestion, here it is shaped as a helper, based on Stephen Toub's Implementing Then with Await pattern:

public static class TaskExt
{
    /// <summary>
    /// Use: await LongRunningTask.Then(DoSomethingMore, cancellationToken)
    /// </summary>
    public static async Task Then(
        this Task antecedent, Action continuation, CancellationToken token)
    {
        await antecedent.When(token);
        continuation();
    }

    /// <summary>
    /// Use: await LongRunningTask.When(cancellationToken)
    /// </summary>
    public static async Task When(
        this Task antecedent, CancellationToken token)
    {
        token.ThrowIfCancellationRequested();

        var tcs = new TaskCompletionSource<Empty>();
        using (token.Register(() => tcs.TrySetCanceled()))
            await Task.WhenAny(antecedent, tcs.Task);

        token.ThrowIfCancellationRequested();
    }

    struct Empty { };
}

Perhaps, the first ThrowIfCancellationRequested() is redundant, but I haven't thoroughly considered all edge cases.

Answer 2

While this answer is conceptually the same as Noseratio's, I am not satisfied by a few details of the implementation, and as such am publishing my proposed implementation of the helper so that it can be commented on by other people on this question.

public static async Task<TResult> WhenNotCanceled<TResult>(this Task<TResult> mainTask, CancellationToken cancellationToken)
{
    if (!cancellationToken.CanBeCanceled) {
        return await mainTask.ConfigureAwait(false);
    }

    cancellationToken.ThrowIfCancellationRequested();

    Task<TResult> completedTask;

    var cancellationTaskSource = new TaskCompletionSource<TResult>();
    using (cancellationToken.Register(() => cancellationTaskSource.TrySetCanceled(), useSynchronizationContext: false)
        completedTask = await Task.WhenAny(mainTask, cancellationTaskSource.Task).ConfigureAwait(false);

    cancellationToken.ThrowIfCancellationRequested();
    return await completedTask.ConfigureAwait(false);
}

public static async Task WhenNotCanceled(this Task mainTask, CancellationToken cancellationToken)
{
    if (!cancellationToken.CanBeCanceled) {
        await mainTask.ConfigureAwait(false);
        return;
    }

    cancellationToken.ThrowIfCancellationRequested();

    Task completedTask;

    var cancellationTaskSource = new TaskCompletionSource<object>();
    using (cancellationToken.Register(() => cancellationTaskSource.TrySetCanceled(), useSynchronizationContext: false)
        completedTask = await Task.WhenAny(mainTask, cancellationTaskSource.Task).ConfigureAwait(false);

    cancellationToken.ThrowIfCancellationRequested();
    await completedTask.ConfigureAwait(false);
}

Async pattern without cancel:

public async Task IndependentlyCancelableSuccessorTask()
{
    await LongRunningTask;
    DoSomethingMore();
}

Async pattern with cancel and WhenNotCanceled:

public async Task IndependentlyCancelableSuccessorTask(CancellationToken cancellationToken)
{
    await LongRunningTask.WhenNotCanceled(cancellationToken);
    DoSomethingMore();
}

Answer 3

My answer is only slightly different than @Jean Hominal's answer and incorporates @Noseratio's approach as well:

public static class TaskExtensionMethods
{
    public static Task<TResult> OrWhenCancelled<TResult>(this Task<TResult> mainTask, CancellationToken cancellationToken)
    {
        if (!cancellationToken.CanBeCanceled)
            return mainTask;

        return OrWhenCancelled_(mainTask, cancellationToken);
    }

    private static async Task<TResult> OrWhenCancelled_<TResult>(this Task<TResult> mainTask, CancellationToken cancellationToken)
    {
        Task cancellationTask = Task.Delay(Timeout.Infinite, cancellationToken);
        await Task.WhenAny(mainTask, cancellationTask).ConfigureAwait(false);

        cancellationToken.ThrowIfCancellationRequested();
        return await mainTask;
    }

    public static Task OrWhenCancelled(this Task mainTask, CancellationToken cancellationToken)
    {
        if (!cancellationToken.CanBeCanceled)
            return mainTask;

        return OrWhenCancelled_(mainTask, cancellationToken);
    }

    private static async Task OrWhenCancelled_(this Task mainTask, CancellationToken cancellationToken)
    {
        Task cancellationTask = Task.Delay(Timeout.Infinite, cancellationToken);
        await Task.WhenAny(mainTask, cancellationTask).ConfigureAwait(false);
        cancellationToken.ThrowIfCancellationRequested();
        await mainTask;
    }
}

Discussion:

• All of the solutions (including this one), do not correctly handle the case where the original ContinueWith specified a TaskScheduler. Specifically, consider a TaskScheduler created TaskScheduler.FromCurrentSynchronizationContext for usage in UI scenarios. In that case, with the original ContinueWith approach you were guaranteed that the cancellation token was checked prior to running the delegate but after already getting on to Main thread (see this answer). That is, the old approach has the nice effect of checking the Cancellation token "one last time" on the main thread prior to considering the result of the task (i.e. trumping whether the main task finished or faulted). This means that in addition to using these extension methods, the new code must wrap its await in a try/finally to do its final check of the CancellationToken :(. See this question.

• @Noseratio's solution could handle the above issue (if needed), but it has the downside of requiring that continuation be placed into a delegate. In my opinion, this defeats one of the big advantages of converting to using await: the code doesn't end up in a delegate, it is just after an await and reads like normal sequential code.

Notes:

• I wish I could have specified that the empty lambda never runs (i.e. instead of only running on cancellation), but the .ContinueWith method doesn't allow that. So, I (mostly arbitrarily chose OnlyOnCancelled)