Skip to content

Commit 1b9e4e3

Browse files
BrennanConroyBrennanConroy
authored
Fix queue count in rate limiters (#90810)
1 parent f868d08 commit 1b9e4e3

File tree

4 files changed

+136
-22
lines changed

4 files changed

+136
-22
lines changed

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

Lines changed: 34 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -156,8 +156,17 @@ protected override ValueTask<RateLimitLease> AcquireAsyncCore(int permitCount, C
156156
Debug.Assert(_queueCount >= 0);
157157
if (!oldestRequest.TrySetResult(FailedLease))
158158
{
159-
// Updating queue count is handled by the cancellation code
160-
_queueCount += oldestRequest.Count;
159+
if (!oldestRequest.QueueCountModified)
160+
{
161+
// We already updated the queue count, the Cancel code is about to run or running and waiting on our lock,
162+
// tell Cancel not to do anything
163+
oldestRequest.QueueCountModified = true;
164+
}
165+
else
166+
{
167+
// Updating queue count was handled by the cancellation code, don't double count
168+
_queueCount += oldestRequest.Count;
169+
}
161170
}
162171
else
163172
{
@@ -277,10 +286,19 @@ private void Release(int releaseCount)
277286
// Check if request was canceled
278287
if (!nextPendingRequest.TrySetResult(lease))
279288
{
280-
// Queued item was canceled so add count back
289+
// Queued item was canceled so add count back, permits weren't acquired
281290
_permitCount += nextPendingRequest.Count;
282-
// Updating queue count is handled by the cancellation code
283-
_queueCount += nextPendingRequest.Count;
291+
if (!nextPendingRequest.QueueCountModified)
292+
{
293+
// We already updated the queue count, the Cancel code is about to run or running and waiting on our lock,
294+
// tell Cancel not to do anything
295+
nextPendingRequest.QueueCountModified = true;
296+
}
297+
else
298+
{
299+
// Updating queue count was handled by the cancellation code, don't double count
300+
_queueCount += nextPendingRequest.Count;
301+
}
284302
}
285303
else
286304
{
@@ -399,6 +417,9 @@ private sealed class RequestRegistration : TaskCompletionSource<RateLimitLease>
399417
private readonly CancellationToken _cancellationToken;
400418
private CancellationTokenRegistration _cancellationTokenRegistration;
401419

420+
// Update under the limiter lock and only if the queue count was updated by the calling code
421+
public bool QueueCountModified { get; set; }
422+
402423
// this field is used only by the disposal mechanics and never shared between threads
403424
private RequestRegistration? _next;
404425

@@ -429,7 +450,14 @@ private static void Cancel(object? state)
429450
var limiter = (ConcurrencyLimiter)registration.Task.AsyncState!;
430451
lock (limiter.Lock)
431452
{
432-
limiter._queueCount -= registration.Count;
453+
// Queuing and replenishing code might modify the _queueCount, since there is no guarantee of when the cancellation
454+
// code runs and we only want to update the _queueCount once, we set a bool (under a lock) so either method
455+
// can update the count and not double count.
456+
if (!registration.QueueCountModified)
457+
{
458+
limiter._queueCount -= registration.Count;
459+
registration.QueueCountModified = true;
460+
}
433461
}
434462
}
435463
}

‎src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/FixedWindowRateLimiter.cs‎

Lines changed: 34 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -173,7 +173,17 @@ protected override ValueTask<RateLimitLease> AcquireAsyncCore(int permitCount, C
173173
Debug.Assert(_queueCount >= 0);
174174
if (!oldestRequest.TrySetResult(FailedLease))
175175
{
176-
_queueCount += oldestRequest.Count;
176+
if (!oldestRequest.QueueCountModified)
177+
{
178+
// We already updated the queue count, the Cancel code is about to run or running and waiting on our lock,
179+
// tell Cancel not to do anything
180+
oldestRequest.QueueCountModified = true;
181+
}
182+
else
183+
{
184+
// Updating queue count was handled by the cancellation code, don't double count
185+
_queueCount += oldestRequest.Count;
186+
}
177187
}
178188
else
179189
{
@@ -330,10 +340,19 @@ private void ReplenishInternal(long nowTicks)
330340

331341
if (!nextPendingRequest.TrySetResult(SuccessfulLease))
332342
{
333-
// Queued item was canceled so add count back
343+
// Queued item was canceled so add count back, permits weren't acquired
334344
_permitCount += nextPendingRequest.Count;
335-
// Updating queue count is handled by the cancellation code
336-
_queueCount += nextPendingRequest.Count;
345+
if (!nextPendingRequest.QueueCountModified)
346+
{
347+
// We already updated the queue count, the Cancel code is about to run or running and waiting on our lock,
348+
// tell Cancel not to do anything
349+
nextPendingRequest.QueueCountModified = true;
350+
}
351+
else
352+
{
353+
// Updating queue count was handled by the cancellation code, don't double count
354+
_queueCount += nextPendingRequest.Count;
355+
}
337356
}
338357
else
339358
{
@@ -435,6 +454,9 @@ private sealed class RequestRegistration : TaskCompletionSource<RateLimitLease>
435454
private readonly CancellationToken _cancellationToken;
436455
private CancellationTokenRegistration _cancellationTokenRegistration;
437456

457+
// Update under the limiter lock and only if the queue count was updated by the calling code
458+
public bool QueueCountModified { get; set; }
459+
438460
// this field is used only by the disposal mechanics and never shared between threads
439461
private RequestRegistration? _next;
440462

@@ -465,7 +487,14 @@ private static void Cancel(object? state)
465487
var limiter = (FixedWindowRateLimiter)registration.Task.AsyncState!;
466488
lock (limiter.Lock)
467489
{
468-
limiter._queueCount -= registration.Count;
490+
// Queuing and replenishing code might modify the _queueCount, since there is no guarantee of when the cancellation
491+
// code runs and we only want to update the _queueCount once, we set a bool (under a lock) so either method
492+
// can update the count and not double count.
493+
if (!registration.QueueCountModified)
494+
{
495+
limiter._queueCount -= registration.Count;
496+
registration.QueueCountModified = true;
497+
}
469498
}
470499
}
471500
}

‎src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/SlidingWindowRateLimiter.cs‎

Lines changed: 34 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -185,7 +185,17 @@ protected override ValueTask<RateLimitLease> AcquireAsyncCore(int permitCount, C
185185
Debug.Assert(_queueCount >= 0);
186186
if (!oldestRequest.TrySetResult(FailedLease))
187187
{
188-
_queueCount += oldestRequest.Count;
188+
if (!oldestRequest.QueueCountModified)
189+
{
190+
// We already updated the queue count, the Cancel code is about to run or running and waiting on our lock,
191+
// tell Cancel not to do anything
192+
oldestRequest.QueueCountModified = true;
193+
}
194+
else
195+
{
196+
// Updating queue count was handled by the cancellation code, don't double count
197+
_queueCount += oldestRequest.Count;
198+
}
189199
}
190200
else
191201
{
@@ -342,11 +352,20 @@ private void ReplenishInternal(long nowTicks)
342352

343353
if (!nextPendingRequest.TrySetResult(SuccessfulLease))
344354
{
345-
// Queued item was canceled so add count back
355+
// Queued item was canceled so add count back, permits weren't acquired
346356
_permitCount += nextPendingRequest.Count;
347357
_requestsPerSegment[_currentSegmentIndex] -= nextPendingRequest.Count;
348-
// Updating queue count is handled by the cancellation code
349-
_queueCount += nextPendingRequest.Count;
358+
if (!nextPendingRequest.QueueCountModified)
359+
{
360+
// We already updated the queue count, the Cancel code is about to run or running and waiting on our lock,
361+
// tell Cancel not to do anything
362+
nextPendingRequest.QueueCountModified = true;
363+
}
364+
else
365+
{
366+
// Updating queue count was handled by the cancellation code, don't double count
367+
_queueCount += nextPendingRequest.Count;
368+
}
350369
}
351370
else
352371
{
@@ -448,6 +467,9 @@ private sealed class RequestRegistration : TaskCompletionSource<RateLimitLease>
448467
private readonly CancellationToken _cancellationToken;
449468
private CancellationTokenRegistration _cancellationTokenRegistration;
450469

470+
// Update under the limiter lock and only if the queue count was updated by the calling code
471+
public bool QueueCountModified { get; set; }
472+
451473
// this field is used only by the disposal mechanics and never shared between threads
452474
private RequestRegistration? _next;
453475

@@ -478,7 +500,14 @@ private static void Cancel(object? state)
478500
var limiter = (SlidingWindowRateLimiter)registration.Task.AsyncState!;
479501
lock (limiter.Lock)
480502
{
481-
limiter._queueCount -= registration.Count;
503+
// Queuing and replenishing code might modify the _queueCount, since there is no guarantee of when the cancellation
504+
// code runs and we only want to update the _queueCount once, we set a bool (under a lock) so either method
505+
// can update the count and not double count.
506+
if (!registration.QueueCountModified)
507+
{
508+
limiter._queueCount -= registration.Count;
509+
registration.QueueCountModified = true;
510+
}
482511
}
483512
}
484513
}

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

Lines changed: 34 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -178,8 +178,17 @@ protected override ValueTask<RateLimitLease> AcquireAsyncCore(int tokenCount, Ca
178178
Debug.Assert(_queueCount >= 0);
179179
if (!oldestRequest.TrySetResult(FailedLease))
180180
{
181-
// Updating queue count is handled by the cancellation code
182-
_queueCount += oldestRequest.Count;
181+
if (!oldestRequest.QueueCountModified)
182+
{
183+
// We already updated the queue count, the Cancel code is about to run or running and waiting on our lock,
184+
// tell Cancel not to do anything
185+
oldestRequest.QueueCountModified = true;
186+
}
187+
else
188+
{
189+
// Updating queue count was handled by the cancellation code, don't double count
190+
_queueCount += oldestRequest.Count;
191+
}
183192
}
184193
else
185194
{
@@ -345,10 +354,19 @@ private void ReplenishInternal(long nowTicks)
345354

346355
if (!nextPendingRequest.TrySetResult(SuccessfulLease))
347356
{
348-
// Queued item was canceled so add count back
357+
// Queued item was canceled so add count back, permits weren't acquired
349358
_tokenCount += nextPendingRequest.Count;
350-
// Updating queue count is handled by the cancellation code
351-
_queueCount += nextPendingRequest.Count;
359+
if (!nextPendingRequest.QueueCountModified)
360+
{
361+
// We already updated the queue count, the Cancel code is about to run or running and waiting on our lock,
362+
// tell Cancel not to do anything
363+
nextPendingRequest.QueueCountModified = true;
364+
}
365+
else
366+
{
367+
// Updating queue count was handled by the cancellation code, don't double count
368+
_queueCount += nextPendingRequest.Count;
369+
}
352370
}
353371
else
354372
{
@@ -450,6 +468,9 @@ private sealed class RequestRegistration : TaskCompletionSource<RateLimitLease>
450468
private readonly CancellationToken _cancellationToken;
451469
private CancellationTokenRegistration _cancellationTokenRegistration;
452470

471+
// Update under the limiter lock and only if the queue count was updated by the calling code
472+
public bool QueueCountModified { get; set; }
473+
453474
// this field is used only by the disposal mechanics and never shared between threads
454475
private RequestRegistration? _next;
455476

@@ -480,7 +501,14 @@ private static void Cancel(object? state)
480501
var limiter = (TokenBucketRateLimiter)registration.Task.AsyncState!;
481502
lock (limiter.Lock)
482503
{
483-
limiter._queueCount -= registration.Count;
504+
// Queuing and replenishing code might modify the _queueCount, since there is no guarantee of when the cancellation
505+
// code runs and we only want to update the _queueCount once, we set a bool (under a lock) so either method
506+
// can update the count and not double count.
507+
if (!registration.QueueCountModified)
508+
{
509+
limiter._queueCount -= registration.Count;
510+
registration.QueueCountModified = true;
511+
}
484512
}
485513
}
486514
}

0 commit comments

Comments
 (0)