diff --git a/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/ConcurrencyLimiter.cs b/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/ConcurrencyLimiter.cs index 6b5a4014990ef..7131b4fe1d799 100644 --- a/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/ConcurrencyLimiter.cs +++ b/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/ConcurrencyLimiter.cs @@ -156,8 +156,17 @@ protected override ValueTask AcquireAsyncCore(int permitCount, C Debug.Assert(_queueCount >= 0); if (!oldestRequest.TrySetResult(FailedLease)) { - // Updating queue count is handled by the cancellation code - _queueCount += oldestRequest.Count; + if (!oldestRequest.QueueCountModified) + { + // We already updated the queue count, the Cancel code is about to run or running and waiting on our lock, + // tell Cancel not to do anything + oldestRequest.QueueCountModified = true; + } + else + { + // Updating queue count was handled by the cancellation code, don't double count + _queueCount += oldestRequest.Count; + } } else { @@ -277,10 +286,19 @@ private void Release(int releaseCount) // Check if request was canceled if (!nextPendingRequest.TrySetResult(lease)) { - // Queued item was canceled so add count back + // Queued item was canceled so add count back, permits weren't acquired _permitCount += nextPendingRequest.Count; - // Updating queue count is handled by the cancellation code - _queueCount += nextPendingRequest.Count; + if (!nextPendingRequest.QueueCountModified) + { + // We already updated the queue count, the Cancel code is about to run or running and waiting on our lock, + // tell Cancel not to do anything + nextPendingRequest.QueueCountModified = true; + } + else + { + // Updating queue count was handled by the cancellation code, don't double count + _queueCount += nextPendingRequest.Count; + } } else { @@ -399,6 +417,9 @@ private sealed class RequestRegistration : TaskCompletionSource private readonly CancellationToken _cancellationToken; private CancellationTokenRegistration _cancellationTokenRegistration; + // Update under the limiter lock and only if the queue count was updated by the calling code + public bool QueueCountModified { get; set; } + // this field is used only by the disposal mechanics and never shared between threads private RequestRegistration? _next; @@ -429,7 +450,14 @@ private static void Cancel(object? state) var limiter = (ConcurrencyLimiter)registration.Task.AsyncState!; lock (limiter.Lock) { - limiter._queueCount -= registration.Count; + // Queuing and replenishing code might modify the _queueCount, since there is no guarantee of when the cancellation + // code runs and we only want to update the _queueCount once, we set a bool (under a lock) so either method + // can update the count and not double count. + if (!registration.QueueCountModified) + { + limiter._queueCount -= registration.Count; + registration.QueueCountModified = true; + } } } } diff --git a/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/FixedWindowRateLimiter.cs b/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/FixedWindowRateLimiter.cs index d09c7973b18aa..daaed9cf5ce42 100644 --- a/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/FixedWindowRateLimiter.cs +++ b/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/FixedWindowRateLimiter.cs @@ -173,7 +173,17 @@ protected override ValueTask AcquireAsyncCore(int permitCount, C Debug.Assert(_queueCount >= 0); if (!oldestRequest.TrySetResult(FailedLease)) { - _queueCount += oldestRequest.Count; + if (!oldestRequest.QueueCountModified) + { + // We already updated the queue count, the Cancel code is about to run or running and waiting on our lock, + // tell Cancel not to do anything + oldestRequest.QueueCountModified = true; + } + else + { + // Updating queue count was handled by the cancellation code, don't double count + _queueCount += oldestRequest.Count; + } } else { @@ -330,10 +340,19 @@ private void ReplenishInternal(long nowTicks) if (!nextPendingRequest.TrySetResult(SuccessfulLease)) { - // Queued item was canceled so add count back + // Queued item was canceled so add count back, permits weren't acquired _permitCount += nextPendingRequest.Count; - // Updating queue count is handled by the cancellation code - _queueCount += nextPendingRequest.Count; + if (!nextPendingRequest.QueueCountModified) + { + // We already updated the queue count, the Cancel code is about to run or running and waiting on our lock, + // tell Cancel not to do anything + nextPendingRequest.QueueCountModified = true; + } + else + { + // Updating queue count was handled by the cancellation code, don't double count + _queueCount += nextPendingRequest.Count; + } } else { @@ -435,6 +454,9 @@ private sealed class RequestRegistration : TaskCompletionSource private readonly CancellationToken _cancellationToken; private CancellationTokenRegistration _cancellationTokenRegistration; + // Update under the limiter lock and only if the queue count was updated by the calling code + public bool QueueCountModified { get; set; } + // this field is used only by the disposal mechanics and never shared between threads private RequestRegistration? _next; @@ -465,7 +487,14 @@ private static void Cancel(object? state) var limiter = (FixedWindowRateLimiter)registration.Task.AsyncState!; lock (limiter.Lock) { - limiter._queueCount -= registration.Count; + // Queuing and replenishing code might modify the _queueCount, since there is no guarantee of when the cancellation + // code runs and we only want to update the _queueCount once, we set a bool (under a lock) so either method + // can update the count and not double count. + if (!registration.QueueCountModified) + { + limiter._queueCount -= registration.Count; + registration.QueueCountModified = true; + } } } } diff --git a/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/SlidingWindowRateLimiter.cs b/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/SlidingWindowRateLimiter.cs index a179720ede33f..23dbf98e0fcde 100644 --- a/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/SlidingWindowRateLimiter.cs +++ b/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/SlidingWindowRateLimiter.cs @@ -185,7 +185,17 @@ protected override ValueTask AcquireAsyncCore(int permitCount, C Debug.Assert(_queueCount >= 0); if (!oldestRequest.TrySetResult(FailedLease)) { - _queueCount += oldestRequest.Count; + if (!oldestRequest.QueueCountModified) + { + // We already updated the queue count, the Cancel code is about to run or running and waiting on our lock, + // tell Cancel not to do anything + oldestRequest.QueueCountModified = true; + } + else + { + // Updating queue count was handled by the cancellation code, don't double count + _queueCount += oldestRequest.Count; + } } else { @@ -342,11 +352,20 @@ private void ReplenishInternal(long nowTicks) if (!nextPendingRequest.TrySetResult(SuccessfulLease)) { - // Queued item was canceled so add count back + // Queued item was canceled so add count back, permits weren't acquired _permitCount += nextPendingRequest.Count; _requestsPerSegment[_currentSegmentIndex] -= nextPendingRequest.Count; - // Updating queue count is handled by the cancellation code - _queueCount += nextPendingRequest.Count; + if (!nextPendingRequest.QueueCountModified) + { + // We already updated the queue count, the Cancel code is about to run or running and waiting on our lock, + // tell Cancel not to do anything + nextPendingRequest.QueueCountModified = true; + } + else + { + // Updating queue count was handled by the cancellation code, don't double count + _queueCount += nextPendingRequest.Count; + } } else { @@ -448,6 +467,9 @@ private sealed class RequestRegistration : TaskCompletionSource private readonly CancellationToken _cancellationToken; private CancellationTokenRegistration _cancellationTokenRegistration; + // Update under the limiter lock and only if the queue count was updated by the calling code + public bool QueueCountModified { get; set; } + // this field is used only by the disposal mechanics and never shared between threads private RequestRegistration? _next; @@ -478,7 +500,14 @@ private static void Cancel(object? state) var limiter = (SlidingWindowRateLimiter)registration.Task.AsyncState!; lock (limiter.Lock) { - limiter._queueCount -= registration.Count; + // Queuing and replenishing code might modify the _queueCount, since there is no guarantee of when the cancellation + // code runs and we only want to update the _queueCount once, we set a bool (under a lock) so either method + // can update the count and not double count. + if (!registration.QueueCountModified) + { + limiter._queueCount -= registration.Count; + registration.QueueCountModified = true; + } } } } 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 5ad7859792ff7..67a3a55a29ad0 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 @@ -178,8 +178,17 @@ protected override ValueTask AcquireAsyncCore(int tokenCount, Ca Debug.Assert(_queueCount >= 0); if (!oldestRequest.TrySetResult(FailedLease)) { - // Updating queue count is handled by the cancellation code - _queueCount += oldestRequest.Count; + if (!oldestRequest.QueueCountModified) + { + // We already updated the queue count, the Cancel code is about to run or running and waiting on our lock, + // tell Cancel not to do anything + oldestRequest.QueueCountModified = true; + } + else + { + // Updating queue count was handled by the cancellation code, don't double count + _queueCount += oldestRequest.Count; + } } else { @@ -345,10 +354,19 @@ private void ReplenishInternal(long nowTicks) if (!nextPendingRequest.TrySetResult(SuccessfulLease)) { - // Queued item was canceled so add count back + // Queued item was canceled so add count back, permits weren't acquired _tokenCount += nextPendingRequest.Count; - // Updating queue count is handled by the cancellation code - _queueCount += nextPendingRequest.Count; + if (!nextPendingRequest.QueueCountModified) + { + // We already updated the queue count, the Cancel code is about to run or running and waiting on our lock, + // tell Cancel not to do anything + nextPendingRequest.QueueCountModified = true; + } + else + { + // Updating queue count was handled by the cancellation code, don't double count + _queueCount += nextPendingRequest.Count; + } } else { @@ -450,6 +468,9 @@ private sealed class RequestRegistration : TaskCompletionSource private readonly CancellationToken _cancellationToken; private CancellationTokenRegistration _cancellationTokenRegistration; + // Update under the limiter lock and only if the queue count was updated by the calling code + public bool QueueCountModified { get; set; } + // this field is used only by the disposal mechanics and never shared between threads private RequestRegistration? _next; @@ -480,7 +501,14 @@ private static void Cancel(object? state) var limiter = (TokenBucketRateLimiter)registration.Task.AsyncState!; lock (limiter.Lock) { - limiter._queueCount -= registration.Count; + // Queuing and replenishing code might modify the _queueCount, since there is no guarantee of when the cancellation + // code runs and we only want to update the _queueCount once, we set a bool (under a lock) so either method + // can update the count and not double count. + if (!registration.QueueCountModified) + { + limiter._queueCount -= registration.Count; + registration.QueueCountModified = true; + } } } }