How to get effect of Task.WhenAny for a Task and CancellationToken?

Question

I have interactive task which in "worst" scenario is not executed at all, thus it is represented by TaskCompletionSource.

I would like to wait for either this task completes, or token which I received is cancelled -- whichever happens first. Perfect tool for such job would be Task.WhenAny, the only problem is it takes only tasks, and I have one Task and one CancellationToken.

How to wait (asynchronously, like Task.WhenAny) for the first event triggered -- completed task, or cancelled token?

async Task MyCodeAsync(CancellationToken token)
{
  var tcs = new TaskCompletionSource<UserData>(); // represents interactive part

  await Task.WhenAny(tcs.Task, token); // imaginary call

  UserData data = tcs.Task.Result; // user interacted, let's continue
  ...
}

I don't create/manage token, so I cannot change it. I have to deal with it.

Update: For such particular case one could use Register method on token to cancel the TaskCompletionSource. For more general method please see Matthew Watson answer.

Answer 1

You could just create an extra task that returns when the cancel token's wait handle is signalled:

var factory = new CancellationTokenSource();
var token   = factory.Token;

await Task.WhenAny(
    Task.Run(() => token.WaitHandle.WaitOne()),
    myTask());

(However, be aware that this - while simple - does use up an extra thread, which is clearly not ideal. See later for an alternate solution which doesn't use an extra thread.)

If you want to check which task completed, you will have to keep a copy of the tasks before calling WhenAny() so you can compare them to the return value, for example:

using System;
using System.Threading;
using System.Threading.Tasks;

namespace ConsoleApp1
{
    class Program
    {
        static async Task Main()
        {
            var factory = new CancellationTokenSource(1000); // Change to 3000 for different result.
            var token   = factory.Token;
            var task    = myTask();

            var result = await Task.WhenAny(
                Task.Run(() => token.WaitHandle.WaitOne()),
                task);

            if (result == task)
                Console.WriteLine("myTask() completed");
            else
                Console.WriteLine("cancel token was signalled");
        }

        static async Task myTask()
        {
            await Task.Delay(2000);
        }
    }
}

---

If you don't want to waste an entire thread waiting for the cancellation token to be signalled, you can use CancellationToken.Register() to register a callback with which you can set the result of a TaskCompletionSource:

(Lifted from here)

public static Task WhenCanceled(CancellationToken cancellationToken)
{
    var tcs = new TaskCompletionSource<bool>();
    cancellationToken.Register(s => ((TaskCompletionSource<bool>) s).SetResult(true), tcs);
    return tcs.Task;
}

You can then use that as follows:

using System;
using System.Threading;
using System.Threading.Tasks;

namespace ConsoleApp1
{
    class Program
    {
        static async Task Main()
        {
            var factory = new CancellationTokenSource(1000);
            var token   = factory.Token;
            var task    = myTask();

            var result = await Task.WhenAny(
                WhenCanceled(token),
                task);

            if (result == task)
                Console.WriteLine("myTask() completed");
            else
                Console.WriteLine("cancel token was signalled");
        }

        public static Task WhenCanceled(CancellationToken cancellationToken)
        {
            var tcs = new TaskCompletionSource<bool>();
            cancellationToken.Register(s => ((TaskCompletionSource<bool>) s).SetResult(true), tcs);
            return tcs.Task;
        }

        static async Task myTask()
        {
            await Task.Delay(2000);
        }
    }
}

This is a preferable approach for the general case.

Answer 2

With this scenario, you have to be extremely careful about leaks. In particular, having objects referenced by delegates that are registered to a long-lived CancellationToken.

The approach that I eventually ended up with in my AsyncEx library looks like this:

public static async Task<T> WaitAsync<T>(this Task<T> task, CancellationToken token)
{
  var tcs = new TaskCompletionSource<T>();
  using (token.Register(() => tcs.TrySetCanceled(token), useSynchronizationContext: false)
    return await await Task.WhenAny(task, tcs.Task).ConfigureAwait(false);
}

The code above ensures that the registration is disposed if the CancellationToken is not canceled.

Usage:

async Task MyCodeAsync(CancellationToken token)
{
  UserData data = await userDataTask.WaitAsync(token);
}