Async method without cancellation support, do it my way.

In the last post, I talked about Dennis Doomen’s LiquidProjects project and the challenge they faced with asynchronous APIs that were not cancelable. Dennis came up with the following solution to the infinitely running Task.Delay operations.

public static class TaskExtensions {
    public static async Task<TResult> WithWaitCancellation<TResult>(this Task<TResult> task,
            CancellationToken cancellationToken) {
        using (var combined = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken))            
        {
            var delayTask = Task.Delay(Timeout.Infinite, combined.Token);
            Task completedTask = await Task.WhenAny(task, delayTask);
            if (completedTask == task)
            {
                combined.Cancel();
                return await task;
            }
            else
            {
                cancellationToken.ThrowIfCancellationRequested();
                throw new InvalidOperationException("Infinite delay task completed.");
            }
        }
    }
}

The code above creates a linked token source, an infinitely delayed task which observes the token referenced by the linked token source. When the outer token source cancels the linked token source will also cancel the token owned by it. When the task returned by Task.WhenAny is the actual work task then the linked token source is canceled. This will automatically cancel the delayed task. The benefit of this implementation is its simplicity. The runtime performance is good but we can do better especially when we’d be using the extension method over and over again we can do a few optimizations. Here is a possible approach

public static Task<TResult> WaitWithCancellation<TResult>(this Task<TResult> task, CancellationToken token = default(CancellationToken))
{
    var tcs = new TaskCompletionSource<TResult>();
    var registration = token.Register(s => {
        var source = (TaskCompletionSource<TResult>) s;
        source.TrySetCanceled();
    }, tcs);

    task.ContinueWith((t, s) => {
        var tcsAndRegistration = (Tuple<TaskCompletionSource<TResult>, CancellationTokenRegistration>) s;

        if (t.IsFaulted && t.Exception!= null) {
            tcsAndRegistration.Item1.TrySetException(t.Exception.GetBaseException());
        }

        if (t.IsCanceled) {
            tcsAndRegistration.Item1.TrySetCanceled();
        }

        if (t.IsCompleted) {
            tcsAndRegistration.Item1.TrySetResult(t.Result);
        }

        tcsAndRegistration.Item2.Dispose();
    }, Tuple.Create(tcs, registration), CancellationToken.None, TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.Default);

    return tcs.Task;
}

Instead of using a delayed task with an infinite time span we use a TaskCompletionSource. The TaskCompletionSource represents a task which is only completed, canceled or faulted when we instruct the source accordingly. So essentially this is an infinitely running task. On the token passed in we register a registration which will get triggered when the token gets canceled. When the token gets canceled, we set the completion source to canceled as well. On the actual task, we register a continuation which sets the state of the task completion source according to the state of the antecedent task (which is the actual worker task). Inside the continuation, we also dispose the token registration not to leak memory. The continuation is scheduled with a runtime hint ExecuteSynchronously to tell the TPL runtime that the work running inside the continuation is short lived and can try to be executed while completing the antecedent task.

To avoid closure allocations we use for the token registration and the continuation the overloads which allow us to pass in a state object. Where we need to pass in multiple things into the delegate we use a tuple to represent those items. Furthermore, we use the TrySet methods to set the task completion source in a safe manner into its final state to avoid raising exceptions that could occur due to races.

Let’s see what the runtime differences of both code snippets are:

You can try it yourself with BenchmarkDotNet, and the tests found on my MicroBenchmark repo. So the second snippet is more complex to understand but it is overall faster and slightly better because it allocates less garbage and omits the asynchronous state machine generation.

Please don’t get me wrong this is slightly esoteric. Depending on your application or system you are building, I would usually go for Dennis’ approach because it is simpler to understand and good enough. If you need a bit more performance and want to save allocations the second approach written by me might come in handy.

Thanks again Dennis for this cool asynchronous challenge.