Re-use cancellation tokens in the https middleware

src/Servers/Kestrel/Core/src/Internal/CancellationTokenSourcePool.cs

@@ -0,0 +1,64 @@
+// 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.Collections.Concurrent;
+using System.Threading;
+
+namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
+{
+ internal class CancellationTokenSourcePool
+ {
+ private const int MaxQueueSize = 1024;
+
+ private readonly ConcurrentQueue<PooledCancellationTokenSource> _queue = new();
+ private int _count;
+
+ public PooledCancellationTokenSource Rent()
+ {
+ if (_queue.TryDequeue(out var cts))
+ {
+ Interlocked.Decrement(ref _count);
+ return cts;
+ }
+ return new PooledCancellationTokenSource(this);
+ }
+
+ private bool Return(PooledCancellationTokenSource cts)
+ {
+ if (Interlocked.Increment(ref _count) > MaxQueueSize || !cts.TryReset())
+ {
+ Interlocked.Decrement(ref _count);
+ return false;
+ }
+
+ _queue.Enqueue(cts);
+ return true;
+ }
+
+ /// <summary>
+ /// A <see cref="CancellationTokenSource"/> with a back pointer to the pool it came from.
+ /// Dispose will return it to the pool.
+ /// </summary>
+ public class PooledCancellationTokenSource : CancellationTokenSource
+ {
+ private readonly CancellationTokenSourcePool _pool;
+
+ public PooledCancellationTokenSource(CancellationTokenSourcePool pool)
+ {
+ _pool = pool;
+ }
+
+ protected override void Dispose(bool disposing)
+ {
+ if (disposing)
+ {
+ // If we failed to return to the pool then dispose
+ if (!_pool.Return(this))
+ {
+ base.Dispose(disposing);
+ }
+ }
+ }
+ }
+ }
+}

src/Servers/Kestrel/Core/src/Middleware/HttpsConnectionMiddleware.cs

@@ -49,6 +49,9 @@ internal class HttpsConnectionMiddleware
  private readonly HttpsOptionsCallback? _httpsOptionsCallback;
  private readonly object? _httpsOptionsCallbackState;
 
+ // Pool for cancellation tokens that cancel the handshake
+ private readonly CancellationTokenSourcePool _ctsPool = new();
+
  public HttpsConnectionMiddleware(ConnectionDelegate next, HttpsConnectionAdapterOptions options)
  : this(next, options, loggerFactory: NullLoggerFactory.Instance)
  {
@@ -150,7 +153,9 @@ public async Task OnConnectionAsync(ConnectionContext context)
 
  try
  {
- using var cancellationTokenSource = new CancellationTokenSource(_handshakeTimeout);
+ using var cancellationTokenSource = _ctsPool.Rent();
+ cancellationTokenSource.CancelAfter(_handshakeTimeout);
+
  if (_httpsOptionsCallback is null)
  {
  await DoOptionsBasedHandshakeAsync(context, sslStream, feature, cancellationTokenSource.Token);