/// <summary>
/// Completes the oldest wait matching the specified <paramref name="key"/> with the specified <paramref name="result"/>.
/// </summary>
/// <typeparam name="T">The wait result type.</typeparam>
/// <param name="key">The unique WaitKey for the wait.</param>
/// <param name="result">The wait result.</param>
public void Complete <T>(WaitKey key, T result)
{
if (Waits.TryGetValue(key, out var queue) && queue.TryDequeue(out var wait))
{
((TaskCompletionSource <T>)wait.TaskCompletionSource).SetResult(result);
}
}
/// <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);
}
/// <summary>
/// Throws the specified <paramref name="exception"/> on the oldest wait matching the specified <paramref name="key"/>.
/// </summary>
/// <param name="key">The unique WaitKey for the wait.</param>
/// <param name="exception">The Exception to throw.</param>
public void Throw(WaitKey key, Exception exception)
{
if (Waits.TryGetValue(key, out var queue) && queue.TryDequeue(out var wait))
{
wait.TaskCompletionSource.SetException(exception);
}
}
private void Disposition(WaitKey key, Action <PendingWait> action)
{
if (Waits.TryGetValue(key, out var queue) && queue.TryDequeue(out var wait))
{
action(wait);
wait.Dispose();
if (Locks.TryGetValue(key, out var recordLock))
{
// enter a read lock first; TryPeek and TryDequeue are atomic so there's no risky operation until later.
recordLock.EnterUpgradeableReadLock();
try
{
// clean up entries in the Waits and Locks dictionaries if the corresponding ConcurrentQueue is empty.
// this is tricky, because we don't want to remove a record if another thread is in the process of
// enqueueing a new wait.
if (queue.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. it's ok and expected that Wait() might add this record
// back to the dictionary as soon as this unblocks; we're preventing new waits from being discarded if
// they are added by another thread just prior to the TryRemove() operation below.
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 and it should be impossible to remove a record containing a non-empty queue
if (queue.IsEmpty)
{
Waits.TryRemove(key, out _);
Locks.TryRemove(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);
}
/// <summary>
/// Adds a new wait for the specified <paramref name="key"/> which does not time out.
/// </summary>
/// <typeparam name="T">The wait result type.</typeparam>
/// <param name="key">A unique WaitKey for the wait.</param>
/// <param name="cancellationToken">The cancellation token for the wait.</param>
/// <returns>A Task representing the wait.</returns>
public Task <T> WaitIndefinitely <T>(WaitKey key, CancellationToken?cancellationToken = null)
{
return(Wait <T>(key, int.MaxValue, cancellationToken));
}
/// <summary>
/// Adds a new wait for the specified <paramref name="key"/> which does not time out.
/// </summary>
/// <param name="key">A unique WaitKey for the wait.</param>
/// <param name="cancellationToken">The cancellation token for the wait.</param>
/// <returns>A Task representing the wait.</returns>
public Task WaitIndefinitely(WaitKey key, CancellationToken?cancellationToken = null)
{
return(WaitIndefinitely <object>(key, cancellationToken));
}
/// <summary>
/// Adds a new wait for the specified <paramref name="key"/> and with the specified <paramref name="timeout"/>.
/// </summary>
/// <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 Wait(WaitKey key, int?timeout = null, CancellationToken?cancellationToken = null)
{
return(Wait <object>(key, timeout, cancellationToken));
}
/// <summary>
/// Completes the oldest wait matching the specified <paramref name="key"/>.
/// </summary>
/// <param name="key">The unique WaitKey for the wait.</param>
public void Complete(WaitKey key)
{
Complete <object>(key, null);
}
/// <summary>
/// Completes the oldest wait matching the specified <paramref name="key"/> with the specified <paramref name="result"/>.
/// </summary>
/// <typeparam name="T">The wait result type.</typeparam>
/// <param name="key">The unique WaitKey for the wait.</param>
/// <param name="result">The wait result.</param>
public void Complete <T>(WaitKey key, T result)
{
Disposition(key, wait =>
((TaskCompletionSource <T>)wait.TaskCompletionSource).TrySetResult(result));
}
/// <summary>
/// Cancels the oldest wait matching the specified <paramref name="key"/>.
/// </summary>
/// <param name="key">The unique WaitKey for the wait.</param>
public void Cancel(WaitKey key)
{
Disposition(key, wait =>
wait.TaskCompletionSource.TrySetCanceled());
}
/// <summary>
/// Causes the oldest wait matching the specified <paramref name="key"/> to time out.
/// </summary>
/// <param name="key">The unique WaitKey for the wait.</param>
public void Timeout(WaitKey key)
{
Disposition(key, wait =>
wait.TaskCompletionSource.TrySetException(new TimeoutException($"The wait timed out after {wait.Timeout} milliseconds")));
}
/// <summary>
/// Throws the specified <paramref name="exception"/> on the oldest wait matching the specified <paramref name="key"/>.
/// </summary>
/// <param name="key">The unique WaitKey for the wait.</param>
/// <param name="exception">The Exception to throw.</param>
public void Throw(WaitKey key, Exception exception)
{
Disposition(key, wait =>
wait.TaskCompletionSource.TrySetException(exception));
}
