Wrap the PipeWriter used by Kestrel's output writing logic (#12081)

Author: halter73

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

@@ -12,6 +12,7 @@
 using Microsoft.AspNetCore.Internal;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Features;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
+using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure.PipeWriterHelpers;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
 {
@@ -45,7 +46,7 @@ internal class Http1OutputProducer : IHttpOutputProducer, IHttpOutputAborter, ID
         private long _unflushedBytes;
         private int _currentMemoryPrefixBytes;
 
-        private readonly PipeWriter _pipeWriter;
+        private readonly ConcurrentPipeWriter _pipeWriter;
         private IMemoryOwner<byte> _fakeMemoryOwner;
 
         // Chunked responses need to be treated uniquely when using GetMemory + Advance.
@@ -81,18 +82,16 @@ public Http1OutputProducer(
             IHttpMinResponseDataRateFeature minResponseDataRateFeature,
             MemoryPool<byte> memoryPool)
         {
-            _pipeWriter = pipeWriter;
+            // Allow appending more data to the PipeWriter when a flush is pending.
+            _pipeWriter = new ConcurrentPipeWriter(pipeWriter, memoryPool);
             _connectionId = connectionId;
             _connectionContext = connectionContext;
             _log = log;
             _minResponseDataRateFeature = minResponseDataRateFeature;
-            _flusher = new TimingPipeFlusher(pipeWriter, timeoutControl, log);
+            _flusher = new TimingPipeFlusher(_pipeWriter, timeoutControl, log);
             _memoryPool = memoryPool;
         }
 
-        // For tests
-        internal PipeWriter PipeWriter => _pipeWriter;
-
         public Task WriteDataAsync(ReadOnlySpan<byte> buffer, CancellationToken cancellationToken = default)
         {
             if (cancellationToken.IsCancellationRequested)
@@ -408,6 +407,8 @@ public void Dispose()
         {
             lock (_contextLock)
             {
+                _pipeWriter.Abort();
+
                 if (_fakeMemoryOwner != null)
                 {
                     _fakeMemoryOwner.Dispose();

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

@@ -87,6 +87,7 @@ public Http2Connection(HttpConnectionContext context)
                 context.TimeoutControl,
                 httpLimits.MinResponseDataRate,
                 context.ConnectionId,
+                context.MemoryPool,
                 context.ServiceContext.Log);
 
             _hpackDecoder = new HPackDecoder(http2Limits.HeaderTableSize, http2Limits.MaxRequestHeaderFieldSize);

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

@@ -15,6 +15,7 @@
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2.FlowControl;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2.HPack;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
+using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure.PipeWriterHelpers;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2
 {
@@ -26,7 +27,7 @@ internal class Http2FrameWriter
         private readonly object _writeLock = new object();
         private readonly Http2Frame _outgoingFrame;
         private readonly HPackEncoder _hpackEncoder = new HPackEncoder();
-        private readonly PipeWriter _outputWriter;
+        private readonly ConcurrentPipeWriter _outputWriter;
         private readonly ConnectionContext _connectionContext;
         private readonly Http2Connection _http2Connection;
         private readonly OutputFlowControl _connectionOutputFlowControl;
@@ -51,9 +52,11 @@ public Http2FrameWriter(
             ITimeoutControl timeoutControl,
             MinDataRate minResponseDataRate,
             string connectionId,
+            MemoryPool<byte> memoryPool,
             IKestrelTrace log)
         {
-            _outputWriter = outputPipeWriter;
+            // Allow appending more data to the PipeWriter when a flush is pending.
+            _outputWriter = new ConcurrentPipeWriter(outputPipeWriter, memoryPool);
             _connectionContext = connectionContext;
             _http2Connection = http2Connection;
             _connectionOutputFlowControl = connectionOutputFlowControl;
@@ -89,6 +92,7 @@ public void Complete()
 
                 _completed = true;
                 _connectionOutputFlowControl.Abort();
+                _outputWriter.Abort();
             }
         }
 

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

@@ -12,6 +12,7 @@
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2.FlowControl;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
+using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure.PipeWriterHelpers;
 using Microsoft.Extensions.Logging;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2
@@ -29,7 +30,8 @@ internal class Http2OutputProducer : IHttpOutputProducer, IHttpOutputAborter
         private readonly MemoryPool<byte> _memoryPool;
         private readonly Http2Stream _stream;
         private readonly object _dataWriterLock = new object();
-        private readonly Pipe _dataPipe;
+        private readonly PipeWriter _pipeWriter;
+        private readonly PipeReader _pipeReader;
         private readonly ValueTask<FlushResult> _dataWriteProcessingTask;
         private bool _startedWritingDataFrames;
         private bool _completed;
@@ -43,7 +45,6 @@ public Http2OutputProducer(
             int streamId,
             Http2FrameWriter frameWriter,
             StreamOutputFlowControl flowControl,
-            ITimeoutControl timeoutControl,
             MemoryPool<byte> pool,
             Http2Stream stream,
             IKestrelTrace log)
@@ -55,8 +56,14 @@ public Http2OutputProducer(
             _stream = stream;
             _log = log;
 
-            _dataPipe = CreateDataPipe(pool);
-            _flusher = new TimingPipeFlusher(_dataPipe.Writer, timeoutControl, log);
+            var pipe = CreateDataPipe(pool);
+
+            _pipeWriter = new ConcurrentPipeWriter(pipe.Writer, pool);
+            _pipeReader = pipe.Reader;
+
+            // No need to pass in timeoutControl here, since no minDataRates are passed to the TimingPipeFlusher.
+            // The minimum output data rate is enforced at the connection level by Http2FrameWriter.
+            _flusher = new TimingPipeFlusher(_pipeWriter, timeoutControl: null, log);
             _dataWriteProcessingTask = ProcessDataWrites();
         }
 
@@ -193,7 +200,7 @@ public Task WriteDataAsync(ReadOnlySpan<byte> data, CancellationToken cancellati
 
                 _startedWritingDataFrames = true;
 
-                _dataPipe.Writer.Write(data);
+                _pipeWriter.Write(data);
                 return _flusher.FlushAsync(this, cancellationToken).GetAsTask();
             }
         }
@@ -210,7 +217,7 @@ public ValueTask<FlushResult> WriteStreamSuffixAsync()
                 _completed = true;
                 _suffixSent = true;
 
-                _dataPipe.Writer.Complete();
+                _pipeWriter.Complete();
                 return _dataWriteProcessingTask;
             }
         }
@@ -239,7 +246,7 @@ public void Advance(int bytes)
 
                 _startedWritingDataFrames = true;
 
-                _dataPipe.Writer.Advance(bytes);
+                _pipeWriter.Advance(bytes);
             }
         }
 
@@ -254,7 +261,7 @@ public Span<byte> GetSpan(int sizeHint = 0)
                     return GetFakeMemory(sizeHint).Span;
                 }
 
-                return _dataPipe.Writer.GetSpan(sizeHint);
+                return _pipeWriter.GetSpan(sizeHint);
             }
         }
 
@@ -269,7 +276,7 @@ public Memory<byte> GetMemory(int sizeHint = 0)
                     return GetFakeMemory(sizeHint);
                 }
 
-                return _dataPipe.Writer.GetMemory(sizeHint);
+                return _pipeWriter.GetMemory(sizeHint);
             }
         }
 
@@ -282,7 +289,7 @@ public void CancelPendingFlush()
                     return;
                 }
 
-                _dataPipe.Writer.CancelPendingFlush();
+                _pipeWriter.CancelPendingFlush();
             }
         }
 
@@ -306,7 +313,7 @@ public ValueTask<FlushResult> WriteDataToPipeAsync(ReadOnlySpan<byte> data, Canc
 
                 _startedWritingDataFrames = true;
 
-                _dataPipe.Writer.Write(data);
+                _pipeWriter.Write(data);
                 return _flusher.FlushAsync(this, cancellationToken);
             }
         }
@@ -345,7 +352,7 @@ public void Stop()
                 // Complete with an exception to prevent an end of stream data frame from being sent without an
                 // explicit call to WriteStreamSuffixAsync. ConnectionAbortedExceptions are swallowed, so the
                 // message doesn't matter
-                _dataPipe.Writer.Complete(new OperationCanceledException());
+                _pipeWriter.Complete(new OperationCanceledException());
 
                 _frameWriter.AbortPendingStreamDataWrites(_flowControl);
             }
@@ -364,7 +371,7 @@ private async ValueTask<FlushResult> ProcessDataWrites()
 
                 do
                 {
-                    readResult = await _dataPipe.Reader.ReadAsync();
+                    readResult = await _pipeReader.ReadAsync();
 
                     if (readResult.IsCompleted && _stream.ResponseTrailers?.Count > 0)
                     {
@@ -393,7 +400,7 @@ private async ValueTask<FlushResult> ProcessDataWrites()
                         flushResult = await _frameWriter.WriteDataAsync(_streamId, _flowControl, readResult.Buffer, endStream: readResult.IsCompleted);
                     }
 
-                    _dataPipe.Reader.AdvanceTo(readResult.Buffer.End);
+                    _pipeReader.AdvanceTo(readResult.Buffer.End);
                 } while (!readResult.IsCompleted);
             }
             catch (OperationCanceledException)
@@ -405,7 +412,7 @@ private async ValueTask<FlushResult> ProcessDataWrites()
                 _log.LogCritical(ex, nameof(Http2OutputProducer) + "." + nameof(ProcessDataWrites) + " observed an unexpected exception.");
             }
 
-            _dataPipe.Reader.Complete();
+            _pipeReader.Complete();
 
             return flushResult;
 

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

@@ -51,7 +51,6 @@ public Http2Stream(Http2StreamContext context)
                 context.StreamId,
                 context.FrameWriter,
                 _outputFlowControl,
-                context.TimeoutControl,
                 context.MemoryPool,
                 this,
                 context.ServiceContext.Log);

src/Servers/Kestrel/Core/src/Internal/Infrastructure/PipeWriterHelpers/BufferSegment.cs

@@ -0,0 +1,130 @@
+// 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.Diagnostics;
+using System.Runtime.CompilerServices;
+
+namespace System.IO.Pipelines
+{
+    // Copied from https://github.com/dotnet/corefx/blob/de3902bb56f1254ec1af4bf7d092fc2c048734cc/src/System.IO.Pipelines/src/System/IO/Pipelines/BufferSegment.cs
+    internal sealed class BufferSegment : ReadOnlySequenceSegment<byte>
+    {
+        private object _memoryOwner;
+        private BufferSegment _next;
+        private int _end;
+
+        /// <summary>
+        /// The End represents the offset into AvailableMemory where the range of "active" bytes ends. At the point when the block is leased
+        /// the End is guaranteed to be equal to Start. The value of Start may be assigned anywhere between 0 and
+        /// Buffer.Length, and must be equal to or less than End.
+        /// </summary>
+        public int End
+        {
+            get => _end;
+            set
+            {
+                Debug.Assert(value <= AvailableMemory.Length);
+
+                _end = value;
+                Memory = AvailableMemory.Slice(0, value);
+            }
+        }
+
+        /// <summary>
+        /// Reference to the next block of data when the overall "active" bytes spans multiple blocks. At the point when the block is
+        /// leased Next is guaranteed to be null. Start, End, and Next are used together in order to create a linked-list of discontiguous
+        /// working memory. The "active" memory is grown when bytes are copied in, End is increased, and Next is assigned. The "active"
+        /// memory is shrunk when bytes are consumed, Start is increased, and blocks are returned to the pool.
+        /// </summary>
+        public BufferSegment NextSegment
+        {
+            get => _next;
+            set
+            {
+                Next = value;
+                _next = value;
+            }
+        }
+
+        public void SetOwnedMemory(IMemoryOwner<byte> memoryOwner)
+        {
+            _memoryOwner = memoryOwner;
+            AvailableMemory = memoryOwner.Memory;
+        }
+
+        public void SetOwnedMemory(byte[] arrayPoolBuffer)
+        {
+            _memoryOwner = arrayPoolBuffer;
+            AvailableMemory = arrayPoolBuffer;
+        }
+
+        public void SetUnownedMemory(Memory<byte> memory)
+        {
+            AvailableMemory = memory;
+        }
+
+        public void ResetMemory()
+        {
+            if (_memoryOwner is IMemoryOwner<byte> owner)
+            {
+                owner.Dispose();
+            }
+            else if (_memoryOwner is byte[] array)
+            {
+                ArrayPool<byte>.Shared.Return(array);
+            }
+
+            // Order of below field clears is significant as it clears in a sequential order
+            // https://github.com/dotnet/corefx/pull/35256#issuecomment-462800477
+            Next = null;
+            RunningIndex = 0;
+            Memory = default;
+            _memoryOwner = null;
+            _next = null;
+            _end = 0;
+            AvailableMemory = default;
+        }
+
+        // Exposed for testing
+        internal object MemoryOwner => _memoryOwner;
+
+        public Memory<byte> AvailableMemory { get; private set; }
+
+        public int Length => End;
+
+        public int WritableBytes
+        {
+            [MethodImpl(MethodImplOptions.AggressiveInlining)]
+            get => AvailableMemory.Length - End;
+        }
+
+        public void SetNext(BufferSegment segment)
+        {
+            Debug.Assert(segment != null);
+            Debug.Assert(Next == null);
+
+            NextSegment = segment;
+
+            segment = this;
+
+            while (segment.Next != null)
+            {
+                segment.NextSegment.RunningIndex = segment.RunningIndex + segment.Length;
+                segment = segment.NextSegment;
+            }
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        internal static long GetLength(BufferSegment startSegment, int startIndex, BufferSegment endSegment, int endIndex)
+        {
+            return (endSegment.RunningIndex + (uint)endIndex) - (startSegment.RunningIndex + (uint)startIndex);
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        internal static long GetLength(long startPosition, BufferSegment endSegment, int endIndex)
+        {
+            return (endSegment.RunningIndex + (uint)endIndex) - startPosition;
+        }
+    }
+}