Implement IdleDuration for RedisTokenBucketRateLimiter.

cristipufu · Apr 1, 2024 · 969cca5 · 969cca5
1 parent e7715fd
commit 969cca5
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 3 deletions.
diff --git a/src/RedisRateLimiting/TokenBucket/RedisTokenBucketRateLimiter.cs b/src/RedisRateLimiting/TokenBucket/RedisTokenBucketRateLimiter.cs
@@ -1,6 +1,7 @@
 using RedisRateLimiting.Concurrency;
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Threading;
 using System.Threading.RateLimiting;
 using System.Threading.Tasks;
@@ -14,7 +15,12 @@ public class RedisTokenBucketRateLimiter<TKey> : RateLimiter
 
         private readonly TokenBucketLease FailedLease = new(isAcquired: false, null);
 
-        public override TimeSpan? IdleDuration => TimeSpan.Zero;
+        private int _activeRequestsCount;
+        private long _lastActiveTimestamp = Stopwatch.GetTimestamp();
+
+        public override TimeSpan? IdleDuration => Interlocked.CompareExchange(ref _activeRequestsCount, 0, 0) > 0
+            ? TimeSpan.Zero
+            : TimeSpan.FromTicks(Stopwatch.GetTimestamp() - _lastActiveTimestamp);
 
         public RedisTokenBucketRateLimiter(TKey partitionKey, RedisTokenBucketRateLimiterOptions options)
         {
@@ -55,18 +61,30 @@ public RedisTokenBucketRateLimiter(TKey partitionKey, RedisTokenBucketRateLimite
             throw new NotImplementedException();
         }
 
-        protected override ValueTask<RateLimitLease> AcquireAsyncCore(int permitCount, CancellationToken cancellationToken)
+        protected override async ValueTask<RateLimitLease> AcquireAsyncCore(int permitCount, CancellationToken cancellationToken)
         {
+            _lastActiveTimestamp = Stopwatch.GetTimestamp();
             if (permitCount > _options.TokenLimit)
             {
                 throw new ArgumentOutOfRangeException(nameof(permitCount), permitCount, string.Format("{0} permit(s) exceeds the permit limit of {1}.", permitCount, _options.TokenLimit));
             }
 
-            return AcquireAsyncCoreInternal();
+            Interlocked.Increment(ref _activeRequestsCount);
+            try
+            {
+                return await AcquireAsyncCoreInternal();
+            }
+            finally
+            {
+                Interlocked.Decrement(ref _activeRequestsCount);
+                _lastActiveTimestamp = Stopwatch.GetTimestamp();
+            }
         }
 
         protected override RateLimitLease AttemptAcquireCore(int permitCount)
         {
+            _lastActiveTimestamp = Stopwatch.GetTimestamp();
+
             // https://github.com/cristipufu/aspnetcore-redis-rate-limiting/issues/66
             return FailedLease;
         }

diff --git a/test/RedisRateLimiting.Tests/UnitTests/TokenBucketUnitTests.cs b/test/RedisRateLimiting.Tests/UnitTests/TokenBucketUnitTests.cs
@@ -100,5 +100,30 @@ public async Task CanAcquireAsyncResource()
             using var lease2 = await limiter.AcquireAsync();
             Assert.False(lease2.IsAcquired);
         }
+
+        [Fact]
+        public async Task IdleDurationIsUpdated()
+        {
+            using var limiter = new RedisTokenBucketRateLimiter<string>(
+                partitionKey: Guid.NewGuid().ToString(),
+                new RedisTokenBucketRateLimiterOptions
+                {
+                    TokenLimit = 1,
+                    TokensPerPeriod = 1,
+                    ReplenishmentPeriod = TimeSpan.FromMinutes(1),
+                    ConnectionMultiplexerFactory = Fixture.ConnectionMultiplexerFactory,
+                });
+            await Task.Delay(TimeSpan.FromMilliseconds(5));
+            Assert.NotEqual(TimeSpan.Zero, limiter.IdleDuration);
+
+            var previousIdleDuration = limiter.IdleDuration;
+            using var lease = await limiter.AcquireAsync();
+            Assert.True(limiter.IdleDuration < previousIdleDuration);
+
+            await Task.Delay(TimeSpan.FromMilliseconds(5));
+            previousIdleDuration = limiter.IdleDuration;
+            using var lease2 = limiter.AttemptAcquire();
+            Assert.True(limiter.IdleDuration < previousIdleDuration);
+        }
     }
 }