RedisTokenBucketRateLimiter now supports permitCount parameter

Author: manuelspezzani

src/RedisRateLimiting/TokenBucket/RedisTokenBucketManager.cs

@@ -56,7 +56,7 @@ public RedisTokenBucketManager(
             RateLimitTimestampKey = new RedisKey($"rl:{{{partitionKey}}}:ts");
         }
 
-        internal async Task<RedisTokenBucketResponse> TryAcquireLeaseAsync()
+        internal async Task<RedisTokenBucketResponse> TryAcquireLeaseAsync(int permitCount)
         {
             var now = DateTimeOffset.UtcNow;
             var nowUnixTimeSeconds = now.ToUnixTimeSeconds();
@@ -73,7 +73,7 @@ internal async Task<RedisTokenBucketResponse> TryAcquireLeaseAsync()
                     tokens_per_period = _options.TokensPerPeriod,
                     token_limit = _options.TokenLimit,
                     replenish_period = _options.ReplenishmentPeriod.TotalSeconds,
-                    permit_count = 1D,
+                    permit_count = permitCount,
                 });
 
             var result = new RedisTokenBucketResponse();
@@ -87,7 +87,7 @@ internal async Task<RedisTokenBucketResponse> TryAcquireLeaseAsync()
             return result;
         }
 
-        internal RedisTokenBucketResponse TryAcquireLease()
+        internal RedisTokenBucketResponse TryAcquireLease(int permitCount)
         {
             var now = DateTimeOffset.UtcNow;
             var nowUnixTimeSeconds = now.ToUnixTimeSeconds();
@@ -104,7 +104,7 @@ internal RedisTokenBucketResponse TryAcquireLease()
                     tokens_per_period = _options.TokensPerPeriod,
                     token_limit = _options.TokenLimit,
                     replenish_period = _options.ReplenishmentPeriod.TotalSeconds,
-                    permit_count = 1D,
+                    permit_count = permitCount,
                 });
 
             var result = new RedisTokenBucketResponse();

src/RedisRateLimiting/TokenBucket/RedisTokenBucketRateLimiter.cs

@@ -60,7 +60,7 @@ protected override ValueTask<RateLimitLease> AcquireAsyncCore(int permitCount, C
                 throw new ArgumentOutOfRangeException(nameof(permitCount), permitCount, string.Format("{0} permit(s) exceeds the permit limit of {1}.", permitCount, _options.TokenLimit));
             }
 
-            return AcquireAsyncCoreInternal();
+            return AcquireAsyncCoreInternal(permitCount);
         }
 
         protected override RateLimitLease AttemptAcquireCore(int permitCount)
@@ -75,7 +75,7 @@ protected override RateLimitLease AttemptAcquireCore(int permitCount)
                 Limit = _options.TokenLimit,
             };
 
-            var response = _redisManager.TryAcquireLease();
+            var response = _redisManager.TryAcquireLease(permitCount);
 
             leaseContext.Allowed = response.Allowed;
             leaseContext.Count = response.Count;
@@ -88,14 +88,14 @@ protected override RateLimitLease AttemptAcquireCore(int permitCount)
             return new TokenBucketLease(isAcquired: false, leaseContext);
         }
 
-        private async ValueTask<RateLimitLease> AcquireAsyncCoreInternal()
+        private async ValueTask<RateLimitLease> AcquireAsyncCoreInternal(int permitCount)
         {
             var leaseContext = new TokenBucketLeaseContext
             {
                 Limit = _options.TokenLimit,
             };
 
-            var response = await _redisManager.TryAcquireLeaseAsync();
+            var response = await _redisManager.TryAcquireLeaseAsync(permitCount);
 
             leaseContext.Allowed = response.Allowed;
             leaseContext.Count = response.Count;

test/RedisRateLimiting.Tests/UnitTests/TokenBucketUnitTests.cs

@@ -107,5 +107,28 @@ public async Task CanAcquireAsyncResource()
             using var lease2 = await limiter.AcquireAsync();
             Assert.False(lease2.IsAcquired);
         }
+
+        [Fact]
+        public async Task SupportsPermitCountFlag()
+        {
+            using var limiter = new RedisTokenBucketRateLimiter<string>(
+                "Test_SupportsPermitCountFlag_TB",
+                new RedisTokenBucketRateLimiterOptions
+                {
+                    TokenLimit = 5,
+                    TokensPerPeriod = 5,
+                    ReplenishmentPeriod = TimeSpan.FromMinutes(1),
+                    ConnectionMultiplexerFactory = Fixture.ConnectionMultiplexerFactory,
+                });
+
+            using var lease = await limiter.AcquireAsync(4);
+            Assert.True(lease.IsAcquired);
+
+            using var lease2 = await limiter.AcquireAsync(3);
+            Assert.False(lease2.IsAcquired);
+
+            using var lease3 = await limiter.AcquireAsync(1);
+            Assert.True(lease3.IsAcquired);
+        }
     }
 }