Improve tracking of connection lifetime (#2646)

Author: halter73

src/Kestrel.Core/Internal/HttpConnection.cs

@@ -11,7 +11,6 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Connections;
-using Microsoft.AspNetCore.Connections.Features;
 using Microsoft.AspNetCore.Hosting.Server;
 using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Adapter.Internal;
@@ -30,6 +29,7 @@ public class HttpConnection : ITimeoutControl, IConnectionTimeoutFeature
 
         private readonly HttpConnectionContext _context;
         private readonly TaskCompletionSource<object> _socketClosedTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+        private readonly TaskCompletionSource<object> _lifetimeTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
 
         private IList<IAdaptedConnection> _adaptedConnections;
         private IDuplexPipe _adaptedTransport;
@@ -53,8 +53,6 @@ public class HttpConnection : ITimeoutControl, IConnectionTimeoutFeature
         private int _writeTimingWrites;
         private long _writeTimingTimeoutTimestamp;
 
-        private Task _lifetimeTask;
-
         public HttpConnection(HttpConnectionContext context)
         {
             _context = context;
@@ -98,12 +96,7 @@ public HttpConnection(HttpConnectionContext context)
 
         private IKestrelTrace Log => _context.ServiceContext.Log;
 
-        public Task StartRequestProcessing<TContext>(IHttpApplication<TContext> application)
-        {
-            return _lifetimeTask = ProcessRequestsAsync(application);
-        }
-
-        private async Task ProcessRequestsAsync<TContext>(IHttpApplication<TContext> httpApplication)
+        public async Task ProcessRequestsAsync<TContext>(IHttpApplication<TContext> httpApplication)
         {
             try
             {
@@ -132,9 +125,6 @@ private async Task ProcessRequestsAsync<TContext>(IHttpApplication<TContext> htt
                     _adaptedTransport = adaptedPipeline;
                 }
 
-                // Do this before the first await so we don't yield control to the transport until we've
-                // added the connection to the connection manager
-                _context.ServiceContext.ConnectionManager.AddConnection(_context.HttpConnectionId, this);
                 _lastTimestamp = _context.ServiceContext.SystemClock.UtcNow.Ticks;
 
                 _context.ConnectionFeatures.Set<IConnectionTimeoutFeature>(this);
@@ -194,7 +184,6 @@ private async Task ProcessRequestsAsync<TContext>(IHttpApplication<TContext> htt
             }
             finally
             {
-                _context.ServiceContext.ConnectionManager.RemoveConnection(_context.HttpConnectionId);
                 DisposeAdaptedConnections();
 
                 if (_http1Connection?.IsUpgraded == true)
@@ -204,6 +193,8 @@ private async Task ProcessRequestsAsync<TContext>(IHttpApplication<TContext> htt
 
                 Log.ConnectionStop(ConnectionId);
                 KestrelEventSource.Log.ConnectionStop(this);
+
+                _lifetimeTcs.SetResult(null);
             }
         }
 
@@ -269,7 +260,7 @@ public Task StopProcessingNextRequestAsync()
                 }
             }
 
-            return _lifetimeTask;
+            return _lifetimeTcs.Task;
         }
 
         public void OnInputOrOutputCompleted()

src/Kestrel.Core/Internal/HttpConnectionMiddleware.cs

@@ -73,22 +73,30 @@ public async Task OnConnectionAsync(ConnectionContext connectionContext)
             }
 
             var connection = new HttpConnection(httpConnectionContext);
+            _serviceContext.ConnectionManager.AddConnection(httpConnectionId, connection);
 
-            var processingTask = connection.StartRequestProcessing(_application);
+            try
+            {
+                var processingTask = connection.ProcessRequestsAsync(_application);
 
-            connectionContext.Transport.Input.OnWriterCompleted(
-                (_, state) => ((HttpConnection)state).OnInputOrOutputCompleted(),
-                connection);
+                connectionContext.Transport.Input.OnWriterCompleted(
+                    (_, state) => ((HttpConnection)state).OnInputOrOutputCompleted(),
+                    connection);
 
-            connectionContext.Transport.Output.OnReaderCompleted(
-                (_, state) => ((HttpConnection)state).OnInputOrOutputCompleted(),
-                connection);
+                connectionContext.Transport.Output.OnReaderCompleted(
+                    (_, state) => ((HttpConnection)state).OnInputOrOutputCompleted(),
+                    connection);
 
-            await AsTask(lifetimeFeature.ConnectionClosed);
+                await AsTask(lifetimeFeature.ConnectionClosed);
 
-            connection.OnConnectionClosed();
+                connection.OnConnectionClosed();
 
-            await processingTask;
+                await processingTask;
+            }
+            finally
+            {
+                _serviceContext.ConnectionManager.RemoveConnection(httpConnectionId);
+            }
         }
 
         private Task AsTask(CancellationToken token)