Eliminate HTTP2 HPack enumerator allocations (#19393)

src/Servers/Kestrel/Core/src/Internal/Http2/Http2FrameWriter.cs

@@ -28,6 +28,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2
         private readonly object _writeLock = new object();
         private readonly Http2Frame _outgoingFrame;
         private readonly HPackEncoder _hpackEncoder = new HPackEncoder();
+        private readonly Http2HeadersEnumerator _headersEnumerator = new Http2HeadersEnumerator();
         private readonly ConcurrentPipeWriter _outputWriter;
         private readonly ConnectionContext _connectionContext;
         private readonly Http2Connection _http2Connection;
@@ -160,7 +161,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2
             |                           Padding (*)                       ...
             +---------------------------------------------------------------+
         */
-        public void WriteResponseHeaders(int streamId, int statusCode, Http2HeadersFrameFlags headerFrameFlags, IHeaderDictionary headers)
+        public void WriteResponseHeaders(int streamId, int statusCode, Http2HeadersFrameFlags headerFrameFlags, HttpResponseHeaders headers)
         {
             lock (_writeLock)
             {
@@ -171,9 +172,10 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2
 
                 try
                 {
+                    _headersEnumerator.Initialize(headers);
                     _outgoingFrame.PrepareHeaders(headerFrameFlags, streamId);
                     var buffer = _headerEncodingBuffer.AsSpan();
-                    var done = _hpackEncoder.BeginEncode(statusCode, EnumerateHeaders(headers), buffer, out var payloadLength);
+                    var done = _hpackEncoder.BeginEncode(statusCode, _headersEnumerator, buffer, out var payloadLength);
                     FinishWritingHeaders(streamId, payloadLength, done);
                 }
                 catch (HPackEncodingException hex)
@@ -196,9 +198,10 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2
 
                 try
                 {
+                    _headersEnumerator.Initialize(headers);
                     _outgoingFrame.PrepareHeaders(Http2HeadersFrameFlags.END_STREAM, streamId);
                     var buffer = _headerEncodingBuffer.AsSpan();
-                    var done = _hpackEncoder.BeginEncode(EnumerateHeaders(headers), buffer, out var payloadLength);
+                    var done = _hpackEncoder.BeginEncode(_headersEnumerator, buffer, out var payloadLength);
                     FinishWritingHeaders(streamId, payloadLength, done);
                 }
                 catch (HPackEncodingException hex)
@@ -662,16 +665,5 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2
                 flowControl.Abort();
             }
         }
-
-        private static IEnumerable<KeyValuePair<string, string>> EnumerateHeaders(IHeaderDictionary headers)
-        {
-            foreach (var header in headers)
-            {
-                foreach (var value in header.Value)
-                {
-                    yield return new KeyValuePair<string, string>(header.Key, value);
-                }
-            }
-        }
     }
 }

src/Servers/Kestrel/Core/src/Internal/Http2/Http2HeaderEnumerator.cs

@@ -0,0 +1,161 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System.Collections;
+using System.Collections.Generic;
+using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http;
+using Microsoft.Extensions.Primitives;
+
+namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2
+{
+    internal sealed class Http2HeadersEnumerator : IEnumerator<KeyValuePair<string, string>>
+    {
+        private bool _isTrailers;
+        private HttpResponseHeaders.Enumerator _headersEnumerator;
+        private HttpResponseTrailers.Enumerator _trailersEnumerator;
+        private IEnumerator<KeyValuePair<string, StringValues>> _genericEnumerator;
+        private StringValues.Enumerator _stringValuesEnumerator;
+
+        public KeyValuePair<string, string> Current { get; private set; }
+        object IEnumerator.Current => Current;
+
+        public Http2HeadersEnumerator()
+        {
+        }
+
+        public void Initialize(HttpResponseHeaders headers)
+        {
+            _headersEnumerator = headers.GetEnumerator();
+            _trailersEnumerator = default;
+            _genericEnumerator = null;
+            _isTrailers = false;
+
+            _stringValuesEnumerator = default;
+            Current = default;
+        }
+
+        public void Initialize(HttpResponseTrailers headers)
+        {
+            _headersEnumerator = default;
+            _trailersEnumerator = headers.GetEnumerator();
+            _genericEnumerator = null;
+            _isTrailers = true;
+
+            _stringValuesEnumerator = default;
+            Current = default;
+        }
+
+        public void Initialize(IDictionary<string, StringValues> headers)
+        {
+            _headersEnumerator = default;
+            _trailersEnumerator = default;
+            _genericEnumerator = headers.GetEnumerator();
+            _isTrailers = false;
+
+            _stringValuesEnumerator = default;
+            Current = default;
+        }
+
+        public bool MoveNext()
+        {
+            if (MoveNextOnStringEnumerator())
+            {
+                return true;
+            }
+
+            if (!TryGetNextStringEnumerator(out _stringValuesEnumerator))
+            {
+                return false;
+            }
+
+            return MoveNextOnStringEnumerator();
+        }
+
+        private string GetCurrentKey()
+        {
+            if (_genericEnumerator != null)
+            {
+                return _genericEnumerator.Current.Key;
+            }
+            else if (_isTrailers)
+            {
+                return _trailersEnumerator.Current.Key;
+            }
+            else
+            {
+                return _headersEnumerator.Current.Key;
+            }
+        }
+
+        private bool MoveNextOnStringEnumerator()
+        {
+            var result = _stringValuesEnumerator.MoveNext();
+            Current = result ? new KeyValuePair<string, string>(GetCurrentKey(), _stringValuesEnumerator.Current) : default;
+            return result;
+        }
+
+        private bool TryGetNextStringEnumerator(out StringValues.Enumerator enumerator)
+        {
+            if (_genericEnumerator != null)
+            {
+                if (!_genericEnumerator.MoveNext())
+                {
+                    enumerator = default;
+                    return false;
+                }
+                else
+                {
+                    enumerator = _genericEnumerator.Current.Value.GetEnumerator();
+                    return true;
+                }
+            }
+            else if (_isTrailers)
+            {
+                if (!_trailersEnumerator.MoveNext())
+                {
+                    enumerator = default;
+                    return false;
+                }
+                else
+                {
+                    enumerator = _trailersEnumerator.Current.Value.GetEnumerator();
+                    return true;
+                }
+            }
+            else
+            {
+                if (!_headersEnumerator.MoveNext())
+                {
+                    enumerator = default;
+                    return false;
+                }
+                else
+                {
+                    enumerator = _headersEnumerator.Current.Value.GetEnumerator();
+                    return true;
+                }
+            }
+        }
+
+        public void Reset()
+        {
+            if (_genericEnumerator != null)
+            {
+                _genericEnumerator.Reset();
+            }
+            else if (_isTrailers)
+            {
+                _trailersEnumerator.Reset();
+            }
+            else
+            {
+                _headersEnumerator.Reset();
+            }
+            _stringValuesEnumerator = default;
+        }
+
+        public void Dispose()
+        {
+        }
+    }
+}

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

@@ -3,7 +3,12 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
+using System.Linq;
 using System.Net.Http.HPack;
+using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http;
+using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2;
+using Microsoft.Extensions.Primitives;
 using Xunit;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
@@ -94,11 +99,11 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             var length = 0;
             if (statusCode.HasValue)
             {
-                Assert.True(encoder.BeginEncode(statusCode.Value, headers, payload, out length));
+                Assert.True(encoder.BeginEncode(statusCode.Value, GetHeadersEnumerator(headers), payload, out length));
             }
             else
             {
-                Assert.True(encoder.BeginEncode(headers, payload, out length));
+                Assert.True(encoder.BeginEncode(GetHeadersEnumerator(headers), payload, out length));
             }
             Assert.Equal(expectedPayload.Length, length);
 
@@ -159,7 +164,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
 
             // When !exactSize, slices are one byte short of fitting the next header
             var sliceLength = expectedStatusCodePayload.Length + (exactSize ? 0 : expectedDateHeaderPayload.Length - 1);
-            Assert.False(encoder.BeginEncode(statusCode, headers, payload.Slice(offset, sliceLength), out var length));
+            Assert.False(encoder.BeginEncode(statusCode, GetHeadersEnumerator(headers), payload.Slice(offset, sliceLength), out var length));
             Assert.Equal(expectedStatusCodePayload.Length, length);
             Assert.Equal(expectedStatusCodePayload, payload.Slice(0, length).ToArray());
 
@@ -184,5 +189,16 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             Assert.Equal(expectedServerHeaderPayload.Length, length);
             Assert.Equal(expectedServerHeaderPayload, payload.Slice(offset, length).ToArray());
         }
+
+        private static Http2HeadersEnumerator GetHeadersEnumerator(IEnumerable<KeyValuePair<string, string>> headers)
+        {
+            var groupedHeaders = headers
+                .GroupBy(k => k.Key)
+                .ToDictionary(g => g.Key, g => new StringValues(g.Select(gg => gg.Value).ToArray()));
+
+            var enumerator = new Http2HeadersEnumerator();
+            enumerator.Initialize(groupedHeaders);
+            return enumerator;
+        }
     }
 }

src/Servers/Kestrel/perf/Kestrel.Performance/Http2FrameWriterBenchmark.cs

@@ -0,0 +1,59 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System.Buffers;
+using System.IO.Pipelines;
+using BenchmarkDotNet.Attributes;
+using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http;
+using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2;
+using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2.FlowControl;
+using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
+using Microsoft.Extensions.Logging.Abstractions;
+
+namespace Microsoft.AspNetCore.Server.Kestrel.Performance
+{
+    public class Http2FrameWriterBenchmark
+    {
+        private MemoryPool<byte> _memoryPool;
+        private Pipe _pipe;
+        private Http2FrameWriter _frameWriter;
+        private HttpResponseHeaders _responseHeaders;
+
+        [GlobalSetup]
+        public void GlobalSetup()
+        {
+            _memoryPool = SlabMemoryPoolFactory.Create();
+
+            var options = new PipeOptions(_memoryPool, readerScheduler: PipeScheduler.Inline, writerScheduler: PipeScheduler.Inline, useSynchronizationContext: false);
+            _pipe = new Pipe(options);
+
+            _frameWriter = new Http2FrameWriter(
+                new NullPipeWriter(),
+                connectionContext: null,
+                http2Connection: null,
+                new OutputFlowControl(initialWindowSize: uint.MaxValue),
+                timeoutControl: null,
+                minResponseDataRate: null,
+                "TestConnectionId",
+                _memoryPool,
+                new KestrelTrace(NullLogger.Instance));
+
+            _responseHeaders = new HttpResponseHeaders();
+            _responseHeaders.HeaderContentType = "application/json";
+            _responseHeaders.HeaderContentLength = "1024";
+        }
+
+        [Benchmark]
+        public void WriteResponseHeaders()
+        {
+            _frameWriter.WriteResponseHeaders(0, 200, Http2HeadersFrameFlags.END_HEADERS, _responseHeaders);
+        }
+
+        [GlobalCleanup]
+        public void Dispose()
+        {
+            _pipe.Writer.Complete();
+            _memoryPool?.Dispose();
+        }
+    }
+}

src/Servers/Kestrel/perf/Kestrel.Performance/Mocks/NullPipeWriter.cs

@@ -0,0 +1,42 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+using System.IO.Pipelines;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.AspNetCore.Server.Kestrel.Performance
+{
+    internal class NullPipeWriter : PipeWriter
+    {
+        private byte[] _buffer = new byte[1024 * 128];
+
+        public override void Advance(int bytes)
+        {
+        }
+
+        public override void CancelPendingFlush()
+        {
+        }
+
+        public override void Complete(Exception exception = null)
+        {
+        }
+
+        public override ValueTask<FlushResult> FlushAsync(CancellationToken cancellationToken = default)
+        {
+            return new ValueTask<FlushResult>(new FlushResult(false, true));
+        }
+
+        public override Memory<byte> GetMemory(int sizeHint = 0)
+        {
+            return _buffer;
+        }
+
+        public override Span<byte> GetSpan(int sizeHint = 0)
+        {
+            return _buffer;
+        }
+    }
+}

src/Servers/Kestrel/samples/http2cat/http2cat.csproj

@@ -38,4 +38,4 @@
   </ItemGroup>
 
 </Project>
- 
+

src/Servers/Kestrel/test/InMemory.FunctionalTests/Http2/Http2TestBase.cs

@@ -26,6 +26,7 @@ using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2.FlowControl;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
 using Microsoft.AspNetCore.Testing;
 using Microsoft.Extensions.Logging;
+using Microsoft.Extensions.Primitives;
 using Microsoft.Net.Http.Headers;
 using Moq;
 using Xunit;
@@ -504,7 +505,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             frame.PrepareHeaders(Http2HeadersFrameFlags.NONE, streamId);
 
             var buffer = _headerEncodingBuffer.AsSpan();
-            var done = _hpackEncoder.BeginEncode(headers, buffer, out var length);
+            var done = _hpackEncoder.BeginEncode(GetHeadersEnumerator(headers), buffer, out var length);
             frame.PayloadLength = length;
 
             if (done)
@@ -539,6 +540,17 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             return FlushAsync(writableBuffer);
         }
 
+        private static Http2HeadersEnumerator GetHeadersEnumerator(IEnumerable<KeyValuePair<string, string>> headers)
+        {
+            var groupedHeaders = headers
+                .GroupBy(k => k.Key)
+                .ToDictionary(g => g.Key, g => new StringValues(g.Select(gg => gg.Value).ToArray()));
+
+            var enumerator = new Http2HeadersEnumerator();
+            enumerator.Initialize(groupedHeaders);
+            return enumerator;
+        }
+
         /* https://tools.ietf.org/html/rfc7540#section-6.2
             +---------------+
             |Pad Length? (8)|
@@ -565,7 +577,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             extendedHeader[0] = padLength;
             var payload = buffer.Slice(extendedHeaderLength, buffer.Length - padLength - extendedHeaderLength);
 
-            _hpackEncoder.BeginEncode(headers, payload, out var length);
+            _hpackEncoder.BeginEncode(GetHeadersEnumerator(headers), payload, out var length);
             var padding = buffer.Slice(extendedHeaderLength + length, padLength);
             padding.Fill(0);
 
@@ -608,7 +620,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             extendedHeader[4] = priority;
             var payload = buffer.Slice(extendedHeaderLength);
 
-            _hpackEncoder.BeginEncode(headers, payload, out var length);
+            _hpackEncoder.BeginEncode(GetHeadersEnumerator(headers), payload, out var length);
 
             frame.PayloadLength = extendedHeaderLength + length;
 
@@ -655,7 +667,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             extendedHeader[5] = priority;
             var payload = buffer.Slice(extendedHeaderLength, buffer.Length - padLength - extendedHeaderLength);
 
-            _hpackEncoder.BeginEncode(headers, payload, out var length);
+            _hpackEncoder.BeginEncode(GetHeadersEnumerator(headers), payload, out var length);
             var padding = buffer.Slice(extendedHeaderLength + length, padLength);
             padding.Fill(0);
 
@@ -776,7 +788,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
 
             frame.PrepareHeaders(flags, streamId);
             var buffer = _headerEncodingBuffer.AsMemory();
-            var done = _hpackEncoder.BeginEncode(headers, buffer.Span, out var length);
+            var done = _hpackEncoder.BeginEncode(GetHeadersEnumerator(headers), buffer.Span, out var length);
             frame.PayloadLength = length;
 
             Http2FrameWriter.WriteHeader(frame, outputWriter);
@@ -869,7 +881,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
 
             frame.PrepareContinuation(flags, streamId);
             var buffer = _headerEncodingBuffer.AsMemory();
-            var done = _hpackEncoder.BeginEncode(headers, buffer.Span, out var length);
+            var done = _hpackEncoder.BeginEncode(GetHeadersEnumerator(headers), buffer.Span, out var length);
             frame.PayloadLength = length;
 
             Http2FrameWriter.WriteHeader(frame, outputWriter);

src/Shared/Http2cat/Http2Utilities.cs

@@ -20,6 +20,7 @@ using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
 using Microsoft.Extensions.Logging;
 using Microsoft.Net.Http.Headers;
+using IHttpHeadersHandler = System.Net.Http.IHttpHeadersHandler;
 
 namespace Microsoft.AspNetCore.Http2Cat
 {
@@ -235,7 +236,7 @@ namespace Microsoft.AspNetCore.Http2Cat
             frame.PrepareHeaders(Http2HeadersFrameFlags.NONE, streamId);
 
             var buffer = _headerEncodingBuffer.AsSpan();
-            var done = _hpackEncoder.BeginEncode(headers, buffer, out var length);
+            var done = _hpackEncoder.BeginEncode(headers.GetEnumerator(), buffer, out var length);
             frame.PayloadLength = length;
 
             if (done)
@@ -334,7 +335,7 @@ namespace Microsoft.AspNetCore.Http2Cat
             extendedHeader[0] = padLength;
             var payload = buffer.Slice(extendedHeaderLength, buffer.Length - padLength - extendedHeaderLength);
 
-            _hpackEncoder.BeginEncode(headers, payload, out var length);
+            _hpackEncoder.BeginEncode(headers.GetEnumerator(), payload, out var length);
             var padding = buffer.Slice(extendedHeaderLength + length, padLength);
             padding.Fill(0);
 
@@ -376,7 +377,7 @@ namespace Microsoft.AspNetCore.Http2Cat
             extendedHeader[4] = priority;
             var payload = buffer.Slice(extendedHeaderLength);
 
-            _hpackEncoder.BeginEncode(headers, payload, out var length);
+            _hpackEncoder.BeginEncode(headers.GetEnumerator(), payload, out var length);
 
             frame.PayloadLength = extendedHeaderLength + length;
 
@@ -422,7 +423,7 @@ namespace Microsoft.AspNetCore.Http2Cat
             extendedHeader[5] = priority;
             var payload = buffer.Slice(extendedHeaderLength, buffer.Length - padLength - extendedHeaderLength);
 
-            _hpackEncoder.BeginEncode(headers, payload, out var length);
+            _hpackEncoder.BeginEncode(headers.GetEnumerator(), payload, out var length);
             var padding = buffer.Slice(extendedHeaderLength + length, padLength);
             padding.Fill(0);
 
@@ -548,7 +549,7 @@ namespace Microsoft.AspNetCore.Http2Cat
 
             frame.PrepareHeaders(flags, streamId);
             var buffer = _headerEncodingBuffer.AsMemory();
-            var done = _hpackEncoder.BeginEncode(headers, buffer.Span, out var length);
+            var done = _hpackEncoder.BeginEncode(headers.GetEnumerator(), buffer.Span, out var length);
             frame.PayloadLength = length;
 
             WriteHeader(frame, outputWriter);
@@ -641,7 +642,7 @@ namespace Microsoft.AspNetCore.Http2Cat
 
             frame.PrepareContinuation(flags, streamId);
             var buffer = _headerEncodingBuffer.AsMemory();
-            var done = _hpackEncoder.BeginEncode(headers, buffer.Span, out var length);
+            var done = _hpackEncoder.BeginEncode(headers.GetEnumerator(), buffer.Span, out var length);
             frame.PayloadLength = length;
 
             WriteHeader(frame, outputWriter);

src/Shared/runtime/Http2/Hpack/HPackEncoder.cs

@@ -2,26 +2,30 @@
 // Licensed under the Apache License, Version 2.0.
 // See THIRD-PARTY-NOTICES.TXT in the project root for license information.
 
-using System.Collections.Generic;
 using System.Diagnostics;
+#if KESTREL
+using HeadersEnumerator = Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2.Http2HeadersEnumerator;
+#else
+using HeadersEnumerator = System.Collections.Generic.IEnumerator<System.Collections.Generic.KeyValuePair<string, string>>;
+#endif
 
 namespace System.Net.Http.HPack
 {
     internal class HPackEncoder
     {
-        private IEnumerator<KeyValuePair<string, string>> _enumerator;
+        private HeadersEnumerator _enumerator;
 
-        public bool BeginEncode(IEnumerable<KeyValuePair<string, string>> headers, Span<byte> buffer, out int length)
+        public bool BeginEncode(HeadersEnumerator enumerator, Span<byte> buffer, out int length)
         {
-            _enumerator = headers.GetEnumerator();
+            _enumerator = enumerator;
             _enumerator.MoveNext();
 
             return Encode(buffer, out length);
         }
 
-        public bool BeginEncode(int statusCode, IEnumerable<KeyValuePair<string, string>> headers, Span<byte> buffer, out int length)
+        public bool BeginEncode(int statusCode, HeadersEnumerator enumerator, Span<byte> buffer, out int length)
         {
-            _enumerator = headers.GetEnumerator();
+            _enumerator = enumerator;
             _enumerator.MoveNext();
 
             int statusCodeLength = EncodeStatusCode(statusCode, buffer);