ILockProvider.cs

using System;
using System.Collections.Generic;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

namespace Polly.Utilities
{
	/// <summary>
	/// Defines operations for locks used by Polly policies.
	/// </summary>
	public interface ILockProviderAsync
	{
		/// <summary>
		/// Waits to acquire the lock.
		/// </summary>
		/// <param name="key">A string key being used by the execution.</param>
		/// <param name="context">The Polly execution context consuming this lock.</param>
		/// <param name="cancellationToken">A cancellation token to cancel waiting to acquire the lock.</param>
		/// <throws>OperationCanceledException, if the passed <paramref name="cancellationToken"/> is signaled before the lock is acquired.</throws>
		/// <throws>InvalidOperationException, invalid lock state</throws>
		ValueTask<IDisposable> AcquireLockAsync(string key, Context context, CancellationToken cancellationToken);
	}

	/// <summary>
	/// Defines operations for locks used by Polly policies.
	/// </summary>
	public interface ILockProvider
	{
		/// <summary>
		/// Waits to acquire the lock.
		/// </summary>
		/// <param name="key">A string key being used by the execution.</param>
		/// <param name="context">The Polly execution context consuming this lock.</param>
		/// <throws>InvalidOperationException, invalid lock state</throws>
		IDisposable AcquireLock(string key, Context context);
	}

	/// <summary>
	/// Lock provider that locks on a key per process. The locking mechanism is designed to be able
	/// to be requested and released on different threads if needed.
	/// </summary>
	public class ProcessLockProviderAsync : ILockProviderAsync
	{
		// TODO: Pass via IOC or other method instead of hard-coding static
		internal static readonly int[] keyLocks = new int[1024];

		private class ProcessLockAsync : IDisposable
		{
			private uint hash;
			private bool gotLock;

			[System.Runtime.CompilerServices.MethodImpl(System.Runtime.CompilerServices.MethodImplOptions.AggressiveInlining)]
			public void Dispose()
			{
				if (gotLock)
				{
					gotLock = false;
					ProcessLockProviderAsync.keyLocks[hash] = 0;
				}
                // else we do not care, it can be disposed in an error case and we will simply ignore that the key locks were not touched
			}

			[System.Runtime.CompilerServices.MethodImpl(System.Runtime.CompilerServices.MethodImplOptions.AggressiveInlining)]
			public async ValueTask<IDisposable> AcquireLockAsync(string key, Context context, CancellationToken cancellationToken)
			{
				// Monitor.Enter and Monitor.Exit are tied to a specific thread and are
				//  slower than this spin lock, which does not care about threads and will execute very
				//  quickly, regardless of lock contention
				// https://stackoverflow.com/questions/11001760/monitor-enter-and-monitor-exit-in-different-threads

				// Get a hash based on the key, use this to lock on a specific int in the array. The array is designed
				// to be small enough to not use very much memory, but large enough to avoid collisions.
				// Even if there is a collision, it will be resolved very quickly.
				hash = (uint)key.GetHashCode() % (uint)ProcessLockProviderAsync.keyLocks.Length;

				// To get the lock, we must change the int at hash index from a 0 to a 1. If the value is
				//  already a 1, we don't get the lock. The return value must be 0 (the original value of the int).
				// it is very unlikely to have any contention here, but if so, the spin cycle should be very short.
				// Parameter index 1 (value of 1) is the value to change to if the existing value (Parameter index 2) is 0.
				while (!cancellationToken.IsCancellationRequested && Interlocked.CompareExchange(ref ProcessLockProviderAsync.keyLocks[hash], 1, 0) == 1)
				{
					// give up a clock cycle, we want to get back and try to get the lock again very quickly
					await Task.Yield();
				}

				if (cancellationToken.IsCancellationRequested)
				{
					throw new OperationCanceledException(cancellationToken);
				}

				gotLock = true;
				return this;
			}
		}

		[System.Runtime.CompilerServices.MethodImpl(System.Runtime.CompilerServices.MethodImplOptions.AggressiveInlining)]
		public ValueTask<IDisposable> AcquireLockAsync(string key, Context context, CancellationToken cancellationToken)
		{
			return new ProcessLockAsync().AcquireLockAsync(key, context, cancellationToken);
		}
	}

	/// <summary>
	/// Lock provider that locks on a key per process. The locking mechanism is designed to be able
	/// to be requested and released on different threads if needed.
	/// </summary>
	public class ProcessLockProvider : ILockProvider
	{
		private class ProcessLock : IDisposable
		{
			private uint hash;
			private bool gotLock;

			[System.Runtime.CompilerServices.MethodImpl(System.Runtime.CompilerServices.MethodImplOptions.AggressiveInlining)]
			public void Dispose()
			{
				if (gotLock)
				{
					gotLock = false;
					ProcessLockProviderAsync.keyLocks[hash] = 0;
				}
                // if constructor had exception, we will not get in Dispose as the object will never be created, if the constructor succeeds, gotLock will always be true
                // we still use the gotLock bool in case of multiple dispose calls
			}

			[System.Runtime.CompilerServices.MethodImpl(System.Runtime.CompilerServices.MethodImplOptions.AggressiveInlining)]
			public ProcessLockAsync(string key, Context context)
			{
				hash = (uint)key.GetHashCode() % (uint)ProcessLockProviderAsync.keyLocks.Length;
				while (Interlocked.CompareExchange(ref ProcessLockProviderAsync.keyLocks[hash], 1, 0) == 1)
				{
					Task.Yield().GetAwaiter().GetResult();
				}
				gotLock = true;
			}
		}

		/// <inheritdoc />
		public IDisposable AcquireLock(string key, Context context)
		{
			return new ProcessLock(key, context);
		}
	}
}

You are welcome!