using System;
using System.Collections.Generic;
using System.Threading;
using System.Threading.Tasks;
namespace AppsensPro.Infrastructure.Threading
{
public sealed class KeyedLock<TKey>
{
/// <summary>
/// Simple wrapper to allow counting how many references an object has
/// </summary>
/// <typeparam name="T"></typeparam>
private sealed class RefCounted<T>
{
public RefCounted(T value)
{
RefCount = 1;
Value = value;
}
public int RefCount { get; set; }
public T Value { get; private set; }
}
/// <summary>
/// Internally maintained dictionary of <see cref="SemaphoreSlim"/>
/// </summary>
private readonly Dictionary<TKey, RefCounted<SemaphoreSlim>> _semaphoreSlims
= new Dictionary<TKey, RefCounted<SemaphoreSlim>>();
/// <summary>
/// Atomically get or create a value, with the associated <see cref="RefCounted{T}"/> incremented
/// </summary>
/// <param name="key"></param>
/// <returns></returns>
private SemaphoreSlim GetOrCreate(TKey key)
{
RefCounted<SemaphoreSlim> item;
lock (_semaphoreSlims)
{
if (_semaphoreSlims.TryGetValue(key, out item))
{
++item.RefCount;
}
else
{
item = new RefCounted<SemaphoreSlim>(new SemaphoreSlim(1, 1));
_semaphoreSlims[key] = item;
}
}
return item.Value;
}
/// <summary>
/// Lock with specified key, returning a disposable <see cref="Releaser"/>
/// </summary>
/// <param name="key"></param>
/// <returns></returns>
public IDisposable Lock(TKey key)
{
GetOrCreate(key).Wait();
return new Releaser(key, _semaphoreSlims);
}
/// <summary>
/// Asynchronously lock with specified key, returning a disposable <see cref="Releaser"/>
/// </summary>
/// <param name="key"></param>
/// <returns></returns>
public async Task<IDisposable> LockAsync(TKey key)
{
await GetOrCreate(key).WaitAsync().ConfigureAwait(false);
return new Releaser(key, _semaphoreSlims);
}
/// <summary>
/// The returned Disposable <see cref="Releaser"/>.
/// The lock will only be removed from dictionary if exists no reference left
/// </summary>
private sealed class Releaser : IDisposable
{
private TKey _key;
private readonly Dictionary<TKey, RefCounted<SemaphoreSlim>> _locks;
public Releaser(TKey key, Dictionary<TKey, RefCounted<SemaphoreSlim>> locks)
{
_key = key;
_locks = locks;
}
public void Dispose()
{
RefCounted<SemaphoreSlim> item;
lock (_locks)
{
item = _locks[_key];
--item.RefCount;
if (item.RefCount == 0)
_locks.Remove(_key);
}
item.Value.Release();
}
}
}
}
