// <copyright file="Waiter.cs" company="JP Dillingham">
// Copyright (c) JP Dillingham. All rights reserved.
//
// This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License
// as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty
// of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.See the GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License along with this program. If not, see https://www.gnu.org/licenses/.
// </copyright>
namespace Soulseek
{
using System;
using System.Collections.Concurrent;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
/// <summary>
/// Enables await-able server messages.
/// </summary>
internal class Waiter : IWaiter
{
private const int DefaultTimeoutValue = 5000;
/// <summary>
/// Initializes a new instance of the <see cref="Waiter"/> class with the default timeout.
/// </summary>
public Waiter()
: this(DefaultTimeoutValue)
{
}
/// <summary>
/// Initializes a new instance of the <see cref="Waiter"/> class with the specified <paramref name="defaultTimeout"/>.
/// </summary>
/// <param name="defaultTimeout">The default timeout duration for message waits.</param>
public Waiter(int defaultTimeout)
{
DefaultTimeout = defaultTimeout;
}
/// <summary>
/// Gets the default timeout duration, in milliseconds.
/// </summary>
public int DefaultTimeout { get; private set; }
private bool Disposed { get; set; }
private ConcurrentDictionary<WaitKey, ReaderWriterLockSlim> Locks { get; } = new ConcurrentDictionary<WaitKey, ReaderWriterLockSlim>();
private ConcurrentDictionary<WaitKey, ConcurrentQueue<PendingWait>> Waits { get; } = new ConcurrentDictionary<WaitKey, ConcurrentQueue<PendingWait>>();
/// <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>
/// Cancels all waits.
/// </summary>
public void CancelAll()
{
var keys = Waits.Keys.ToList();
foreach (var key in keys)
{
Cancel(key);
}
}
/// <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>
/// 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>
/// Disposes this instance.
/// </summary>
public void Dispose()
{
Dispose(true);
GC.SuppressFinalize(this);
}
/// <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));
}
/// <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>
/// 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, in milliseconds.</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>
/// 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 MonitorWaits(), 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>
/// <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"/> 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>
/// Disposes this instance.
/// </summary>
/// <param name="disposing">A value indicating whether disposal is in progress.</param>
protected virtual void Dispose(bool disposing)
{
if (!Disposed)
{
if (disposing)
{
CancelAll();
}
Disposed = true;
}
}
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>
/// The composite value for the wait dictionary.
/// </summary>
internal class PendingWait : IDisposable
{
/// <summary>
/// Initializes a new instance of the <see cref="PendingWait"/> class.
/// </summary>
/// <param name="taskCompletionSource">The task completion source for the wait task.</param>
/// <param name="timeout">The number of milliseconds after which the wait is to time out.</param>
/// <param name="cancelAction">The action to invoke when the task is cancelled.</param>
/// <param name="timeoutAction">The action to invoke when the task times out.</param>
/// <param name="cancellationToken">The cancellation token for the wait.</param>
public PendingWait(dynamic taskCompletionSource, int timeout, Action cancelAction, Action timeoutAction, CancellationToken cancellationToken)
{
TaskCompletionSource = taskCompletionSource;
Timeout = timeout;
CancelAction = cancelAction;
TimeoutAction = timeoutAction;
CancellationToken = cancellationToken;
}
/// <summary>
/// Gets the task completion source for the wait task.
/// </summary>
public dynamic TaskCompletionSource { get; }
/// <summary>
/// Gets the number of milliseconds after which the wait is to time out.
/// </summary>
public int Timeout { get; }
private Action CancelAction { get; set; }
private CancellationToken CancellationToken { get; set; }
private CancellationTokenRegistration CancellationTokenRegistration { get; set; }
private bool Disposed { get; set; }
private Action TimeoutAction { get; set; }
private CancellationTokenRegistration TimeoutTokenRegistration { get; set; }
private CancellationTokenSource TimeoutTokenSource { get; set; }
/// <summary>
/// Releases the managed and unmanaged resources used by the <see cref="PendingWait"/>.
/// </summary>
public void Dispose()
{
Dispose(true);
GC.SuppressFinalize(this);
}
/// <summary>
/// Register cancellation and timeout actions.
/// </summary>
public void Register()
{
CancellationTokenRegistration = CancellationToken.Register(() => CancelAction());
TimeoutTokenSource = new CancellationTokenSource(Timeout);
TimeoutTokenRegistration = TimeoutTokenSource.Token.Register(() => TimeoutAction());
}
/// <summary>
/// Releases the managed and unmanaged resources used by the <see cref="PendingWait"/>.
/// </summary>
/// <param name="disposing">A value indicating whether the object is being disposed.</param>
protected virtual void Dispose(bool disposing)
{
if (!Disposed)
{
if (disposing)
{
CancellationTokenRegistration.Dispose();
TimeoutTokenSource.Dispose();
TimeoutTokenRegistration.Dispose();
}
Disposed = true;
}
}
}
}
}