Limit pool size and unit test

Author: JamesNK

src/Servers/Kestrel/Transport.Quic/src/Internal/QuicConnectionContext.cs

@@ -168,7 +168,7 @@ public override ValueTask<ConnectionContext> ConnectAsync(IFeatureCollection? fe
             return new ValueTask<ConnectionContext>(context);
         }
 
-        internal void ReturnStream(QuicStreamContext stream)
+        internal bool TryReturnStream(QuicStreamContext stream)
         {
             lock (_poolLock)
             {
@@ -191,9 +191,15 @@ internal void ReturnStream(QuicStreamContext stream)
                     _streamPoolHeartbeatInitialized = true;
                 }
 
-                stream.PoolExpirationTicks = Volatile.Read(ref _heartbeatTicks) + StreamPoolExpiryTicks;
-                StreamPool.Push(stream);
+                if (stream.CanReuse && StreamPool.Count < MaxStreamPoolSize)
+                {
+                    stream.PoolExpirationTicks = Volatile.Read(ref _heartbeatTicks) + StreamPoolExpiryTicks;
+                    StreamPool.Push(stream);
+                    return true;
+                }
             }
+
+            return false;
         }
 
         private void RemoveExpiredStreams()

src/Servers/Kestrel/Transport.Quic/src/Internal/QuicConnectionListener.cs

@@ -41,6 +41,8 @@ public QuicConnectionListener(QuicTransportOptions options, IQuicTrace log, EndP
             quicListenerOptions.ServerAuthenticationOptions = sslServerAuthenticationOptions;
             quicListenerOptions.ListenEndPoint = endpoint as IPEndPoint;
             quicListenerOptions.IdleTimeout = options.IdleTimeout;
+            quicListenerOptions.MaxBidirectionalStreams = options.MaxBidirectionalStreamCount;
+            quicListenerOptions.MaxUnidirectionalStreams = options.MaxUnidirectionalStreamCount;
 
             _listener = new QuicListener(QuicImplementationProviders.MsQuic, quicListenerOptions);
 

src/Servers/Kestrel/Transport.Quic/src/Internal/QuicStreamContext.cs

@@ -393,12 +393,12 @@ public override async ValueTask DisposeAsync()
             _stream.Dispose();
             _stream = null!;
 
-            if (CanReuse)
-            {
-                _connection.ReturnStream(this);
-            }
-            else
+            if (!_connection.TryReturnStream(this))
             {
+                // Dispose when one of:
+                // - Stream is not bidirection
+                // - Stream didn't complete gracefully
+                // - Pool is full
                 DisposeCore();
             }
         }

src/Servers/Kestrel/Transport.Quic/test/QuicConnectionContextTests.cs

@@ -7,6 +7,7 @@
 using System.Net.Quic;
 using System.Text;
 using System.Threading.Tasks;
+using Microsoft.AspNetCore.Connections;
 using Microsoft.AspNetCore.Connections.Features;
 using Microsoft.AspNetCore.Server.Kestrel.FunctionalTests;
 using Microsoft.AspNetCore.Server.Kestrel.Transport.Quic.Internal;
@@ -223,6 +224,94 @@ public async Task StreamPool_Heartbeat_ExpiredStreamRemoved()
             Assert.Equal(0, quicConnectionContext.StreamPool.Count);
         }
 
+        [ConditionalFact]
+        [MsQuicSupported]
+        public async Task StreamPool_ManyConcurrentStreams_StreamPoolFull()
+        {
+            // Arrange
+            await using var connectionListener = await QuicTestHelpers.CreateConnectionListenerFactory(LoggerFactory);
+
+            var options = QuicTestHelpers.CreateClientConnectionOptions(connectionListener.EndPoint);
+            using var clientConnection = new QuicConnection(QuicImplementationProviders.MsQuic, options);
+            await clientConnection.ConnectAsync().DefaultTimeout();
+
+            await using var serverConnection = await connectionListener.AcceptAsync().DefaultTimeout();
+
+            var testHeartbeatFeature = new TestHeartbeatFeature();
+            serverConnection.Features.Set<IConnectionHeartbeatFeature>(testHeartbeatFeature);
+
+            // Act
+            var quicConnectionContext = Assert.IsType<QuicConnectionContext>(serverConnection);
+            Assert.Equal(0, quicConnectionContext.StreamPool.Count);
+
+            var pauseCompleteTcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+            var allConnectionsOnServerTcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+            var streamTasks = new List<Task>();
+            var requestState = new RequestState(clientConnection, serverConnection, allConnectionsOnServerTcs, pauseCompleteTcs.Task);
+
+            const int StreamsSent = 101;
+            for (var i = 0; i < StreamsSent; i++)
+            {
+                // TODO: Race condition in QUIC library.
+                // Delay between sending streams to avoid
+                // https://github.com/dotnet/runtime/issues/55249
+                await Task.Delay(50);
+                streamTasks.Add(SendStream(requestState));
+            }
+
+            await allConnectionsOnServerTcs.Task.DefaultTimeout();
+            pauseCompleteTcs.SetResult();
+
+            await Task.WhenAll(streamTasks).DefaultTimeout();
+
+            // Assert
+            // Up to 100 streams are pooled.
+            Assert.Equal(100, quicConnectionContext.StreamPool.Count);
+
+            static async Task SendStream(RequestState requestState)
+            {
+                var clientStream = requestState.QuicConnection.OpenBidirectionalStream();
+                await clientStream.WriteAsync(TestData, endStream: true).DefaultTimeout();
+                var serverStream = await requestState.ServerConnection.AcceptAsync().DefaultTimeout();
+                var readResult = await serverStream.Transport.Input.ReadAtLeastAsync(TestData.Length).DefaultTimeout();
+                serverStream.Transport.Input.AdvanceTo(readResult.Buffer.End);
+
+                // Input should be completed.
+                readResult = await serverStream.Transport.Input.ReadAsync();
+                Assert.True(readResult.IsCompleted);
+
+                lock (requestState)
+                {
+                    requestState.ActiveConcurrentConnections++;
+                    if (requestState.ActiveConcurrentConnections == StreamsSent)
+                    {
+                        requestState.AllConnectionsOnServerTcs.SetResult();
+                    }
+                }
+
+                await requestState.PauseCompleteTask;
+
+                // Complete reading and writing.
+                await serverStream.Transport.Input.CompleteAsync();
+                await serverStream.Transport.Output.CompleteAsync();
+
+                var quicStreamContext = Assert.IsType<QuicStreamContext>(serverStream);
+
+                // Both send and receive loops have exited.
+                await quicStreamContext._processingTask.DefaultTimeout();
+                await quicStreamContext.DisposeAsync();
+            }
+        }
+
+        private record RequestState(
+            QuicConnection QuicConnection,
+            MultiplexedConnectionContext ServerConnection,
+            TaskCompletionSource AllConnectionsOnServerTcs,
+            Task PauseCompleteTask)
+        {
+            public int ActiveConcurrentConnections { get; set; }
+        };
+
         private class TestSystemClock : ISystemClock
         {
             public DateTimeOffset UtcNow { get; set; }

src/Servers/Kestrel/Transport.Quic/test/QuicTestHelpers.cs

@@ -32,6 +32,8 @@ public static QuicTransportFactory CreateTransportFactory(ILoggerFactory loggerF
             var quicTransportOptions = new QuicTransportOptions();
             quicTransportOptions.Alpn = Alpn;
             quicTransportOptions.IdleTimeout = TimeSpan.FromMinutes(1);
+            quicTransportOptions.MaxBidirectionalStreamCount = 200;
+            quicTransportOptions.MaxUnidirectionalStreamCount = 200;
             if (systemClock != null)
             {
                 quicTransportOptions.SystemClock = systemClock;
@@ -74,8 +76,8 @@ public static QuicClientConnectionOptions CreateClientConnectionOptions(EndPoint
         {
             return new QuicClientConnectionOptions
             {
-                MaxBidirectionalStreams = 10,
-                MaxUnidirectionalStreams = 20,
+                MaxBidirectionalStreams = 200,
+                MaxUnidirectionalStreams = 200,
                 RemoteEndPoint = remoteEndPoint,
                 ClientAuthenticationOptions = new SslClientAuthenticationOptions
                 {
@@ -98,14 +100,12 @@ public static async Task<QuicStreamContext> CreateAndCompleteBidirectionalStream
 
             // Input should be completed.
             readResult = await serverStream.Transport.Input.ReadAsync();
+            Assert.True(readResult.IsCompleted);
 
             // Complete reading and writing.
             await serverStream.Transport.Input.CompleteAsync();
             await serverStream.Transport.Output.CompleteAsync();
 
-            // Assert
-            Assert.True(readResult.IsCompleted);
-
             var quicStreamContext = Assert.IsType<QuicStreamContext>(serverStream);
 
             // Both send and receive loops have exited.