Merge pull request #50 from kamilslusarczyk/issue/6

Added support for RedisSlidingWindow statistics
cristipufu · Apr 21, 2023 · d5151b7 · d5151b7
2 parents 27cfdf5 + f1ae6c1
commit d5151b7
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 1 deletion.
diff --git a/src/RedisRateLimiting/SlidingWindow/RedisSlidingWindowManager.cs b/src/RedisRateLimiting/SlidingWindow/RedisSlidingWindowManager.cs
@@ -1,5 +1,6 @@
 using StackExchange.Redis;
 using System;
+using System.Threading.RateLimiting;
 using System.Threading.Tasks;
 
 namespace RedisRateLimiting.Concurrency
@@ -9,6 +10,7 @@ internal class RedisSlidingWindowManager
         private readonly IConnectionMultiplexer _connectionMultiplexer;
         private readonly RedisSlidingWindowRateLimiterOptions _options;
         private readonly RedisKey RateLimitKey;
+        private readonly RedisKey StatsRateLimitKey;
 
         private static readonly LuaScript _redisScript = LuaScript.Prepare(
           @"local limit = tonumber(@permit_limit)
@@ -28,8 +30,22 @@ internal class RedisSlidingWindowManager
 
             redis.call(""expireat"", @rate_limit_key, timestamp + window + 1)
 
+            if allowed
+            then
+                redis.call(""hincrby"", @stats_key, 'total_successful', 1)
+            else
+                redis.call(""hincrby"", @stats_key, 'total_failed', 1)
+            end
+
             return { allowed, count }");
 
+        private static readonly LuaScript StatisticsScript = LuaScript.Prepare(
+            @"local count = redis.call(""zcard"", @rate_limit_key)
+            local total_successful_count = redis.call(""hget"", @stats_key, 'total_successful')
+            local total_failed_count = redis.call(""hget"", @stats_key, 'total_failed')
+
+            return { count, total_successful_count, total_failed_count }");
+
         public RedisSlidingWindowManager(
             string partitionKey,
             RedisSlidingWindowRateLimiterOptions options)
@@ -38,6 +54,7 @@ public RedisSlidingWindowManager(
             _connectionMultiplexer = options.ConnectionMultiplexerFactory!.Invoke();
 
             RateLimitKey = new RedisKey($"rl:{{{partitionKey}}}");
+            StatsRateLimitKey = new RedisKey($"rl:{{{partitionKey}}}:stats");
         }
 
         internal async Task<RedisSlidingWindowResponse> TryAcquireLeaseAsync(string requestId)
@@ -54,6 +71,7 @@ internal async Task<RedisSlidingWindowResponse> TryAcquireLeaseAsync(string requ
                     rate_limit_key = RateLimitKey,
                     permit_limit = _options.PermitLimit,
                     window = _options.Window.TotalSeconds,
+                    stats_key = StatsRateLimitKey,
                     current_time = nowUnixTimeSeconds,
                     unique_id = requestId,
                 });
@@ -83,6 +101,7 @@ internal RedisSlidingWindowResponse TryAcquireLease(string requestId)
                    rate_limit_key = RateLimitKey,
                    permit_limit = _options.PermitLimit,
                    window = _options.Window.TotalSeconds,
+                   stats_key = StatsRateLimitKey,
                    current_time = nowUnixTimeSeconds,
                    unique_id = requestId,
                });
@@ -97,6 +116,31 @@ internal RedisSlidingWindowResponse TryAcquireLease(string requestId)
 
             return result;
         }
+
+        internal RateLimiterStatistics? GetStatistics()
+        {
+            var database = _connectionMultiplexer.GetDatabase();
+
+            var response = (RedisValue[]?)database.ScriptEvaluate(
+                StatisticsScript,
+                new
+                {
+                    rate_limit_key = RateLimitKey,
+                    stats_key = StatsRateLimitKey,
+                });
+
+            if (response == null)
+            {
+                return null;
+            }
+
+            return new RateLimiterStatistics
+            {
+                CurrentAvailablePermits = _options.PermitLimit - (long)response[0],
+                TotalSuccessfulLeases = (long)response[1],
+                TotalFailedLeases = (long)response[2],
+            };
+        }
     }
 
     internal class RedisSlidingWindowResponse

diff --git a/src/RedisRateLimiting/SlidingWindow/RedisSlidingWindowRateLimiter.cs b/src/RedisRateLimiting/SlidingWindow/RedisSlidingWindowRateLimiter.cs
@@ -45,7 +45,7 @@ public RedisSlidingWindowRateLimiter(TKey partitionKey, RedisSlidingWindowRateLi
 
         public override RateLimiterStatistics? GetStatistics()
         {
-            throw new NotImplementedException();
+            return _redisManager.GetStatistics();
         }
 
         protected override ValueTask<RateLimitLease> AcquireAsyncCore(int permitCount, CancellationToken cancellationToken)

diff --git a/test/RedisRateLimiting.Tests/UnitTests/SlidingWindowUnitTests.cs b/test/RedisRateLimiting.Tests/UnitTests/SlidingWindowUnitTests.cs
@@ -78,6 +78,8 @@ public async Task ThrowsWhenAcquiringMoreThanLimit()
         [Fact]
         public async Task CanAcquireAsyncResource()
         {
+            await Fixture.ClearStatisticsAsync("Test_CanAcquireAsyncResource_SW");
+
             using var limiter = new RedisSlidingWindowRateLimiter<string>(
                 "Test_CanAcquireAsyncResource_SW",
                 new RedisSlidingWindowRateLimiterOptions
@@ -92,6 +94,17 @@ public async Task CanAcquireAsyncResource()
 
             using var lease2 = await limiter.AcquireAsync();
             Assert.False(lease2.IsAcquired);
+
+            var stats = limiter.GetStatistics()!;
+            Assert.Equal(1, stats.TotalSuccessfulLeases);
+            Assert.Equal(1, stats.TotalFailedLeases);
+            Assert.Equal(0, stats.CurrentAvailablePermits);
+
+            lease.Dispose();
+            lease2.Dispose();
+
+            stats = limiter.GetStatistics()!;
+            Assert.Equal(0, stats.CurrentAvailablePermits);
         }
     }
 }