Skip to content

Fix AttemptAcquire(0) when token isn't available#123841

Open
reedz wants to merge 2 commits intodotnet:mainfrom
reedz:feature/token-bucket-rate-limiter-fix
Open

Fix AttemptAcquire(0) when token isn't available#123841
reedz wants to merge 2 commits intodotnet:mainfrom
reedz:feature/token-bucket-rate-limiter-fix

Conversation

@reedz
Copy link
Contributor

@reedz reedz commented Jan 31, 2026

Fixes #118192

Updated the _tokenCount checks to be fraction-aware (i.e. no longer >0, but >=1) and added unit tests to reproduce the issue.

Copilot AI review requested due to automatic review settings January 31, 2026 14:09
@dotnet-policy-service dotnet-policy-service bot added the community-contribution Indicates that the PR has been added by a community member label Jan 31, 2026
@dotnet-policy-service
Copy link
Contributor

Tagging subscribers to this area: @agocke, @VSadov
See info in area-owners.md if you want to be subscribed.

Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Fixes TokenBucketRateLimiter behavior where AttemptAcquire(0)/AcquireAsync(0) could succeed with only fractional tokens available (inconsistent with CurrentAvailablePermits truncation), and adds tests covering the fractional-token scenarios.

Changes:

  • Update zero-token acquisition fast paths to require at least one whole token (_tokenCount >= 1) rather than any positive fractional value.
  • Ensure queued zero-token requests are only fulfilled once at least one whole token is available.
  • Add unit tests reproducing and validating the fractional-token behavior for both sync and async zero-token acquisition.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File Description
src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/TokenBucketRateLimiter.cs Tightens zero-token acquisition/queue fulfillment conditions to treat fractional _tokenCount as unavailable.
src/libraries/System.Threading.RateLimiting/tests/TokenBucketRateLimiterTests.cs Adds targeted tests to cover fractional-token edge cases for AttemptAcquire(0) and AcquireAsync(0).

Comment on lines +516 to +522
var acquireTask = limiter.AcquireAsync(0);
Assert.False(acquireTask.IsCompleted);

Replenish(limiter, 1);
using var lease2 = await acquireTask;
Assert.True(lease2.IsAcquired);
}
Copy link

Copilot AI Jan 31, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In this test, awaiting acquireTask without any timeout can hang the test run if the replenishment math/rounding doesn’t quite reach the >= 1 threshold (especially since the test intentionally works with fractional tokens). Consider awaiting with a bounded timeout (e.g., convert to Task and use WaitAsync/WhenAny) and/or replenishing with a larger elapsed time on the second replenish to guarantee completion.

Copilot uses AI. Check for mistakes.
Comment on lines +551 to +552
Replenish(limiter, 1);
Assert.Equal(1, limiter.GetStatistics()!.CurrentAvailablePermits);
Copy link

Copilot AI Jan 31, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This assertion/await sequence is sensitive to timing conversion and truncation: if the final replenishment leaves _tokenCount just under 1.0, CurrentAvailablePermits will still be 0 and the awaited task may not complete. To make the test robust, consider replenishing with a slightly larger elapsed time before asserting/awaiting (or assert availability using a threshold rather than relying on hitting exactly 1.0).

Suggested change
Replenish(limiter, 1);
Assert.Equal(1, limiter.GetStatistics()!.CurrentAvailablePermits);
Replenish(limiter, 2);
Assert.True(limiter.GetStatistics()!.CurrentAvailablePermits >= 1);

Copilot uses AI. Check for mistakes.
Comment on lines +339 to 340
else if (_tokenCount >= nextPendingRequest.Count && (nextPendingRequest.Count > 0 || _tokenCount >= 1))
{
Copy link

Copilot AI Jan 31, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This fulfillment condition works, but it’s hard to read due to redundant comparisons (for Count > 0, _tokenCount >= Count already implies >= 1). Consider rewriting it in a more direct form (e.g., branch on nextPendingRequest.Count == 0) to reduce cognitive load and help avoid future mistakes around the tokenCount == 0 special-case.

Copilot uses AI. Check for mistakes.
@stephentoub
Copy link
Member

🤖 Copilot Code Review — PR #123841

Holistic Assessment

Motivation: This PR addresses a valid inconsistency identified in issue #118192 where AttemptAcquire(0) returns a successful lease when _tokenCount is in the range (0, 1), while GetStatistics().CurrentAvailablePermits truncates to long and reports 0. The goal of aligning these behaviors is reasonable.

Approach: The fix changes the condition from _tokenCount > 0 to _tokenCount >= 1 in several places. This ensures zero-token acquisitions (which are used to probe "are permits exhausted") only succeed when at least one whole token is available, matching the truncated statistics.

Summary: ⚠️ Needs Human Review. The code changes appear logically consistent and the fix aligns the semantics of AttemptAcquire(0) with GetStatistics(). However, there are some considerations that a maintainer should evaluate:

  1. Semantic clarification needed: The API documentation says "Set permitCount to 0 to get whether permits are exhausted." The question is whether "exhausted" means _tokenCount == 0 (no tokens at all) or (long)_tokenCount == 0 (no whole tokens). This PR chooses the latter interpretation, which is consistent with what statistics reports, but this is a behavioral change that some callers may depend on.

  2. The tests are good: Contrary to initial concerns, the PR does include tests that create fractional token states (e.g., Replenish(limiter, 1) with ReplenishmentPeriod=2ms and TokensPerPeriod=1 adds 0.5 tokens).


Detailed Findings

✅ Code Changes — Logically consistent

The four changes are consistent with each other:

  • AttemptAcquireCore (line 115): Changes fast-path zero-token check from > 0 to >= 1
  • AcquireAsyncCore (line 149): Same change for async path ✓
  • TryLeaseUnsynchronized (line 230): Changes != 0 to >= 1, which is equivalent for non-negative values and more consistent ✓
  • ReplenishInternal (line 339): Adds (nextPendingRequest.Count > 0 || _tokenCount >= 1) to prevent fulfilling queued zero-token requests with fractional tokens ✓

The logic in ReplenishInternal correctly handles both cases:

  • For nextPendingRequest.Count > 0: The existing _tokenCount >= nextPendingRequest.Count is sufficient (if we have enough for N tokens, we have at least N >= 1)
  • For nextPendingRequest.Count == 0: The additional _tokenCount >= 1 check ensures we don't fulfill a zero-token probe with fractional tokens

✅ Tests — Cover the fractional token scenario

The three new tests properly exercise the fix:

  1. AcquireZero_WithFractionalTokens_Fails: Verifies AttemptAcquire(0) fails with 0.5 tokens
  2. AcquireAsyncZero_WithFractionalTokens_Queues: Verifies AcquireAsync(0) queues and waits for a whole token
  3. AcquireAsyncZero_QueuedRequest_NotFulfilledWithFractionalTokens: Verifies queued zero-token requests aren't fulfilled until a whole token is available

⚠️ Thread Safety — Pre-existing pattern, but worth noting

In AttemptAcquireCore and AcquireAsyncCore, _tokenCount (a double) is read outside the lock before the comparison. On 32-bit platforms, reading a 64-bit double is not atomic and could theoretically result in a torn read. However:

  • This pattern existed before this PR (the check was just > 0 instead of >= 1)
  • These are optimistic fast-path checks; the lock path will re-verify
  • This is not a regression introduced by this PR

💡 Suggestion — Consider alternative approach using integer tokens

An alternative design would be to store _tokenCount as an integer (or use separate integer/fractional components) to avoid the truncation mismatch entirely. However, this would be a larger refactor and the current fix is minimal and correct.


Cross-Model Review Notes

This review was synthesized from multiple AI models (Gemini, GPT). Key points of agreement:

  • The logic changes are consistent across all paths
  • The thread-safety concern is pre-existing and not introduced by this PR
  • The fix correctly aligns AttemptAcquire(0) behavior with GetStatistics()

One model incorrectly claimed the tests don't cover fractional tokens, but review of the actual test code shows they do create fractional token states via partial replenishment.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area-System.Threading community-contribution Indicates that the PR has been added by a community member

Projects

None yet

Development

Successfully merging this pull request may close these issues.

TokenBucketRateLimiter: AttemptAcquire(0) succeeds while CurrentAvailablePermits is 0 due to fractional _tokenCount

2 participants