Refactoring & extensibility

src/AspNetCoreRateLimit/AspNetCoreRateLimit.csproj

@@ -3,15 +3,16 @@
   <PropertyGroup>
     <TargetFramework>netstandard2.0</TargetFramework>
     <Description>ASP.NET Core rate limiting middleware</Description>
-    <VersionPrefix>1.0.6</VersionPrefix>
-    <Authors>Stefan Prodan</Authors>
+    <Authors>Stefan Prodan, Cristi Pufu</Authors>
     <AssemblyName>AspNetCoreRateLimit</AssemblyName>
     <PackageId>AspNetCoreRateLimit</PackageId>
     <PackageTags>aspnetcore;rate-limit;throttle</PackageTags>
     <PackageProjectUrl>https://github.com/stefanprodan/AspNetCoreRateLimit</PackageProjectUrl>
     <PackageLicenseUrl>http://opensource.org/licenses/MIT</PackageLicenseUrl>
     <RepositoryType>git</RepositoryType>
     <RepositoryUrl>https://github.com/stefanprodan/AspNetCoreRateLimit</RepositoryUrl>
+    <LangVersion>7.3</LangVersion>
+    <Version>3.0.0</Version>
   </PropertyGroup>
 
   <ItemGroup>

src/AspNetCoreRateLimit/AsyncKeyLock/AsyncKeyLock.cs

@@ -0,0 +1,133 @@
+// Copyright (c) Six Labors and contributors.
+// Licensed under the Apache License, Version 2.0.
+// Thanks to https://github.com/SixLabors/ImageSharp.Web/
+using System;
+using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace AspNetCoreRateLimit
+{
+    /// <summary>
+    /// The async key lock prevents multiple asynchronous threads acting upon the same object with the given key at the same time.
+    /// It is designed so that it does not block unique requests allowing a high throughput.
+    /// </summary>
+    internal sealed class AsyncKeyLock
+    {
+        /// <summary>
+        /// A collection of doorman counters used for tracking references to the same key.
+        /// </summary>
+        private static readonly Dictionary<string, AsyncKeyLockDoorman> Keys = new Dictionary<string, AsyncKeyLockDoorman>();
+
+        /// <summary>
+        /// A pool of unused doorman counters that can be re-used to avoid allocations.
+        /// </summary>
+        private static readonly Stack<AsyncKeyLockDoorman> Pool = new Stack<AsyncKeyLockDoorman>(MaxPoolSize);
+
+        /// <summary>
+        /// Maximum size of the doorman pool. If the pool is already full when releasing
+        /// a doorman, it is simply left for garbage collection.
+        /// </summary>
+        private const int MaxPoolSize = 20;
+
+        /// <summary>
+        /// SpinLock used to protect access to the Keys and Pool collections.
+        /// </summary>
+        private static SpinLock _spinLock = new SpinLock(false);
+
+        /// <summary>
+        /// Locks the current thread in read mode asynchronously.
+        /// </summary>
+        /// <param name="key">The key identifying the specific object to lock against.</param>
+        /// <returns>
+        /// The <see cref="Task{IDisposable}"/> that will release the lock.
+        /// </returns>
+        public async Task<IDisposable> ReaderLockAsync(string key)
+        {
+            AsyncKeyLockDoorman doorman = GetDoorman(key);
+
+            return await doorman.ReaderLockAsync().ConfigureAwait(false);
+        }
+
+        /// <summary>
+        /// Locks the current thread in write mode asynchronously.
+        /// </summary>
+        /// <param name="key">The key identifying the specific object to lock against.</param>
+        /// <returns>
+        /// The <see cref="Task{IDisposable}"/> that will release the lock.
+        /// </returns>
+        public async Task<IDisposable> WriterLockAsync(string key)
+        {
+            AsyncKeyLockDoorman doorman = GetDoorman(key);
+
+            return await doorman.WriterLockAsync().ConfigureAwait(false);
+        }
+
+        /// <summary>
+        /// Gets the doorman for the specified key. If no such doorman exists, an unused doorman
+        /// is obtained from the pool (or a new one is allocated if the pool is empty), and it's
+        /// assigned to the requested key.
+        /// </summary>
+        /// <param name="key">The key for the desired doorman.</param>
+        /// <returns>The <see cref="Doorman"/>.</returns>
+        private static AsyncKeyLockDoorman GetDoorman(string key)
+        {
+            AsyncKeyLockDoorman doorman;
+            bool lockTaken = false;
+            try
+            {
+                _spinLock.Enter(ref lockTaken);
+
+                if (!Keys.TryGetValue(key, out doorman))
+                {
+                    doorman = (Pool.Count > 0) ? Pool.Pop() : new AsyncKeyLockDoorman(ReleaseDoorman);
+                    doorman.Key = key;
+                    Keys.Add(key, doorman);
+                }
+
+                doorman.RefCount++;
+            }
+            finally
+            {
+                if (lockTaken)
+                {
+                    _spinLock.Exit();
+                }
+            }
+
+            return doorman;
+        }
+
+        /// <summary>
+        /// Releases a reference to a doorman. If the ref-count hits zero, then the doorman is
+        /// returned to the pool (or is simply left for the garbage collector to cleanup if the
+        /// pool is already full).
+        /// </summary>
+        /// <param name="doorman">The <see cref="Doorman"/>.</param>
+        private static void ReleaseDoorman(AsyncKeyLockDoorman doorman)
+        {
+            bool lockTaken = false;
+            try
+            {
+                _spinLock.Enter(ref lockTaken);
+
+                if (--doorman.RefCount == 0)
+                {
+                    Keys.Remove(doorman.Key);
+                    if (Pool.Count < MaxPoolSize)
+                    {
+                        doorman.Key = null;
+                        Pool.Push(doorman);
+                    }
+                }
+            }
+            finally
+            {
+                if (lockTaken)
+                {
+                    _spinLock.Exit();
+                }
+            }
+        }
+    }
+}

src/AspNetCoreRateLimit/AsyncKeyLock/AsyncKeyLockDoorman.cs

@@ -0,0 +1,171 @@
+// Copyright (c) Six Labors and contributors.
+// Licensed under the Apache License, Version 2.0.
+// Thanks to https://github.com/SixLabors/ImageSharp.Web/
+using System;
+using System.Collections.Generic;
+using System.Threading.Tasks;
+
+namespace AspNetCoreRateLimit
+{
+    /// <summary>
+    /// An asynchronous locker that provides read and write locking policies.
+    /// </summary>
+    internal sealed class AsyncKeyLockDoorman
+    {
+        private readonly Queue<TaskCompletionSource<Releaser>> _waitingWriters;
+        private readonly Task<Releaser> _readerReleaser;
+        private readonly Task<Releaser> _writerReleaser;
+        private readonly Action<AsyncKeyLockDoorman> _reset;
+        private TaskCompletionSource<Releaser> _waitingReader;
+        private int _readersWaiting;
+        private int _status;
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="AsyncKeyLockDoorman"/> class.
+        /// </summary>
+        /// <param name="reset">The reset action.</param>
+        public AsyncKeyLockDoorman(Action<AsyncKeyLockDoorman> reset)
+        {
+            _waitingWriters = new Queue<TaskCompletionSource<Releaser>>();
+            _waitingReader = new TaskCompletionSource<Releaser>();
+            _status = 0;
+
+            _readerReleaser = Task.FromResult(new Releaser(this, false));
+            _writerReleaser = Task.FromResult(new Releaser(this, true));
+            _reset = reset;
+        }
+
+        /// <summary>
+        /// Gets or sets the key that this doorman is mapped to.
+        /// </summary>
+        public string Key { get; set; }
+
+        /// <summary>
+        /// Gets or sets the current reference count on this doorman.
+        /// </summary>
+        public int RefCount { get; set; }
+
+        /// <summary>
+        /// Locks the current thread in read mode asynchronously.
+        /// </summary>
+        /// <returns>The <see cref="Task{Releaser}"/>.</returns>
+        public Task<Releaser> ReaderLockAsync()
+        {
+            lock (_waitingWriters)
+            {
+                if (_status >= 0 && _waitingWriters.Count == 0)
+                {
+                    ++_status;
+                    return _readerReleaser;
+                }
+                else
+                {
+                    ++_readersWaiting;
+                    return _waitingReader.Task.ContinueWith(t => t.Result);
+                }
+            }
+        }
+
+        /// <summary>
+        /// Locks the current thread in write mode asynchronously.
+        /// </summary>
+        /// <returns>The <see cref="Task{Releaser}"/>.</returns>
+        public Task<Releaser> WriterLockAsync()
+        {
+            lock (_waitingWriters)
+            {
+                if (_status == 0)
+                {
+                    _status = -1;
+                    return _writerReleaser;
+                }
+                else
+                {
+                    var waiter = new TaskCompletionSource<Releaser>();
+                    _waitingWriters.Enqueue(waiter);
+                    return waiter.Task;
+                }
+            }
+        }
+
+        private void ReaderRelease()
+        {
+            TaskCompletionSource<Releaser> toWake = null;
+
+            lock (_waitingWriters)
+            {
+                --_status;
+
+                if (_status == 0)
+                {
+                    if (_waitingWriters.Count > 0)
+                    {
+                        _status = -1;
+                        toWake = _waitingWriters.Dequeue();
+                    }
+                }
+            }
+
+            _reset(this);
+
+            toWake?.SetResult(new Releaser(this, true));
+        }
+
+        private void WriterRelease()
+        {
+            TaskCompletionSource<Releaser> toWake = null;
+            bool toWakeIsWriter = false;
+
+            lock (_waitingWriters)
+            {
+                if (_waitingWriters.Count > 0)
+                {
+                    toWake = _waitingWriters.Dequeue();
+                    toWakeIsWriter = true;
+                }
+                else if (_readersWaiting > 0)
+                {
+                    toWake = _waitingReader;
+                    _status = _readersWaiting;
+                    _readersWaiting = 0;
+                    _waitingReader = new TaskCompletionSource<Releaser>();
+                }
+                else
+                {
+                    _status = 0;
+                }
+            }
+
+            _reset(this);
+
+            toWake?.SetResult(new Releaser(this, toWakeIsWriter));
+        }
+
+        public readonly struct Releaser : IDisposable
+        {
+            private readonly AsyncKeyLockDoorman toRelease;
+            private readonly bool writer;
+
+            internal Releaser(AsyncKeyLockDoorman toRelease, bool writer)
+            {
+                this.toRelease = toRelease;
+                this.writer = writer;
+            }
+
+            public void Dispose()
+            {
+                if (toRelease != null)
+                {
+                    if (writer)
+                    {
+                        toRelease.WriterRelease();
+                    }
+                    else
+                    {
+                        toRelease.ReaderRelease();
+                    }
+                }
+            }
+        }
+    }
+}