General connection management (#2834)

This change makes the handling of graceful shutdown work for more than just http scenarios. This should allow us to move TLS further out and should also allow us to start moving things to connection middleware instead of connection adapters.
Summary of the things changed/added:
- Added IConnectionLifetimeNotificationFeature that represents an attempt to gracefully close the connection that isn't being aborted. This feels pretty awful but we may have to do it.
- Moved connection management to the ConnectionDispatcher and out of the HttpConnectionMiddleware
- Removed Http from the names of the ConnectionManager and Heartbeat

src/Connections.Abstractions/Features/IConnectionLifetimeNotificationFeature.cs

@@ -0,0 +1,14 @@
+// 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.Threading;
+
+namespace Microsoft.AspNetCore.Connections.Features
+{
+    public interface IConnectionLifetimeNotificationFeature
+    {
+        CancellationToken ConnectionClosedRequested { get; set; }
+
+        void RequestClose();
+    }
+}

src/Kestrel.Core/Internal/ConnectionDispatcher.cs

@@ -17,6 +17,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
 {
     public class ConnectionDispatcher : IConnectionDispatcher
     {
+        private static long _lastConnectionId = long.MinValue;
+
         private readonly ServiceContext _serviceContext;
         private readonly ConnectionDelegate _connectionDelegate;
 
@@ -44,28 +46,51 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
             // This *must* be set before returning from OnConnection
             connection.Application = pair.Application;
 
-            return Execute(connection);
+            return Execute(new KestrelConnection(connection));
         }
 
-        private async Task Execute(ConnectionContext connectionContext)
+        private async Task Execute(KestrelConnection connection)
         {
-            using (BeginConnectionScope(connectionContext))
+            var id = Interlocked.Increment(ref _lastConnectionId);
+            var connectionContext = connection.TransportConnection;
+
+            try
             {
-                try
+                _serviceContext.ConnectionManager.AddConnection(id, connection);
+
+                Log.ConnectionStart(connectionContext.ConnectionId);
+                KestrelEventSource.Log.ConnectionStart(connectionContext);
+
+                using (BeginConnectionScope(connectionContext))
                 {
-                    await _connectionDelegate(connectionContext);
+                    try
-                }
+                    {
-                catch (Exception ex)
+                        await _connectionDelegate(connectionContext);
-                {
+                    }
-                    Log.LogCritical(0, ex, $"{nameof(ConnectionDispatcher)}.{nameof(Execute)}() {connectionContext.ConnectionId}");
+                    catch (Exception ex)
-                }
+                    {
-                finally
+                        Log.LogCritical(0, ex, $"{nameof(ConnectionDispatcher)}.{nameof(Execute)}() {connectionContext.ConnectionId}");
-                {
+                    }
-                    // Complete the transport PipeReader and PipeWriter after calling into application code
+                    finally
-                    connectionContext.Transport.Input.Complete();
+                    {
-                    connectionContext.Transport.Output.Complete();
+                        // Complete the transport PipeReader and PipeWriter after calling into application code
+                        connectionContext.Transport.Input.Complete();
+                        connectionContext.Transport.Output.Complete();
+                    }
+
+                    // Wait for the transport to close
+                    await CancellationTokenAsTask(connectionContext.ConnectionClosed);
                 }
             }
+            finally
+            {
+                connection.Complete();
+
+                _serviceContext.ConnectionManager.RemoveConnection(id);
+
+                Log.ConnectionStop(connectionContext.ConnectionId);
+                KestrelEventSource.Log.ConnectionStop(connectionContext);
+            }
         }
 
         private IDisposable BeginConnectionScope(ConnectionContext connectionContext)
@@ -78,6 +103,20 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
             return null;
         }
 
+        private static Task CancellationTokenAsTask(CancellationToken token)
+        {
+            if (token.IsCancellationRequested)
+            {
+                return Task.CompletedTask;
+            }
+
+            // Transports already dispatch prior to tripping ConnectionClosed
+            // since application code can register to this token.
+            var tcs = new TaskCompletionSource<object>();
+            token.Register(state => ((TaskCompletionSource<object>)state).SetResult(null), tcs);
+            return tcs.Task;
+        }
+
         // Internal for testing
         internal static PipeOptions GetInputPipeOptions(ServiceContext serviceContext, MemoryPool<byte> memoryPool, PipeScheduler writerScheduler) => new PipeOptions
         (

src/Kestrel.Core/Internal/HttpConnection.cs

@@ -11,6 +11,7 @@ using System.Net;
 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;
@@ -28,8 +29,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
         private static readonly ReadOnlyMemory<byte> Http2Id = new[] { (byte)'h', (byte)'2' };
 
         private readonly HttpConnectionContext _context;
-        private readonly TaskCompletionSource<object> _socketClosedTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+        private readonly ISystemClock _systemClock;
-        private readonly TaskCompletionSource<object> _lifetimeTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
 
         private IList<IAdaptedConnection> _adaptedConnections;
         private IDuplexPipe _adaptedTransport;
@@ -56,6 +56,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
         public HttpConnection(HttpConnectionContext context)
         {
             _context = context;
+            _systemClock = _context.ServiceContext.SystemClock;
         }
 
         // For testing
@@ -100,20 +101,13 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
         {
             try
             {
-                // TODO: When we start tracking all connection middleware for shutdown, go back
-                // to logging connections tart and stop in ConnectionDispatcher so we get these
-                // logs for all connection middleware.
-                Log.ConnectionStart(ConnectionId);
-                KestrelEventSource.Log.ConnectionStart(this);
-
                 AdaptedPipeline adaptedPipeline = null;
                 var adaptedPipelineTask = Task.CompletedTask;
 
-                // _adaptedTransport must be set prior to adding the connection to the manager in order
+                // _adaptedTransport must be set prior to wiring up callbacks
-                // to allow the connection to be aported prior to protocol selection.
+                // to allow the connection to be aborted prior to protocol selection.
                 _adaptedTransport = _context.Transport;
 
-
                 if (_context.ConnectionAdapters.Count > 0)
                 {
                     adaptedPipeline = new AdaptedPipeline(_adaptedTransport,
@@ -124,58 +118,79 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
                     _adaptedTransport = adaptedPipeline;
                 }
 
-                _lastTimestamp = _context.ServiceContext.SystemClock.UtcNow.Ticks;
+                // This feature should never be null in Kestrel
+                var connectionHeartbeatFeature = _context.ConnectionFeatures.Get<IConnectionHeartbeatFeature>();
 
-                _context.ConnectionFeatures.Set<IConnectionTimeoutFeature>(this);
+                Debug.Assert(connectionHeartbeatFeature != null, nameof(IConnectionHeartbeatFeature) + " is missing!");
 
-                if (adaptedPipeline != null)
+                connectionHeartbeatFeature?.OnHeartbeat(state => ((HttpConnection)state).Tick(), this);
+
+                var connectionLifetimeNotificationFeature = _context.ConnectionFeatures.Get<IConnectionLifetimeNotificationFeature>();
+
+                Debug.Assert(connectionLifetimeNotificationFeature != null, nameof(IConnectionLifetimeNotificationFeature) + " is missing!");
+
+                using (connectionLifetimeNotificationFeature?.ConnectionClosedRequested.Register(state => ((HttpConnection)state).StopProcessingNextRequest(), this))
                 {
-                    // Stream can be null here and run async will close the connection in that case
+                    _lastTimestamp = _context.ServiceContext.SystemClock.UtcNow.Ticks;
-                    var stream = await ApplyConnectionAdaptersAsync();
-                    adaptedPipelineTask = adaptedPipeline.RunAsync(stream);
-                }
 
-                IRequestProcessor requestProcessor = null;
+                    _context.ConnectionFeatures.Set<IConnectionTimeoutFeature>(this);
 
-                lock (_protocolSelectionLock)
+                    if (adaptedPipeline != null)
-                {
-                    // Ensure that the connection hasn't already been stopped.
-                    if (_protocolSelectionState == ProtocolSelectionState.Initializing)
                     {
-                        switch (SelectProtocol())
+                        // Stream can be null here and run async will close the connection in that case
-                        {
+                        var stream = await ApplyConnectionAdaptersAsync();
-                            case HttpProtocols.Http1:
+                        adaptedPipelineTask = adaptedPipeline.RunAsync(stream);
-                                // _http1Connection must be initialized before adding the connection to the connection manager
-                                requestProcessor = _http1Connection = CreateHttp1Connection(_adaptedTransport);
-                                _protocolSelectionState = ProtocolSelectionState.Selected;
-                                break;
-                            case HttpProtocols.Http2:
-                                // _http2Connection must be initialized before yielding control to the transport thread,
-                                // to prevent a race condition where _http2Connection.Abort() is called just as
-                                // _http2Connection is about to be initialized.
-                                requestProcessor = CreateHttp2Connection(_adaptedTransport);
-                                _protocolSelectionState = ProtocolSelectionState.Selected;
-                                break;
-                            case HttpProtocols.None:
-                                // An error was already logged in SelectProtocol(), but we should close the connection.
-                                Abort(ex: null);
-                                break;
-                            default:
-                                // SelectProtocol() only returns Http1, Http2 or None.
-                                throw new NotSupportedException($"{nameof(SelectProtocol)} returned something other than Http1, Http2 or None.");
-                        }
-
-                        _requestProcessor = requestProcessor;
                     }
-                }
 
-                if (requestProcessor != null)
+                    IRequestProcessor requestProcessor = null;
-                {
-                    await requestProcessor.ProcessRequestsAsync(httpApplication);
-                }
 
-                await adaptedPipelineTask;
+                    lock (_protocolSelectionLock)
-                await _socketClosedTcs.Task;
+                    {
+                        // Ensure that the connection hasn't already been stopped.
+                        if (_protocolSelectionState == ProtocolSelectionState.Initializing)
+                        {
+                            switch (SelectProtocol())
+                            {
+                                case HttpProtocols.Http1:
+                                    // _http1Connection must be initialized before adding the connection to the connection manager
+                                    requestProcessor = _http1Connection = CreateHttp1Connection(_adaptedTransport);
+                                    _protocolSelectionState = ProtocolSelectionState.Selected;
+                                    break;
+                                case HttpProtocols.Http2:
+                                    // _http2Connection must be initialized before yielding control to the transport thread,
+                                    // to prevent a race condition where _http2Connection.Abort() is called just as
+                                    // _http2Connection is about to be initialized.
+                                    requestProcessor = CreateHttp2Connection(_adaptedTransport);
+                                    _protocolSelectionState = ProtocolSelectionState.Selected;
+                                    break;
+                                case HttpProtocols.None:
+                                    // An error was already logged in SelectProtocol(), but we should close the connection.
+                                    Abort(ex: null);
+                                    break;
+                                default:
+                                    // SelectProtocol() only returns Http1, Http2 or None.
+                                    throw new NotSupportedException($"{nameof(SelectProtocol)} returned something other than Http1, Http2 or None.");
+                            }
+
+                            _requestProcessor = requestProcessor;
+                        }
+                    }
+
+                    _context.Transport.Input.OnWriterCompleted(
+                        (_, state) => ((HttpConnection)state).OnInputOrOutputCompleted(),
+                        this);
+
+                    _context.Transport.Output.OnReaderCompleted(
+                        (_, state) => ((HttpConnection)state).OnInputOrOutputCompleted(),
+                        this);
+
+                    if (requestProcessor != null)
+                    {
+                        await requestProcessor.ProcessRequestsAsync(httpApplication);
+                    }
+
+                    await adaptedPipelineTask;
+                }
             }
             catch (Exception ex)
             {
@@ -189,11 +204,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
                 {
                     _context.ServiceContext.ConnectionManager.UpgradedConnectionCount.ReleaseOne();
                 }
-
-                Log.ConnectionStop(ConnectionId);
-                KestrelEventSource.Log.ConnectionStop(this);
-
-                _lifetimeTcs.SetResult(null);
             }
         }
 
@@ -235,12 +245,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
             });
         }
 
-        public void OnConnectionClosed()
+        private void StopProcessingNextRequest()
-        {
-            _socketClosedTcs.TrySetResult(null);
-        }
-
-        public Task StopProcessingNextRequestAsync()
         {
             lock (_protocolSelectionLock)
             {
@@ -257,11 +262,9 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
                         break;
                 }
             }
-
-            return _lifetimeTcs.Task;
         }
 
-        public void OnInputOrOutputCompleted()
+        private void OnInputOrOutputCompleted()
         {
             lock (_protocolSelectionLock)
             {
@@ -281,7 +284,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
             }
         }
 
-        public void Abort(ConnectionAbortedException ex)
+        private void Abort(ConnectionAbortedException ex)
         {
             lock (_protocolSelectionLock)
             {
@@ -301,13 +304,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
             }
         }
 
-        public Task AbortAsync(ConnectionAbortedException ex)
-        {
-            Abort(ex);
-
-            return _socketClosedTcs.Task;
-        }
-
         private async Task<Stream> ApplyConnectionAdaptersAsync()
         {
             var connectionAdapters = _context.ConnectionAdapters;
@@ -381,6 +377,11 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
             return http2Enabled && (!hasTls || Http2Id.Span.SequenceEqual(applicationProtocol.Span)) ? HttpProtocols.Http2 : HttpProtocols.Http1;
         }
 
+        private void Tick()
+        {
+            Tick(_systemClock.UtcNow);
+        }
+
         public void Tick(DateTimeOffset now)
         {
             if (_protocolSelectionState == ProtocolSelectionState.Aborted)

src/Kestrel.Core/Internal/HttpConnectionContext.cs

@@ -14,7 +14,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
     public class HttpConnectionContext
     {
         public string ConnectionId { get; set; }
-        public long HttpConnectionId { get; set; }
         public HttpProtocols Protocols { get; set; }
         public ConnectionContext ConnectionContext { get; set; }
         public ServiceContext ServiceContext { get; set; }

src/Kestrel.Core/Internal/HttpConnectionMiddleware.cs

@@ -16,8 +16,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
 {
     public class HttpConnectionMiddleware<TContext>
     {
-        private static long _lastHttpConnectionId = long.MinValue;
-
         private readonly IList<IConnectionAdapter> _connectionAdapters;
         private readonly ServiceContext _serviceContext;
         private readonly IHttpApplication<TContext> _application;
@@ -33,19 +31,16 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
             _connectionAdapters = adapters;
         }
 
-        public async Task OnConnectionAsync(ConnectionContext connectionContext)
+        public Task OnConnectionAsync(ConnectionContext connectionContext)
         {
             // We need the transport feature so that we can cancel the output reader that the transport is using
             // This is a bit of a hack but it preserves the existing semantics
             var memoryPoolFeature = connectionContext.Features.Get<IMemoryPoolFeature>();
 
-            var httpConnectionId = Interlocked.Increment(ref _lastHttpConnectionId);
-
             var httpConnectionContext = new HttpConnectionContext
             {
                 ConnectionId = connectionContext.ConnectionId,
                 ConnectionContext = connectionContext,
-                HttpConnectionId = httpConnectionId,
                 Protocols = _protocols,
                 ServiceContext = _serviceContext,
                 ConnectionFeatures = connectionContext.Features,
@@ -55,7 +50,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
             };
 
             var connectionFeature = connectionContext.Features.Get<IHttpConnectionFeature>();
-            var lifetimeFeature = connectionContext.Features.Get<IConnectionLifetimeFeature>();
 
             if (connectionFeature != null)
             {
@@ -71,44 +65,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
             }
 
             var connection = new HttpConnection(httpConnectionContext);
-            _serviceContext.ConnectionManager.AddConnection(httpConnectionId, connection);
 
-            try
+            return connection.ProcessRequestsAsync(_application);
-            {
-                var processingTask = connection.ProcessRequestsAsync(_application);
-
-                connectionContext.Transport.Input.OnWriterCompleted(
-                    (_, state) => ((HttpConnection)state).OnInputOrOutputCompleted(),
-                    connection);
-
-                connectionContext.Transport.Output.OnReaderCompleted(
-                    (_, state) => ((HttpConnection)state).OnInputOrOutputCompleted(),
-                    connection);
-
-                await CancellationTokenAsTask(lifetimeFeature.ConnectionClosed);
-
-                connection.OnConnectionClosed();
-
-                await processingTask;
-            }
-            finally
-            {
-                _serviceContext.ConnectionManager.RemoveConnection(httpConnectionId);
-            }
-        }
-
-        private static Task CancellationTokenAsTask(CancellationToken token)
-        {
-            if (token.IsCancellationRequested)
-            {
-                return Task.CompletedTask;
-            }
-
-            // Transports already dispatch prior to tripping ConnectionClosed
-            // since application code can register to this token.
-            var tcs = new TaskCompletionSource<object>();
-            token.Register(() => tcs.SetResult(null));
-            return tcs.Task;
         }
     }
 }

src/Kestrel.Core/Internal/Infrastructure/ConnectionManager.cs

@@ -6,17 +6,17 @@ using System.Collections.Concurrent;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
 {
-    public class HttpConnectionManager
+    public class ConnectionManager
     {
-        private readonly ConcurrentDictionary<long, HttpConnectionReference> _connectionReferences = new ConcurrentDictionary<long, HttpConnectionReference>();
+        private readonly ConcurrentDictionary<long, ConnectionReference> _connectionReferences = new ConcurrentDictionary<long, ConnectionReference>();
         private readonly IKestrelTrace _trace;
 
-        public HttpConnectionManager(IKestrelTrace trace, long? upgradedConnectionLimit)
+        public ConnectionManager(IKestrelTrace trace, long? upgradedConnectionLimit)
             : this(trace, GetCounter(upgradedConnectionLimit))
         {
         }
 
-        public HttpConnectionManager(IKestrelTrace trace, ResourceCounter upgradedConnections)
+        public ConnectionManager(IKestrelTrace trace, ResourceCounter upgradedConnections)
         {
             UpgradedConnectionCount = upgradedConnections;
             _trace = trace;
@@ -27,9 +27,9 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
         /// </summary>
         public ResourceCounter UpgradedConnectionCount { get; }
 
-        public void AddConnection(long id, HttpConnection connection)
+        public void AddConnection(long id, KestrelConnection connection)
         {
-            if (!_connectionReferences.TryAdd(id, new HttpConnectionReference(connection)))
+            if (!_connectionReferences.TryAdd(id, new ConnectionReference(connection)))
             {
                 throw new ArgumentException(nameof(id));
             }
@@ -43,7 +43,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
             }
         }
 
-        public void Walk(Action<HttpConnection> callback)
+        public void Walk(Action<KestrelConnection> callback)
         {
             foreach (var kvp in _connectionReferences)
             {

src/Kestrel.Core/Internal/Infrastructure/ConnectionManagerShutdownExtensions.cs

@@ -8,29 +8,31 @@ using Microsoft.AspNetCore.Connections;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
 {
-    public static class HttpConnectionManagerShutdownExtensions
+    public static class ConnectionManagerShutdownExtensions
     {
-        public static async Task<bool> CloseAllConnectionsAsync(this HttpConnectionManager connectionManager, CancellationToken token)
+        public static async Task<bool> CloseAllConnectionsAsync(this ConnectionManager connectionManager, CancellationToken token)
         {
             var closeTasks = new List<Task>();
 
             connectionManager.Walk(connection =>
             {
-                closeTasks.Add(connection.StopProcessingNextRequestAsync());
+                connection.TransportConnection.RequestClose();
+                closeTasks.Add(connection.ExecutionTask);
             });
 
             var allClosedTask = Task.WhenAll(closeTasks.ToArray());
             return await Task.WhenAny(allClosedTask, CancellationTokenAsTask(token)).ConfigureAwait(false) == allClosedTask;
         }
 
-        public static async Task<bool> AbortAllConnectionsAsync(this HttpConnectionManager connectionManager)
+        public static async Task<bool> AbortAllConnectionsAsync(this ConnectionManager connectionManager)
         {
             var abortTasks = new List<Task>();
             var canceledException = new ConnectionAbortedException(CoreStrings.ConnectionAbortedDuringServerShutdown);
 
             connectionManager.Walk(connection =>
             {
-                abortTasks.Add(connection.AbortAsync(canceledException));
+                connection.TransportConnection.Abort(canceledException);
+                abortTasks.Add(connection.ExecutionTask);
             });
 
             var allAbortedTask = Task.WhenAll(abortTasks.ToArray());

src/Kestrel.Core/Internal/Infrastructure/ConnectionReference.cs

@@ -0,0 +1,26 @@
+// 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;
+using Microsoft.AspNetCore.Connections;
+
+namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
+{
+    public class ConnectionReference
+    {
+        private readonly WeakReference<KestrelConnection> _weakReference;
+
+        public ConnectionReference(KestrelConnection connection)
+        {
+            _weakReference = new WeakReference<KestrelConnection>(connection);
+            ConnectionId = connection.TransportConnection.ConnectionId;
+        }
+
+        public string ConnectionId { get; }
+
+        public bool TryGetConnection(out KestrelConnection connection)
+        {
+            return _weakReference.TryGetTarget(out connection);
+        }
+    }
+}

src/Kestrel.Core/Internal/Infrastructure/HeartbeatManager.cs

@@ -5,27 +5,29 @@ using System;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
 {
-    public class HttpHeartbeatManager : IHeartbeatHandler
+    public class HeartbeatManager : IHeartbeatHandler, ISystemClock
     {
-        private readonly HttpConnectionManager _connectionManager;
+        private readonly ConnectionManager _connectionManager;
-        private readonly Action<HttpConnection> _walkCallback;
+        private readonly Action<KestrelConnection> _walkCallback;
         private DateTimeOffset _now;
 
-        public HttpHeartbeatManager(HttpConnectionManager connectionManager)
+        public HeartbeatManager(ConnectionManager connectionManager)
         {
             _connectionManager = connectionManager;
             _walkCallback = WalkCallback;
         }
 
+        public DateTimeOffset UtcNow => _now;
+
         public void OnHeartbeat(DateTimeOffset now)
         {
             _now = now;
             _connectionManager.Walk(_walkCallback);
         }
 
-        private void WalkCallback(HttpConnection connection)
+        private void WalkCallback(KestrelConnection connection)
         {
-            connection.Tick(_now);
+            connection.TransportConnection.TickHeartbeat();
         }
     }
 }

src/Kestrel.Core/Internal/Infrastructure/HttpConnectionReference.cs

@@ -1,25 +0,0 @@
-// 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;
-
-namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
-{
-    public class HttpConnectionReference
-    {
-        private readonly WeakReference<HttpConnection> _weakReference;
-
-        public HttpConnectionReference(HttpConnection connection)
-        {
-            _weakReference = new WeakReference<HttpConnection>(connection);
-            ConnectionId = connection.ConnectionId;
-        }
-
-        public string ConnectionId { get; }
-
-        public bool TryGetConnection(out HttpConnection connection)
-        {
-            return _weakReference.TryGetTarget(out connection);
-        }
-    }
-}

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

@@ -0,0 +1,28 @@
+// 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.Diagnostics.Tracing;
+using System.Runtime.CompilerServices;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Connections;
+using Microsoft.AspNetCore.Server.Kestrel.Transport.Abstractions.Internal;
+
+namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
+{
+    public class KestrelConnection
+    {
+        private TaskCompletionSource<object> _executionTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        public KestrelConnection(TransportConnection transportConnection)
+        {
+            TransportConnection = transportConnection;
+            ExecutionTask = _executionTcs.Task;
+        }
+
+        public TransportConnection TransportConnection { get; }
+
+        public Task ExecutionTask { get; }
+
+        internal void Complete() => _executionTcs.TrySetResult(null);
+    }
+}

src/Kestrel.Core/Internal/Infrastructure/KestrelEventSource.cs

@@ -2,8 +2,11 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System.Diagnostics.Tracing;
+using System.Net;
 using System.Runtime.CompilerServices;
+using Microsoft.AspNetCore.Connections;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http;
+using Microsoft.AspNetCore.Server.Kestrel.Transport.Abstractions.Internal;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
 {
@@ -25,15 +28,15 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
         // - Avoid renaming methods or parameters marked with EventAttribute. EventSource uses these to form the event object.
 
         [NonEvent]
-        public void ConnectionStart(HttpConnection connection)
+        public void ConnectionStart(TransportConnection connection)
         {
             // avoid allocating strings unless this event source is enabled
             if (IsEnabled())
             {
                 ConnectionStart(
                     connection.ConnectionId,
-                    connection.LocalEndPoint?.ToString(),
+                    connection.LocalAddress != null ? new IPEndPoint(connection.LocalAddress, connection.LocalPort).ToString() : null,
-                    connection.RemoteEndPoint?.ToString());
+                    connection.RemoteAddress != null ? new IPEndPoint(connection.RemoteAddress, connection.RemotePort).ToString() : null);
             }
         }
 
@@ -52,7 +55,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
         }
 
         [NonEvent]
-        public void ConnectionStop(HttpConnection connection)
+        public void ConnectionStop(TransportConnection connection)
         {
             if (IsEnabled())
             {

src/Kestrel.Core/Internal/ServiceContext.cs

@@ -19,7 +19,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
 
         public DateHeaderValueManager DateHeaderValueManager { get; set; }
 
-        public HttpConnectionManager ConnectionManager { get; set; }
+        public ConnectionManager ConnectionManager { get; set; }
 
         public Heartbeat Heartbeat { get; set; }
 

src/Kestrel.Core/KestrelServer.cs

@@ -67,17 +67,15 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core
             var serverOptions = options.Value ?? new KestrelServerOptions();
             var logger = loggerFactory.CreateLogger("Microsoft.AspNetCore.Server.Kestrel");
             var trace = new KestrelTrace(logger);
-            var connectionManager = new HttpConnectionManager(
+            var connectionManager = new ConnectionManager(
                 trace,
                 serverOptions.Limits.MaxConcurrentUpgradedConnections);
 
-            var systemClock = new SystemClock();
+            var heartbeatManager = new HeartbeatManager(connectionManager);
-            var dateHeaderValueManager = new DateHeaderValueManager(systemClock);
+            var dateHeaderValueManager = new DateHeaderValueManager(heartbeatManager);
-
-            var httpHeartbeatManager = new HttpHeartbeatManager(connectionManager);
             var heartbeat = new Heartbeat(
-                new IHeartbeatHandler[] { dateHeaderValueManager, httpHeartbeatManager },
+                new IHeartbeatHandler[] { dateHeaderValueManager, heartbeatManager },
-                systemClock,
+                new SystemClock(),
                 DebuggerWrapper.Singleton,
                 trace);
 
@@ -102,7 +100,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core
                 Log = trace,
                 HttpParser = new HttpParser<Http1ParsingHandler>(trace.IsEnabled(LogLevel.Information)),
                 Scheduler = scheduler,
-                SystemClock = systemClock,
+                SystemClock = heartbeatManager,
                 DateHeaderValueManager = dateHeaderValueManager,
                 ConnectionManager = connectionManager,
                 Heartbeat = heartbeat,
@@ -118,7 +116,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core
 
         private IKestrelTrace Trace => ServiceContext.Log;
 
-        private HttpConnectionManager ConnectionManager => ServiceContext.ConnectionManager;
+        private ConnectionManager ConnectionManager => ServiceContext.ConnectionManager;
 
         public async Task StartAsync<TContext>(IHttpApplication<TContext> application, CancellationToken cancellationToken)
         {

src/Kestrel.Transport.Abstractions/Internal/TransportConnection.FeatureCollection.cs

@@ -20,6 +20,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Abstractions.Internal
                                                IApplicationTransportFeature,
                                                ITransportSchedulerFeature,
                                                IConnectionLifetimeFeature,
+                                               IConnectionHeartbeatFeature,
+                                               IConnectionLifetimeNotificationFeature,
                                                IBytesWrittenFeature
     {
         // NOTE: When feature interfaces are added to or removed from this TransportConnection class implementation,
@@ -85,8 +87,21 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Abstractions.Internal
             set => ConnectionClosed = value;
         }
 
+        CancellationToken IConnectionLifetimeNotificationFeature.ConnectionClosedRequested
+        {
+            get => ConnectionClosedRequested;
+            set => ConnectionClosedRequested = value;
+        }
+
         void IConnectionLifetimeFeature.Abort() => Abort(abortReason: null);
 
+        void IConnectionLifetimeNotificationFeature.RequestClose() => RequestClose();
+
+        void IConnectionHeartbeatFeature.OnHeartbeat(System.Action<object> action, object state)
+        {
+            OnHeartbeat(action, state);
+        }
+
         long IBytesWrittenFeature.TotalBytesWritten => TotalBytesWritten;
     }
 }

src/Kestrel.Transport.Abstractions/Internal/TransportConnection.Generated.cs

@@ -20,6 +20,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Abstractions.Internal
         private static readonly Type IApplicationTransportFeatureType = typeof(IApplicationTransportFeature);
         private static readonly Type ITransportSchedulerFeatureType = typeof(ITransportSchedulerFeature);
         private static readonly Type IConnectionLifetimeFeatureType = typeof(IConnectionLifetimeFeature);
+        private static readonly Type IConnectionHeartbeatFeatureType = typeof(IConnectionHeartbeatFeature);
+        private static readonly Type IConnectionLifetimeNotificationFeatureType = typeof(IConnectionLifetimeNotificationFeature);
         private static readonly Type IBytesWrittenFeatureType = typeof(IBytesWrittenFeature);
 
         private object _currentIHttpConnectionFeature;
@@ -30,6 +32,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Abstractions.Internal
         private object _currentIApplicationTransportFeature;
         private object _currentITransportSchedulerFeature;
         private object _currentIConnectionLifetimeFeature;
+        private object _currentIConnectionHeartbeatFeature;
+        private object _currentIConnectionLifetimeNotificationFeature;
         private object _currentIBytesWrittenFeature;
 
         private int _featureRevision;
@@ -46,6 +50,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Abstractions.Internal
             _currentIApplicationTransportFeature = this;
             _currentITransportSchedulerFeature = this;
             _currentIConnectionLifetimeFeature = this;
+            _currentIConnectionHeartbeatFeature = this;
+            _currentIConnectionLifetimeNotificationFeature = this;
             _currentIBytesWrittenFeature = this;
 
         }
@@ -134,6 +140,14 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Abstractions.Internal
                 {
                     feature = _currentIConnectionLifetimeFeature;
                 }
+                else if (key == IConnectionHeartbeatFeatureType)
+                {
+                    feature = _currentIConnectionHeartbeatFeature;
+                }
+                else if (key == IConnectionLifetimeNotificationFeatureType)
+                {
+                    feature = _currentIConnectionLifetimeNotificationFeature;
+                }
                 else if (key == IBytesWrittenFeatureType)
                 {
                     feature = _currentIBytesWrittenFeature;
@@ -182,6 +196,14 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Abstractions.Internal
                 {
                     _currentIConnectionLifetimeFeature = value;
                 }
+                else if (key == IConnectionHeartbeatFeatureType)
+                {
+                    _currentIConnectionHeartbeatFeature = value;
+                }
+                else if (key == IConnectionLifetimeNotificationFeatureType)
+                {
+                    _currentIConnectionLifetimeNotificationFeature = value;
+                }
                 else if (key == IBytesWrittenFeatureType)
                 {
                     _currentIBytesWrittenFeature = value;
@@ -228,6 +250,14 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Abstractions.Internal
             {
                 feature = (TFeature)_currentIConnectionLifetimeFeature;
             }
+            else if (typeof(TFeature) == typeof(IConnectionHeartbeatFeature))
+            {
+                feature = (TFeature)_currentIConnectionHeartbeatFeature;
+            }
+            else if (typeof(TFeature) == typeof(IConnectionLifetimeNotificationFeature))
+            {
+                feature = (TFeature)_currentIConnectionLifetimeNotificationFeature;
+            }
             else if (typeof(TFeature) == typeof(IBytesWrittenFeature))
             {
                 feature = (TFeature)_currentIBytesWrittenFeature;
@@ -275,6 +305,14 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Abstractions.Internal
             {
                 _currentIConnectionLifetimeFeature = feature;
             }
+            else if (typeof(TFeature) == typeof(IConnectionHeartbeatFeature))
+            {
+                _currentIConnectionHeartbeatFeature = feature;
+            }
+            else if (typeof(TFeature) == typeof(IConnectionLifetimeNotificationFeature))
+            {
+                _currentIConnectionLifetimeNotificationFeature = feature;
+            }
             else if (typeof(TFeature) == typeof(IBytesWrittenFeature))
             {
                 _currentIBytesWrittenFeature = feature;
@@ -319,6 +357,14 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Abstractions.Internal
             {
                 yield return new KeyValuePair<Type, object>(IConnectionLifetimeFeatureType, _currentIConnectionLifetimeFeature);
             }
+            if (_currentIConnectionHeartbeatFeature != null)
+            {
+                yield return new KeyValuePair<Type, object>(IConnectionHeartbeatFeatureType, _currentIConnectionHeartbeatFeature);
+            }
+            if (_currentIConnectionLifetimeNotificationFeature != null)
+            {
+                yield return new KeyValuePair<Type, object>(IConnectionLifetimeNotificationFeatureType, _currentIConnectionLifetimeNotificationFeature);
+            }
             if (_currentIBytesWrittenFeature != null)
             {
                 yield return new KeyValuePair<Type, object>(IBytesWrittenFeatureType, _currentIBytesWrittenFeature);

src/Kestrel.Transport.Abstractions/Internal/TransportConnection.cs

@@ -1,6 +1,7 @@
 // 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;
 using System.Buffers;
 using System.Collections.Generic;
 using System.IO.Pipelines;
@@ -14,10 +15,15 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Abstractions.Internal
     public abstract partial class TransportConnection : ConnectionContext
     {
         private IDictionary<object, object> _items;
+        private List<(Action<object> handler, object state)> _heartbeatHandlers;
+        private readonly object _heartbeatLock = new object();
+        protected readonly CancellationTokenSource _connectionClosingCts = new CancellationTokenSource();
 
         public TransportConnection()
         {
             FastReset();
+
+            ConnectionClosedRequested = _connectionClosingCts.Token;
         }
 
         public IPAddress RemoteAddress { get; set; }
@@ -55,6 +61,37 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Abstractions.Internal
 
         public CancellationToken ConnectionClosed { get; set; }
 
+        public CancellationToken ConnectionClosedRequested { get; set; }
+
+        public void TickHeartbeat()
+        {
+            lock (_heartbeatLock)
+            {
+                if (_heartbeatHandlers == null)
+                {
+                    return;
+                }
+
+                foreach (var (handler, state) in _heartbeatHandlers)
+                {
+                    handler(state);
+                }
+            }
+        }
+
+        public void OnHeartbeat(Action<object> action, object state)
+        {
+            lock (_heartbeatLock)
+            {
+                if (_heartbeatHandlers == null)
+                {
+                    _heartbeatHandlers = new List<(Action<object> handler, object state)>();
+                }
+
+                _heartbeatHandlers.Add((action, state));
+            }
+        }
+
         // DO NOT remove this override to ConnectionContext.Abort. Doing so would cause
         // any TransportConnection that does not override Abort or calls base.Abort
         // to stack overflow when IConnectionLifetimeFeature.Abort() is called.
@@ -65,5 +102,18 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Abstractions.Internal
         {
             Output.CancelPendingRead();
         }
+
+        public void RequestClose()
+        {
+            try
+            {
+                _connectionClosingCts.Cancel();
+            }
+            catch (ObjectDisposedException)
+            {
+                // There's a race where the token could be disposed
+                // swallow the exception and no-op
+            }
+        }
     }
 }

src/Kestrel.Transport.Libuv/Internal/ILibuvTrace.cs

@@ -27,5 +27,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Libuv.Internal
         void ConnectionPause(string connectionId);
 
         void ConnectionResume(string connectionId);
+
+        void ConnectionAborted(string connectionId);
     }
 }

src/Kestrel.Transport.Libuv/Internal/LibuvConnection.cs

@@ -120,7 +120,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Libuv.Internal
         {
             _abortReason = abortReason;
             Output.CancelPendingRead();
-            
+
             // This cancels any pending I/O.
             Thread.Post(s => s.Dispose(), _socket);
         }
@@ -129,6 +129,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Libuv.Internal
         public void Dispose()
         {
             _connectionClosedTokenSource.Dispose();
+            _connectionClosingCts.Dispose();
         }
 
         // Called on Libuv thread

src/Kestrel.Transport.Libuv/Internal/LibuvTrace.cs

@@ -35,6 +35,9 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Libuv.Internal
         private static readonly Action<ILogger, string, Exception> _connectionReset =
             LoggerMessage.Define<string>(LogLevel.Debug, 19, @"Connection id ""{ConnectionId}"" reset.");
 
+        private static readonly Action<ILogger, string, Exception> _connectionAborted =
+            LoggerMessage.Define<string>(LogLevel.Debug, new EventId(20, nameof(ConnectionAborted)), @"Connection id ""{ConnectionId}"" aborted.");
+
         private readonly ILogger _logger;
 
         public LibuvTrace(ILogger logger)
@@ -95,6 +98,11 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Libuv.Internal
             _connectionResume(_logger, connectionId, null);
         }
 
+        public void ConnectionAborted(string connectionId)
+        {
+            _connectionAborted(_logger, connectionId, null);
+        }
+
         public IDisposable BeginScope<TState>(TState state) => _logger.BeginScope(state);
 
         public bool IsEnabled(LogLevel logLevel) => _logger.IsEnabled(logLevel);

src/Kestrel.Transport.Sockets/Internal/ISocketsTrace.cs

@@ -19,5 +19,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.Internal
         void ConnectionPause(string connectionId);
 
         void ConnectionResume(string connectionId);
+
+        void ConnectionAborted(string connectionId);
     }
 }

src/Kestrel.Transport.Sockets/Internal/SocketConnection.cs

@@ -95,6 +95,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.Internal
 
         public override void Abort(ConnectionAbortedException abortReason)
         {
+            _trace.ConnectionAborted(ConnectionId);
+
             _abortReason = abortReason;
             Output.CancelPendingRead();
 
@@ -106,6 +108,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.Internal
         public void Dispose()
         {
             _connectionClosedTokenSource.Dispose();
+            _connectionClosingCts.Dispose();
         }
 
         private async Task DoReceive()
@@ -188,17 +191,21 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.Internal
 
                 var flushTask = Input.FlushAsync();
 
-                if (!flushTask.IsCompleted)
+                var paused = !flushTask.IsCompleted;
+
+                if (paused)
                 {
                     _trace.ConnectionPause(ConnectionId);
+                }
 
-                    await flushTask;
+                var result = await flushTask;
 
+                if (paused)
+                {
                     _trace.ConnectionResume(ConnectionId);
                 }
 
-                var result = flushTask.GetAwaiter().GetResult();
+                if (result.IsCompleted || result.IsCanceled)
-                if (result.IsCompleted)
                 {
                     // Pipe consumer is shut down, do we stop writing
                     break;
@@ -244,6 +251,9 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.Internal
 
                 // Complete the output after disposing the socket
                 Output.Complete(error);
+
+                // Cancel any pending flushes so that the input loop is un-paused
+                Input.CancelPendingFlush();
             }
         }
 

src/Kestrel.Transport.Sockets/Internal/SocketsTrace.cs

@@ -28,6 +28,9 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.Internal
         private static readonly Action<ILogger, string, Exception> _connectionReset =
             LoggerMessage.Define<string>(LogLevel.Debug, new EventId(19, nameof(ConnectionReset)), @"Connection id ""{ConnectionId}"" reset.");
 
+        private static readonly Action<ILogger, string, Exception> _connectionAborted =
+            LoggerMessage.Define<string>(LogLevel.Debug, new EventId(20, nameof(ConnectionAborted)), @"Connection id ""{ConnectionId}"" aborted.");
+
         private readonly ILogger _logger;
 
         public SocketsTrace(ILogger logger)
@@ -83,6 +86,11 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.Internal
             _connectionResume(_logger, connectionId, null);
         }
 
+        public void ConnectionAborted(string connectionId)
+        {
+            _connectionAborted(_logger, connectionId, null);
+        }
+
         public IDisposable BeginScope<TState>(TState state) => _logger.BeginScope(state);
 
         public bool IsEnabled(LogLevel logLevel) => _logger.IsEnabled(logLevel);

test/Kestrel.Core.Tests/ConnectionDispatcherTests.cs

@@ -5,6 +5,7 @@ using System;
 using System.Collections.Generic;
 using System.IO.Pipelines;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal;
 using Microsoft.AspNetCore.Server.Kestrel.Transport.Abstractions.Internal;
@@ -23,7 +24,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             var tcs = new TaskCompletionSource<object>();
             var dispatcher = new ConnectionDispatcher(serviceContext, _ => tcs.Task);
 
-            var connection = Mock.Of<TransportConnection>();
+            var connection = new Mock<TransportConnection>() { CallBase = true }.Object;
+            connection.ConnectionClosed = new CancellationToken(canceled: true);
 
             dispatcher.OnConnection(connection);
 
@@ -51,7 +53,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             var serviceContext = new TestServiceContext();
             var dispatcher = new ConnectionDispatcher(serviceContext, _ => Task.CompletedTask);
 
-            var mockConnection = new Mock<TransportConnection>();
+            var mockConnection = new Mock<TransportConnection>() { CallBase = true };
+            mockConnection.Object.ConnectionClosed = new CancellationToken(canceled: true);
             var mockPipeReader = new Mock<PipeReader>();
             var mockPipeWriter = new Mock<PipeWriter>();
             var mockPipe = new Mock<IDuplexPipe>();

test/Kestrel.Core.Tests/HttpConnectionManagerTests.cs

@@ -3,8 +3,10 @@
 
 using System;
 using System.Runtime.CompilerServices;
+using Microsoft.AspNetCore.Connections;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
+using Microsoft.AspNetCore.Server.Kestrel.Transport.Abstractions.Internal;
 using Microsoft.AspNetCore.Testing;
 using Moq;
 using Xunit;
@@ -18,7 +20,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
         {
             var connectionId = "0";
             var trace = new Mock<IKestrelTrace>();
-            var httpConnectionManager = new HttpConnectionManager(trace.Object, ResourceCounter.Unlimited);
+            var httpConnectionManager = new ConnectionManager(trace.Object, ResourceCounter.Unlimited);
 
             // Create HttpConnection in inner scope so it doesn't get rooted by the current frame.
             UnrootedConnectionsGetRemovedFromHeartbeatInnerScope(connectionId, httpConnectionManager, trace);
@@ -36,14 +38,12 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
         [MethodImpl(MethodImplOptions.NoInlining)]
         private void UnrootedConnectionsGetRemovedFromHeartbeatInnerScope(
             string connectionId,
-            HttpConnectionManager httpConnectionManager,
+            ConnectionManager httpConnectionManager,
             Mock<IKestrelTrace> trace)
         {
-            var httpConnection = new HttpConnection(new HttpConnectionContext
+            var mock = new Mock<TransportConnection>();
-            {
+            mock.Setup(m => m.ConnectionId).Returns(connectionId);
-                ServiceContext = new TestServiceContext(),
+            var httpConnection = new KestrelConnection(mock.Object);
-                ConnectionId = connectionId
-            });
 
             httpConnectionManager.AddConnection(0, httpConnection);
 

test/Kestrel.Core.Tests/HttpConnectionTests.cs

@@ -43,7 +43,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
                 ConnectionAdapters = new List<IConnectionAdapter>(),
                 ConnectionFeatures = connectionFeatures,
                 MemoryPool = _memoryPool,
-                HttpConnectionId = long.MinValue,
                 Transport = pair.Transport,
                 ServiceContext = new TestServiceContext
                 {

test/Kestrel.InMemory.FunctionalTests/HttpsTests.cs

@@ -312,7 +312,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests
             LoggerFactory.AddProvider(loggerProvider);
 
             var testContext = new TestServiceContext(LoggerFactory);
-            var heartbeatManager = new HttpHeartbeatManager(testContext.ConnectionManager);
+            var heartbeatManager = new HeartbeatManager(testContext.ConnectionManager);
 
             var handshakeStartedTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
             TimeSpan handshakeTimeout = default;

test/Kestrel.InMemory.FunctionalTests/KeepAliveTimeoutTests.cs

@@ -27,7 +27,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests
         public async Task ConnectionClosedWhenKeepAliveTimeoutExpires()
         {
             var testContext = new TestServiceContext(LoggerFactory);
-            var heartbeatManager = new HttpHeartbeatManager(testContext.ConnectionManager);
+            var heartbeatManager = new HeartbeatManager(testContext.ConnectionManager);
 
             using (var server = CreateServer(testContext))
             using (var connection = server.CreateConnection())
@@ -51,7 +51,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests
         public async Task ConnectionKeptAliveBetweenRequests()
         {
             var testContext = new TestServiceContext(LoggerFactory);
-            var heartbeatManager = new HttpHeartbeatManager(testContext.ConnectionManager);
+            var heartbeatManager = new HeartbeatManager(testContext.ConnectionManager);
 
             using (var server = CreateServer(testContext))
             using (var connection = server.CreateConnection())
@@ -76,7 +76,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests
         public async Task ConnectionNotTimedOutWhileRequestBeingSent()
         {
             var testContext = new TestServiceContext(LoggerFactory);
-            var heartbeatManager = new HttpHeartbeatManager(testContext.ConnectionManager);
+            var heartbeatManager = new HeartbeatManager(testContext.ConnectionManager);
 
             using (var server = CreateServer(testContext))
             using (var connection = server.CreateConnection())
@@ -113,7 +113,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests
         private async Task ConnectionNotTimedOutWhileAppIsRunning()
         {
             var testContext = new TestServiceContext(LoggerFactory);
-            var heartbeatManager = new HttpHeartbeatManager(testContext.ConnectionManager);
+            var heartbeatManager = new HeartbeatManager(testContext.ConnectionManager);
             var cts = new CancellationTokenSource();
 
             using (var server = CreateServer(testContext, longRunningCt: cts.Token))
@@ -150,7 +150,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests
         private async Task ConnectionTimesOutWhenOpenedButNoRequestSent()
         {
             var testContext = new TestServiceContext(LoggerFactory);
-            var heartbeatManager = new HttpHeartbeatManager(testContext.ConnectionManager);
+            var heartbeatManager = new HeartbeatManager(testContext.ConnectionManager);
 
             using (var server = CreateServer(testContext))
             using (var connection = server.CreateConnection())
@@ -167,7 +167,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests
         private async Task KeepAliveTimeoutDoesNotApplyToUpgradedConnections()
         {
             var testContext = new TestServiceContext(LoggerFactory);
-            var heartbeatManager = new HttpHeartbeatManager(testContext.ConnectionManager);
+            var heartbeatManager = new HeartbeatManager(testContext.ConnectionManager);
             var cts = new CancellationTokenSource();
 
             using (var server = CreateServer(testContext, upgradeCt: cts.Token))

test/Kestrel.InMemory.FunctionalTests/RequestBodyTimeoutTests.cs

@@ -22,7 +22,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests
         {
             var gracePeriod = TimeSpan.FromSeconds(5);
             var serviceContext = new TestServiceContext(LoggerFactory);
-            var heartbeatManager = new HttpHeartbeatManager(serviceContext.ConnectionManager);
+            var heartbeatManager = new HeartbeatManager(serviceContext.ConnectionManager);
 
             var appRunningEvent = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
 
@@ -136,7 +136,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests
         {
             var gracePeriod = TimeSpan.FromSeconds(5);
             var serviceContext = new TestServiceContext(LoggerFactory);
-            var heartbeatManager = new HttpHeartbeatManager(serviceContext.ConnectionManager);
+            var heartbeatManager = new HeartbeatManager(serviceContext.ConnectionManager);
 
             var appRunningTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
             var exceptionSwallowedTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);

test/Kestrel.InMemory.FunctionalTests/RequestHeadersTimeoutTests.cs

@@ -26,7 +26,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests
         public async Task ConnectionAbortedWhenRequestHeadersNotReceivedInTime(string headers)
         {
             var testContext = new TestServiceContext(LoggerFactory);
-            var heartbeatManager = new HttpHeartbeatManager(testContext.ConnectionManager);
+            var heartbeatManager = new HeartbeatManager(testContext.ConnectionManager);
 
             using (var server = CreateServer(testContext))
             using (var connection = server.CreateConnection())
@@ -47,7 +47,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests
         public async Task RequestHeadersTimeoutCanceledAfterHeadersReceived()
         {
             var testContext = new TestServiceContext(LoggerFactory);
-            var heartbeatManager = new HttpHeartbeatManager(testContext.ConnectionManager);
+            var heartbeatManager = new HeartbeatManager(testContext.ConnectionManager);
 
             using (var server = CreateServer(testContext))
             using (var connection = server.CreateConnection())
@@ -76,7 +76,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests
         public async Task ConnectionAbortedWhenRequestLineNotReceivedInTime(string requestLine)
         {
             var testContext = new TestServiceContext(LoggerFactory);
-            var heartbeatManager = new HttpHeartbeatManager(testContext.ConnectionManager);
+            var heartbeatManager = new HeartbeatManager(testContext.ConnectionManager);
 
             using (var server = CreateServer(testContext))
             using (var connection = server.CreateConnection())
@@ -95,7 +95,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests
         public async Task TimeoutNotResetOnEachRequestLineCharacterReceived()
         {
             var testContext = new TestServiceContext(LoggerFactory);
-            var heartbeatManager = new HttpHeartbeatManager(testContext.ConnectionManager);
+            var heartbeatManager = new HeartbeatManager(testContext.ConnectionManager);
 
             using (var server = CreateServer(testContext))
             using (var connection = server.CreateConnection())

test/Kestrel.InMemory.FunctionalTests/RequestTests.cs

@@ -877,7 +877,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests
             var appEvent = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
             var delayEvent = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
             var serviceContext = new TestServiceContext(LoggerFactory);
-            var heartbeatManager = new HttpHeartbeatManager(serviceContext.ConnectionManager);
+            var heartbeatManager = new HeartbeatManager(serviceContext.ConnectionManager);
 
             using (var server = new TestServer(async context =>
             {

test/Kestrel.InMemory.FunctionalTests/UpgradeTests.cs

@@ -254,7 +254,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests
             const int limit = 10;
             var upgradeTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
             var serviceContext = new TestServiceContext(LoggerFactory);
-            serviceContext.ConnectionManager = new HttpConnectionManager(serviceContext.Log, ResourceCounter.Quota(limit));
+            serviceContext.ConnectionManager = new ConnectionManager(serviceContext.Log, ResourceCounter.Quota(limit));
 
             using (var server = new TestServer(async context =>
             {

test/Kestrel.Transport.FunctionalTests/Http2/ShutdownTests.cs

@@ -109,6 +109,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests.Http2
                 MemoryPoolFactory = memoryPoolFactory.Create
             };
 
+            TestApplicationErrorLogger.ThrowOnUngracefulShutdown = false;
+
             // Abortive shutdown leaves one request hanging
             using (var server = new TestServer(async context =>
             {

test/shared/TestApplicationErrorLogger.cs

@@ -20,6 +20,8 @@ namespace Microsoft.AspNetCore.Testing
 
         public bool ThrowOnCriticalErrors { get; set; } = true;
 
+        public bool ThrowOnUngracefulShutdown { get; set; } = true;
+
         public ConcurrentQueue<LogMessage> Messages { get; } = new ConcurrentQueue<LogMessage>();
 
         public ConcurrentQueue<object> Scopes { get; } = new ConcurrentQueue<object>();
@@ -59,7 +61,9 @@ namespace Microsoft.AspNetCore.Testing
             }
 
             // Fail tests where not all the connections close during server shutdown.
-            if (eventId.Id == 21 && eventId.Name == nameof(KestrelTrace.NotAllConnectionsAborted))
+            if (ThrowOnUngracefulShutdown &&
+                ((eventId.Id == 16 && eventId.Name == nameof(KestrelTrace.NotAllConnectionsClosedGracefully)) ||
+                 (eventId.Id == 21 && eventId.Name == nameof(KestrelTrace.NotAllConnectionsAborted))))
             {
                 var log = $"Log {logLevel}[{eventId}]: {formatter(state, exception)} {exception?.Message}";
                 throw new Exception($"Shutdown failure. {log}");

test/shared/TestServiceContext.cs

@@ -45,7 +45,7 @@ namespace Microsoft.AspNetCore.Testing
             SystemClock = new SystemClock();
             DateHeaderValueManager = new DateHeaderValueManager(SystemClock);
 
-            var heartbeatManager = new HttpHeartbeatManager(ConnectionManager);
+            var heartbeatManager = new HeartbeatManager(ConnectionManager);
             Heartbeat = new Heartbeat(
                 new IHeartbeatHandler[] { DateHeaderValueManager, heartbeatManager },
                 SystemClock,
@@ -61,7 +61,7 @@ namespace Microsoft.AspNetCore.Testing
             MockSystemClock = new MockSystemClock();
             SystemClock = MockSystemClock;
             DateHeaderValueManager = new DateHeaderValueManager(MockSystemClock);
-            ConnectionManager = new HttpConnectionManager(Log, ResourceCounter.Unlimited);
+            ConnectionManager = new ConnectionManager(Log, ResourceCounter.Unlimited);
             HttpParser = new HttpParser<Http1ParsingHandler>(Log.IsEnabled(LogLevel.Information));
             ServerOptions = new KestrelServerOptions
             {

tools/CodeGenerator/TransportConnectionFeatureCollection.cs

@@ -19,6 +19,8 @@ namespace CodeGenerator
                 "IApplicationTransportFeature",
                 "ITransportSchedulerFeature",
                 "IConnectionLifetimeFeature",
+                "IConnectionHeartbeatFeature",
+                "IConnectionLifetimeNotificationFeature",
                 "IBytesWrittenFeature",
             };