Http3FramingStreamWriter perf improvements (10%-500%)

Author: ladeak

src/CHttpServer/CHttpServer/Http3/Http3FramingStreamWriter.cs

@@ -13,17 +13,22 @@ private readonly struct Segment()
 
         public byte[] Reference { get; init; } = [];
         public Memory<byte> Used { get; init; } = Memory<byte>.Empty;
+
+        public bool IsEmpty => Used.IsEmpty;
+
+        public bool IsAllocated => Reference.Length > 0;
     }
 
     private readonly byte[] _buffer = new byte[9];
     private readonly Lock _lockObject = new();
-    private readonly List<Segment> _segments = new List<Segment>(128) { new Segment() };
+    private readonly List<Segment> _segments = new List<Segment>(128);
     private readonly ArrayPool<byte> _memoryPool = memoryPool ?? ArrayPool<byte>.Shared;
     private Stream _responseStream = responseStream;
     private readonly byte _frameType = frameType;
     private CancellationTokenSource? _cts;
     private bool _isCompleted = false;
     private long _unflushedBytes;
+    private Segment _currentSegment = Segment.Empty;
     private TaskCompletionSource _tcs = new(TaskCreationOptions.RunContinuationsAsynchronously);
     private Func<CancellationToken, Task>? _onResponseStartingCallback = onResponseStartingCallback;
 
@@ -60,12 +65,11 @@ public void Reset(Stream responseStream, Func<CancellationToken, Task>? onRespon
     public override void Advance(int bytes)
     {
         ThrowIfCompleted();
-        var current = _segments[^1];
+        var current = _currentSegment;
         var usedLength = current.Used.Length;
         var available = current.Reference.Length - usedLength;
         ArgumentOutOfRangeException.ThrowIfLessThan(available, bytes);
-        current = current with { Used = current.Reference.AsMemory(0, usedLength + bytes) };
-        _segments[^1] = current;
+        _currentSegment = current with { Used = current.Reference.AsMemory(0, usedLength + bytes) };
         _unflushedBytes += bytes;
     }
 
@@ -111,23 +115,19 @@ public override Memory<byte> GetMemory(int sizeHint = 0)
     {
         ThrowIfCompleted();
         ArgumentOutOfRangeException.ThrowIfLessThan(sizeHint, 0);
-        var current = _segments[^1];
+        var current = _currentSegment;
         var available = current.Reference.Length - current.Used.Length;
         if (sizeHint <= available)
             return current.Reference.AsMemory(current.Used.Length);
 
-        var segment = new Segment() { Reference = _memoryPool.Rent(sizeHint) }; // Minimum size left for ArrayPool.
+        // Not enough memory left in the current segment.
+        if (!current.IsEmpty)
+            _segments.Add(current);
+        else if (current.IsAllocated)
+            _memoryPool.Return(current.Reference, true);
 
-        // If it is unused segment, but too small, return the currently rented array and replace the segment.
-        if (current.Used.Length == 0)
-        {
-            if (current.Reference.Length != 0)
-                _memoryPool.Return(current.Reference, true);
-            _segments[^1] = segment;
-        }
-        else
-            _segments.Add(segment);
-        return segment.Reference.AsMemory();
+        _currentSegment = new Segment() { Reference = _memoryPool.Rent(sizeHint) }; // Minimum size left for ArrayPool.
+        return _currentSegment.Reference.AsMemory();
     }
 
     public override Span<byte> GetSpan(int sizeHint = 0) => GetMemory(sizeHint).Span;
@@ -189,21 +189,20 @@ public override async ValueTask<FlushResult> FlushAsync(CancellationToken cancel
             var dataFrameHeaderLength = PrepareFrameHeader(_unflushedBytes);
             await _responseStream.WriteAsync(_buffer.AsMemory(0, dataFrameHeaderLength), localToken);
 
-            int i = 0;
-            var emptySegment = Segment.Empty;
-            for (; i < _segments.Count; i++)
+            // Shortcut: if flushed after every write, no need to address segments
+            if (_segments.Count == 0)
             {
-                var memory = _segments[i];
-                if (memory.Reference.Length == 0)
-                    break;
-                if (memory.Used.Length > 0)
-                    await _responseStream.WriteAsync(memory.Used, localToken);
-                _memoryPool.Return(memory.Reference, true);
-                _segments[i] = emptySegment;
+                if (!_currentSegment.IsEmpty)
+                    await _responseStream.WriteAsync(_currentSegment.Used, localToken);
+                if (_currentSegment.IsAllocated)
+                    _memoryPool.Return(_currentSegment.Reference, true);
+                _currentSegment = Segment.Empty;
+                _unflushedBytes = 0;
+                await _responseStream.FlushAsync();
+                return new FlushResult(false, false);
             }
-            _unflushedBytes = 0;
-            await _responseStream.FlushAsync(localToken);
-            return new FlushResult(isCanceled: false, isCompleted: false);
+
+            return await FlushAllSegmentsAsync(localToken);
         }
         catch (OperationCanceledException)
         {
@@ -233,6 +232,27 @@ public override async ValueTask<FlushResult> FlushAsync(CancellationToken cancel
         }
     }
 
+    private async Task<FlushResult> FlushAllSegmentsAsync(CancellationToken localToken)
+    {
+        if (!_currentSegment.IsEmpty)
+            _segments.Add(_currentSegment);
+        else if (_currentSegment.IsAllocated)
+            _memoryPool.Return(_currentSegment.Reference, true);
+        _currentSegment = Segment.Empty;
+
+        for (int i = 0; i < _segments.Count; i++)
+        {
+            var memory = _segments[i];
+            if (!memory.IsEmpty)
+                await _responseStream.WriteAsync(memory.Used, localToken);
+            _memoryPool.Return(memory.Reference, true);
+        }
+        _segments.Clear();
+        _unflushedBytes = 0;
+        await _responseStream.FlushAsync(localToken);
+        return new FlushResult(isCanceled: false, isCompleted: false);
+    }
+
     private void Flush()
     {
         if (_unflushedBytes == 0)
@@ -245,19 +265,40 @@ private void Flush()
         }
         var dataFrameHeaderLength = PrepareFrameHeader(_unflushedBytes);
         _responseStream.Write(_buffer.AsSpan(0, dataFrameHeaderLength));
+
+        // Shortcut: if flushed after every write, no need to address segments
+        if (_segments.Count == 0)
+        {
+            if (!_currentSegment.IsEmpty)
+                _responseStream.Write(_currentSegment.Used.Span);
+            if (_currentSegment.IsAllocated)
+                _memoryPool.Return(_currentSegment.Reference, true);
+            _currentSegment = Segment.Empty;
+            _unflushedBytes = 0;
+            _responseStream.Flush();
+            return;
+        }
+
+        FlushAllSegments();
+    }
+
+    private void FlushAllSegments()
+    {
+        if (!_currentSegment.IsEmpty)
+            _segments.Add(_currentSegment);
+        else if (_currentSegment.IsAllocated)
+            _memoryPool.Return(_currentSegment.Reference, true);
+        _currentSegment = Segment.Empty;
+
         var source = CollectionsMarshal.AsSpan(_segments);
-        int i = 0;
-        var emptySegment = Segment.Empty;
-        for (; i < _segments.Count; i++)
+        for (int i = 0; i < _segments.Count; i++)
         {
             ref var memory = ref source[i];
-            if (memory.Reference.Length == 0)
-                break;
             if (memory.Used.Length > 0)
                 _responseStream.Write(memory.Used.Span);
             _memoryPool.Return(memory.Reference, true);
-            source[i] = emptySegment;
         }
+        _segments.Clear();
         _unflushedBytes = 0;
         _responseStream.Flush();
     }
@@ -290,7 +331,9 @@ private void ClearSegments(Span<Segment> source)
         }
         _unflushedBytes = 0;
         _segments.Clear();
-        _segments.Add(Segment.Empty);
+        if (_currentSegment.IsAllocated)
+            _memoryPool.Return(_currentSegment.Reference);
+        _currentSegment = Segment.Empty;
     }
 
     private void ThrowIfCompleted()

tests/CHttp.Benchmarks/Program.cs

@@ -206,3 +206,24 @@ public bool EncodeLong()
         return QPackIntegerEncoder.TryEncode(buffer, Value, 7, out _);
     }
 }
+
+[SimpleJob, DisassemblyDiagnoser]
+public class Http3FramingStreamWriterBenchmarks
+{
+    private Http3FramingStreamWriter _writer = new Http3FramingStreamWriter(Stream.Null, 0);
+
+    [Benchmark]
+    public async Task FlushAsync()
+    {
+        Span<byte> data = [0, 1, 2, 3, 4];
+        var memory0 = _writer.GetMemory(data.Length);
+        data.CopyTo(memory0.Span);
+        _writer.Advance(data.Length);
+
+        var memory1 = _writer.GetMemory(data.Length);
+        data.CopyTo(memory1.Span);
+        _writer.Advance(data.Length);
+
+        await _writer.FlushAsync();
+    }
+}

tests/CHttpServer.Tests/Http3/Http3FramingStreamWriterTests.cs

@@ -1,4 +1,5 @@
 using System.Buffers;
+using System.Formats.Asn1;
 using System.Net.Quic;
 using System.Runtime.InteropServices;
 using CHttpServer.Http3;
@@ -362,6 +363,24 @@ public async Task WriteAsync_Into_ClosedStream()
         Assert.Equal(0, arrayPool.OutstandingBytes);
     }
 
+    [Fact]
+    public async Task GetMemory_WhenNoDataAvailable_InSegment()
+    {
+        var ms = new MemoryStream();
+        var sut = new Http3FramingStreamWriter(ms, 0);
+        Span<byte> data = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+        var memory0 = sut.GetMemory(data.Length);
+        data.CopyTo(memory0.Span);
+        sut.Advance(data.Length);
+
+        var memory1 = sut.GetMemory(data.Length);
+        data.CopyTo(memory1.Span);
+        sut.Advance(data.Length);
+
+        await sut.FlushAsync(TestContext.Current.CancellationToken);
+        Assert.Equal(24, ms.Length);
+    }
+
     [Fact]
     public async Task GetMemory_WhenDataAvailable_InSegment()
     {
@@ -581,58 +600,58 @@ public async Task WriteAsync_Writes_FrameType()
     }
 
     private class WaitBeforeWriteStream(TaskCompletionSource tcs) : MemoryStream
-    {
-        public int WrittenBytes { get; private set; }
+{
+    public int WrittenBytes { get; private set; }
 
-        public override async ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
-        {
-            await tcs.Task;
-            await base.WriteAsync(buffer, cancellationToken);
-            WrittenBytes += buffer.Length;
-        }
+    public override async ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
+    {
+        await tcs.Task;
+        await base.WriteAsync(buffer, cancellationToken);
+        WrittenBytes += buffer.Length;
     }
+}
 
-    private class WaitAfterWriteStream(TaskCompletionSource tcs) : MemoryStream
-    {
-        public int WrittenBytes { get; private set; }
+private class WaitAfterWriteStream(TaskCompletionSource tcs) : MemoryStream
+{
+    public int WrittenBytes { get; private set; }
 
-        public override async ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
-        {
-            await base.WriteAsync(buffer, cancellationToken);
-            WrittenBytes += buffer.Length;
-            await tcs.Task;
-        }
+    public override async ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
+    {
+        await base.WriteAsync(buffer, cancellationToken);
+        WrittenBytes += buffer.Length;
+        await tcs.Task;
     }
+}
 
-    private class TestArrayPool : ArrayPool<byte>
-    {
-        private readonly ArrayPool<byte> _internalPool;
-        private readonly List<byte[]> _rentedArrays = new();
+private class TestArrayPool : ArrayPool<byte>
+{
+    private readonly ArrayPool<byte> _internalPool;
+    private readonly List<byte[]> _rentedArrays = new();
 
-        public TestArrayPool()
-        {
-            _internalPool = Shared;
-        }
+    public TestArrayPool()
+    {
+        _internalPool = Shared;
+    }
 
-        public TestArrayPool(int maxArrayLength, int maxArrayPerBucket)
-        {
-            _internalPool = Create(maxArrayLength, maxArrayPerBucket);
-        }
+    public TestArrayPool(int maxArrayLength, int maxArrayPerBucket)
+    {
+        _internalPool = Create(maxArrayLength, maxArrayPerBucket);
+    }
 
-        public int OutstandingBytes => _rentedArrays.Sum(x => x.Length);
+    public int OutstandingBytes => _rentedArrays.Sum(x => x.Length);
 
-        public override byte[] Rent(int minimumLength)
-        {
-            var buffer = _internalPool.Rent(minimumLength);
-            _rentedArrays.Add(buffer);
-            return buffer;
-        }
+    public override byte[] Rent(int minimumLength)
+    {
+        var buffer = _internalPool.Rent(minimumLength);
+        _rentedArrays.Add(buffer);
+        return buffer;
+    }
 
-        public override void Return(byte[] array, bool clearArray = false)
-        {
-            _internalPool.Return(array, clearArray);
-            _rentedArrays.Remove(array);
-        }
+    public override void Return(byte[] array, bool clearArray = false)
+    {
+        _internalPool.Return(array, clearArray);
+        _rentedArrays.Remove(array);
     }
 }
+}