TokenBucketRateLimiter.WaitAsync() does not immediately reduce the _queueCount on cancellation

[halter73]

Description

TokenBucketRateLimiter.WaitAsync(int, CancellationToken) still counts the requested tokens against the QueueLimit after they've been canceled by the CancellationToken after an incomplete ValueTask is returned. This is eventually fixed when the request is dequeued from the Deque<RequestRegistration> by ReplenishInternal, but that could be too late.

Reproduction Steps

var limiter = new TokenBucketRateLimiter(new TokenBucketRateLimiterOptions(1, QueueProcessingOrder.OldestFirst, 1,
    TimeSpan.Zero, 1, autoReplenishment: false));

limiter.Acquire(1);

var cts = new CancellationTokenSource();
var firstWaitTask = limiter.WaitAsync(1, cts.Token);

cts.Cancel();
await Assert.ThrowsAsync<OperationCanceledException>(() => firstWaitTask.AsTask());

// This incorrectly returns a failed lease immediately because ctc.Cancel() did not reduce the _queueCount.
var secondWaitTask = limiter.WaitAsync(1, cts.Token);
//Assert.False(secondWaitTask.IsCompleted)

Assert.True(limiter.TryReplenish());

// WaitAsync() no longer returns a failed lease because we called TryReplenish(), but that shouldn't be necessary
var thirdWaitTask = limiter.WaitAsync(1, cts.Token);
Assert.False(thirdWaitTask.IsCompleted)

Expected behavior

The secondWaitTask should initially be incomplete and later complete successfully when the TokenBucketRateLimiter is replenished.

Actual behavior

The secondWaitTask immediately returns a non-acquired lease because the QueueLimit has supposedly been exceeded (but not really).

Regression?

No.

Known Workarounds

No response

Configuration

No response

Other information

No response

[ghost]

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

Issue Details

---

Description

TokenBucketRateLimiter.WaitAsync(int, CancellationToken) still counts the requested tokens against the QueueLimit after they've been canceled by the CancellationToken after an incomplete ValueTask is returned. This is eventually fixed when the request is dequeued from the Deque<RequestRegistration> by ReplenishInternal, but that could be too late.

Reproduction Steps

var limiter = new TokenBucketRateLimiter(new TokenBucketRateLimiterOptions(1, QueueProcessingOrder.OldestFirst, 1,
    TimeSpan.Zero, 1, autoReplenishment: false));

limiter.Acquire(1);

var cts = new CancellationTokenSource();
var firstWaitTask = limiter.WaitAsync(1, cts.Token);

cts.Cancel();
await Assert.ThrowsAsync<OperationCanceledException>(() => firstWaitTask.AsTask());

// This incorrectly returns a failed lease immediately because ctc.Cancel() did not reduce the _queueCount.
var secondWaitTask = limiter.WaitAsync(1, cts.Token);
//Assert.False(secondWaitTask.IsCompleted)

Assert.True(limiter.TryReplenish());

// WaitAsync() no longer returns a failed lease because we called TryReplenish(), but that shouldn't be necessary
var thirdWaitTask = limiter.WaitAsync(1, cts.Token);
Assert.False(thirdWaitTask.IsCompleted)

Expected behavior

The secondWaitTask should initially be incomplete and later complete successfully when the TokenBucketRateLimiter is replenished.

Actual behavior

The secondWaitTask immediately returns a non-acquired lease because the QueueLimit has supposedly been exceeded (but not really).

Regression?

No.

Known Workarounds

No response

Configuration

No response

Other information

No response

Author:	halter73
Assignees:	-
Labels:	area-System.Threading, untriaged
Milestone:	-

[stephentoub]

Sounds similar to #761. They're using the same helper type, right? We might want to address them together. I haven't fixed that one yet as I've been concerned about the perf impact.

[halter73]

I'm not very familiar with the Channel code. What helper type are you referring to? AsyncOperation? I don't think were sharing anything. I think the following code will need to be changed to decrement _queueCount:

runtime/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/TokenBucketRateLimiter.cs

Lines 132 to 135 in 0b8204f

	ctr = cancellationToken.Register(static obj =>
	{
	((TaskCompletionSource<RateLimitLease>)obj!).TrySetException(new OperationCanceledException());
	}, tcs);

And this logic will need to be updated to not decrement _queueCount in the cancellation case:

runtime/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/TokenBucketRateLimiter.cs

Lines 277 to 286 in 0b8204f

	_queueCount -= nextPendingRequest.Count;
	_tokenCount -= nextPendingRequest.Count;
	Debug.Assert(_queueCount >= 0);
	Debug.Assert(_tokenCount >= 0);
	if (!nextPendingRequest.Tcs.TrySetResult(SuccessfulLease))
	{
	// Queued item was canceled so add count back
	_tokenCount += nextPendingRequest.Count;
	}

We're also somewhat concerned about perf because the CancellationToken.Register callback will need a reference to the TokenBucketRateLimiter instance and the tcs forcing an extra allocation. It's far from the only allocation on this code path though, and we've discussed pooling a custom IValueTaskSource in the future which we could then use to store the extra reference avoiding the extra allocation.

[stephentoub]

What helper type are you referring to? AsyncOperation? I don't think were sharing anything.

Ah, ok. I must be misremembering. I thought it was also using AsyncOperation, but I must have been thinking of this:

runtime/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/ConcurrencyLimiter.cs

Line 311 in 67848f2

// Perf: Use AsyncOperation<TResult> instead