Reduce HttpRequestHeader enumerator allocations (#32236)

Author: benaadams

src/Servers/Kestrel/Core/src/Internal/Http/HttpRequestHeaders.cs

@@ -16,6 +16,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
 {
     internal sealed partial class HttpRequestHeaders : HttpHeaders
     {
+        private EnumeratorCache? _enumeratorCache;
         private long _previousBits;
 
         public bool ReuseHeaderValues { get; set; }
@@ -65,6 +66,7 @@ protected override void ClearFast()
             // Clear ContentLength and any unknown headers as we will never reuse them
             _contentLength = null;
             MaybeUnknown?.Clear();
+            _enumeratorCache?.Reset();
         }
 
         private static long ParseContentLength(string value)
@@ -148,7 +150,73 @@ public Enumerator GetEnumerator()
 
         protected override IEnumerator<KeyValuePair<string, StringValues>> GetEnumeratorFast()
         {
-            return GetEnumerator();
+            // Get or create the cache.
+            var cache = _enumeratorCache ??= new();
+
+            EnumeratorBox enumerator;
+            if (cache.CachedEnumerator is not null)
+            {
+                // Previous enumerator, reuse that one.
+                enumerator = cache.InUseEnumerator = cache.CachedEnumerator;
+                // Set previous to null so if there is a second enumerator call
+                // during the same request it doesn't get the same one.
+                cache.CachedEnumerator = null;
+            }
+            else
+            {
+                // Create new enumerator box and store as in use.
+                enumerator = cache.InUseEnumerator = new();
+            }
+
+            // Set the underlying struct enumerator to a new one.
+            enumerator.Enumerator = new Enumerator(this);
+            return enumerator;
+        }
+
+        private class EnumeratorCache
+        {
+            /// <summary>
+            /// Enumerator created from previous request
+            /// </summary>
+            public EnumeratorBox? CachedEnumerator { get; set; }
+            /// <summary>
+            /// Enumerator used on this request
+            /// </summary>
+            public EnumeratorBox? InUseEnumerator { get; set; }
+
+            /// <summary>
+            /// Moves InUseEnumerator to CachedEnumerator
+            /// </summary>
+            public void Reset()
+            {
+                var enumerator = InUseEnumerator;
+                if (enumerator is not null)
+                {
+                    InUseEnumerator = null;
+                    enumerator.Enumerator = default;
+                    CachedEnumerator = enumerator;
+                }
+            }
+        }
+
+        /// <summary>
+        /// Strong box enumerator for the IEnumerator interface to cache and amortizate the
+        /// IEnumerator allocations across requests if the header collection is commonly
+        /// enumerated for forwarding in a reverse-proxy type situation.
+        /// </summary>
+        private class EnumeratorBox : IEnumerator<KeyValuePair<string, StringValues>>
+        {
+            public Enumerator Enumerator;
+
+            public KeyValuePair<string, StringValues> Current => Enumerator.Current;
+
+            public bool MoveNext() => Enumerator.MoveNext();
+
+            object IEnumerator.Current => Current;
+
+            public void Dispose() { }
+
+            public void Reset() => throw new NotSupportedException();
         }
 
         public partial struct Enumerator : IEnumerator<KeyValuePair<string, StringValues>>

src/Servers/Kestrel/Core/test/HttpRequestHeadersTests.cs

@@ -99,9 +99,64 @@ public void SameExceptionThrownForMissingKey()
         }
 
         [Fact]
-        public void EntriesCanBeEnumerated()
+        public void EntriesCanBeEnumeratedAfterResets()
+        {
+            HttpRequestHeaders headers = new HttpRequestHeaders();
+
+            EnumerateEntries((IHeaderDictionary)headers);
+            headers.Reset();
+            EnumerateEntries((IDictionary<string, StringValues>)headers);
+            headers.Reset();
+            EnumerateEntries((IHeaderDictionary)headers);
+            headers.Reset();
+            EnumerateEntries((IDictionary<string, StringValues>)headers);
+        }
+
+        [Fact]
+        public void EnumeratorNotReusedBeforeReset()
+        {
+            HttpRequestHeaders headers = new HttpRequestHeaders();
+            IEnumerable<KeyValuePair<string, StringValues>> enumerable = headers;
+
+            var enumerator0 = enumerable.GetEnumerator();
+            var enumerator1 = enumerable.GetEnumerator();
+
+            Assert.NotSame(enumerator0, enumerator1);
+        }
+
+        [Fact]
+        public void EnumeratorReusedAfterReset()
+        {
+            HttpRequestHeaders headers = new HttpRequestHeaders();
+            IEnumerable<KeyValuePair<string, StringValues>> enumerable = headers;
+
+            var enumerator0 = enumerable.GetEnumerator();
+            headers.Reset();
+            var enumerator1 = enumerable.GetEnumerator();
+
+            Assert.Same(enumerator0, enumerator1);
+        }
+
+        private static void EnumerateEntries(IHeaderDictionary headers)
+        {
+            var v1 = new[] { "localhost" };
+            var v2 = new[] { "0" };
+            var v3 = new[] { "value" };
+            headers.Host = v1;
+            headers.ContentLength = 0;
+            headers["custom"] = v3;
+
+            Assert.Equal(
+                new[] {
+                    new KeyValuePair<string, StringValues>("Host", v1),
+                    new KeyValuePair<string, StringValues>("Content-Length", v2),
+                    new KeyValuePair<string, StringValues>("custom", v3),
+                },
+                headers);
+        }
+
+        private static void EnumerateEntries(IDictionary<string, StringValues> headers)
         {
-            IDictionary<string, StringValues> headers = new HttpRequestHeaders();
             var v1 = new[] { "localhost" };
             var v2 = new[] { "0" };
             var v3 = new[] { "value" };