|
// Copyright (c) .NET Foundation. All rights reserved.
// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
using System;
using System.Collections.Generic;
using System.Threading;
using System.Threading.Tasks;
namespace NuGet.Common
{
/// <summary>A class that allows one thread at a time synchronization, but is fine grained per key.
/// For example, the key can be a resource name that must only be accessed one at a time,
/// but different resources can be accessed concurrently.</summary>
internal sealed class KeyedLock : IDisposable
{
/// <summary>The dictionary that contains a <see cref="LockState"/> for each key.</summary>
/// <remarks>Both reading and modifying this dictionary must be synchronized though <see cref="_dictionaryLock"/>.</remarks>
private readonly Dictionary<string, LockState> _locks;
/// <summary>A lock to synchronize reading and modifying <see cref="_locks"/>.</summary>
private readonly SemaphoreSlim _dictionaryLock;
internal KeyedLock()
{
_locks = new Dictionary<string, LockState>();
_dictionaryLock = new SemaphoreSlim(initialCount: 1);
}
internal async Task EnterAsync(string key, CancellationToken token)
{
if (key == null)
{
throw new ArgumentNullException(nameof(key));
}
LockState lockState;
// Get the dictionary lock, so no other call to Enter[Async] or Exit[Async] can modify
// _locks while the current call does.
await _dictionaryLock.WaitAsync(token);
try
{
lockState = GetOrCreate(key);
}
finally
{
_dictionaryLock.Release();
}
try
{
// wait on the key-specific lock
await lockState.Semaphore.WaitAsync(token);
}
catch
{
// GetOrCreate(key) increments the lock state counter. Since this task failed to obtain the lock
// and is no longer waiting for it, decrease the count so that it will be eligible for cleanup
// when the code holding the lock releases it.
// If Exit[Async] ran for this key after the above WaitAsync failed, but before this catch block runs,
// it means that the LockState won't be removed from the dictionary, so technically a memory leak.
// But there isn't a multi-threaded correctness problem, the key still only allows one thread at a time
// to access it.
Interlocked.Decrement(ref lockState.Count);
throw;
}
}
internal void Enter(string key)
{
if (key == null)
{
throw new ArgumentNullException(nameof(key));
}
LockState lockState;
// Get the dictionary lock, so no other call to Enter[Async] or Exit[Async] can modify
// _locks while the current call does.
_dictionaryLock.Wait(CancellationToken.None);
try
{
lockState = GetOrCreate(key);
}
finally
{
_dictionaryLock.Release();
}
try
{
// wait on the key-specific lock
lockState.Semaphore.Wait(CancellationToken.None);
}
catch
{
// GetOrCreate(key) increments the lock state counter. Since this task failed to obtain the lock
// and is no longer waiting for it, decrease the count so that it will be eligible for cleanup
// when the code holding the lock releases it.
// If Exit[Async] ran for this key after the above WaitAsync failed, but before this catch block runs,
// it means that the LockState won't be removed from the dictionary, so technically a memory leak.
// But there isn't a multi-threaded correctness problem, the key still only allows one thread at a time
// to access it.
Interlocked.Decrement(ref lockState.Count);
throw;
}
}
private LockState GetOrCreate(string key)
{
// The caller holds the dictionary lock, so we know no other call to Enter[Async] or Exit[Async] can be
// running at this instant.
if (_locks.TryGetValue(key, out var lockState))
{
// LockState.Count is a reference counter for the number of threads/tasks that want access to this key's lock
// Since the current thread wants access, increment the counter. Doing this while we hold the dictionary lock
// ensures that Exit[Async] trying to run at the same time won't have timing issues with regards to changing
// and reading the value at the same time.
Interlocked.Increment(ref lockState.Count);
}
else
{
// Key doesn't yet exist, so create one and set the initial count to 1, indicating that the current thread
// has interest in the key's LockState.
lockState = new LockState();
_locks[key] = lockState;
}
return lockState;
}
internal async Task ExitAsync(string key)
{
if (key == null)
{
throw new ArgumentNullException(nameof(key));
}
// Get the dictionary lock, so no other call to Enter[Async] or Exit[Async] can modify
// _locks while the current call does.
await _dictionaryLock.WaitAsync();
try
{
Cleanup(key);
}
finally
{
_dictionaryLock.Release();
}
}
internal void Exit(string key)
{
if (key == null)
{
throw new ArgumentNullException(nameof(key));
}
// Get the dictionary lock, so no other call to Enter[Async] or Exit[Async] can modify
// _locks while the current call does.
_dictionaryLock.Wait();
try
{
Cleanup(key);
}
finally
{
_dictionaryLock.Release();
}
}
private void Cleanup(string key)
{
// The caller holds the dictionary lock, so we know no other call to Enter[Async] or Exit[Async] can be
// running at this instant.
var lockState = _locks[key];
// Release the per-key lock, allowing any Enter[Async] to now obtain the lock
lockState.Semaphore.Release();
// Decrement the counter while still holding the dictionary lock to ensure that no other Enter[Async]
// or Exit[Async] can cause timing issues.
var count = Interlocked.Decrement(ref lockState.Count);
if (count == 0)
{
// count == 0 means that this was the only/last thread accessing this key's LockState. Therefore, we
// can dispose of it and remove the key from the dictionary.
lockState.Dispose();
_locks.Remove(key);
}
}
/// <summary>Nested class to hold the state of per-key locks.</summary>
private sealed class LockState : IDisposable
{
/// <summary>The synchronization object used to ensure only 1 thread at a time can obtain the key's lock.</summary>
public SemaphoreSlim Semaphore { get; }
/// <summary>A counter of how many threads/tasks have interest in the key's lock. When this reaches zero,
/// it means no more threads or tasks want access to the resource, and the dictionary can clear its memory
/// of the key and this state.</summary>
public int Count;
public LockState()
{
Semaphore = new SemaphoreSlim(initialCount: 1);
Count = 1;
}
public void Dispose()
{
Semaphore.Dispose();
GC.SuppressFinalize(this);
}
}
public void Dispose()
{
_dictionaryLock.Dispose();
foreach (var kvp in _locks)
{
kvp.Value.Semaphore.Dispose();
}
_locks.Clear();
}
}
}
|