Skip to content

Commit bb1dd3a

Browse files
cristipufucristipufu
authored
Merge pull request #157 from dlxeon/feat/idleduration-for-bucketratelimiter
Implement IdleDuration for RedisTokenBucketRateLimiter
2 parents c1522a8 + 24c8000 commit bb1dd3a

File tree

8 files changed

+161
-11
lines changed

8 files changed

+161
-11
lines changed

‎src/RedisRateLimiting/Concurrency/RedisConcurrencyRateLimiter.cs

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
using System;
33
using System.Collections.Concurrent;
44
using System.Collections.Generic;
5+
using System.Diagnostics;
56
using System.Threading;
67
using System.Threading.RateLimiting;
78
using System.Threading.Tasks;
@@ -10,6 +11,8 @@ namespace RedisRateLimiting
1011
{
1112
public class RedisConcurrencyRateLimiter<TKey> : RateLimiter
1213
{
14+
private static readonly double TickFrequency = (double)TimeSpan.TicksPerSecond / Stopwatch.Frequency;
15+
1316
private readonly RedisConcurrencyManager _redisManager;
1417
private readonly RedisConcurrencyRateLimiterOptions _options;
1518
private readonly ConcurrentQueue<Request> _queue = new();
@@ -20,7 +23,12 @@ public class RedisConcurrencyRateLimiter<TKey> : RateLimiter
2023

2124
private readonly ConcurrencyLease FailedLease = new(false, null, null);
2225

23-
public override TimeSpan? IdleDuration => TimeSpan.Zero;
26+
private int _activeRequestsCount;
27+
private long _idleSince = Stopwatch.GetTimestamp();
28+
29+
public override TimeSpan? IdleDuration => Interlocked.CompareExchange(ref _activeRequestsCount, 0, 0) > 0
30+
? null
31+
: new TimeSpan((long)((Stopwatch.GetTimestamp() - _idleSince) * TickFrequency));
2432

2533
public RedisConcurrencyRateLimiter(TKey partitionKey, RedisConcurrencyRateLimiterOptions options)
2634
{
@@ -64,14 +72,24 @@ public RedisConcurrencyRateLimiter(TKey partitionKey, RedisConcurrencyRateLimite
6472
return _redisManager.GetStatistics();
6573
}
6674

67-
protected override ValueTask<RateLimitLease> AcquireAsyncCore(int permitCount, CancellationToken cancellationToken)
75+
protected override async ValueTask<RateLimitLease> AcquireAsyncCore(int permitCount, CancellationToken cancellationToken)
6876
{
77+
_idleSince = Stopwatch.GetTimestamp();
6978
if (permitCount > _options.PermitLimit)
7079
{
7180
throw new ArgumentOutOfRangeException(nameof(permitCount), permitCount, string.Format("{0} permit(s) exceeds the permit limit of {1}.", permitCount, _options.PermitLimit));
7281
}
7382

74-
return AcquireAsyncCoreInternal(cancellationToken);
83+
Interlocked.Increment(ref _activeRequestsCount);
84+
try
85+
{
86+
return await AcquireAsyncCoreInternal(cancellationToken);
87+
}
88+
finally
89+
{
90+
Interlocked.Decrement(ref _activeRequestsCount);
91+
_idleSince = Stopwatch.GetTimestamp();
92+
}
7593
}
7694

7795
protected override RateLimitLease AttemptAcquireCore(int permitCount)

‎src/RedisRateLimiting/FixedWindow/RedisFixedWindowRateLimiter.cs

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
using RedisRateLimiting.Concurrency;
22
using System;
33
using System.Collections.Generic;
4+
using System.Diagnostics;
45
using System.Threading;
56
using System.Threading.RateLimiting;
67
using System.Threading.Tasks;
@@ -9,12 +10,19 @@ namespace RedisRateLimiting
910
{
1011
public class RedisFixedWindowRateLimiter<TKey> : RateLimiter
1112
{
13+
private static readonly double TickFrequency = (double)TimeSpan.TicksPerSecond / Stopwatch.Frequency;
14+
1215
private readonly RedisFixedWindowManager _redisManager;
1316
private readonly RedisFixedWindowRateLimiterOptions _options;
1417

1518
private readonly FixedWindowLease FailedLease = new(isAcquired: false, null);
1619

17-
public override TimeSpan? IdleDuration => TimeSpan.Zero;
20+
private int _activeRequestsCount;
21+
private long _idleSince = Stopwatch.GetTimestamp();
22+
23+
public override TimeSpan? IdleDuration => Interlocked.CompareExchange(ref _activeRequestsCount, 0, 0) > 0
24+
? null
25+
: new TimeSpan((long)((Stopwatch.GetTimestamp() - _idleSince) * TickFrequency));
1826

1927
public RedisFixedWindowRateLimiter(TKey partitionKey, RedisFixedWindowRateLimiterOptions options)
2028
{
@@ -74,7 +82,17 @@ private async ValueTask<RateLimitLease> AcquireAsyncCoreInternal(int permitCount
7482
Window = _options.Window,
7583
};
7684

77-
var response = await _redisManager.TryAcquireLeaseAsync(permitCount);
85+
RedisFixedWindowResponse response;
86+
Interlocked.Increment(ref _activeRequestsCount);
87+
try
88+
{
89+
response = await _redisManager.TryAcquireLeaseAsync(permitCount);
90+
}
91+
finally
92+
{
93+
Interlocked.Decrement(ref _activeRequestsCount);
94+
_idleSince = Stopwatch.GetTimestamp();
95+
}
7896

7997
leaseContext.Count = response.Count;
8098
leaseContext.RetryAfter = response.RetryAfter;

‎src/RedisRateLimiting/SlidingWindow/RedisSlidingWindowRateLimiter.cs

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
using RedisRateLimiting.Concurrency;
22
using System;
33
using System.Collections.Generic;
4+
using System.Diagnostics;
45
using System.Threading;
56
using System.Threading.RateLimiting;
67
using System.Threading.Tasks;
@@ -9,12 +10,19 @@ namespace RedisRateLimiting
910
{
1011
public class RedisSlidingWindowRateLimiter<TKey> : RateLimiter
1112
{
13+
private static readonly double TickFrequency = (double)TimeSpan.TicksPerSecond / Stopwatch.Frequency;
14+
1215
private readonly RedisSlidingWindowManager _redisManager;
1316
private readonly RedisSlidingWindowRateLimiterOptions _options;
1417

1518
private readonly SlidingWindowLease FailedLease = new(isAcquired: false, null);
1619

17-
public override TimeSpan? IdleDuration => TimeSpan.Zero;
20+
private int _activeRequestsCount;
21+
private long _idleSince = Stopwatch.GetTimestamp();
22+
23+
public override TimeSpan? IdleDuration => Interlocked.CompareExchange(ref _activeRequestsCount, 0, 0) > 0
24+
? null
25+
: new TimeSpan((long)((Stopwatch.GetTimestamp() - _idleSince) * TickFrequency));
1826

1927
public RedisSlidingWindowRateLimiter(TKey partitionKey, RedisSlidingWindowRateLimiterOptions options)
2028
{
@@ -50,14 +58,24 @@ public RedisSlidingWindowRateLimiter(TKey partitionKey, RedisSlidingWindowRateLi
5058
return _redisManager.GetStatistics();
5159
}
5260

53-
protected override ValueTask<RateLimitLease> AcquireAsyncCore(int permitCount, CancellationToken cancellationToken)
61+
protected override async ValueTask<RateLimitLease> AcquireAsyncCore(int permitCount, CancellationToken cancellationToken)
5462
{
63+
_idleSince = Stopwatch.GetTimestamp();
5564
if (permitCount > _options.PermitLimit)
5665
{
5766
throw new ArgumentOutOfRangeException(nameof(permitCount), permitCount, string.Format("{0} permit(s) exceeds the permit limit of {1}.", permitCount, _options.PermitLimit));
5867
}
5968

60-
return AcquireAsyncCoreInternal();
69+
Interlocked.Increment(ref _activeRequestsCount);
70+
try
71+
{
72+
return await AcquireAsyncCoreInternal();
73+
}
74+
finally
75+
{
76+
Interlocked.Decrement(ref _activeRequestsCount);
77+
_idleSince = Stopwatch.GetTimestamp();
78+
}
6179
}
6280

6381
protected override RateLimitLease AttemptAcquireCore(int permitCount)

‎src/RedisRateLimiting/TokenBucket/RedisTokenBucketRateLimiter.cs

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
using RedisRateLimiting.Concurrency;
22
using System;
33
using System.Collections.Generic;
4+
using System.Diagnostics;
45
using System.Threading;
56
using System.Threading.RateLimiting;
67
using System.Threading.Tasks;
@@ -9,12 +10,19 @@ namespace RedisRateLimiting
910
{
1011
public class RedisTokenBucketRateLimiter<TKey> : RateLimiter
1112
{
13+
private static readonly double TickFrequency = (double)TimeSpan.TicksPerSecond / Stopwatch.Frequency;
14+
1215
private readonly RedisTokenBucketManager _redisManager;
1316
private readonly RedisTokenBucketRateLimiterOptions _options;
1417

1518
private readonly TokenBucketLease FailedLease = new(isAcquired: false, null);
1619

17-
public override TimeSpan? IdleDuration => TimeSpan.Zero;
20+
private int _activeRequestsCount;
21+
private long _idleSince = Stopwatch.GetTimestamp();
22+
23+
public override TimeSpan? IdleDuration => Interlocked.CompareExchange(ref _activeRequestsCount, 0, 0) > 0
24+
? null
25+
: new TimeSpan((long)((Stopwatch.GetTimestamp() - _idleSince) * TickFrequency));
1826

1927
public RedisTokenBucketRateLimiter(TKey partitionKey, RedisTokenBucketRateLimiterOptions options)
2028
{
@@ -55,14 +63,24 @@ public RedisTokenBucketRateLimiter(TKey partitionKey, RedisTokenBucketRateLimite
5563
throw new NotImplementedException();
5664
}
5765

58-
protected override ValueTask<RateLimitLease> AcquireAsyncCore(int permitCount, CancellationToken cancellationToken)
66+
protected override async ValueTask<RateLimitLease> AcquireAsyncCore(int permitCount, CancellationToken cancellationToken)
5967
{
68+
_idleSince = Stopwatch.GetTimestamp();
6069
if (permitCount > _options.TokenLimit)
6170
{
6271
throw new ArgumentOutOfRangeException(nameof(permitCount), permitCount, string.Format("{0} permit(s) exceeds the permit limit of {1}.", permitCount, _options.TokenLimit));
6372
}
6473

65-
return AcquireAsyncCoreInternal(permitCount);
74+
Interlocked.Increment(ref _activeRequestsCount);
75+
try
76+
{
77+
return await AcquireAsyncCoreInternal(permitCount);
78+
}
79+
finally
80+
{
81+
Interlocked.Decrement(ref _activeRequestsCount);
82+
_idleSince = Stopwatch.GetTimestamp();
83+
}
6684
}
6785

6886
protected override RateLimitLease AttemptAcquireCore(int permitCount)

‎test/RedisRateLimiting.Tests/UnitTests/ConcurrencyUnitTests.cs

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -339,6 +339,26 @@ public async Task GetPermitWhilePermitEmptyQueueNotEmptyGetsQueued()
339339
using var lease3 = await wait3;
340340
Assert.True(lease3.IsAcquired);
341341
}
342+
343+
[Fact]
344+
public async Task IdleDurationIsUpdated()
345+
{
346+
await using var limiter = new RedisConcurrencyRateLimiter<string>(
347+
partitionKey: Guid.NewGuid().ToString(),
348+
new RedisConcurrencyRateLimiterOptions
349+
{
350+
PermitLimit = 1,
351+
QueueLimit = 1,
352+
TryDequeuePeriod = TimeSpan.FromHours(1),
353+
ConnectionMultiplexerFactory = Fixture.ConnectionMultiplexerFactory,
354+
});
355+
await Task.Delay(TimeSpan.FromMilliseconds(5));
356+
Assert.NotEqual(TimeSpan.Zero, limiter.IdleDuration);
357+
358+
var previousIdleDuration = limiter.IdleDuration;
359+
using var lease = await limiter.AcquireAsync();
360+
Assert.True(limiter.IdleDuration < previousIdleDuration);
361+
}
342362

343363
static internal void ForceDequeue(RedisConcurrencyRateLimiter<string> limiter)
344364
{

‎test/RedisRateLimiting.Tests/UnitTests/FixedWindowUnitTests.cs

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,5 +101,24 @@ public async Task CanAcquireMultiplePermits()
101101
using var lease3 = await limiter.AcquireAsync(permitCount: 2);
102102
Assert.True(lease3.IsAcquired);
103103
}
104+
105+
[Fact]
106+
public async Task IdleDurationIsUpdated()
107+
{
108+
await using var limiter = new RedisFixedWindowRateLimiter<string>(
109+
partitionKey: Guid.NewGuid().ToString(),
110+
new RedisFixedWindowRateLimiterOptions
111+
{
112+
PermitLimit = 1,
113+
Window = TimeSpan.FromMinutes(1),
114+
ConnectionMultiplexerFactory = Fixture.ConnectionMultiplexerFactory,
115+
});
116+
await Task.Delay(TimeSpan.FromMilliseconds(5));
117+
Assert.NotEqual(TimeSpan.Zero, limiter.IdleDuration);
118+
119+
var previousIdleDuration = limiter.IdleDuration;
120+
using var lease = await limiter.AcquireAsync();
121+
Assert.True(limiter.IdleDuration < previousIdleDuration);
122+
}
104123
}
105124
}

‎test/RedisRateLimiting.Tests/UnitTests/SlidingWindowUnitTests.cs

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,5 +124,24 @@ public async Task CanAcquireAsyncResourceWithSmallWindow()
124124
using var lease4 = await limiter.AcquireAsync();
125125
Assert.False(lease4.IsAcquired);
126126
}
127+
128+
[Fact]
129+
public async Task IdleDurationIsUpdated()
130+
{
131+
await using var limiter = new RedisSlidingWindowRateLimiter<string>(
132+
partitionKey: Guid.NewGuid().ToString(),
133+
new RedisSlidingWindowRateLimiterOptions
134+
{
135+
PermitLimit = 1,
136+
Window = TimeSpan.FromMilliseconds(600),
137+
ConnectionMultiplexerFactory = Fixture.ConnectionMultiplexerFactory,
138+
});
139+
await Task.Delay(TimeSpan.FromMilliseconds(5));
140+
Assert.NotEqual(TimeSpan.Zero, limiter.IdleDuration);
141+
142+
var previousIdleDuration = limiter.IdleDuration;
143+
using var lease = await limiter.AcquireAsync();
144+
Assert.True(limiter.IdleDuration < previousIdleDuration);
145+
}
127146
}
128147
}

‎test/RedisRateLimiting.Tests/UnitTests/TokenBucketUnitTests.cs

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -123,5 +123,25 @@ public async Task CanAcquireMultiPermits()
123123
using var lease3 = await limiter.AcquireAsync(1);
124124
Assert.True(lease3.IsAcquired);
125125
}
126+
127+
[Fact]
128+
public async Task IdleDurationIsUpdated()
129+
{
130+
await using var limiter = new RedisTokenBucketRateLimiter<string>(
131+
partitionKey: Guid.NewGuid().ToString(),
132+
new RedisTokenBucketRateLimiterOptions
133+
{
134+
TokenLimit = 1,
135+
TokensPerPeriod = 1,
136+
ReplenishmentPeriod = TimeSpan.FromMinutes(1),
137+
ConnectionMultiplexerFactory = Fixture.ConnectionMultiplexerFactory,
138+
});
139+
await Task.Delay(TimeSpan.FromMilliseconds(5));
140+
Assert.NotEqual(TimeSpan.Zero, limiter.IdleDuration);
141+
142+
var previousIdleDuration = limiter.IdleDuration;
143+
using var lease = await limiter.AcquireAsync();
144+
Assert.True(limiter.IdleDuration < previousIdleDuration);
145+
}
126146
}
127147
}

0 commit comments

Comments
 (0)