/// <summary>
/// Cancels a <see cref="TaskCompletionSource{TResult}.Task"/> if a given <see cref="CancellationToken"/> is canceled.
/// </summary>
/// <typeparam name="T">The type of value returned by a successfully completed <see cref="Task{TResult}"/>.</typeparam>
/// <param name="taskCompletionSource">The <see cref="TaskCompletionSource{TResult}"/> to cancel.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/>.</param>
/// <param name="cancellationCallback">A callback to invoke when cancellation occurs.</param>
internal static void AttachCancellation <T>(this TaskCompletionSource <T> taskCompletionSource, CancellationToken cancellationToken, ICancellationNotification cancellationCallback = null)
{
Requires.NotNull(taskCompletionSource, nameof(taskCompletionSource));
if (cancellationToken.CanBeCanceled)
{
if (cancellationToken.IsCancellationRequested)
{
taskCompletionSource.TrySetCanceled(cancellationToken);
}
else
{
var tuple = new CancelableTaskCompletionSource <T>(taskCompletionSource, cancellationToken, cancellationCallback);
tuple.CancellationTokenRegistration = cancellationToken.Register(
s =>
{
var t = (CancelableTaskCompletionSource <T>)s;
t.TaskCompletionSource.TrySetCanceled(t.CancellationToken);
t.CancellationCallback?.OnCanceled();
},
tuple,
useSynchronizationContext: false);
taskCompletionSource.Task.ContinueWith(
(_, s) =>
{
var t = (CancelableTaskCompletionSource <T>)s;
t.CancellationTokenRegistration.Dispose();
},
tuple,
CancellationToken.None,
TaskContinuationOptions.ExecuteSynchronously,
TaskScheduler.Default);
}
}
}
/// <summary>
/// Initializes a new instance of the <see cref="CancelableTaskCompletionSource{T}"/> class.
/// </summary>
/// <param name="taskCompletionSource">The task completion source.</param>
/// <param name="cancellationToken">The cancellation token.</param>
/// <param name="cancellationCallback">A callback to invoke when cancellation occurs.</param>
internal CancelableTaskCompletionSource(TaskCompletionSource <T> taskCompletionSource, CancellationToken cancellationToken, ICancellationNotification cancellationCallback)
{
this.TaskCompletionSource = taskCompletionSource ?? throw new ArgumentNullException(nameof(taskCompletionSource));
this.CancellationToken = cancellationToken;
this.CancellationCallback = cancellationCallback;
}
/// <summary>
/// Cancels a <see cref="TaskCompletionSource{TResult}.Task"/> if a given <see cref="CancellationToken"/> is canceled.
/// </summary>
/// <typeparam name="T">The type of value returned by a successfully completed <see cref="Task{TResult}"/>.</typeparam>
/// <param name="taskCompletionSource">The <see cref="TaskCompletionSource{TResult}"/> to cancel.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/>.</param>
/// <param name="cancellationCallback">A callback to invoke when cancellation occurs.</param>
internal static void AttachCancellation <T>(this TaskCompletionSource <T> taskCompletionSource, CancellationToken cancellationToken, ICancellationNotification cancellationCallback = null)
{
Requires.NotNull(taskCompletionSource, nameof(taskCompletionSource));
if (cancellationToken.CanBeCanceled && !taskCompletionSource.Task.IsCompleted)
{
if (cancellationToken.IsCancellationRequested)
{
taskCompletionSource.TrySetCanceled(cancellationToken);
}
else
{
var tuple = new CancelableTaskCompletionSource <T>(taskCompletionSource, cancellationCallback, cancellationToken);
tuple.CancellationTokenRegistration = cancellationToken.Register(
s =>
{
var t = (CancelableTaskCompletionSource <T>)s;
if (t.TaskCompletionSource.TrySetCanceled(t.CancellationToken))
{
t.CancellationCallback?.OnCanceled();
}
},
tuple,
useSynchronizationContext: false);
// In certain race conditions, our continuation could execute inline. We could force it to always run
// asynchronously, but then in the common case it becomes less efficient.
// Instead, we will optimize for the common (no-race) case and detect if we were inlined, and if so, defer the work
// to avoid making our caller block for arbitrary code since CTR.Dispose blocks for in-progress cancellation notification to complete.
taskCompletionSource.Task.ContinueWith(
(_, s) =>
{
var t = (CancelableTaskCompletionSource <T>)s;
if (t.ContinuationScheduled || !t.OnOwnerThread)
{
// We're not executing inline... Go ahead and do the work.
t.CancellationTokenRegistration.Dispose();
}
else if (!t.CancellationToken.IsCancellationRequested) // If the CT is canceled, the CTR is implicitly disposed.
{
// We hit the race where the task is already completed another way,
// and our continuation is executing inline with our caller.
// Dispose our CTR from the threadpool to avoid blocking on 3rd party code.
ThreadPool.QueueUserWorkItem(
s2 =>
{
try
{
var t2 = (CancelableTaskCompletionSource <T>)s2;
t2.CancellationTokenRegistration.Dispose();
}
catch (Exception ex)
{
// Swallow any exception.
Report.Fail(ex.Message);
}
},
s);
}
},
tuple,
CancellationToken.None,
TaskContinuationOptions.ExecuteSynchronously,
TaskScheduler.Default);
tuple.ContinuationScheduled = true;
}
}
}
