Added FrameConnection.OnConnectionClosed back (#2028)

Author: davidfowl

src/Kestrel.Core/Internal/FrameConnection.cs

@@ -28,6 +28,7 @@ public class FrameConnection : ITimeoutControl
 
         private readonly FrameConnectionContext _context;
         private IList<IAdaptedConnection> _adaptedConnections;
+        private readonly TaskCompletionSource<object> _socketClosedTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
         private Frame _frame;
         private Http2Connection _http2Connection;
         private volatile int _http2ConnectionState;
@@ -151,6 +152,7 @@ private async Task ProcessRequestsAsync<TContext>(IHttpApplication<TContext> htt
                 }
 
                 await adaptedPipelineTask;
+                await _socketClosedTcs.Task;
             }
             catch (Exception ex)
             {
@@ -189,6 +191,13 @@ internal void CreateFrame<TContext>(IHttpApplication<TContext> httpApplication,
             });
         }
 
+        public void OnConnectionClosed(Exception ex)
+        {
+            Abort(ex);
+
+            _socketClosedTcs.TrySetResult(null);
+        }
+
         public Task StopAsync()
         {
             Debug.Assert(_frame != null, $"{nameof(_frame)} is null");

src/Kestrel.Core/Internal/HttpConnectionMiddleware.cs

@@ -13,6 +13,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
 {
     public class HttpConnectionMiddleware<TContext>
     {
+        private static Action<Exception, object> _completeTcs = CompleteTcs;
+
         private static long _lastFrameConnectionId = long.MinValue;
 
         private readonly IList<IConnectionAdapter> _connectionAdapters;
@@ -67,30 +69,40 @@ public Task OnConnectionAsync(ConnectionContext connectionContext)
             var processingTask = connection.StartRequestProcessing(_application);
 
             var inputTcs = new TaskCompletionSource<object>();
+            var outputTcs = new TaskCompletionSource<object>();
 
-            // Abort the frame when the transport writer completes
-            connectionContext.Transport.Input.OnWriterCompleted((error, state) =>
-            {
-                var tcs = (TaskCompletionSource<object>)state;
-
-                if (error != null)
-                {
-                    tcs.TrySetException(error);
-                }
-                else
-                {
-                    tcs.TrySetResult(null);
-                }
-            },
-            inputTcs);
+            // The reason we don't fire events directly from these callbacks is because it seems
+            // like the transport callbacks root the state object (even after it fires)
+            connectionContext.Transport.Input.OnWriterCompleted(_completeTcs, inputTcs);
+            connectionContext.Transport.Output.OnReaderCompleted(_completeTcs, outputTcs);
 
             inputTcs.Task.ContinueWith((task, state) =>
             {
                 ((FrameConnection)state).Abort(task.Exception?.InnerException);
             },
             connection, TaskContinuationOptions.ExecuteSynchronously);
 
+            outputTcs.Task.ContinueWith((task, state) =>
+            {
+                ((FrameConnection)state).OnConnectionClosed(task.Exception?.InnerException);
+            },
+            connection, TaskContinuationOptions.ExecuteSynchronously);
+
             return processingTask;
         }
+
+        private static void CompleteTcs(Exception error, object state)
+        {
+            var tcs = (TaskCompletionSource<object>)state;
+
+            if (error != null)
+            {
+                tcs.TrySetException(error);
+            }
+            else
+            {
+                tcs.TrySetResult(null);
+            }
+        }
     }
 }