Skip to content

Commit

Permalink
Re-use cancellation tokens in the https middleware (#31528)
Browse files Browse the repository at this point in the history
* Re-use cancellation tokens in the https middleware
- These tokens are usually short lived and exist for the purposes of cancelling the handshake.
  • Loading branch information
davidfowl authored Apr 6, 2021
1 parent 3143d95 commit 5de9b7d
Show file tree
Hide file tree
Showing 2 changed files with 70 additions and 1 deletion.
Original file line number Diff line number Diff line change
@@ -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);
}
}
}
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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)
{
Expand Down Expand Up @@ -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);
Expand Down

0 comments on commit 5de9b7d

Please sign in to comment.