Added support for RedisSlidingWindow statistics

cristipufu · Apr 14, 2023 · 5ef7f16 · 5ef7f16
1 parent a77709a
commit 5ef7f16
Show file tree

Hide file tree

Showing 2 changed files with 45 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 string 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 = $"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)