Make implementation aligned with FixedWindowRateLimiter: more precise…

… implementation of IdleDuration, similar behavior
cristipufu · Jun 8, 2024 · cb22a88 · cb22a88
1 parent 969cca5
commit cb22a88
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 12 deletions.
diff --git a/src/RedisRateLimiting/TokenBucket/RedisTokenBucketRateLimiter.cs b/src/RedisRateLimiting/TokenBucket/RedisTokenBucketRateLimiter.cs
@@ -10,17 +10,19 @@ namespace RedisRateLimiting
 {
     public class RedisTokenBucketRateLimiter<TKey> : RateLimiter
     {
+        private static readonly double TickFrequency = (double)TimeSpan.TicksPerSecond / Stopwatch.Frequency;
+
         private readonly RedisTokenBucketManager _redisManager;
         private readonly RedisTokenBucketRateLimiterOptions _options;
 
         private readonly TokenBucketLease FailedLease = new(isAcquired: false, null);
 
         private int _activeRequestsCount;
-        private long _lastActiveTimestamp = Stopwatch.GetTimestamp();
+        private long _idleSince = Stopwatch.GetTimestamp();
 
         public override TimeSpan? IdleDuration => Interlocked.CompareExchange(ref _activeRequestsCount, 0, 0) > 0
-            ? TimeSpan.Zero
-            : TimeSpan.FromTicks(Stopwatch.GetTimestamp() - _lastActiveTimestamp);
+            ? null
+            : new TimeSpan((long)((Stopwatch.GetTimestamp() - _idleSince) * TickFrequency));
 
         public RedisTokenBucketRateLimiter(TKey partitionKey, RedisTokenBucketRateLimiterOptions options)
         {
@@ -63,7 +65,7 @@ public RedisTokenBucketRateLimiter(TKey partitionKey, RedisTokenBucketRateLimite
 
         protected override async ValueTask<RateLimitLease> AcquireAsyncCore(int permitCount, CancellationToken cancellationToken)
         {
-            _lastActiveTimestamp = Stopwatch.GetTimestamp();
+            _idleSince = 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));
@@ -77,14 +79,12 @@ protected override async ValueTask<RateLimitLease> AcquireAsyncCore(int permitCo
             finally
             {
                 Interlocked.Decrement(ref _activeRequestsCount);
-                _lastActiveTimestamp = Stopwatch.GetTimestamp();
+                _idleSince = 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
@@ -119,11 +119,6 @@ public async Task IdleDurationIsUpdated()
             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);
         }
     }
 }