From 447f4bc29824131beadd3d7c26497893f9c88f32 Mon Sep 17 00:00:00 2001
From: David Fowler <davidfowl@gmail.com>
Date: Mon, 28 Jan 2019 20:37:11 +0000
Subject: [PATCH] Reuse the stack for OnCompleted and OnStarting calls (#6833)

- Reviewing profiles for MVC and other things that use RegisterForDispose, this shows up as a significant amount of allocations.
---
 src/Servers/Kestrel/Core/src/CoreStrings.resx |   4 +-
 .../Core/src/Internal/Http/HttpProtocol.cs    | 104 ++++++++++--------
 .../InMemory.FunctionalTests/ResponseTests.cs |  93 ++++++++++++++++
 3 files changed, 151 insertions(+), 50 deletions(-)

diff --git a/src/Servers/Kestrel/Core/src/CoreStrings.resx b/src/Servers/Kestrel/Core/src/CoreStrings.resx
index b17228eec4..774305483a 100644
--- a/src/Servers/Kestrel/Core/src/CoreStrings.resx
+++ b/src/Servers/Kestrel/Core/src/CoreStrings.resx
@@ -1,4 +1,4 @@
-<?xml version="1.0" encoding="utf-8"?>
+<?xml version="1.0" encoding="utf-8"?>
 <root>
   <!-- 
     Microsoft ResX Schema 
@@ -590,4 +590,4 @@ For more information on configuring HTTPS see https://go.microsoft.com/fwlink/?l
   <data name="Http2ErrorStreamAborted" xml:space="preserve">
     <value>A frame of type {frameType} was received after stream {streamId} was reset or aborted.</value>
   </data>
-</root>
\ No newline at end of file
+</root>
diff --git a/src/Servers/Kestrel/Core/src/Internal/Http/HttpProtocol.cs b/src/Servers/Kestrel/Core/src/Internal/Http/HttpProtocol.cs
index 1edbc54a8e..8a8926b176 100644
--- a/src/Servers/Kestrel/Core/src/Internal/Http/HttpProtocol.cs
+++ b/src/Servers/Kestrel/Core/src/Internal/Http/HttpProtocol.cs
@@ -17,7 +17,6 @@ using Microsoft.AspNetCore.Connections;
 using Microsoft.AspNetCore.Hosting.Server;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Http.Features;
-using Microsoft.AspNetCore.Http.Internal;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Primitives;
@@ -31,9 +30,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
         private static readonly byte[] _bytesTransferEncodingChunked = Encoding.ASCII.GetBytes("\r\nTransfer-Encoding: chunked");
         private static readonly byte[] _bytesServer = Encoding.ASCII.GetBytes("\r\nServer: " + Constants.ServerName);
 
-        private readonly object _onStartingSync = new Object();
-        private readonly object _onCompletedSync = new Object();
-
         protected Streams _streams;
 
         private Stack<KeyValuePair<Func<object, Task>, object>> _onStarting;
@@ -313,8 +309,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
 
         public void Reset()
         {
-            _onStarting = null;
-            _onCompleted = null;
+            _onStarting?.Clear();
+            _onCompleted?.Clear();
 
             _requestProcessingStatus = RequestProcessingStatus.RequestPending;
             _autoChunk = false;
@@ -588,7 +584,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
                 // already failed. If an OnStarting callback throws we can go through
                 // our normal error handling in ProduceEnd.
                 // https://github.com/aspnet/KestrelHttpServer/issues/43
-                if (!HasResponseStarted && _applicationException == null && _onStarting != null)
+                if (!HasResponseStarted && _applicationException == null && _onStarting?.Count > 0)
                 {
                     await FireOnStarting();
                 }
@@ -624,7 +620,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
                     }
                 }
 
-                if (_onCompleted != null)
+                if (_onCompleted?.Count > 0)
                 {
                     await FireOnCompleted();
                 }
@@ -649,43 +645,32 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
 
         public void OnStarting(Func<object, Task> callback, object state)
         {
-            lock (_onStartingSync)
+            if (HasResponseStarted)
             {
-                if (HasResponseStarted)
-                {
-                    ThrowResponseAlreadyStartedException(nameof(OnStarting));
-                }
-
-                if (_onStarting == null)
-                {
-                    _onStarting = new Stack<KeyValuePair<Func<object, Task>, object>>();
-                }
-                _onStarting.Push(new KeyValuePair<Func<object, Task>, object>(callback, state));
+                ThrowResponseAlreadyStartedException(nameof(OnStarting));
             }
+
+            if (_onStarting == null)
+            {
+                _onStarting = new Stack<KeyValuePair<Func<object, Task>, object>>();
+            }
+            _onStarting.Push(new KeyValuePair<Func<object, Task>, object>(callback, state));
         }
 
         public void OnCompleted(Func<object, Task> callback, object state)
         {
-            lock (_onCompletedSync)
+            if (_onCompleted == null)
             {
-                if (_onCompleted == null)
-                {
-                    _onCompleted = new Stack<KeyValuePair<Func<object, Task>, object>>();
-                }
-                _onCompleted.Push(new KeyValuePair<Func<object, Task>, object>(callback, state));
+                _onCompleted = new Stack<KeyValuePair<Func<object, Task>, object>>();
             }
+            _onCompleted.Push(new KeyValuePair<Func<object, Task>, object>(callback, state));
         }
 
         protected Task FireOnStarting()
         {
-            Stack<KeyValuePair<Func<object, Task>, object>> onStarting;
-            lock (_onStartingSync)
-            {
-                onStarting = _onStarting;
-                _onStarting = null;
-            }
+            var onStarting = _onStarting;
 
-            if (onStarting == null)
+            if (onStarting == null || onStarting.Count == 0)
             {
                 return Task.CompletedTask;
             }
@@ -700,10 +685,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
         {
             try
             {
-                var count = onStarting.Count;
-                for (var i = 0; i < count; i++)
+                while (onStarting.TryPop(out var entry))
                 {
-                    var entry = onStarting.Pop();
                     var task = entry.Key.Invoke(entry.Value);
                     if (!ReferenceEquals(task, Task.CompletedTask))
                     {
@@ -725,10 +708,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
             {
                 await currentTask;
 
-                var count = onStarting.Count;
-                for (var i = 0; i < count; i++)
+                while (onStarting.TryPop(out var entry))
                 {
-                    var entry = onStarting.Pop();
                     await entry.Key.Invoke(entry.Value);
                 }
             }
@@ -740,24 +721,50 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
 
         protected Task FireOnCompleted()
         {
-            Stack<KeyValuePair<Func<object, Task>, object>> onCompleted;
-            lock (_onCompletedSync)
-            {
-                onCompleted = _onCompleted;
-                _onCompleted = null;
-            }
+            var onCompleted = _onCompleted;
 
-            if (onCompleted == null)
+            if (onCompleted == null || onCompleted.Count == 0)
             {
                 return Task.CompletedTask;
             }
 
-            return FireOnCompletedAwaited(onCompleted);
+            return FireOnCompletedMayAwait(onCompleted);
         }
 
-        private async Task FireOnCompletedAwaited(Stack<KeyValuePair<Func<object, Task>, object>> onCompleted)
+        private Task FireOnCompletedMayAwait(Stack<KeyValuePair<Func<object, Task>, object>> onCompleted)
         {
-            foreach (var entry in onCompleted)
+
+            while (onCompleted.TryPop(out var entry))
+            {
+                try
+                {
+                    var task = entry.Key.Invoke(entry.Value);
+                    if (!ReferenceEquals(task, Task.CompletedTask))
+                    {
+                        return FireOnCompletedAwaited(task, onCompleted);
+                    }
+                }
+                catch (Exception ex)
+                {
+                    ReportApplicationError(ex);
+                }
+            }
+
+            return Task.CompletedTask;
+        }
+
+        private async Task FireOnCompletedAwaited(Task currentTask, Stack<KeyValuePair<Func<object, Task>, object>> onCompleted)
+        {
+            try
+            {
+                await currentTask;
+            }
+            catch (Exception ex)
+            {
+                Log.ApplicationError(ConnectionId, TraceIdentifier, ex);
+            }
+
+            while (onCompleted.TryPop(out var entry))
             {
                 try
                 {
@@ -768,6 +775,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
                     Log.ApplicationError(ConnectionId, TraceIdentifier, ex);
                 }
             }
+
         }
 
         public Task FlushAsync(CancellationToken cancellationToken = default(CancellationToken))
diff --git a/src/Servers/Kestrel/test/InMemory.FunctionalTests/ResponseTests.cs b/src/Servers/Kestrel/test/InMemory.FunctionalTests/ResponseTests.cs
index a098277e68..7b8ae6b955 100644
--- a/src/Servers/Kestrel/test/InMemory.FunctionalTests/ResponseTests.cs
+++ b/src/Servers/Kestrel/test/InMemory.FunctionalTests/ResponseTests.cs
@@ -2001,6 +2001,99 @@ namespace Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests
             Assert.Equal(2, TestApplicationErrorLogger.Messages.Where(message => message.LogLevel == LogLevel.Error).Count());
         }
 
+        [Fact]
+        public async Task OnStartingThrowsInsideOnStartingCallbacksRuns()
+        {
+            var testContext = new TestServiceContext(LoggerFactory);
+            var tcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+
+            using (var server = new TestServer(async httpContext =>
+            {
+                var response = httpContext.Response;
+                response.OnStarting(state1 =>
+                {
+                    response.OnStarting(state2 =>
+                    {
+                        tcs.TrySetResult(null);
+                        return Task.CompletedTask;
+                    },
+                    null);
+
+                    return Task.CompletedTask;
+
+                }, null);
+
+                response.Headers["Content-Length"] = new[] { "11" };
+
+                await response.Body.WriteAsync(Encoding.ASCII.GetBytes("Hello World"), 0, 11);
+            }, testContext))
+            {
+                using (var connection = server.CreateConnection())
+                {
+                    await connection.Send(
+                        "GET / HTTP/1.1",
+                        "Host:",
+                        "",
+                        "");
+                    await connection.Receive(
+                        "HTTP/1.1 200 OK",
+                        $"Date: {testContext.DateHeaderValue}",
+                        "");
+
+                    await tcs.Task.DefaultTimeout();
+                }
+
+                await server.StopAsync();
+            }
+        }
+
+        [Fact]
+        public async Task OnCompletedThrowsInsideOnCompletedCallbackRuns()
+        {
+            var testContext = new TestServiceContext(LoggerFactory);
+            var tcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+
+            using (var server = new TestServer(async httpContext =>
+            {
+                var response = httpContext.Response;
+                response.OnCompleted(state1 =>
+                {
+                    response.OnCompleted(state2 =>
+                    {
+                        tcs.TrySetResult(null);
+
+                        return Task.CompletedTask;
+                    },
+                    null);
+
+                    return Task.CompletedTask;
+
+                }, null);
+
+                response.Headers["Content-Length"] = new[] { "11" };
+
+                await response.Body.WriteAsync(Encoding.ASCII.GetBytes("Hello World"), 0, 11);
+            }, testContext))
+            {
+                using (var connection = server.CreateConnection())
+                {
+                    await connection.Send(
+                        "GET / HTTP/1.1",
+                        "Host:",
+                        "",
+                        "");
+                    await connection.Receive(
+                        "HTTP/1.1 200 OK",
+                        $"Date: {testContext.DateHeaderValue}",
+                        "");
+
+                    await tcs.Task.DefaultTimeout();
+                }
+
+                await server.StopAsync();
+            }
+        }
+
         [Fact]
         public async Task ThrowingInOnCompletedIsLogged()
         {