Update

Author: BrennanConroy

src/RateLimiting/src/ConcurrencyLimiter.cs

@@ -82,7 +82,7 @@ protected override ValueTask<RateLimitLease> WaitAsyncCore(int permitCount, Canc
                 throw new ArgumentOutOfRangeException(nameof(permitCount), $"{permitCount} permits exceeds the permit limit of {_options.PermitLimit}.");
             }
 
-            // Return SuccessfulAcquisition if requestedCount is 0 and resources are available
+            // Return SuccessfulLease if requestedCount is 0 and resources are available
             if (permitCount == 0 && _permitCount > 0)
             {
                 return new ValueTask<RateLimitLease>(SuccessfulLease);
@@ -192,14 +192,6 @@ private void Release(int releaseCount)
             }
         }
 
-        private void CancellationRequested(TaskCompletionSource<RateLimitLease> tcs, CancellationToken token)
-        {
-            lock (_lock)
-            {
-                tcs.TrySetException(new OperationCanceledException(token));
-            }
-        }
-
         private class ConcurrencyLease : RateLimitLease
         {
             private bool _disposed;
@@ -270,8 +262,6 @@ public RequestRegistration(int requestedCount, TaskCompletionSource<RateLimitLea
             public TaskCompletionSource<RateLimitLease> Tcs { get; }
 
             public CancellationTokenRegistration CancellationTokenRegistration { get; }
-
-            public CancellationTokenRegistration CancellationTokenRegistration { get; }
         }
     }
 }

src/RateLimiting/src/RateLimiter.cs

@@ -31,7 +31,6 @@ public abstract class RateLimiter
         /// <remarks>
         /// Set <paramref name="permitCount"/> to 0 to get whether permits are exhausted.
         /// </remarks>
-        /// <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(int permitCount = 1)

src/RateLimiting/src/TokenBucketRateLimiter.cs

@@ -5,6 +5,7 @@
 
 using System.Collections.Generic;
 using System.Diagnostics;
+using System.Diagnostics.CodeAnalysis;
 using System.Threading.Tasks;
 
 namespace System.Threading.RateLimiting
@@ -19,16 +20,18 @@ public sealed class TokenBucketRateLimiter : RateLimiter
         private long _lastReplenishmentTick = Environment.TickCount;
 
         private readonly Timer? _renewTimer;
-        private readonly object _lock = new object();
         private readonly TokenBucketRateLimiterOptions _options;
         private readonly Deque<RequestRegistration> _queue = new Deque<RequestRegistration>();
 
+        // Use the queue as the lock field so we don't need to allocate another object for a lock and have another field in the object
+        private object Lock => _queue;
+
         private static readonly RateLimitLease SuccessfulLease = new TokenBucketLease(true, null);
 
         /// <summary>
         /// Initializes the <see cref="TokenBucketRateLimiter"/>.
         /// </summary>
-        /// <param name="options"></param>
+        /// <param name="options">Options to specify the behavior of the <see cref="TokenBucketRateLimiter"/>.</param>
         public TokenBucketRateLimiter(TokenBucketRateLimiterOptions options)
         {
             _tokenCount = options.TokenLimit;
@@ -40,10 +43,7 @@ public TokenBucketRateLimiter(TokenBucketRateLimiterOptions options)
             }
         }
 
-        /// <summary>
-        /// An estimated count of available tokens.
-        /// </summary>
-        /// <returns></returns>
+        /// <inheritdoc/>
         public override int GetAvailablePermits() => _tokenCount;
 
         /// <inheritdoc/>
@@ -52,26 +52,25 @@ protected override RateLimitLease AcquireCore(int tokenCount)
             // These amounts of resources can never be acquired
             if (tokenCount > _options.TokenLimit)
             {
-                throw new InvalidOperationException($"{tokenCount} tokens exceeds the token limit of {_options.TokenLimit}.");
+                throw new ArgumentOutOfRangeException(nameof(tokenCount), $"{tokenCount} tokens exceeds the token limit of {_options.TokenLimit}.");
             }
 
-            // Return SuccessfulAcquisition or FailedAcquisition depending to indicate limiter state
+            // Return SuccessfulLease or FailedLease depending to indicate limiter state
             if (tokenCount == 0)
             {
-                if (GetAvailablePermits() > 0)
+                if (_tokenCount > 0)
                 {
                     return SuccessfulLease;
                 }
 
                 return CreateFailedTokenLease(tokenCount);
             }
 
-            lock (_lock)
+            lock (Lock)
             {
-                if (GetAvailablePermits() >= tokenCount)
+                if (TryLeaseUnsynchronized(tokenCount, out RateLimitLease? lease))
                 {
-                    _tokenCount -= tokenCount;
-                    return SuccessfulLease;
+                    return lease;
                 }
 
                 return CreateFailedTokenLease(tokenCount);
@@ -86,23 +85,20 @@ protected override ValueTask<RateLimitLease> WaitAsyncCore(int tokenCount, Cance
             // These amounts of resources can never be acquired
             if (tokenCount > _options.TokenLimit)
             {
-                throw new InvalidOperationException($"{tokenCount} token(s) exceeds the permit limit of {_options.TokenLimit}.");
+                throw new ArgumentOutOfRangeException(nameof(tokenCount), $"{tokenCount} token(s) exceeds the permit limit of {_options.TokenLimit}.");
             }
 
             // Return SuccessfulAcquisition if requestedCount is 0 and resources are available
-            if (tokenCount == 0 && GetAvailablePermits() > 0)
+            if (tokenCount == 0 && _tokenCount > 0)
             {
-                // Perf: static failed/successful value tasks?
                 return new ValueTask<RateLimitLease>(SuccessfulLease);
             }
 
-            lock (_lock)
+            lock (Lock)
             {
-                if (GetAvailablePermits() >= tokenCount && GetAvailablePermits() != 0)
+                if (TryLeaseUnsynchronized(tokenCount, out RateLimitLease? lease))
                 {
-                    _tokenCount -= tokenCount;
-                    // Perf: static failed/successful value tasks?
-                    return new ValueTask<RateLimitLease>(SuccessfulLease);
+                    return new ValueTask<RateLimitLease>(lease);
                 }
 
                 // Don't queue if queue limit reached
@@ -116,8 +112,12 @@ protected override ValueTask<RateLimitLease> WaitAsyncCore(int tokenCount, Cance
                 CancellationTokenRegistration ctr;
                 if (cancellationToken.CanBeCanceled)
                 {
-                    ctr = cancellationToken.Register(obj => CancellationRequested((TaskCompletionSource<RateLimitLease>)obj, cancellationToken), tcs);
+                    ctr = cancellationToken.Register(obj =>
+                    {
+                        ((TaskCompletionSource<RateLimitLease>)obj).TrySetException(new OperationCanceledException(cancellationToken));
+                    }, tcs);
                 }
+
                 RequestRegistration registration = new RequestRegistration(tokenCount, tcs, ctr);
                 _queue.EnqueueTail(registration);
                 _queueCount += tokenCount;
@@ -130,12 +130,39 @@ protected override ValueTask<RateLimitLease> WaitAsyncCore(int tokenCount, Cance
 
         private RateLimitLease CreateFailedTokenLease(int tokenCount)
         {
-            int replenishAmount = tokenCount - GetAvailablePermits() + _queueCount;
+            int replenishAmount = tokenCount - _tokenCount + _queueCount;
             // can't have 0 replenish periods, that would mean it should be a successful lease
             // if TokensPerPeriod is larger than the replenishAmount needed then it would be 0
             int replenishPeriods = Math.Max(replenishAmount / _options.TokensPerPeriod, 1);
 
-            return new TokenBucketLease(false, TimeSpan.FromTicks(_options.ReplenishmentPeriod.Ticks*replenishPeriods));
+            return new TokenBucketLease(false, TimeSpan.FromTicks(_options.ReplenishmentPeriod.Ticks * replenishPeriods));
+        }
+
+        private bool TryLeaseUnsynchronized(int tokenCount, [NotNullWhen(true)] out RateLimitLease? lease)
+        {
+            // if permitCount is 0 we want to queue it if there are no available permits
+            if (_tokenCount >= tokenCount && _tokenCount != 0)
+            {
+                if (tokenCount == 0)
+                {
+                    // Edge case where the check before the lock showed 0 available permits but when we got the lock some permits were now available
+                    lease = SuccessfulLease;
+                    return true;
+                }
+
+                // a. if there are no items queued we can lease
+                // b. if there are items queued but the processing order is newest first, then we can lease the incoming request since it is the newest
+                if (_queueCount == 0 || (_queueCount > 0 && _options.QueueProcessingOrder == QueueProcessingOrder.NewestFirst))
+                {
+                    _tokenCount -= tokenCount;
+                    Debug.Assert(_tokenCount >= 0);
+                    lease = SuccessfulLease;
+                    return true;
+                }
+            }
+
+            lease = null;
+            return false;
         }
 
         /// <summary>
@@ -165,7 +192,7 @@ private static void Replenish(object? state)
             long nowTicks = Environment.TickCount * TimeSpan.TicksPerMillisecond;
 
             // method is re-entrant (from Timer), lock to avoid multiple simultaneous replenishes
-            lock (limiter!._lock)
+            lock (limiter!.Lock)
             {
                 if (nowTicks - limiter._lastReplenishmentTick < limiter._options.ReplenishmentPeriod.Ticks)
                 {
@@ -214,7 +241,6 @@ private static void Replenish(object? state)
                         Debug.Assert(limiter._queueCount >= 0);
                         Debug.Assert(limiter._tokenCount >= 0);
 
-                        // requestToFulfill == request
                         if (!nextPendingRequest.Tcs.TrySetResult(SuccessfulLease))
                         {
                             // Queued item was canceled so add count back
@@ -231,14 +257,6 @@ private static void Replenish(object? state)
             }
         }
 
-        private void CancellationRequested(TaskCompletionSource<RateLimitLease> tcs, CancellationToken token)
-        {
-            lock (_lock)
-            {
-                tcs.TrySetException(new OperationCanceledException(token));
-            }
-        }
-
         private class TokenBucketLease : RateLimitLease
         {
             private readonly TimeSpan? _retryAfter;
@@ -255,12 +273,10 @@ public TokenBucketLease(bool isAcquired, TimeSpan? retryAfter)
 
             private IEnumerable<string> Enumerable()
             {
-                if (_retryAfter is null)
+                if (_retryAfter is not null)
                 {
-                    yield break;
+                    yield return MetadataName.RetryAfter.Name;
                 }
-
-                yield return MetadataName.RetryAfter.Name;
             }
 
             public override bool TryGetMetadata(string metadataName, out object? metadata)
@@ -271,14 +287,14 @@ public override bool TryGetMetadata(string metadataName, out object? metadata)
                     return true;
                 }
 
-                metadata = null;
+                metadata = default;
                 return false;
             }
 
             protected override void Dispose(bool disposing) { }
         }
 
-        private struct RequestRegistration
+        private readonly struct RequestRegistration
         {
             public RequestRegistration(int tokenCount, TaskCompletionSource<RateLimitLease> tcs, CancellationTokenRegistration cancellationTokenRegistration)
             {

src/RateLimiting/src/TokenBucketRateLimiterOptions.cs

@@ -78,7 +78,7 @@ public TokenBucketRateLimiterOptions(
         /// Determines the behaviour of <see cref="RateLimiter.WaitAsync"/> when not enough resources can be leased.
         /// </summary>
         /// <value>
-        /// <see cref="QueueProcessingOrder.OldestFirst"/>
+        /// <see cref="QueueProcessingOrder.OldestFirst"/> by default.
         /// </value>
         public QueueProcessingOrder QueueProcessingOrder { get; }
 

src/RateLimiting/test/BaseRateLimiterTests.cs

@@ -0,0 +1,80 @@
+// 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.Test
+{
+    public abstract class BaseRateLimiterTests
+    {
+        [Fact]
+        public abstract void CanAcquireResource();
+
+        [Fact]
+        public abstract void InvalidOptionsThrows();
+
+        [Fact]
+        public abstract Task CanAcquireResourceAsync();
+
+        [Fact]
+        public abstract Task CanAcquireResourceAsync_QueuesAndGrabsOldest();
+
+        [Fact]
+        public abstract Task CanAcquireResourceAsync_QueuesAndGrabsNewest();
+
+        [Fact]
+        public abstract Task FailsWhenQueuingMoreThanLimit();
+
+        [Fact]
+        public abstract Task QueueAvailableAfterQueueLimitHitAndResources_BecomeAvailable();
+
+        [Fact]
+        public abstract void ThrowsWhenAcquiringMoreThanLimit();
+
+        [Fact]
+        public abstract Task ThrowsWhenWaitingForMoreThanLimit();
+
+        [Fact]
+        public abstract void ThrowsWhenAcquiringLessThanZero();
+
+        [Fact]
+        public abstract Task ThrowsWhenWaitingForLessThanZero();
+
+        [Fact]
+        public abstract void AcquireZero_WithAvailability();
+
+        [Fact]
+        public abstract void AcquireZero_WithoutAvailability();
+
+        [Fact]
+        public abstract Task WaitAsyncZero_WithAvailability();
+
+        [Fact]
+        public abstract Task WaitAsyncZero_WithoutAvailabilityWaitsForAvailability();
+
+        [Fact]
+        public abstract Task CanDequeueMultipleResourcesAtOnce();
+
+        [Fact]
+        public abstract Task CanAcquireResourcesWithWaitAsyncWithQueuedItemsIfNewestFirst();
+
+        [Fact]
+        public abstract Task CannotAcquireResourcesWithWaitAsyncWithQueuedItemsIfOldestFirst();
+
+        [Fact]
+        public abstract Task CanCancelWaitAsyncAfterQueuing();
+
+        [Fact]
+        public abstract Task CanCancelWaitAsyncBeforeQueuing();
+
+        [Fact]
+        public abstract Task CanAcquireResourcesWithAcquireWithQueuedItemsIfNewestFirst();
+
+        [Fact]
+        public abstract Task CannotAcquireResourcesWithAcquireWithQueuedItemsIfOldestFirst();
+
+        [Fact]
+        public abstract void NoMetadataOnAcquiredLease();
+    }
+}