Fix queue count in rate limiters (#90878)

Co-authored-by: Brennan <brecon@microsoft.com>

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 6b5a401..7131b4f 100644 (file)
--- 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 @@ namespace System.Threading.RateLimiting
                             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 @@ namespace System.Threading.RateLimiting
                         // 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 @@ namespace System.Threading.RateLimiting
             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 @@ namespace System.Threading.RateLimiting
                     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 d09c797..daaed9c 100644 (file)
--- 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 @@ namespace System.Threading.RateLimiting
                             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 @@ namespace System.Threading.RateLimiting
 
                         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 @@ namespace System.Threading.RateLimiting
             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 @@ namespace System.Threading.RateLimiting
                     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 a179720..23dbf98 100644 (file)
--- 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 @@ namespace System.Threading.RateLimiting
                             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 @@ namespace System.Threading.RateLimiting
 
                         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 @@ namespace System.Threading.RateLimiting
             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 @@ namespace System.Threading.RateLimiting
                     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 5ad7859..67a3a55 100644 (file)
--- 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 @@ namespace System.Threading.RateLimiting
                             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 @@ namespace System.Threading.RateLimiting
 
                         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 @@ namespace System.Threading.RateLimiting
             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 @@ namespace System.Threading.RateLimiting
                     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;
+                        }
                     }
                 }
             }