RedisFixedWindowRateLimiter now supports permitCount parameter

(cherry picked from commit bb6baed024a4595a639085db09df954c90b064bf)

Author: manuelspezzani

src/RedisRateLimiting/FixedWindow/RedisFixedWindowManager.cs

@@ -13,6 +13,9 @@ internal class RedisFixedWindowManager
 
         private static readonly LuaScript _redisScript = LuaScript.Prepare(
           @"local expires_at = tonumber(redis.call(""get"", @expires_at_key))
+            local current = tonumber(redis.call(""get"", @rate_limit_key))
+            local requested = tonumber(@increment_amount)
+            local limit = tonumber(@permit_limit)
 
             if not expires_at or expires_at < tonumber(@current_time) then
                 -- this is either a brand new window,
@@ -27,13 +30,19 @@ internal class RedisFixedWindowManager
                 redis.call(""expireat"", @expires_at_key, @next_expires_at + 1)
                 -- since the database was updated, return the new value
                 expires_at = @next_expires_at
+                current = 0
             end
 
-            -- now that the window either already exists or it was freshly initialized,
-            -- increment the counter(`incrby` returns a number)
-            local current = redis.call(""incrby"", @rate_limit_key, @increment_amount)
+            local allowed = current + requested <= limit
 
-            return { current, expires_at }");
+            if allowed
+            then
+                -- now that the window either already exists or it was freshly initialized,
+                -- increment the counter(`incrby` returns a number)
+                current = redis.call(""incrby"", @rate_limit_key, @increment_amount)
+            end
+
+            return { current, expires_at, allowed }");
 
         public RedisFixedWindowManager(
             string partitionKey,
@@ -46,7 +55,7 @@ public RedisFixedWindowManager(
             RateLimitExpireKey = new RedisKey($"rl:{{{partitionKey}}}:exp");
         }
 
-        internal async Task<RedisFixedWindowResponse> TryAcquireLeaseAsync()
+        internal async Task<RedisFixedWindowResponse> TryAcquireLeaseAsync(int permitCount)
         {
             var now = DateTimeOffset.UtcNow;
             var nowUnixTimeSeconds = now.ToUnixTimeSeconds();
@@ -59,9 +68,10 @@ internal async Task<RedisFixedWindowResponse> TryAcquireLeaseAsync()
                 {
                     rate_limit_key = RateLimitKey,
                     expires_at_key = RateLimitExpireKey,
+                    permit_limit = _options.PermitLimit,
                     next_expires_at = now.Add(_options.Window).ToUnixTimeSeconds(),
                     current_time = nowUnixTimeSeconds,
-                    increment_amount = 1D,
+                    increment_amount = permitCount,
                 });
 
             var result = new RedisFixedWindowResponse();
@@ -70,13 +80,14 @@ internal async Task<RedisFixedWindowResponse> TryAcquireLeaseAsync()
             {
                 result.Count = (long)response[0];
                 result.ExpiresAt = (long)response[1];
+                result.Allowed = (bool)response[2];
                 result.RetryAfter = TimeSpan.FromSeconds(result.ExpiresAt - nowUnixTimeSeconds);
             }
 
             return result;
         }
 
-        internal RedisFixedWindowResponse TryAcquireLease()
+        internal RedisFixedWindowResponse TryAcquireLease(int permitCount)
         {
             var now = DateTimeOffset.UtcNow;
             var nowUnixTimeSeconds = now.ToUnixTimeSeconds();
@@ -89,9 +100,10 @@ internal RedisFixedWindowResponse TryAcquireLease()
                 {
                     rate_limit_key = RateLimitKey,
                     expires_at_key = RateLimitExpireKey,
+                    permit_limit = _options.PermitLimit,
                     next_expires_at = now.Add(_options.Window).ToUnixTimeSeconds(),
                     current_time = nowUnixTimeSeconds,
-                    increment_amount = 1D,
+                    increment_amount = permitCount,
                 });
 
             var result = new RedisFixedWindowResponse();
@@ -100,6 +112,7 @@ internal RedisFixedWindowResponse TryAcquireLease()
             {
                 result.Count = (long)response[0];
                 result.ExpiresAt = (long)response[1];
+                result.Allowed = (bool)response[2];
                 result.RetryAfter = TimeSpan.FromSeconds(result.ExpiresAt - nowUnixTimeSeconds);
             }
 
@@ -112,5 +125,6 @@ internal class RedisFixedWindowResponse
         internal long ExpiresAt { get; set; }
         internal TimeSpan RetryAfter { get; set; }
         internal long Count { get; set; }
+        internal bool Allowed { get; set; }
     }
 }

src/RedisRateLimiting/FixedWindow/RedisFixedWindowRateLimiter.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)
@@ -71,38 +71,28 @@ protected override RateLimitLease AttemptAcquireCore(int permitCount)
                 Window = _options.Window,
             };
 
-            var response = _redisManager.TryAcquireLease();
+            var response = _redisManager.TryAcquireLease(permitCount);
 
             leaseContext.Count = response.Count;
             leaseContext.RetryAfter = response.RetryAfter;
-
-            if (leaseContext.Count > _options.PermitLimit)
-            {
-                return new FixedWindowLease(isAcquired: false, leaseContext);
-            }
-
-            return new FixedWindowLease(isAcquired: true, leaseContext);
+            
+            return new FixedWindowLease(isAcquired: response.Allowed, leaseContext);
         }
 
-        private async ValueTask<RateLimitLease> AcquireAsyncCoreInternal()
+        private async ValueTask<RateLimitLease> AcquireAsyncCoreInternal(int permitCount)
         {
             var leaseContext = new FixedWindowLeaseContext
             {
                 Limit = _options.PermitLimit,
                 Window = _options.Window,
             };
 
-            var response = await _redisManager.TryAcquireLeaseAsync();
+            var response = await _redisManager.TryAcquireLeaseAsync(permitCount);
 
             leaseContext.Count = response.Count;
             leaseContext.RetryAfter = response.RetryAfter;
 
-            if (leaseContext.Count > _options.PermitLimit)
-            {
-                return new FixedWindowLease(isAcquired: false, leaseContext);
-            }
-
-            return new FixedWindowLease(isAcquired: true, leaseContext);
+            return new FixedWindowLease(isAcquired: response.Allowed, leaseContext);
         }
 
         private sealed class FixedWindowLeaseContext

test/RedisRateLimiting.Tests/UnitTests/FixedWindowUnitTests.cs

@@ -86,5 +86,27 @@ public async Task CanAcquireAsyncResource()
             using var lease2 = await limiter.AcquireAsync();
             Assert.False(lease2.IsAcquired);
         }
+
+        [Fact]
+        public async Task SupportsPermitCountFlag()
+        {
+            using var limiter = new RedisFixedWindowRateLimiter<string>(
+                "Test_SupportsPermitCountFlag_FW",
+                new RedisFixedWindowRateLimiterOptions
+                {
+                    PermitLimit = 5,
+                    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: 3);
+            Assert.False(lease2.IsAcquired);
+
+            using var lease3 = await limiter.AcquireAsync(permitCount: 2);
+            Assert.True(lease3.IsAcquired);
+        }
     }
 }