How to combine TaskCompletionSource and CancellationTokenSource?
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.
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.
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.