Use inline struct to avoid extra stack referencedavidfowl/event-stack

3 files changed, 99 insertions, 53 deletions

diff --git a/src/Servers/Kestrel/Core/src/Internal/EventStack.cs b/src/Servers/Kestrel/Core/src/Internal/EventStack.cs
new file mode 100644
index 0000000000..0fa5784a1c
--- /dev/null
+++ b/src/Servers/Kestrel/Core/src/Internal/EventStack.cs
@@ -0,0 +1,70 @@
+using System;
+using System.Collections.Generic;
+using System.Runtime.CompilerServices;
+using System.Threading.Tasks;
+
+namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
+{
+    internal struct EventStack
+    {
+        private KeyValuePair<Func<object, Task>, object>[] _array;
+        private int _size;
+
+        public EventStack(int size)
+        {
+            _array = size == 0 ? Array.Empty<KeyValuePair<Func<object, Task>, object>>() : new KeyValuePair<Func<object, Task>, object>[size];
+            _size = 0;
+        }
+
+        public int Count => _size;
+
+        public bool TryPop(out KeyValuePair<Func<object, Task>, object> result)
+        {
+            int size = _size - 1;
+            var array = _array;
+
+            if ((uint)size >= (uint)array.Length)
+            {
+                result = default;
+                return false;
+            }
+
+            _size = size;
+            result = array[size];
+            array[size] = default;
+            return true;
+        }
+
+        // Pushes an item to the top of the stack.
+        public void Push(KeyValuePair<Func<object, Task>, object> item)
+        {
+            int size = _size;
+            var array = _array;
+
+            if ((uint)size < (uint)array.Length)
+            {
+                array[size] = item;
+                _size = size + 1;
+            }
+            else
+            {
+                PushWithResize(item);
+            }
+        }
+
+        // Non-inline from Stack.Push to improve its code quality as uncommon path
+        [MethodImpl(MethodImplOptions.NoInlining)]
+        private void PushWithResize(KeyValuePair<Func<object, Task>, object> item)
+        {
+            Array.Resize(ref _array, Math.Max(2 * _array.Length, 1));
+            _array[_size] = item;
+            _size++;
+        }
+
+        public void Clear()
+        {
+            _array.AsSpan().Clear();
+            _size = 0;
+        }
+    }
+}
diff --git a/src/Servers/Kestrel/Core/src/Internal/Http/HttpProtocol.cs b/src/Servers/Kestrel/Core/src/Internal/Http/HttpProtocol.cs
index bf856a53c4..aa8d48f155 100644
--- a/src/Servers/Kestrel/Core/src/Internal/Http/HttpProtocol.cs
+++ b/src/Servers/Kestrel/Core/src/Internal/Http/HttpProtocol.cs
@@ -38,8 +38,9 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
         internal const string SchemeHttps = "https";
 
         protected BodyControl? _bodyControl;
-        private Stack<KeyValuePair<Func<object, Task>, object>>? _onStarting;
-        private Stack<KeyValuePair<Func<object, Task>, object>>? _onCompleted;
+        // Mutable structs, don't mark as readonly!
+        private EventStack _onStarting = new(0);
+        private EventStack _onCompleted = new(0);
 
         private readonly object _abortLock = new object();
         protected volatile bool _connectionAborted;
@@ -333,8 +334,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
 
         public void Reset()
         {
-            _onStarting?.Clear();
-            _onCompleted?.Clear();
+            _onStarting.Clear();
+            _onCompleted.Clear();
             _routeValues?.Clear();
 
             _requestProcessingStatus = RequestProcessingStatus.RequestPending;
@@ -666,7 +667,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?.Count > 0)
+                    if (!HasResponseStarted && _applicationException == null && _onStarting.Count > 0)
                     {
                         await FireOnStarting();
                     }
@@ -733,7 +734,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
                     }
                 }
 
-                if (_onCompleted?.Count > 0)
+                if (_onCompleted.Count > 0)
                 {
                     await FireOnCompleted();
                 }
@@ -760,39 +761,30 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
                 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)
         {
-            if (_onCompleted == null)
-            {
-                _onCompleted = new Stack<KeyValuePair<Func<object, Task>, object>>();
-            }
             _onCompleted.Push(new KeyValuePair<Func<object, Task>, object>(callback, state));
         }
 
         protected Task FireOnStarting()
         {
-            var onStarting = _onStarting;
-            if (onStarting?.Count > 0)
+            if (_onStarting.Count > 0)
             {
-                return ProcessEvents(this, onStarting);
+                return ProcessEvents(this);
             }
 
             return Task.CompletedTask;
 
-            static async Task ProcessEvents(HttpProtocol protocol, Stack<KeyValuePair<Func<object, Task>, object>> events)
+            static async Task ProcessEvents(HttpProtocol protocol)
             {
                 // Try/Catch is outside the loop as any error that occurs is before the request starts.
                 // So we want to report it as an ApplicationError to fail the request and not process more events.
                 try
                 {
-                    while (events.TryPop(out var entry))
+                    while (protocol._onStarting.TryPop(out var entry))
                     {
                         await entry.Key.Invoke(entry.Value);
                     }
@@ -804,30 +796,19 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
             }
         }
 
-        protected Task FireOnCompleted()
+        protected async Task FireOnCompleted()
         {
-            var onCompleted = _onCompleted;
-            if (onCompleted?.Count > 0)
-            {
-                return ProcessEvents(this, onCompleted);
-            }
-
-            return Task.CompletedTask;
-
-            static async Task ProcessEvents(HttpProtocol protocol, Stack<KeyValuePair<Func<object, Task>, object>> events)
+            // Try/Catch is inside the loop as any error that occurs is after the request has finished.
+            // So we will just log it and keep processing the events, as the completion has already happened.
+            while (_onCompleted.TryPop(out var entry))
             {
-                // Try/Catch is inside the loop as any error that occurs is after the request has finished.
-                // So we will just log it and keep processing the events, as the completion has already happened.
-                while (events.TryPop(out var entry))
+                try
                 {
-                    try
-                    {
-                        await entry.Key.Invoke(entry.Value);
-                    }
-                    catch (Exception ex)
-                    {
-                        protocol.Log.ApplicationError(protocol.ConnectionId, protocol.TraceIdentifier, ex);
-                    }
+                    await entry.Key.Invoke(entry.Value);
+                }
+                catch (Exception ex)
+                {
+                    Log.ApplicationError(ConnectionId, TraceIdentifier, ex);
                 }
             }
         }
diff --git a/src/Servers/Kestrel/Core/src/Internal/Infrastructure/KestrelConnection.cs b/src/Servers/Kestrel/Core/src/Internal/Infrastructure/KestrelConnection.cs
index fb7506ab8c..927094c3c5 100644
--- a/src/Servers/Kestrel/Core/src/Internal/Infrastructure/KestrelConnection.cs
+++ b/src/Servers/Kestrel/Core/src/Internal/Infrastructure/KestrelConnection.cs
@@ -16,7 +16,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
         private List<(Action<object> handler, object state)>? _heartbeatHandlers;
         private readonly object _heartbeatLock = new object();
 
-        private Stack<KeyValuePair<Func<object, Task>, object>>? _onCompleted;
+        private EventStack _onCompleted = new(0);
         private bool _completed;
 
         private readonly CancellationTokenSource _connectionClosingCts = new CancellationTokenSource();
@@ -81,10 +81,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
                 throw new InvalidOperationException("The connection is already complete.");
             }
 
-            if (_onCompleted == null)
-            {
-                _onCompleted = new Stack<KeyValuePair<Func<object, Task>, object>>();
-            }
             _onCompleted.Push(new KeyValuePair<Func<object, Task>, object>(callback, state));
         }
 
@@ -96,26 +92,25 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
             }
 
             _completed = true;
-            var onCompleted = _onCompleted;
 
-            if (onCompleted == null || onCompleted.Count == 0)
+            if (_onCompleted.Count == 0)
             {
                 return Task.CompletedTask;
             }
 
-            return CompleteAsyncMayAwait(onCompleted);
+            return CompleteAsyncMayAwait();
         }
 
-        private Task CompleteAsyncMayAwait(Stack<KeyValuePair<Func<object, Task>, object>> onCompleted)
+        private Task CompleteAsyncMayAwait()
         {
-            while (onCompleted.TryPop(out var entry))
+            while (_onCompleted.TryPop(out var entry))
             {
                 try
                 {
                     var task = entry.Key.Invoke(entry.Value);
                     if (!task.IsCompletedSuccessfully)
                     {
-                        return CompleteAsyncAwaited(task, onCompleted);
+                        return CompleteAsyncAwaited(task);
                     }
                 }
                 catch (Exception ex)
@@ -127,7 +122,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
             return Task.CompletedTask;
         }
 
-        private async Task CompleteAsyncAwaited(Task currentTask, Stack<KeyValuePair<Func<object, Task>, object>> onCompleted)
+        private async Task CompleteAsyncAwaited(Task currentTask)
         {
             try
             {
@@ -138,7 +133,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
                 Logger.LogError(ex, "An error occurred running an IConnectionCompleteFeature.OnCompleted callback.");
             }
 
-            while (onCompleted.TryPop(out var entry))
+            while (_onCompleted.TryPop(out var entry))
             {
                 try
                 {