RedisConcurrencyRateLimiter now supports permitCount parameter

Author: manuelspezzani

src/RedisRateLimiting/Concurrency/RedisConcurrencyManager.cs

@@ -1,5 +1,6 @@
 using StackExchange.Redis;
 using System;
+using System.Collections.Generic;
 using System.Threading.RateLimiting;
 using System.Threading.Tasks;
 
@@ -18,6 +19,7 @@ internal class RedisConcurrencyManager
             local queue_limit = tonumber(@queue_limit)
             local try_enqueue = tonumber(@try_enqueue)
             local timestamp = tonumber(@current_time)
+            local requested = tonumber(@permit_count)
             -- max seconds it takes to complete a request
             local ttl = 60
 
@@ -29,10 +31,19 @@ internal class RedisConcurrencyManager
             end
 
             local count = redis.call(""zcard"", @rate_limit_key)
-            local allowed = count < limit
+            local allowed = count + requested <= limit
             local queued = false
             local queue_count = 0
 
+            local addparams = {}
+            local remparams = {}
+            for i=1,requested do
+                local index = i*2
+                addparams[index-1]=timestamp
+                addparams[index]=@unique_id..':'..tostring(i)
+                remparams[i]=addparams[index]
+            end
+
             if allowed
             then
 
@@ -45,23 +56,23 @@ internal class RedisConcurrencyManager
                 if queue_count == 0 or try_enqueue == 0
                 then
 
-                    redis.call(""zadd"", @rate_limit_key, timestamp, @unique_id)
+                    redis.call(""zadd"", @rate_limit_key, unpack(addparams))
 
                     if queue_limit > 0
                     then
                         -- remove from pending queue
-                        redis.call(""zrem"", @queue_key, @unique_id)
+                        redis.call(""zrem"", @queue_key, unpack(remparams))
                     end
                 
                 else
                     -- queue the current request next in line if we have any requests in the pending queue
                     allowed = false
 
-                    queued = queue_count + count < limit + queue_limit
+                    queued = queue_count + count + requested <= limit + queue_limit
 
                     if queued
                     then
-                        redis.call(""zadd"", @queue_key, timestamp, @unique_id)
+                        redis.call(""zadd"", @queue_key, unpack(addparams))
                     end
 
                 end
@@ -72,23 +83,23 @@ internal class RedisConcurrencyManager
                 then
 
                     queue_count = redis.call(""zcard"", @queue_key)
-                    queued = queue_count < queue_limit
+                    queued = queue_count + requested <= queue_limit
 
                     if queued
                     then
-                        redis.call(""zadd"", @queue_key, timestamp, @unique_id)
+                        redis.call(""zadd"", @queue_key, unpack(addparams))
                     end
 
                 end
             end
 
             if allowed
             then
-                redis.call(""hincrby"", @stats_key, 'total_successful', 1)
+                redis.call(""hincrby"", @stats_key, 'total_successful', requested)
             else
                 if queued == false and try_enqueue == 1
                 then
-                    redis.call(""hincrby"", @stats_key, 'total_failed', 1)
+                    redis.call(""hincrby"", @stats_key, 'total_failed', requested)
                 end
             end
 
@@ -114,7 +125,7 @@ public RedisConcurrencyManager(
             StatsRateLimitKey = new RedisKey($"rl:{{{partitionKey}}}:stats");
         }
 
-        internal async Task<RedisConcurrencyResponse> TryAcquireLeaseAsync(string requestId, bool tryEnqueue = false)
+        internal async Task<RedisConcurrencyResponse> TryAcquireLeaseAsync(string requestId, int permitCount, bool tryEnqueue = false)
         {
             var nowUnixTimeSeconds = DateTimeOffset.UtcNow.ToUnixTimeSeconds();
 
@@ -132,6 +143,7 @@ internal async Task<RedisConcurrencyResponse> TryAcquireLeaseAsync(string reques
                     stats_key = StatsRateLimitKey,
                     current_time = nowUnixTimeSeconds,
                     unique_id = requestId,
+                    permit_count = permitCount
                 });
 
             var result = new RedisConcurrencyResponse();
@@ -147,7 +159,7 @@ internal async Task<RedisConcurrencyResponse> TryAcquireLeaseAsync(string reques
             return result;
         }
 
-        internal RedisConcurrencyResponse TryAcquireLease(string requestId, bool tryEnqueue = false)
+        internal RedisConcurrencyResponse TryAcquireLease(string requestId, int permitCount, bool tryEnqueue = false)
         {
             var nowUnixTimeSeconds = DateTimeOffset.UtcNow.ToUnixTimeSeconds();
 
@@ -165,6 +177,7 @@ internal RedisConcurrencyResponse TryAcquireLease(string requestId, bool tryEnqu
                     stats_key = StatsRateLimitKey,
                     current_time = nowUnixTimeSeconds,
                     unique_id = requestId,
+                    permit_count = permitCount
                 });
 
             var result = new RedisConcurrencyResponse();
@@ -180,22 +193,41 @@ internal RedisConcurrencyResponse TryAcquireLease(string requestId, bool tryEnqu
             return result;
         }
 
-        internal void ReleaseLease(string requestId)
+        internal void ReleaseLease(string requestId, int permitCount)
         {
             var database = _connectionMultiplexer.GetDatabase();
-            database.SortedSetRemove(RateLimitKey, requestId);
+
+            for (var i = 1; i <= permitCount; i++)
+            {
+                database.SortedSetRemove(RateLimitKey, $"{requestId}:{i}");
+            }
         }
 
-        internal async Task ReleaseLeaseAsync(string requestId)
+        internal Task ReleaseLeaseAsync(string requestId, int permitCount)
         {
             var database = _connectionMultiplexer.GetDatabase();
-            await database.SortedSetRemoveAsync(RateLimitKey, requestId);
+            var tasks = new List<Task>(permitCount);
+
+            for (var i = 1; i <= permitCount; i++)
+            {
+                tasks.Add(database.SortedSetRemoveAsync(RateLimitKey, $"{requestId}:{i}"));
+            }
+
+            return Task.WhenAll(tasks);
         }
 
-        internal async Task ReleaseQueueLeaseAsync(string requestId)
+        internal Task ReleaseQueueLeaseAsync(string requestId, int permitCount)
         {
             var database = _connectionMultiplexer.GetDatabase();
-            await database.SortedSetRemoveAsync(QueueRateLimitKey, requestId);
+
+            var tasks = new List<Task>(permitCount);
+
+            for (var i = 1; i <= permitCount; i++)
+            {
+                tasks.Add(database.SortedSetRemoveAsync(QueueRateLimitKey, $"{requestId}:{i}"));
+            }
+
+            return Task.WhenAll(tasks);
         }
 
         internal RateLimiterStatistics? GetStatistics()

src/RedisRateLimiting/Concurrency/RedisConcurrencyRateLimiter.cs

@@ -71,7 +71,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(cancellationToken);
+            return AcquireAsyncCoreInternal(permitCount, cancellationToken);
         }
 
         protected override RateLimitLease AttemptAcquireCore(int permitCount)
@@ -85,9 +85,10 @@ protected override RateLimitLease AttemptAcquireCore(int permitCount)
             {
                 Limit = _options.PermitLimit,
                 RequestId = Guid.NewGuid().ToString(),
+                PermitCount = permitCount
             };
 
-            var response = _redisManager.TryAcquireLease(leaseContext.RequestId);
+            var response = _redisManager.TryAcquireLease(leaseContext.RequestId, permitCount);
 
             leaseContext.Count = response.Count;
 
@@ -99,15 +100,16 @@ protected override RateLimitLease AttemptAcquireCore(int permitCount)
             return new ConcurrencyLease(isAcquired: false, this, leaseContext);
         }
 
-        private async ValueTask<RateLimitLease> AcquireAsyncCoreInternal(CancellationToken cancellationToken)
+        private async ValueTask<RateLimitLease> AcquireAsyncCoreInternal(int permitCount, CancellationToken cancellationToken)
         {
             var leaseContext = new ConcurencyLeaseContext
             {
                 Limit = _options.PermitLimit,
                 RequestId = Guid.NewGuid().ToString(),
+                PermitCount = permitCount
             };
 
-            var response = await _redisManager.TryAcquireLeaseAsync(leaseContext.RequestId, tryEnqueue: true);
+            var response = await _redisManager.TryAcquireLeaseAsync(leaseContext.RequestId, permitCount, tryEnqueue: true);
 
             leaseContext.Count = response.Count;
 
@@ -148,7 +150,7 @@ private void Release(ConcurencyLeaseContext leaseContext)
         {
             if (leaseContext.RequestId is null) return;
 
-            _redisManager.ReleaseLease(leaseContext.RequestId);
+            _redisManager.ReleaseLease(leaseContext.RequestId, leaseContext.PermitCount);
         }
 
         private async Task StartDequeueTimerAsync(PeriodicTimer periodicTimer)
@@ -170,7 +172,7 @@ private async Task TryDequeueRequestsAsync()
                         try
                         {
                             // The request was canceled while in the pending queue
-                            await _redisManager.ReleaseQueueLeaseAsync(request.LeaseContext!.RequestId!);
+                            await _redisManager.ReleaseQueueLeaseAsync(request.LeaseContext!.RequestId!, request.LeaseContext!.PermitCount);
                         }
                         finally
                         {
@@ -182,7 +184,7 @@ private async Task TryDequeueRequestsAsync()
                         continue;
                     }
 
-                    var response = await _redisManager.TryAcquireLeaseAsync(request.LeaseContext!.RequestId!);
+                    var response = await _redisManager.TryAcquireLeaseAsync(request.LeaseContext!.RequestId!, request.LeaseContext!.PermitCount);
 
                     request.LeaseContext.Count = response.Count;
 
@@ -195,7 +197,7 @@ private async Task TryDequeueRequestsAsync()
                             if (request.TaskCompletionSource?.TrySetResult(pendingLease) == false)
                             {
                                 // The request was canceled after we acquired the lease
-                                await _redisManager.ReleaseLeaseAsync(request.LeaseContext!.RequestId!);
+                                await _redisManager.ReleaseLeaseAsync(request.LeaseContext!.RequestId!, request.LeaseContext!.PermitCount);
                             }
                         }
                         finally
@@ -257,6 +259,8 @@ private sealed class ConcurencyLeaseContext
             public long Count { get; set; }
 
             public long Limit { get; set; }
+
+            public int PermitCount { get; set; }
         }
 
         private sealed class ConcurrencyLease : RateLimitLease

test/RedisRateLimiting.Tests/UnitTests/ConcurrencyUnitTests.cs

@@ -120,6 +120,59 @@ public void CanAcquireResource()
             lease.Dispose();
         }
 
+        [Fact]
+        public async Task SupportsPermitCountFlag_NoQueue()
+        {
+            using var limiter = new RedisConcurrencyRateLimiter<string>(
+                "Test_SupportsPermitCountFlag_NoQueue_Concurrency",
+                new RedisConcurrencyRateLimiterOptions
+                {
+                    PermitLimit = 2,
+                    QueueLimit = 0,
+                    ConnectionMultiplexerFactory = Fixture.ConnectionMultiplexerFactory,
+                });
+
+            var lease1 = await limiter.AcquireAsync(permitCount: 2);
+            var lease2 = await limiter.AcquireAsync(permitCount: 1);
+
+            Assert.True(lease1.IsAcquired);
+            Assert.False(lease2.IsAcquired);
+
+            lease1.Dispose();
+            lease2.Dispose();
+        }
+
+        [Fact]
+        public async Task SupportsPermitCountFlag_WithQueue()
+        {
+            using var limiter = new RedisConcurrencyRateLimiter<string>(
+                "Test_SupportsPermitCountFlag_WithQueue_Concurrency",
+                new RedisConcurrencyRateLimiterOptions
+                {
+                    PermitLimit = 2,
+                    QueueLimit = 2,
+                    ConnectionMultiplexerFactory = Fixture.ConnectionMultiplexerFactory,
+                });
+
+            var lease1 = await limiter.AcquireAsync(permitCount: 2);
+            Assert.True(lease1.IsAcquired);
+
+            var wait2 = limiter.AcquireAsync(permitCount: 2);
+            var wait3 = limiter.AcquireAsync(permitCount: 1);
+
+            await Task.Delay(1000);
+            lease1.Dispose();
+
+            var lease2 = await wait2;
+            Assert.True(lease2.IsAcquired);
+
+            var lease3 = await wait3;
+            Assert.False(lease3.IsAcquired);
+
+            lease2.Dispose();
+            lease3.Dispose();
+        }
+
         [Fact]
         public async Task CanAcquireResourceAsyncQueuesAndGrabsOldest()
         {