Introduce abstract PartitionedRateLimiter (#67241)

src/libraries/System.Threading.RateLimiting/ref/System.Threading.RateLimiting.cs

@@ -39,6 +39,19 @@ public sealed partial class MetadataName<T> : System.IEquatable<System.Threading
         public static bool operator !=(System.Threading.RateLimiting.MetadataName<T> left, System.Threading.RateLimiting.MetadataName<T> right) { throw null; }
         public override string ToString() { throw null; }
     }
+    public abstract partial class PartitionedRateLimiter<TResource> : System.IAsyncDisposable, System.IDisposable
+    {
+        protected PartitionedRateLimiter() { }
+        public System.Threading.RateLimiting.RateLimitLease Acquire(TResource resourceID, int permitCount = 1) { throw null; }
+        protected abstract System.Threading.RateLimiting.RateLimitLease AcquireCore(TResource resourceID, int permitCount);
+        public void Dispose() { }
+        protected virtual void Dispose(bool disposing) { }
+        public System.Threading.Tasks.ValueTask DisposeAsync() { throw null; }
+        protected virtual System.Threading.Tasks.ValueTask DisposeAsyncCore() { throw null; }
+        public abstract int GetAvailablePermits(TResource resourceID);
+        public System.Threading.Tasks.ValueTask<System.Threading.RateLimiting.RateLimitLease> WaitAsync(TResource resourceID, int permitCount = 1, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
+        protected abstract System.Threading.Tasks.ValueTask<System.Threading.RateLimiting.RateLimitLease> WaitAsyncCore(TResource resourceID, int permitCount, System.Threading.CancellationToken cancellationToken);
+    }
     public enum QueueProcessingOrder
     {
         OldestFirst = 0,

src/libraries/System.Threading.RateLimiting/src/System.Threading.RateLimiting.csproj

-<Project Sdk="Microsoft.NET.Sdk">
+<Project Sdk="Microsoft.NET.Sdk">
   <PropertyGroup>
     <TargetFrameworks>$(NetCoreAppCurrent);$(NetCoreAppMinimum);netstandard2.0;$(NetFrameworkMinimum)</TargetFrameworks>
     <Nullable>enable</Nullable>
@@ -18,13 +18,13 @@ System.Threading.RateLimiting.RateLimitLease</PackageDescription>
     <Compile Include="System\Threading\RateLimiting\ConcurrencyLimiterOptions.cs" />
     <Compile Include="System\Threading\RateLimiting\MetadataName.cs" />
     <Compile Include="System\Threading\RateLimiting\MetadataName.T.cs" />
+    <Compile Include="System\Threading\RateLimiting\PartitionedRateLimiter.T.cs" />
     <Compile Include="System\Threading\RateLimiting\QueueProcessingOrder.cs" />
     <Compile Include="System\Threading\RateLimiting\RateLimiter.cs" />
     <Compile Include="System\Threading\RateLimiting\RateLimitLease.cs" />
     <Compile Include="System\Threading\RateLimiting\TokenBucketRateLimiter.cs" />
     <Compile Include="System\Threading\RateLimiting\TokenBucketRateLimiterOptions.cs" />
-    <Compile Include="$(CommonPath)System\Collections\Generic\Deque.cs"
-             Link="Common\System\Collections\Generic\Deque.cs" />
+    <Compile Include="$(CommonPath)System\Collections\Generic\Deque.cs" Link="Common\System\Collections\Generic\Deque.cs" />
   </ItemGroup>
   <ItemGroup Condition="'$(TargetFrameworkIdentifier)' == '.NETCoreApp'">
     <Reference Include="System.Runtime" />

src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/PartitionedRateLimiter.T.cs

+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Threading.Tasks;
+
+namespace System.Threading.RateLimiting
+{
+    /// <summary>
+    /// Represents a limiter type that users interact with to determine if an operation can proceed given a specific <typeparamref name="TResource"/>.
+    /// </summary>
+    /// <typeparam name="TResource">The resource type that is being limited.</typeparam>
+    public abstract class PartitionedRateLimiter<TResource> : IAsyncDisposable, IDisposable
+    {
+        /// <summary>
+        /// An estimated count of available permits.
+        /// </summary>
+        /// <returns></returns>
+        public abstract int GetAvailablePermits(TResource resourceID);
+
+        /// <summary>
+        /// Fast synchronous attempt to acquire permits.
+        /// </summary>
+        /// <remarks>
+        /// Set <paramref name="permitCount"/> to 0 to get whether permits are exhausted.
+        /// </remarks>
+        /// <param name="resourceID">The resource to limit.</param>
+        /// <param name="permitCount">Number of permits to try and acquire.</param>
+        /// <returns>A successful or failed lease.</returns>
+        /// <exception cref="ArgumentOutOfRangeException"></exception>
+        public RateLimitLease Acquire(TResource resourceID, int permitCount = 1)
+        {
+            if (permitCount < 0)
+            {
+                throw new ArgumentOutOfRangeException(nameof(permitCount));
+            }
+
+            return AcquireCore(resourceID, permitCount);
+        }
+
+        /// <summary>
+        /// Method that <see cref="PartitionedRateLimiter{TResource}"/> implementations implement for <see cref="Acquire"/>.
+        /// </summary>
+        /// <param name="resourceID">The resource to limit.</param>
+        /// <param name="permitCount">Number of permits to try and acquire.</param>
+        /// <returns></returns>
+        protected abstract RateLimitLease AcquireCore(TResource resourceID, int permitCount);
+
+        /// <summary>
+        /// Wait until the requested permits are available or permits can no longer be acquired.
+        /// </summary>
+        /// <remarks>
+        /// Set <paramref name="permitCount"/> to 0 to wait until permits are replenished.
+        /// </remarks>
+        /// <param name="resourceID">The resource to limit.</param>
+        /// <param name="permitCount">Number of permits to try and acquire.</param>
+        /// <param name="cancellationToken">Optional token to allow canceling a queued request for permits.</param>
+        /// <returns>A task that completes when the requested permits are acquired or when the requested permits are denied.</returns>
+        /// <exception cref="ArgumentOutOfRangeException"></exception>
+        public ValueTask<RateLimitLease> WaitAsync(TResource resourceID, int permitCount = 1, CancellationToken cancellationToken = default)
+        {
+            if (permitCount < 0)
+            {
+                throw new ArgumentOutOfRangeException(nameof(permitCount));
+            }
+
+            if (cancellationToken.IsCancellationRequested)
+            {
+                return new ValueTask<RateLimitLease>(Task.FromCanceled<RateLimitLease>(cancellationToken));
+            }
+
+            return WaitAsyncCore(resourceID, permitCount, cancellationToken);
+        }
+
+        /// <summary>
+        /// Method that <see cref="PartitionedRateLimiter{TResource}"/> implementations implement for <see cref="WaitAsync"/>.
+        /// </summary>
+        /// <param name="resourceID">The resource to limit.</param>
+        /// <param name="permitCount">Number of permits to try and acquire.</param>
+        /// <param name="cancellationToken">Optional token to allow canceling a queued request for permits.</param>
+        /// <returns>A task that completes when the requested permits are acquired or when the requested permits are denied.</returns>
+        protected abstract ValueTask<RateLimitLease> WaitAsyncCore(TResource resourceID, int permitCount, CancellationToken cancellationToken);
+
+        /// <summary>
+        /// Dispose method for implementations to write.
+        /// </summary>
+        /// <param name="disposing"></param>
+        protected virtual void Dispose(bool disposing) { }
+
+        /// <summary>
+        /// Disposes the RateLimiter. This completes any queued acquires with a failed lease.
+        /// </summary>
+        public void Dispose()
+        {
+            // Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
+            Dispose(disposing: true);
+            GC.SuppressFinalize(this);
+        }
+
+        /// <summary>
+        /// DisposeAsync method for implementations to write.
+        /// </summary>
+        protected virtual ValueTask DisposeAsyncCore()
+        {
+            return default;
+        }
+
+        /// <summary>
+        /// Disposes the RateLimiter asynchronously.
+        /// </summary>
+        /// <returns>ValueTask representing the completion of the disposal.</returns>
+        public async ValueTask DisposeAsync()
+        {
+            // Perform async cleanup.
+            await DisposeAsyncCore().ConfigureAwait(false);
+
+            // Dispose of unmanaged resources.
+            Dispose(false);
+
+            // Suppress finalization.
+            GC.SuppressFinalize(this);
+        }
+    }
+}

src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/RateLimiter.cs

@@ -102,7 +102,7 @@ protected virtual ValueTask DisposeAsyncCore()
         /// <summary>
         /// Disposes the RateLimiter asynchronously.
         /// </summary>
-        /// <returns>ValueTask representin the completion of the disposal.</returns>
+        /// <returns>ValueTask representing the completion of the disposal.</returns>
         public async ValueTask DisposeAsync()
         {
             // Perform async cleanup.

src/libraries/System.Threading.RateLimiting/tests/PartitionedRateLimiterTests.cs

+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Threading.Tasks;
+using Xunit;
+
+namespace System.Threading.RateLimiting.Tests
+{
+    public class PartitionedRateLimiterTests
+    {
+        [Fact]
+        public void ThrowsWhenAcquiringLessThanZero()
+        {
+            using var limiter = new NotImplementedPartitionedRateLimiter<string>();
+            Assert.Throws<ArgumentOutOfRangeException>(() => limiter.Acquire(string.Empty, -1));
+        }
+
+        [Fact]
+        public async Task ThrowsWhenWaitingForLessThanZero()
+        {
+            using var limiter = new NotImplementedPartitionedRateLimiter<string>();
+            await Assert.ThrowsAsync<ArgumentOutOfRangeException>(async () => await limiter.WaitAsync(string.Empty, -1));
+        }
+
+        [Fact]
+        public async Task WaitAsyncThrowsWhenPassedACanceledToken()
+        {
+            using var limiter = new NotImplementedPartitionedRateLimiter<string>();
+            await Assert.ThrowsAsync<TaskCanceledException>(
+                async () => await limiter.WaitAsync(string.Empty, 1, new CancellationToken(true)));
+        }
+
+        internal class NotImplementedPartitionedRateLimiter<T> : PartitionedRateLimiter<T>
+        {
+            public override int GetAvailablePermits(T resourceID) => throw new NotImplementedException();
+            protected override RateLimitLease AcquireCore(T resourceID, int permitCount) => throw new NotImplementedException();
+            protected override ValueTask<RateLimitLease> WaitAsyncCore(T resourceID, int permitCount, CancellationToken cancellationToken) => throw new NotImplementedException();
+        }
+    }
+}

src/libraries/System.Threading.RateLimiting/tests/System.Threading.RateLimiting.Tests.csproj

@@ -5,6 +5,7 @@
   <ItemGroup>
     <Compile Include="BaseRateLimiterTests.cs" />
     <Compile Include="ConcurrencyLimiterTests.cs" />
+    <Compile Include="PartitionedRateLimiterTests.cs" />
     <Compile Include="TokenBucketRateLimiterTests.cs" />
   </ItemGroup>
   <ItemGroup>