/// <summary>
/// Handles one waiting task can be cancelled.
/// </summary>
/// <param name="status">The status of the waiting task being cancelled.</param>
private void OnWaitingTaskCancelled(WaitingCancellationStatus status)
{
CancellationTokenSource?overallCancellationSource = null;
lock (this.syncObject)
{
if (--this.outstandingWaitingCount != 0 || !this.isCancellationAllowed)
{
// when overall cancellation is not allowed, we cancel this single waiting task.
status.TrySetCanceled(status.CancellationToken);
}
else
{
// otherwise, we cancel the overall computation, when it is done, it will cancel the current waiting task.
overallCancellationSource = this.combinedCancellationTokenSource;
if (overallCancellationSource is not null)
{
this.combinedCancellationTokenSource = null;
this.isCancellationRequested = true;
}
}
}
if (overallCancellationSource is not null)
{
overallCancellationSource.Cancel();
overallCancellationSource.Dispose();
}
}
/// <summary>
/// Initializes a new instance of the <see cref="CancellableJoinComputation"/> class.
/// </summary>
/// <param name="taskFactory">A callback to create the task.</param>
/// <param name="allowCancelled">Whether the inner task can be cancelled.</param>
internal CancellableJoinComputation(Func <CancellationToken, Task> taskFactory, bool allowCancelled)
{
Requires.NotNull(taskFactory, nameof(taskFactory));
if (allowCancelled)
{
this.isCancellationAllowed = true;
this.combinedCancellationTokenSource = new CancellationTokenSource();
this.joinedWaitingList = new List <WaitingCancellationStatus>(capacity: 2);
}
this.InnerTask = taskFactory(this.combinedCancellationTokenSource?.Token ?? CancellationToken.None);
if (allowCancelled)
{
// Note: this continuation is chained asynchronously to prevent being inlined when we trigger the combined cancellation token.
this.InnerTask.ContinueWith(
(t, s) =>
{
var me = (CancellableJoinComputation)s !;
List <WaitingCancellationStatus> allWaitingTasks;
CancellationTokenSource?combinedCancellationTokenSource;
lock (me.syncObject)
{
Assumes.NotNull(me.joinedWaitingList);
allWaitingTasks = me.joinedWaitingList;
combinedCancellationTokenSource = me.combinedCancellationTokenSource;
me.joinedWaitingList = null;
me.combinedCancellationTokenSource = null;
}
combinedCancellationTokenSource?.Dispose();
if (t.IsCanceled)
{
for (int i = 0; i < allWaitingTasks.Count; i++)
{
WaitingCancellationStatus status = allWaitingTasks[i];
if (status.CancellationToken.IsCancellationRequested)
{
status.TrySetCanceled(status.CancellationToken);
}
else
{
status.TrySetCanceled();
}
status.Dispose();
}
}
else if (t.IsFaulted)
{
System.Collections.ObjectModel.ReadOnlyCollection <Exception> exceptions = t.Exception !.InnerExceptions;
for (int i = 0; i < allWaitingTasks.Count; i++)
{
WaitingCancellationStatus status = allWaitingTasks[i];
status.TrySetException(exceptions);
status.Dispose();
}
}
else
{
for (int i = 0; i < allWaitingTasks.Count; i++)
{
WaitingCancellationStatus status = allWaitingTasks[i];
status.TrySetResult(true);
status.Dispose();
}
}
},
this,
CancellationToken.None,
TaskContinuationOptions.RunContinuationsAsynchronously,
TaskScheduler.Default).Forget();
}
}
/// <summary>
/// Try to join the computation.
/// </summary>
/// <param name="isInitialTask">It is true for the initial task starting the computation. This must be called once right after the constructor.</param>
/// <param name="task">Returns a task which can be waited on.</param>
/// <param name="cancellationToken">A cancellation token to abort this waiting.</param>
/// <returns>It returns false, if the inner task is aborted. In which case, no way to join the existing computation.</returns>
internal bool TryJoinComputation(bool isInitialTask, [NotNullWhen(true)] out Task?task, CancellationToken cancellationToken)
{
if (!this.isCancellationAllowed)
{
task = this.JoinNotCancellableTaskAsync(isInitialTask, cancellationToken);
return(true);
}
if (cancellationToken.IsCancellationRequested)
{
if (isInitialTask)
{
// It is a corner case the cancellation token is triggered right after the first task starts. It may need cancel the inner task.
CancellationTokenSource?cancellationTokenSource = null;
lock (this.syncObject)
{
if (this.isCancellationAllowed && this.outstandingWaitingCount == 0 && this.combinedCancellationTokenSource is not null)
{
this.isCancellationRequested = true;
cancellationTokenSource = this.combinedCancellationTokenSource;
this.combinedCancellationTokenSource = null;
}
}
if (cancellationTokenSource is not null)
{
cancellationTokenSource.Cancel();
cancellationTokenSource.Dispose();
}
task = this.InnerTask;
return(true);
}
else
{
task = Task.FromCanceled(cancellationToken);
return(true);
}
}
// if the inner task is joined by a new uncancellable task, we will abandone the cancellation token source because we will never use it anymore.
// we do it outside of our lock.
CancellationTokenSource?combinedCancellationTokenSourceToDispose = null;
try
{
lock (this.syncObject)
{
if (this.isCancellationRequested)
{
// If the earlier computation is aborted, we cannot join it anymore.
task = null;
return(false);
}
if (this.InnerTask.IsCompleted)
{
task = this.InnerTask;
return(true);
}
if (!cancellationToken.CanBeCanceled)
{
// A single joined client which doesn't allow cancellation would turn the entire computation not cancellable.
combinedCancellationTokenSourceToDispose = this.combinedCancellationTokenSource;
this.combinedCancellationTokenSource = null;
this.isCancellationAllowed = false;
task = this.JoinNotCancellableTaskAsync(isInitialTask, CancellationToken.None);
}
else if (!this.isCancellationAllowed)
{
task = this.JoinNotCancellableTaskAsync(isInitialTask, cancellationToken);
}
else
{
Assumes.NotNull(this.joinedWaitingList);
WaitingCancellationStatus status;
// we need increase the outstanding count before creating WiatingCancellationStatus.
// Under a rare race condition the cancellation token can be trigger with this time frame, and lead OnWaitingTaskCancelled to be called recursively
// within this lock. It would be critical to make sure the outstandingWaitingCount to increase before decreasing there.
this.outstandingWaitingCount++;
try
{
status = new WaitingCancellationStatus(this, cancellationToken);
}
catch
{
this.outstandingWaitingCount--;
throw;
}
this.joinedWaitingList.Add(status);
task = status.Task;
}
}
}
finally
{
combinedCancellationTokenSourceToDispose?.Dispose();
}
return(true);
}
