diff --git a/src/Servers/Kestrel/Core/src/Internal/CancellationTokenSourcePool.cs b/src/Servers/Kestrel/Core/src/Internal/CancellationTokenSourcePool.cs new file mode 100644 index 000000000000..ffb3a30e558a --- /dev/null +++ b/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 _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; + } + + /// + /// A with a back pointer to the pool it came from. + /// Dispose will return it to the pool. + /// + 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); + } + } + } + } + } +} diff --git a/src/Servers/Kestrel/Core/src/Middleware/HttpsConnectionMiddleware.cs b/src/Servers/Kestrel/Core/src/Middleware/HttpsConnectionMiddleware.cs index 6742972ead45..501fa13a0a52 100644 --- a/src/Servers/Kestrel/Core/src/Middleware/HttpsConnectionMiddleware.cs +++ b/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);