From 9d8556e7c4af2b4039b8ef8ae44104aca87a44f5 Mon Sep 17 00:00:00 2001
From: David Fowler <davidfowl@gmail.com>
Date: Sat, 26 Aug 2017 21:20:06 -0700
Subject: [PATCH] Added FrameConnection.OnConnectionClosed back (#2028)

---
 src/Kestrel.Core/Internal/FrameConnection.cs  |  9 ++++
 .../Internal/HttpConnectionMiddleware.cs      | 42 ++++++++++++-------
 2 files changed, 36 insertions(+), 15 deletions(-)

diff --git a/src/Kestrel.Core/Internal/FrameConnection.cs b/src/Kestrel.Core/Internal/FrameConnection.cs
index 8171817b7a..f2e3a3e172 100644
--- a/src/Kestrel.Core/Internal/FrameConnection.cs
+++ b/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");
diff --git a/src/Kestrel.Core/Internal/HttpConnectionMiddleware.cs b/src/Kestrel.Core/Internal/HttpConnectionMiddleware.cs
index 221371535b..b47ea9d603 100644
--- a/src/Kestrel.Core/Internal/HttpConnectionMiddleware.cs
+++ b/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,22 +69,12 @@ 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
-            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) =>
             {
@@ -90,7 +82,27 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
             },
             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);
+            }
+        }
     }
 }