Limit pool capacity, add LinkTo helper

microsoft · Oct 20, 2021 · 733e986 · 733e986
1 parent 2a50622
commit 733e986
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 6 deletions.
diff --git a/src/ReverseProxy/Forwarder/StreamCopyHttpContent.cs b/src/ReverseProxy/Forwarder/StreamCopyHttpContent.cs
@@ -143,7 +143,7 @@ async Task SerializeToStreamAsync(Stream stream, TransportContext? context, Canc
             if (_activityToken.Token != cancellationToken)
             {
                 Debug.Assert(cancellationToken.CanBeCanceled);
-                registration = cancellationToken.UnsafeRegister(ActivityCancellationTokenSource.LinkedTokenCancelDelegate, _activityToken);
+                registration = _activityToken.LinkTo(cancellationToken);
             }
 #else
             // On .NET Core 3.1, cancellationToken will always be CancellationToken.None

diff --git a/src/ReverseProxy/Utilities/ActivityCancellationTokenSource.cs b/src/ReverseProxy/Utilities/ActivityCancellationTokenSource.cs
@@ -10,10 +10,12 @@ namespace Yarp.ReverseProxy.Utilities
     internal sealed class ActivityCancellationTokenSource : CancellationTokenSource
     {
 #if NET6_0_OR_GREATER
+        private const int MaxQueueSize = 1024;
         private static readonly ConcurrentQueue<ActivityCancellationTokenSource> _sharedSources = new();
+        private static int _count;
 #endif
 
-        public static readonly Action<object?> LinkedTokenCancelDelegate = static s =>
+        private static readonly Action<object?> _linkedTokenCancelDelegate = static s =>
         {
             ((ActivityCancellationTokenSource)s!).Cancel(throwOnFirstException: false);
         };
@@ -31,7 +33,11 @@ public void ResetTimeout()
         public static ActivityCancellationTokenSource Rent(TimeSpan activityTimeout, CancellationToken linkedToken)
         {
 #if NET6_0_OR_GREATER
-            if (!_sharedSources.TryDequeue(out var cts))
+            if (_sharedSources.TryDequeue(out var cts))
+            {
+                Interlocked.Decrement(ref _count);
+            }
+            else
             {
                 cts = new ActivityCancellationTokenSource();
             }
@@ -40,7 +46,7 @@ public static ActivityCancellationTokenSource Rent(TimeSpan activityTimeout, Can
 #endif
 
             cts._activityTimeoutMs = (int)activityTimeout.TotalMilliseconds;
-            cts._linkedRegistration = linkedToken.UnsafeRegister(LinkedTokenCancelDelegate, cts);
+            cts._linkedRegistration = cts.LinkTo(linkedToken);
             cts.ResetTimeout();
 
             return cts;
@@ -54,12 +60,22 @@ public void Return()
 #if NET6_0_OR_GREATER
             if (TryReset())
             {
-                _sharedSources.Enqueue(this);
-                return;
+                if (Interlocked.Increment(ref _count) <= MaxQueueSize)
+                {
+                    _sharedSources.Enqueue(this);
+                    return;
+                }
+
+                Interlocked.Decrement(ref _count);
             }
 #endif
 
             Dispose();
         }
+
+        public CancellationTokenRegistration LinkTo(CancellationToken linkedToken)
+        {
+            return linkedToken.UnsafeRegister(_linkedTokenCancelDelegate, this);
+        }
     }
 }