RedisSlidingWindowRateLimiter now supports permitCount parameter

Author: manuelspezzani

src/RedisRateLimiting/SlidingWindow/RedisSlidingWindowManager.cs

@@ -16,16 +16,24 @@ internal class RedisSlidingWindowManager
           @"local limit = tonumber(@permit_limit)
             local timestamp = tonumber(@current_time)
             local window = tonumber(@window)
+            local requested = tonumber(@permit_count)
+
+            local zaddparams = {}
+            for i=1,requested do
+                local index = i*2
+                zaddparams[index-1]=timestamp
+                zaddparams[index]=@unique_id..':'..tostring(i)
+            end
 
             -- remove all requests outside current window
             redis.call(""zremrangebyscore"", @rate_limit_key, '-inf', timestamp - window)
 
             local count = redis.call(""zcard"", @rate_limit_key)
-            local allowed = count < limit
+            local allowed = count + requested <= limit
 
             if allowed
             then
-                redis.call(""zadd"", @rate_limit_key, timestamp, @unique_id)
+                redis.call(""zadd"", @rate_limit_key, unpack(zaddparams))
             end
 
             redis.call(""expireat"", @rate_limit_key, timestamp + window + 1)
@@ -57,7 +65,7 @@ public RedisSlidingWindowManager(
             StatsRateLimitKey = new RedisKey($"rl:{{{partitionKey}}}:stats");
         }
 
-        internal async Task<RedisSlidingWindowResponse> TryAcquireLeaseAsync(string requestId)
+        internal async Task<RedisSlidingWindowResponse> TryAcquireLeaseAsync(string requestId, int permitCount)
         {
             var now = DateTimeOffset.UtcNow;
             var nowUnixTimeSeconds = now.ToUnixTimeSeconds();
@@ -74,6 +82,7 @@ internal async Task<RedisSlidingWindowResponse> TryAcquireLeaseAsync(string requ
                     stats_key = StatsRateLimitKey,
                     current_time = nowUnixTimeSeconds,
                     unique_id = requestId,
+                    permit_count = permitCount
                 });
 
             var result = new RedisSlidingWindowResponse();
@@ -87,7 +96,7 @@ internal async Task<RedisSlidingWindowResponse> TryAcquireLeaseAsync(string requ
             return result;
         }
 
-        internal RedisSlidingWindowResponse TryAcquireLease(string requestId)
+        internal RedisSlidingWindowResponse TryAcquireLease(string requestId, int permitCount)
         {
             var now = DateTimeOffset.UtcNow;
             var nowUnixTimeSeconds = now.ToUnixTimeSeconds();
@@ -104,6 +113,7 @@ internal RedisSlidingWindowResponse TryAcquireLease(string requestId)
                    stats_key = StatsRateLimitKey,
                    current_time = nowUnixTimeSeconds,
                    unique_id = requestId,
+                   permit_count = permitCount
                });
 
             var result = new RedisSlidingWindowResponse();

src/RedisRateLimiting/SlidingWindow/RedisSlidingWindowRateLimiter.cs

@@ -55,7 +55,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.PermitLimit));
             }
 
-            return AcquireAsyncCoreInternal();
+            return AcquireAsyncCoreInternal(permitCount);
         }
 
         protected override RateLimitLease AttemptAcquireCore(int permitCount)
@@ -72,7 +72,7 @@ protected override RateLimitLease AttemptAcquireCore(int permitCount)
                 RequestId = Guid.NewGuid().ToString(),
             };
 
-            var response = _redisManager.TryAcquireLease(leaseContext.RequestId);
+            var response = _redisManager.TryAcquireLease(leaseContext.RequestId, permitCount);
 
             leaseContext.Count = response.Count;
             leaseContext.Allowed = response.Allowed;
@@ -85,7 +85,7 @@ protected override RateLimitLease AttemptAcquireCore(int permitCount)
             return new SlidingWindowLease(isAcquired: false, leaseContext);
         }
 
-        private async ValueTask<RateLimitLease> AcquireAsyncCoreInternal()
+        private async ValueTask<RateLimitLease> AcquireAsyncCoreInternal(int permitCount)
         {
             var leaseContext = new SlidingWindowLeaseContext
             {
@@ -94,7 +94,7 @@ private async ValueTask<RateLimitLease> AcquireAsyncCoreInternal()
                 RequestId = Guid.NewGuid().ToString(),
             };
 
-            var response = await _redisManager.TryAcquireLeaseAsync(leaseContext.RequestId);
+            var response = await _redisManager.TryAcquireLeaseAsync(leaseContext.RequestId, permitCount);
 
             leaseContext.Count = response.Count;
             leaseContext.Allowed = response.Allowed;

test/RedisRateLimiting.Tests/UnitTests/SlidingWindowUnitTests.cs

@@ -106,5 +106,24 @@ public async Task CanAcquireAsyncResource()
             stats = limiter.GetStatistics()!;
             Assert.Equal(0, stats.CurrentAvailablePermits);
         }
+
+        [Fact]
+        public async Task SupportsPermitCountFlag()
+        {
+            using var limiter = new RedisSlidingWindowRateLimiter<string>(
+                "Test_SupportsPermitCountFlag_SW",
+                new RedisSlidingWindowRateLimiterOptions
+                {
+                    PermitLimit = 3,
+                    Window = TimeSpan.FromMinutes(1),
+                    ConnectionMultiplexerFactory = Fixture.ConnectionMultiplexerFactory,
+                });
+
+            using var lease = await limiter.AcquireAsync(permitCount: 3);
+            Assert.True(lease.IsAcquired);
+
+            using var lease2 = await limiter.AcquireAsync(permitCount: 1);
+            Assert.False(lease2.IsAcquired);
+        }
     }
 }