Eliminate HTTP2 HPack enumerator allocations

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

@@ -28,6 +28,7 @@ internal class Http2FrameWriter
         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 @@ public ValueTask<FlushResult> Write100ContinueAsync(int streamId)
             |                           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 @@ public void WriteResponseHeaders(int streamId, int statusCode, Http2HeadersFrame
 
                 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 @@ public ValueTask<FlushResult> WriteResponseTrailers(int streamId, HttpResponseTr
 
                 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 @@ public void AbortPendingStreamDataWrites(StreamOutputFlowControl flowControl)
                 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 @@ public void EncodesHeadersInSinglePayloadWhenSpaceAvailable(KeyValuePair<string,
             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 @@ public void EncodesHeadersInMultiplePayloadsWhenSpaceNotAvailable(bool exactSize
 
             // 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 @@ public void EncodesHeadersInMultiplePayloadsWhenSpaceNotAvailable(bool exactSize
             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>
- 
+