From 25605d228abfc7c7531f4c7543cda0819ebeb82c Mon Sep 17 00:00:00 2001 From: reedz Date: Sat, 31 Jan 2026 15:03:03 +0100 Subject: [PATCH 1/2] fix attempt acquire with fractional token availability --- .../RateLimiting/TokenBucketRateLimiter.cs | 8 +- .../tests/TokenBucketRateLimiterTests.cs | 101 ++++++++++++++++++ 2 files changed, 105 insertions(+), 4 deletions(-) diff --git a/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/TokenBucketRateLimiter.cs b/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/TokenBucketRateLimiter.cs index 430a87978e6f2f..0d31324ea68f81 100644 --- a/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/TokenBucketRateLimiter.cs +++ b/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/TokenBucketRateLimiter.cs @@ -112,7 +112,7 @@ protected override RateLimitLease AttemptAcquireCore(int tokenCount) // Return SuccessfulLease or FailedLease depending to indicate limiter state if (tokenCount == 0 && !_disposed) { - if (_tokenCount > 0) + if (_tokenCount >= 1) { Interlocked.Increment(ref _successfulLeasesCount); return SuccessfulLease; @@ -146,7 +146,7 @@ protected override ValueTask AcquireAsyncCore(int tokenCount, Ca ThrowIfDisposed(); // Return SuccessfulAcquisition if requestedCount is 0 and resources are available - if (tokenCount == 0 && _tokenCount > 0) + if (tokenCount == 0 && _tokenCount >= 1) { Interlocked.Increment(ref _successfulLeasesCount); return new ValueTask(SuccessfulLease); @@ -227,7 +227,7 @@ private bool TryLeaseUnsynchronized(int tokenCount, [NotNullWhen(true)] out Rate ThrowIfDisposed(); // if permitCount is 0 we want to queue it if there are no available permits - if (_tokenCount >= tokenCount && _tokenCount != 0) + if (_tokenCount >= tokenCount && _tokenCount >= 1) { if (tokenCount == 0) { @@ -336,7 +336,7 @@ private void ReplenishInternal(long nowTicks) : queue.DequeueTail(); disposer.Add(nextPendingRequest); } - else if (_tokenCount >= nextPendingRequest.Count) + else if (_tokenCount >= nextPendingRequest.Count && (nextPendingRequest.Count > 0 || _tokenCount >= 1)) { // Request can be fulfilled nextPendingRequest = diff --git a/src/libraries/System.Threading.RateLimiting/tests/TokenBucketRateLimiterTests.cs b/src/libraries/System.Threading.RateLimiting/tests/TokenBucketRateLimiterTests.cs index 44388481dec53f..ac0852e73b069f 100644 --- a/src/libraries/System.Threading.RateLimiting/tests/TokenBucketRateLimiterTests.cs +++ b/src/libraries/System.Threading.RateLimiting/tests/TokenBucketRateLimiterTests.cs @@ -464,6 +464,107 @@ public override void AcquireZero_WithoutAvailability() lease2.Dispose(); } + [Fact] + public void AcquireZero_WithFractionalTokens_Fails() + { + var limiter = new TokenBucketRateLimiter(new TokenBucketRateLimiterOptions + { + TokenLimit = 3, + QueueProcessingOrder = QueueProcessingOrder.OldestFirst, + QueueLimit = 1, + ReplenishmentPeriod = TimeSpan.FromMilliseconds(2), + TokensPerPeriod = 1, + AutoReplenishment = false + }); + + using var lease1 = limiter.AttemptAcquire(3); + Assert.True(lease1.IsAcquired); + Assert.Equal(0, limiter.GetStatistics()!.CurrentAvailablePermits); + + // Replenish a fractional amount (less than 1 token) + // With ReplenishmentPeriod=2ms and TokensPerPeriod=1, replenishing 1ms adds 0.5 tokens + Replenish(limiter, 1); + + // Statistics should still report 0 available permits (fractional tokens truncated) + Assert.Equal(0, limiter.GetStatistics()!.CurrentAvailablePermits); + + // AttemptAcquire(0) should fail when there are only fractional tokens + var lease2 = limiter.AttemptAcquire(0); + Assert.False(lease2.IsAcquired); + lease2.Dispose(); + } + + [Fact] + public async Task AcquireAsyncZero_WithFractionalTokens_Queues() + { + var limiter = new TokenBucketRateLimiter(new TokenBucketRateLimiterOptions + { + TokenLimit = 3, + QueueProcessingOrder = QueueProcessingOrder.OldestFirst, + QueueLimit = 1, + ReplenishmentPeriod = TimeSpan.FromMilliseconds(2), + TokensPerPeriod = 1, + AutoReplenishment = false + }); + + using var lease1 = limiter.AttemptAcquire(3); + Assert.True(lease1.IsAcquired); + + // Replenish a fractional amount (less than 1 token) + Replenish(limiter, 1); + + Assert.Equal(0, limiter.GetStatistics()!.CurrentAvailablePermits); + + // AcquireAsync(0) should not immediately succeed with fractional tokens + var acquireTask = limiter.AcquireAsync(0); + Assert.False(acquireTask.IsCompleted); + + // After replenishing enough to have at least 1 token, the task should complete + Replenish(limiter, 1); + using var lease2 = await acquireTask; + Assert.True(lease2.IsAcquired); + } + + [Fact] + public async Task AcquireAsyncZero_QueuedRequest_NotFulfilledWithFractionalTokens() + { + // ReplenishmentPeriod = 4ms, TokensPerPeriod = 1 means 0.25 tokens per ms + var limiter = new TokenBucketRateLimiter(new TokenBucketRateLimiterOptions + { + TokenLimit = 3, + QueueProcessingOrder = QueueProcessingOrder.OldestFirst, + QueueLimit = 1, + ReplenishmentPeriod = TimeSpan.FromMilliseconds(4), + TokensPerPeriod = 1, + AutoReplenishment = false + }); + + // Drain all tokens + using var lease1 = limiter.AttemptAcquire(3); + Assert.True(lease1.IsAcquired); + Assert.Equal(0, limiter.GetStatistics()!.CurrentAvailablePermits); + + // Replenish 0.5 tokens (2ms of 4ms period) + Replenish(limiter, 2); + Assert.Equal(0, limiter.GetStatistics()!.CurrentAvailablePermits); + + // Queue an AcquireAsync(0) - should not immediately succeed with fractional tokens + var acquireTask = limiter.AcquireAsync(0); + Assert.False(acquireTask.IsCompleted, "AcquireAsync(0) should be queued when only fractional tokens available"); + + // Replenish another 0.25 tokens (1ms) - total is now 0.75 tokens, still < 1 + // This triggers queue processing but should NOT fulfill the queued request + Replenish(limiter, 1); + Assert.Equal(0, limiter.GetStatistics()!.CurrentAvailablePermits); + Assert.False(acquireTask.IsCompleted, "Queued AcquireAsync(0) should not be fulfilled with only 0.75 fractional tokens"); + + // Replenish to reach 1 full token - now it should complete + Replenish(limiter, 1); + Assert.Equal(1, limiter.GetStatistics()!.CurrentAvailablePermits); + using var lease2 = await acquireTask; + Assert.True(lease2.IsAcquired); + } + [Fact] public override async Task AcquireAsyncZero_WithAvailability() { From e5ffdcbe5fccc720e0f0d4a08f19654569640714 Mon Sep 17 00:00:00 2001 From: reedz Date: Sat, 31 Jan 2026 15:07:59 +0100 Subject: [PATCH 2/2] remove redundant comments --- .../tests/TokenBucketRateLimiterTests.cs | 15 ++------------- 1 file changed, 2 insertions(+), 13 deletions(-) diff --git a/src/libraries/System.Threading.RateLimiting/tests/TokenBucketRateLimiterTests.cs b/src/libraries/System.Threading.RateLimiting/tests/TokenBucketRateLimiterTests.cs index ac0852e73b069f..b3ae3d7600b268 100644 --- a/src/libraries/System.Threading.RateLimiting/tests/TokenBucketRateLimiterTests.cs +++ b/src/libraries/System.Threading.RateLimiting/tests/TokenBucketRateLimiterTests.cs @@ -510,16 +510,12 @@ public async Task AcquireAsyncZero_WithFractionalTokens_Queues() using var lease1 = limiter.AttemptAcquire(3); Assert.True(lease1.IsAcquired); - // Replenish a fractional amount (less than 1 token) Replenish(limiter, 1); - Assert.Equal(0, limiter.GetStatistics()!.CurrentAvailablePermits); - // AcquireAsync(0) should not immediately succeed with fractional tokens var acquireTask = limiter.AcquireAsync(0); Assert.False(acquireTask.IsCompleted); - // After replenishing enough to have at least 1 token, the task should complete Replenish(limiter, 1); using var lease2 = await acquireTask; Assert.True(lease2.IsAcquired); @@ -528,7 +524,6 @@ public async Task AcquireAsyncZero_WithFractionalTokens_Queues() [Fact] public async Task AcquireAsyncZero_QueuedRequest_NotFulfilledWithFractionalTokens() { - // ReplenishmentPeriod = 4ms, TokensPerPeriod = 1 means 0.25 tokens per ms var limiter = new TokenBucketRateLimiter(new TokenBucketRateLimiterOptions { TokenLimit = 3, @@ -539,26 +534,20 @@ public async Task AcquireAsyncZero_QueuedRequest_NotFulfilledWithFractionalToken AutoReplenishment = false }); - // Drain all tokens using var lease1 = limiter.AttemptAcquire(3); Assert.True(lease1.IsAcquired); Assert.Equal(0, limiter.GetStatistics()!.CurrentAvailablePermits); - // Replenish 0.5 tokens (2ms of 4ms period) Replenish(limiter, 2); Assert.Equal(0, limiter.GetStatistics()!.CurrentAvailablePermits); - // Queue an AcquireAsync(0) - should not immediately succeed with fractional tokens var acquireTask = limiter.AcquireAsync(0); - Assert.False(acquireTask.IsCompleted, "AcquireAsync(0) should be queued when only fractional tokens available"); + Assert.False(acquireTask.IsCompleted); - // Replenish another 0.25 tokens (1ms) - total is now 0.75 tokens, still < 1 - // This triggers queue processing but should NOT fulfill the queued request Replenish(limiter, 1); Assert.Equal(0, limiter.GetStatistics()!.CurrentAvailablePermits); - Assert.False(acquireTask.IsCompleted, "Queued AcquireAsync(0) should not be fulfilled with only 0.75 fractional tokens"); + Assert.False(acquireTask.IsCompleted); - // Replenish to reach 1 full token - now it should complete Replenish(limiter, 1); Assert.Equal(1, limiter.GetStatistics()!.CurrentAvailablePermits); using var lease2 = await acquireTask;