HTTP/3: Pool QuicStreamContext instances

Author: JamesNK

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

@@ -0,0 +1,23 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+
+namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Quic.Internal
+{
+    /// <summary>
+    /// Abstracts the system clock to facilitate testing.
+    /// </summary>
+    internal interface ISystemClock
+    {
+        /// <summary>
+        /// Retrieves the current system time in UTC.
+        /// </summary>
+        DateTimeOffset UtcNow { get; }
+    }
+
+    internal class SystemClock : ISystemClock
+    {
+        public DateTimeOffset UtcNow => DateTimeOffset.UtcNow;
+    }
+}

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

@@ -2,6 +2,7 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
+using System.Diagnostics;
 using System.Net.Quic;
 using System.Threading;
 using System.Threading.Tasks;
@@ -14,6 +15,14 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Quic.Internal
 {
     internal class QuicConnectionContext : TransportMultiplexedConnection, IProtocolErrorCodeFeature
     {
+        // Internal for testing.
+        internal QuicStreamStack StreamPool;
+
+        private bool _streamPoolHeartbeatInitialized;
+        // Ticks updated once per-second in heartbeat event.
+        private long _heartbeatTicks;
+        private readonly object _poolLock = new object();
+
         private readonly QuicConnection _connection;
         private readonly QuicTransportContext _context;
         private readonly IQuicTrace _log;
@@ -23,6 +32,10 @@ internal class QuicConnectionContext : TransportMultiplexedConnection, IProtocol
 
         public long Error { get; set; }
 
+        internal const int InitialStreamPoolSize = 5;
+        internal const int MaxStreamPoolSize = 100;
+        internal const long StreamPoolExpiryTicks = TimeSpan.TicksPerSecond * 5;
+
         public QuicConnectionContext(QuicConnection connection, QuicTransportContext context)
         {
             _log = context.Log;
@@ -31,6 +44,8 @@ public QuicConnectionContext(QuicConnection connection, QuicTransportContext con
             ConnectionClosed = _connectionClosedTokenSource.Token;
             Features.Set<ITlsConnectionFeature>(new FakeTlsConnectionFeature());
             Features.Set<IProtocolErrorCodeFeature>(this);
+
+            StreamPool = new QuicStreamStack(InitialStreamPoolSize);
         }
 
         public override async ValueTask DisposeAsync()
@@ -62,7 +77,25 @@ public override void Abort(ConnectionAbortedException abortReason)
             try
             {
                 var stream = await _connection.AcceptStreamAsync(cancellationToken);
-                var context = new QuicStreamContext(stream, this, _context);
+
+                QuicStreamContext? context = null;
+
+                // Only use pool for bidirectional streams. Just a handful of unidirecitonal
+                // streams are created for a connection and they live for the lifetime of the connection.
+                if (stream.CanRead && stream.CanWrite)
+                {
+                    lock (_poolLock)
+                    {
+                        StreamPool.TryPop(out context);
+                    }
+                }
+
+                if (context == null)
+                {
+                    context = new QuicStreamContext(this, _context);
+                }
+
+                context.Initialize(stream);
                 context.Start();
 
                 _log.AcceptedStream(context);
@@ -124,12 +157,61 @@ public override ValueTask<ConnectionContext> ConnectAsync(IFeatureCollection? fe
                 quicStream = _connection.OpenBidirectionalStream();
             }
 
-            var context = new QuicStreamContext(quicStream, this, _context);
+            // Only a handful of control streams are created by the server and they last for the
+            // lifetime of the connection. No value in pooling them.
+            QuicStreamContext? context = new QuicStreamContext(this, _context);
+            context.Initialize(quicStream);
             context.Start();
 
             _log.ConnectedStream(context);
 
             return new ValueTask<ConnectionContext>(context);
         }
+
+        internal bool TryReturnStream(QuicStreamContext stream)
+        {
+            lock (_poolLock)
+            {
+                if (!_streamPoolHeartbeatInitialized)
+                {
+                    // Heartbeat feature is added to connection features by Kestrel.
+                    // No event is on the context is raised between feature being added and serving
+                    // connections so initialize heartbeat the first time a stream is added to
+                    // the connection's stream pool.
+                    var heartbeatFeature = Features.Get<IConnectionHeartbeatFeature>();
+                    if (heartbeatFeature != null)
+                    {
+                        heartbeatFeature.OnHeartbeat(static state => ((QuicConnectionContext)state).RemoveExpiredStreams(), this);
+                    }
+
+                    // Set ticks for the first time. Ticks are then updated in heartbeat.
+                    var now = _context.Options.SystemClock.UtcNow.Ticks;
+                    Volatile.Write(ref _heartbeatTicks, now);
+
+                    _streamPoolHeartbeatInitialized = true;
+                }
+
+                if (stream.CanReuse && StreamPool.Count < MaxStreamPoolSize)
+                {
+                    stream.PoolExpirationTicks = Volatile.Read(ref _heartbeatTicks) + StreamPoolExpiryTicks;
+                    StreamPool.Push(stream);
+                    return true;
+                }
+            }
+
+            return false;
+        }
+
+        private void RemoveExpiredStreams()
+        {
+            lock (_poolLock)
+            {
+                // Update ticks on heartbeat. A precise value isn't necessary.
+                var now = _context.Options.SystemClock.UtcNow.Ticks;
+                Volatile.Write(ref _heartbeatTicks, now);
+
+                StreamPool.RemoveExpired(now);
+            }
+        }
     }
 }

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);