/// <summary>
/// Adds a new wait for the specified <paramref name="key"/> and with the specified <paramref name="timeout"/>.
/// </summary>
/// <typeparam name="T">The wait result type.</typeparam>
/// <param name="key">A unique WaitKey for the wait.</param>
/// <param name="timeout">The wait timeout.</param>
/// <param name="cancellationToken">The cancellation token for the wait.</param>
/// <returns>A Task representing the wait.</returns>
public Task <T> Wait <T>(WaitKey key, int?timeout = null, CancellationToken?cancellationToken = null)
{
timeout = timeout ?? DefaultTimeout;
var wait = new PendingWait()
{
TaskCompletionSource = new TaskCompletionSource <T>(TaskCreationOptions.RunContinuationsAsynchronously),
DateTime = DateTime.UtcNow,
TimeoutAfter = (int)timeout,
CancellationToken = cancellationToken,
};
var recordLock = Locks.GetOrAdd(key, new ReaderWriterLockSlim());
recordLock.EnterReadLock();
try
{
Waits.AddOrUpdate(key, new ConcurrentQueue <PendingWait>(new[] { wait }), (_, queue) =>
{
queue.Enqueue(wait);
return(queue);
});
}
finally
{
recordLock.ExitReadLock();
}
return(((TaskCompletionSource <T>)wait.TaskCompletionSource).Task);
}
/// <remarks>
/// Not thread safe; ensure this is invoked only by the timer within this class.
/// </remarks>
private void MonitorWaits(object sender, object e)
{
foreach (var record in Waits)
{
// it shouldn't be possible for a record to make it into Waits without a corresponding lock in Locks,
// but use GetOrAdd here anyway.
var recordLock = Locks.GetOrAdd(record.Key, new ReaderWriterLockSlim());
// enter a read lock first; TryPeek and TryDequeue are atomic so there's no risky operation until later.
recordLock.EnterUpgradeableReadLock();
try
{
if (record.Value.TryPeek(out var nextPendingWait))
{
if (nextPendingWait.CancellationToken != null && ((CancellationToken)nextPendingWait.CancellationToken).IsCancellationRequested)
{
if (record.Value.TryDequeue(out var cancelledWait))
{
cancelledWait.TaskCompletionSource.SetException(new OperationCanceledException("The wait was cancelled."));
}
}
else if (nextPendingWait.DateTime.AddSeconds(nextPendingWait.TimeoutAfter) < DateTime.UtcNow && record.Value.TryDequeue(out var timedOutWait))
{
timedOutWait.TaskCompletionSource.SetException(new TimeoutException($"The wait timed out after {timedOutWait.TimeoutAfter} seconds."));
}
}
if (record.Value.IsEmpty)
{
// enter the write lock to prevent Wait() (which obtains a read lock) from enqueing any more waits
// before we can delete the dictionary record
recordLock.EnterWriteLock();
try
{
// check the queue again to ensure Wait() didn't enqueue anything between the last check and when we
// entered the write lock. this is guarateed to be safe since we now have exclusive access to the record
if (record.Value.IsEmpty)
{
Waits.TryRemove(record.Key, out _);
Locks.TryRemove(record.Key, out _);
}
}
finally
{
recordLock.ExitWriteLock();
}
}
}
finally
{
recordLock.ExitUpgradeableReadLock();
}
}
}
/// <summary>
/// Adds a new wait for the specified <paramref name="key"/> and with the specified <paramref name="timeout"/>.
/// </summary>
/// <typeparam name="T">The wait result type.</typeparam>
/// <param name="key">A unique WaitKey for the wait.</param>
/// <param name="timeout">The wait timeout, in milliseconds.</param>
/// <param name="cancellationToken">The cancellation token for the wait.</param>
/// <returns>A Task representing the wait.</returns>
public Task <T> Wait <T>(WaitKey key, int?timeout = null, CancellationToken?cancellationToken = null)
{
timeout ??= DefaultTimeout;
cancellationToken ??= CancellationToken.None;
var taskCompletionSource = new TaskCompletionSource <T>(TaskCreationOptions.RunContinuationsAsynchronously);
var wait = new PendingWait(
taskCompletionSource,
timeout.Value,
cancelAction: () => Cancel(key),
timeoutAction: () => Timeout(key),
cancellationToken.Value);
// obtain a read lock for the key. this is necessary to prevent this code from adding a wait to the ConcurrentQueue
// while the containing dictionary entry is being cleaned up in Disposition(), effectively discarding the new wait.
#pragma warning disable IDE0067, CA2000 // Dispose objects before losing scope
var recordLock = Locks.GetOrAdd(key, new ReaderWriterLockSlim());
#pragma warning restore IDE0067, CA2000 // Dispose objects before losing scope
recordLock.EnterReadLock();
try
{
Waits.AddOrUpdate(key, new ConcurrentQueue <PendingWait>(new[] { wait }), (_, queue) =>
{
queue.Enqueue(wait);
return(queue);
});
}
finally
{
recordLock.ExitReadLock();
}
// defer registration to prevent the wait from being dispositioned prior to being successfully queued this is a
// concern if we are given a timeout of 0, or a cancellation token which is already cancelled
wait.Register();
return(((TaskCompletionSource <T>)wait.TaskCompletionSource).Task);
}
