Merge pull request #157 from dlxeon/feat/idleduration-for-bucketratel…

…imiter Implement IdleDuration for RedisTokenBucketRateLimiter
cristipufu · Jun 13, 2024 · bb1dd3a · bb1dd3a
2 parents c1522a8 + 24c8000
commit bb1dd3a
Show file tree

Hide file tree

Showing 8 changed files with 161 additions and 11 deletions.
diff --git a/src/RedisRateLimiting/Concurrency/RedisConcurrencyRateLimiter.cs b/src/RedisRateLimiting/Concurrency/RedisConcurrencyRateLimiter.cs
@@ -2,6 +2,7 @@
 using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Threading;
 using System.Threading.RateLimiting;
 using System.Threading.Tasks;
@@ -10,6 +11,8 @@ namespace RedisRateLimiting
 {
     public class RedisConcurrencyRateLimiter<TKey> : RateLimiter
     {
+        private static readonly double TickFrequency = (double)TimeSpan.TicksPerSecond / Stopwatch.Frequency;
+
         private readonly RedisConcurrencyManager _redisManager;
         private readonly RedisConcurrencyRateLimiterOptions _options;
         private readonly ConcurrentQueue<Request> _queue = new();
@@ -20,7 +23,12 @@ public class RedisConcurrencyRateLimiter<TKey> : RateLimiter
 
         private readonly ConcurrencyLease FailedLease = new(false, null, null);
 
-        public override TimeSpan? IdleDuration => TimeSpan.Zero;
+        private int _activeRequestsCount;
+        private long _idleSince = Stopwatch.GetTimestamp();
+
+        public override TimeSpan? IdleDuration => Interlocked.CompareExchange(ref _activeRequestsCount, 0, 0) > 0
+            ? null
+            : new TimeSpan((long)((Stopwatch.GetTimestamp() - _idleSince) * TickFrequency));
 
         public RedisConcurrencyRateLimiter(TKey partitionKey, RedisConcurrencyRateLimiterOptions options)
         {
@@ -64,14 +72,24 @@ public RedisConcurrencyRateLimiter(TKey partitionKey, RedisConcurrencyRateLimite
             return _redisManager.GetStatistics();
         }
 
-        protected override ValueTask<RateLimitLease> AcquireAsyncCore(int permitCount, CancellationToken cancellationToken)
+        protected override async ValueTask<RateLimitLease> AcquireAsyncCore(int permitCount, CancellationToken cancellationToken)
         {
+            _idleSince = Stopwatch.GetTimestamp();
             if (permitCount > _options.PermitLimit)
             {
                 throw new ArgumentOutOfRangeException(nameof(permitCount), permitCount, string.Format("{0} permit(s) exceeds the permit limit of {1}.", permitCount, _options.PermitLimit));
             }
 
-            return AcquireAsyncCoreInternal(cancellationToken);
+            Interlocked.Increment(ref _activeRequestsCount);
+            try
+            {
+                return await AcquireAsyncCoreInternal(cancellationToken);
+            }
+            finally
+            {
+                Interlocked.Decrement(ref _activeRequestsCount);
+                _idleSince = Stopwatch.GetTimestamp();
+            }
         }
 
         protected override RateLimitLease AttemptAcquireCore(int permitCount)

diff --git a/src/RedisRateLimiting/FixedWindow/RedisFixedWindowRateLimiter.cs b/src/RedisRateLimiting/FixedWindow/RedisFixedWindowRateLimiter.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;
@@ -9,12 +10,19 @@ namespace RedisRateLimiting
 {
     public class RedisFixedWindowRateLimiter<TKey> : RateLimiter
     {
+        private static readonly double TickFrequency = (double)TimeSpan.TicksPerSecond / Stopwatch.Frequency;
+
         private readonly RedisFixedWindowManager _redisManager;
         private readonly RedisFixedWindowRateLimiterOptions _options;
 
         private readonly FixedWindowLease FailedLease = new(isAcquired: false, null);
 
-        public override TimeSpan? IdleDuration => TimeSpan.Zero;
+        private int _activeRequestsCount;
+        private long _idleSince = Stopwatch.GetTimestamp();
+
+        public override TimeSpan? IdleDuration => Interlocked.CompareExchange(ref _activeRequestsCount, 0, 0) > 0
+            ? null
+            : new TimeSpan((long)((Stopwatch.GetTimestamp() - _idleSince) * TickFrequency));
 
         public RedisFixedWindowRateLimiter(TKey partitionKey, RedisFixedWindowRateLimiterOptions options)
         {
@@ -74,7 +82,17 @@ private async ValueTask<RateLimitLease> AcquireAsyncCoreInternal(int permitCount
                 Window = _options.Window,
             };
 
-            var response = await _redisManager.TryAcquireLeaseAsync(permitCount);
+            RedisFixedWindowResponse response;
+            Interlocked.Increment(ref _activeRequestsCount);
+            try
+            {
+                response = await _redisManager.TryAcquireLeaseAsync(permitCount);
+            }
+            finally
+            {
+                Interlocked.Decrement(ref _activeRequestsCount);
+                _idleSince = Stopwatch.GetTimestamp();
+            }
 
             leaseContext.Count = response.Count;
             leaseContext.RetryAfter = response.RetryAfter;

diff --git a/src/RedisRateLimiting/SlidingWindow/RedisSlidingWindowRateLimiter.cs b/src/RedisRateLimiting/SlidingWindow/RedisSlidingWindowRateLimiter.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;
@@ -9,12 +10,19 @@ namespace RedisRateLimiting
 {
     public class RedisSlidingWindowRateLimiter<TKey> : RateLimiter
     {
+        private static readonly double TickFrequency = (double)TimeSpan.TicksPerSecond / Stopwatch.Frequency;
+
         private readonly RedisSlidingWindowManager _redisManager;
         private readonly RedisSlidingWindowRateLimiterOptions _options;
 
         private readonly SlidingWindowLease FailedLease = new(isAcquired: false, null);
 
-        public override TimeSpan? IdleDuration => TimeSpan.Zero;
+        private int _activeRequestsCount;
+        private long _idleSince = Stopwatch.GetTimestamp();
+
+        public override TimeSpan? IdleDuration => Interlocked.CompareExchange(ref _activeRequestsCount, 0, 0) > 0
+            ? null
+            : new TimeSpan((long)((Stopwatch.GetTimestamp() - _idleSince) * TickFrequency));
 
         public RedisSlidingWindowRateLimiter(TKey partitionKey, RedisSlidingWindowRateLimiterOptions options)
         {
@@ -50,14 +58,24 @@ public RedisSlidingWindowRateLimiter(TKey partitionKey, RedisSlidingWindowRateLi
             return _redisManager.GetStatistics();
         }
 
-        protected override ValueTask<RateLimitLease> AcquireAsyncCore(int permitCount, CancellationToken cancellationToken)
+        protected override async ValueTask<RateLimitLease> AcquireAsyncCore(int permitCount, CancellationToken cancellationToken)
         {
+            _idleSince = Stopwatch.GetTimestamp();
             if (permitCount > _options.PermitLimit)
             {
                 throw new ArgumentOutOfRangeException(nameof(permitCount), permitCount, string.Format("{0} permit(s) exceeds the permit limit of {1}.", permitCount, _options.PermitLimit));
             }
 
-            return AcquireAsyncCoreInternal();
+            Interlocked.Increment(ref _activeRequestsCount);
+            try
+            {
+                return await AcquireAsyncCoreInternal();
+            }
+            finally
+            {
+                Interlocked.Decrement(ref _activeRequestsCount);
+                _idleSince = Stopwatch.GetTimestamp();
+            }
         }
 
         protected override RateLimitLease AttemptAcquireCore(int permitCount)

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;
@@ -9,12 +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);
 
-        public override TimeSpan? IdleDuration => TimeSpan.Zero;
+        private int _activeRequestsCount;
+        private long _idleSince = Stopwatch.GetTimestamp();
+
+        public override TimeSpan? IdleDuration => Interlocked.CompareExchange(ref _activeRequestsCount, 0, 0) > 0
+            ? null
+            : new TimeSpan((long)((Stopwatch.GetTimestamp() - _idleSince) * TickFrequency));
 
         public RedisTokenBucketRateLimiter(TKey partitionKey, RedisTokenBucketRateLimiterOptions options)
         {
@@ -55,14 +63,24 @@ 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)
         {
+            _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));
             }
 
-            return AcquireAsyncCoreInternal(permitCount);
+            Interlocked.Increment(ref _activeRequestsCount);
+            try
+            {
+                return await AcquireAsyncCoreInternal(permitCount);
+            }
+            finally
+            {
+                Interlocked.Decrement(ref _activeRequestsCount);
+                _idleSince = Stopwatch.GetTimestamp();
+            }
         }
 
         protected override RateLimitLease AttemptAcquireCore(int permitCount)

diff --git a/test/RedisRateLimiting.Tests/UnitTests/ConcurrencyUnitTests.cs b/test/RedisRateLimiting.Tests/UnitTests/ConcurrencyUnitTests.cs
@@ -339,6 +339,26 @@ public async Task GetPermitWhilePermitEmptyQueueNotEmptyGetsQueued()
             using var lease3 = await wait3;
             Assert.True(lease3.IsAcquired);
         }
+
+        [Fact]
+        public async Task IdleDurationIsUpdated()
+        {
+            await using var limiter = new RedisConcurrencyRateLimiter<string>(
+                partitionKey: Guid.NewGuid().ToString(),
+                new RedisConcurrencyRateLimiterOptions
+                {
+                    PermitLimit = 1,
+                    QueueLimit = 1,
+                    TryDequeuePeriod = TimeSpan.FromHours(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);
+        }
 
         static internal void ForceDequeue(RedisConcurrencyRateLimiter<string> limiter)
         {

diff --git a/test/RedisRateLimiting.Tests/UnitTests/FixedWindowUnitTests.cs b/test/RedisRateLimiting.Tests/UnitTests/FixedWindowUnitTests.cs
@@ -101,5 +101,24 @@ public async Task CanAcquireMultiplePermits()
             using var lease3 = await limiter.AcquireAsync(permitCount: 2);
             Assert.True(lease3.IsAcquired);
         }
+
+        [Fact]
+        public async Task IdleDurationIsUpdated()
+        {
+            await using var limiter = new RedisFixedWindowRateLimiter<string>(
+                partitionKey: Guid.NewGuid().ToString(),
+                new RedisFixedWindowRateLimiterOptions
+                {
+                    PermitLimit = 1,
+                    Window = 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);
+        }        
     }
 }
diff --git a/test/RedisRateLimiting.Tests/UnitTests/SlidingWindowUnitTests.cs b/test/RedisRateLimiting.Tests/UnitTests/SlidingWindowUnitTests.cs
@@ -124,5 +124,24 @@ public async Task CanAcquireAsyncResourceWithSmallWindow()
             using var lease4 = await limiter.AcquireAsync();
             Assert.False(lease4.IsAcquired);
         }
+
+        [Fact]
+        public async Task IdleDurationIsUpdated()
+        {
+            await using var limiter = new RedisSlidingWindowRateLimiter<string>(
+                partitionKey: Guid.NewGuid().ToString(),
+                new RedisSlidingWindowRateLimiterOptions
+                {
+                    PermitLimit = 1,
+                    Window = TimeSpan.FromMilliseconds(600),
+                    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);
+        }
     }
 }
diff --git a/test/RedisRateLimiting.Tests/UnitTests/TokenBucketUnitTests.cs b/test/RedisRateLimiting.Tests/UnitTests/TokenBucketUnitTests.cs
@@ -123,5 +123,25 @@ public async Task CanAcquireMultiPermits()
             using var lease3 = await limiter.AcquireAsync(1);
             Assert.True(lease3.IsAcquired);
         }
+
+        [Fact]
+        public async Task IdleDurationIsUpdated()
+        {
+            await 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);
+        }
     }
 }