Added FrameConnection.OnConnectionClosed back (#2028)

src/Kestrel.Core/Internal/FrameConnection.cs

@@ -28,6 +28,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
 
         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 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
                 }
 
                 await adaptedPipelineTask;
+                await _socketClosedTcs.Task;
             }
             catch (Exception ex)
             {
@@ -189,6 +191,13 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
             });
         }
 
+        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,9 +69,29 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
             var processingTask = connection.StartRequestProcessing(_application);
 
             var inputTcs = new TaskCompletionSource<object>();
+            var outputTcs = new TaskCompletionSource<object>();
 
-            // Abort the frame when the transport writer completes
+            // The reason we don't fire events directly from these callbacks is because it seems
-            connectionContext.Transport.Input.OnWriterCompleted((error, state) =>
+            // 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;
 
@@ -81,16 +103,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
             {
                 tcs.TrySetResult(null);
             }
-            },
-            inputTcs);
-
-            inputTcs.Task.ContinueWith((task, state) =>
-            {
-                ((FrameConnection)state).Abort(task.Exception?.InnerException);
-            },
-            connection, TaskContinuationOptions.ExecuteSynchronously);
-
-            return processingTask;
         }
     }
 }