How to combine TaskCompletionSource and CancellationTokenSource?

Question

I have such code (simplified here) which awaits finishing task:

var task_completion_source = new TaskCompletionSource<bool>();
observable.Subscribe(b => 
   { 
      if (b) 
          task_completion_source.SetResult(true); 
   });
await task_completion_source.Task;    

The idea is to subscribe and wait for the true in the stream of booleans. This finishes the "task" and I can move on beyond the await.

However I would like to cancel -- but not subscription, but awaiting. I would like to pass cancel token (somehow) to task_completion_source so when I cancel the token source, the await will move on.

How to do it?

Update: CancellationTokenSource is external to this code, all I have here is the token from it.

Answer 1

If I understand you correctly, you can do it like this:

using (cancellationToken.Register(() => {
    // this callback will be executed when token is cancelled
    task_comletion_source.TrySetCanceled();
})) {
    // ...
    await task_comletion_source.Task;
}

Note that it will throw an exception on your await, which you have to handle.

Answer 2

I recommend that you do not build this yourself. There are a number of edge cases around cancellation tokens that are tedious to get right. For example, if the registration returned from Register is never disposed, you can end up with a resource leak.

Instead, you can use the Task.WaitAsync extension method from my AsyncEx.Tasks library:

var task_completion_source = new TaskCompletionSource<bool>();
observable.Subscribe(b => 
{ 
  if (b) 
    task_completion_source.SetResult(true); 
});
await task_completion_source.Task.WaitAsync(cancellationToken);

On a side note, I'd strongly encourage you to use ToTask rather than an explicit TaskCompletionSource. Again, ToTask handles edge cases nicely for you.

Answer 3

Here was my stab at writing this myself. I almost made the mistake for not disposing the Register (thanks to Stephen Cleary)

    /// <summary>
    /// This allows a TaskCompletionSource to be await with a cancellation token and timeout.
    /// 
    /// Example usable:
    /// 
    ///     var tcs = new TaskCompletionSource<bool>();
    ///           ...
    ///     var result = await tcs.WaitAsync(timeoutTokenSource.Token);
    /// 
    /// A TaskCanceledException will be thrown if the given cancelToken is canceled before the tcs completes or errors. 
    /// </summary>
    /// <typeparam name="TResult">Result type of the TaskCompletionSource</typeparam>
    /// <param name="tcs">The task completion source to be used  </param>
    /// <param name="cancelToken">This method will throw an OperationCanceledException if the cancelToken is canceled</param>
    /// <param name="timeoutMs">This method will throw a TimeoutException if it doesn't complete within the given timeout, unless the timeout is less then or equal to 0 or Timeout.Infinite</param>
    /// <param name="updateTcs">If this is true and the given cancelToken is canceled then the underlying tcs will also be canceled.  If this is true a timeout occurs the underlying tcs will be faulted with a TimeoutException.</param>
    /// <returns>The tcs.Task</returns>
    public static async Task<TResult> WaitAsync<TResult>(this TaskCompletionSource<TResult> tcs, CancellationToken cancelToken, int timeoutMs = Timeout.Infinite, bool updateTcs = false)
    {
        // The overrideTcs is used so we can wait for either the give tcs to complete or the overrideTcs.  We do this using the Task.WhenAny method.
        // one issue with WhenAny is that it won't return when a task is canceled, it only returns when a task completes so we complete the
        // overrideTcs when either the cancelToken is canceled or the timeoutMs is reached.
        //
        var overrideTcs = new TaskCompletionSource<TResult>();
        using( var timeoutCancelTokenSource = (timeoutMs <= 0 || timeoutMs == Timeout.Infinite) ? null : new CancellationTokenSource(timeoutMs) )
        {
            var timeoutToken = timeoutCancelTokenSource?.Token ?? CancellationToken.None;
            using( var linkedTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancelToken, timeoutToken) )
            {
                // This method is called when either the linkedTokenSource is canceled.  This lets us assign a value to the overrideTcs so that
                // We can break out of the await WhenAny below.
                //
                void CancelTcs()
                {
                    if( updateTcs && !tcs.Task.IsCompleted )
                    {
                        // ReSharper disable once AccessToDisposedClosure (in this case, CancelTcs will never be called outside the using)
                        if( timeoutCancelTokenSource?.IsCancellationRequested ?? false )
                            tcs.TrySetException(new TimeoutException($"WaitAsync timed out after {timeoutMs}ms"));
                        else
                            tcs.TrySetCanceled();
                    }

                    overrideTcs.TrySetResult(default(TResult));
                }

                using( linkedTokenSource.Token.Register(CancelTcs) )
                {
                    try
                    {
                        await Task.WhenAny(tcs.Task, overrideTcs.Task);
                    }
                    catch { /* ignore */ }

                    // We always favor the result from the given tcs task if it has completed.
                    //
                    if( tcs.Task.IsCompleted )
                    {
                        // We do another await here so that if the tcs.Task has faulted or has been canceled we won't wrap those exceptions
                        // in a nested exception.  While technically accessing the tcs.Task.Result will generate the same exception the
                        // exception will be wrapped in a nested exception.  We don't want that nesting so we just await.
                        await tcs.Task;
                        return tcs.Task.Result;
                    }

                    // It wasn't the tcs.Task that got us our of the above WhenAny so go ahead and timeout or cancel the operation.
                    //
                    if( timeoutCancelTokenSource?.IsCancellationRequested ?? false )
                        throw new TimeoutException($"WaitAsync timed out after {timeoutMs}ms");

                    throw new OperationCanceledException();
                }
            }
        }
    }

This with throw a TaskCanceledException if the cancelToken is canceled before the tcs gets a result or errors.