Track connection earlier (#12387)

- We made a change to dispatch connection execution but that also ended up dispatching the tracking of those connections in the connection manager. While it's not a huge deal most of the time it can affect graceful shutdown as there can be queued connections that are delayed in the thread pool queue which are untracked. This change makes it so we track the KestrelConnection before dispatching.

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

@@ -51,8 +51,14 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
                             break;
                         }
 
+                        // Add the connection to the connection manager before we queue it for execution
                         var id = Interlocked.Increment(ref _lastConnectionId);
-                        var kestrelConnection = new KestrelConnection(id, _serviceContext, _connectionDelegate, connection, _serviceContext.Log);
+                        var kestrelConnection = new KestrelConnection(id, _serviceContext, _connectionDelegate, connection, Log);
+
+                        _serviceContext.ConnectionManager.AddConnection(id, kestrelConnection);
+
+                        Log.ConnectionAccepted(connection.ConnectionId);
+
                         ThreadPool.UnsafeQueueUserWorkItem(kestrelConnection, preferLocal: false);
                     }
                 }

src/Servers/Kestrel/Core/src/Internal/Infrastructure/IKestrelTrace.cs

@@ -11,6 +11,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
 {
     internal interface IKestrelTrace : ILogger
     {
+        void ConnectionAccepted(string connectionId);
+
         void ConnectionStart(string connectionId);
 
         void ConnectionStop(string connectionId);

src/Servers/Kestrel/Core/src/Internal/Infrastructure/KestrelConnection.cs

@@ -186,8 +186,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
 
             try
             {
-                _serviceContext.ConnectionManager.AddConnection(_id, this);
-
                 Logger.ConnectionStart(connectionContext.ConnectionId);
                 KestrelEventSource.Log.ConnectionStart(connectionContext);
 

src/Servers/Kestrel/Core/src/Internal/Infrastructure/KestrelTrace.cs

@@ -5,7 +5,6 @@ using System;
 using Microsoft.AspNetCore.Connections;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2.HPack;
-using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
 using Microsoft.Extensions.Logging;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
@@ -111,6 +110,9 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
             LoggerMessage.Define<string, int>(LogLevel.Information, new EventId(38, nameof(HPackEncodingError)),
                 @"Connection id ""{ConnectionId}"": HPACK encoding error while encoding headers for stream ID {StreamId}.");
 
+        private static readonly Action<ILogger, string, Exception> _connectionAccepted =
+            LoggerMessage.Define<string>(LogLevel.Debug, new EventId(39, nameof(ConnectionAccepted)), @"Connection id ""{ConnectionId}"" accepted.");
+
         protected readonly ILogger _logger;
 
         public KestrelTrace(ILogger logger)
@@ -118,6 +120,11 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
             _logger = logger;
         }
 
+        public virtual void ConnectionAccepted(string connectionId)
+        {
+            _connectionAccepted(_logger, connectionId, null);
+        }
+
         public virtual void ConnectionStart(string connectionId)
         {
             _connectionStart(_logger, connectionId, null);

src/Servers/Kestrel/Core/test/ConnectionDispatcherTests.cs

@@ -30,6 +30,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             var connection = new Mock<DefaultConnectionContext> { CallBase = true }.Object;
             connection.ConnectionClosed = new CancellationToken(canceled: true);
             var kestrelConnection = new KestrelConnection(0, serviceContext, _ => tcs.Task, connection, serviceContext.Log);
+            serviceContext.ConnectionManager.AddConnection(0, kestrelConnection);
 
             var task = kestrelConnection.ExecuteAsync();
 
@@ -79,6 +80,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             var connection = new Mock<DefaultConnectionContext> { CallBase = true }.Object;
             connection.ConnectionClosed = new CancellationToken(canceled: true);
             var kestrelConnection = new KestrelConnection(0, serviceContext, _ => Task.CompletedTask, connection, serviceContext.Log);
+            serviceContext.ConnectionManager.AddConnection(0, kestrelConnection);
             var completeFeature = kestrelConnection.TransportConnection.Features.Get<IConnectionCompleteFeature>();
 
             Assert.NotNull(completeFeature);
@@ -99,6 +101,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             var connection = new Mock<DefaultConnectionContext> { CallBase = true }.Object;
             connection.ConnectionClosed = new CancellationToken(canceled: true);
             var kestrelConnection = new KestrelConnection(0, serviceContext, _ => Task.CompletedTask, connection, serviceContext.Log);
+            serviceContext.ConnectionManager.AddConnection(0, kestrelConnection);
             var completeFeature = kestrelConnection.TransportConnection.Features.Get<IConnectionCompleteFeature>();
 
             Assert.NotNull(completeFeature);

src/Servers/Kestrel/perf/Kestrel.Performance/Mocks/MockTrace.cs

@@ -15,6 +15,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Performance
     {
         public void ApplicationError(string connectionId, string requestId, Exception ex) { }
         public IDisposable BeginScope<TState>(TState state) => null;
+        public void ConnectionAccepted(string connectionId) { }
         public void ConnectionBadRequest(string connectionId, BadHttpRequestException ex) { }
         public void ConnectionDisconnect(string connectionId) { }
         public void ConnectionError(string connectionId, Exception ex) { }

src/Servers/Kestrel/shared/test/CompositeKestrelTrace.cs

@@ -11,7 +11,7 @@ using Microsoft.Extensions.Logging;
 
 namespace Microsoft.AspNetCore.Testing
 {
-    internal class CompositeKestrelTrace: IKestrelTrace
+    internal class CompositeKestrelTrace : IKestrelTrace
     {
         private readonly IKestrelTrace _trace1;
         private readonly IKestrelTrace _trace2;
@@ -33,6 +33,12 @@ namespace Microsoft.AspNetCore.Testing
             return _trace1.IsEnabled(logLevel) || _trace2.IsEnabled(logLevel);
         }
 
+        public void ConnectionAccepted(string connectionId)
+        {
+            _trace1.ConnectionAccepted(connectionId);
+            _trace2.ConnectionAccepted(connectionId);
+        }
+
         public IDisposable BeginScope<TState>(TState state)
         {
             return _trace1.BeginScope(state);