Expose response PipeWriter in Kestrel

Author: jkotalik

src/Http/Http.Abstractions/src/Extensions/HttpResponseWritingExtensions.cs

@@ -61,7 +61,7 @@ public static class HttpResponseWritingExtensions
             }
 
             byte[] data = encoding.GetBytes(text);
-            return response.Body.WriteAsync(data, 0, data.Length, cancellationToken);
+            return response.BodyPipe.WriteAsync(new Memory<byte>(data, 0, data.Length), cancellationToken).AsTask();
         }
     }
-}
+}

src/Http/Http.Features/src/FeatureReferences.cs

@@ -95,4 +95,4 @@ private TFeature UpdateCached<TFeature, TState>(ref TFeature cached, TState stat
         public TFeature Fetch<TFeature>(ref TFeature cached, Func<IFeatureCollection, TFeature> factory)
             where TFeature : class => Fetch(ref cached, Collection, factory);
     }
-}
+}

src/Http/Http/src/StreamPipeWriter.cs

@@ -93,7 +93,7 @@ public override Memory<byte> GetMemory(int sizeHint = 0)
         {
             EnsureCapacity(sizeHint);
 
-            return _currentSegment;
+            return _currentSegment.Slice(_position);
         }
 
         /// <inheritdoc />
@@ -265,7 +265,7 @@ private void AddSegment(int sizeHint = 0)
             }
 
             // Get a new buffer using the minimum segment size, unless the size hint is larger than a single segment.
-            _currentSegmentOwner = _pool.Rent(Math.Max(_minimumSegmentSize, sizeHint));
+            _currentSegmentOwner = _pool.Rent(Math.Min(Math.Max(_minimumSegmentSize, sizeHint), _pool.MaxBufferSize));
             _currentSegment = _currentSegmentOwner.Memory;
             _position = 0;
         }

src/Http/Http/src/WriteOnlyPipeStream.cs

@@ -60,6 +60,8 @@ public override int ReadTimeout
             set => throw new NotSupportedException();
         }
 
+        public PipeWriter InnerPipeWriter => _pipeWriter;
+
         /// <inheritdoc />
         public override int Read(byte[] buffer, int offset, int count)
             => throw new NotSupportedException();

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

@@ -44,13 +44,14 @@ public static int BeginChunkBytes(int dataCount, Span<byte> span)
             return count;
         }
 
-        internal static void WriteBeginChunkBytes(this ref BufferWriter<PipeWriter> start, int dataCount)
+        internal static int WriteBeginChunkBytes(this ref BufferWriter<PipeWriter> start, int dataCount)
         {
             // 10 bytes is max length + \r\n
             start.Ensure(10);
 
             var count = BeginChunkBytes(dataCount, start.Span);
             start.Advance(count);
+            return count;
         }
 
         internal static void WriteEndChunkBytes(this ref BufferWriter<PipeWriter> start)

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

@@ -3,6 +3,7 @@
 
 using System;
 using System.Buffers;
+using System.Diagnostics;
 using System.IO.Pipelines;
 using System.Text;
 using System.Threading;
@@ -26,15 +27,26 @@ public class Http1OutputProducer : IHttpOutputProducer, IHttpOutputAborter, IDis
         private readonly IHttpMinResponseDataRateFeature _minResponseDataRateFeature;
         private readonly TimingPipeFlusher _flusher;
 
-        // This locks access to to all of the below fields
+        // This locks access to all of the below fields
         private readonly object _contextLock = new object();
 
         private bool _completed = false;
         private bool _aborted;
         private long _unflushedBytes;
-
+        private bool _autoChunk;
         private readonly PipeWriter _pipeWriter;
 
+        private const int BeginChunkLengthMax = 5;
+        private const int EndChunkLength = 2;
+
+        // Chunked responses need to be treated uniquely when using GetMemory + Advance.
+        // We need to know the size of the data written to the chunk before calling Advance on the
+        // PipeWriter, meaning we internally track how far we have advanced through a current chunk (_advancedBytesForChunk).
+        // Once write or flush is called, we modify the _currentChunkMemory to prepend the size of data written
+        // and append the end terminator.
+        private int _advancedBytesForChunk;
+        private Memory<byte> _currentChunkMemory;
+
         public Http1OutputProducer(
             PipeWriter pipeWriter,
             string connectionId,
@@ -58,28 +70,92 @@ public Task WriteDataAsync(ReadOnlySpan<byte> buffer, CancellationToken cancella
                 return Task.FromCanceled(cancellationToken);
             }
 
+            return WriteAsync(buffer, cancellationToken).AsTask();
+        }
+
+        public ValueTask<FlushResult> WriteDataToPipeAsync(ReadOnlySpan<byte> buffer, CancellationToken cancellationToken = default)
+        {
+            if (cancellationToken.IsCancellationRequested)
+            {
+                return new ValueTask<FlushResult>(Task.FromCanceled<FlushResult>(cancellationToken));
+            }
+
             return WriteAsync(buffer, cancellationToken);
         }
 
-        public Task WriteStreamSuffixAsync()
+        public ValueTask<FlushResult> WriteStreamSuffixAsync()
         {
             return WriteAsync(_endChunkedResponseBytes.Span);
         }
 
-        public Task FlushAsync(CancellationToken cancellationToken = default)
+        public ValueTask<FlushResult> FlushAsync(CancellationToken cancellationToken = default)
         {
             return WriteAsync(Constants.EmptyData, cancellationToken);
         }
 
-        public Task WriteChunkAsync(ReadOnlySpan<byte> buffer, CancellationToken cancellationToken)
+        public Memory<byte> GetMemory(int sizeHint = 0)
+        {
+            if (_autoChunk)
+            {
+                return GetChunkedMemory(sizeHint);
+            }
+            else
+            {
+                return _pipeWriter.GetMemory(sizeHint);
+            }
+        }
+
+        public Span<byte> GetSpan(int sizeHint = 0)
+        {
+            if (_autoChunk)
+            {
+                return GetChunkedMemory(sizeHint).Span;
+            }
+            else
+            {
+                return _pipeWriter.GetMemory(sizeHint).Span;
+            }
+        }
+
+        public void Advance(int bytes)
+        {
+            if (_autoChunk)
+            {
+                if (bytes < 0)
+                {
+                    throw new ArgumentOutOfRangeException(nameof(bytes));
+                }
+
+                if (bytes + _advancedBytesForChunk > _currentChunkMemory.Length - BeginChunkLengthMax - EndChunkLength)
+                {
+                    throw new InvalidOperationException("Can't advance past buffer size.");
+                }
+                _advancedBytesForChunk += bytes;
+            }
+            else
+            {
+                _pipeWriter.Advance(bytes);
+            }
+        }
+
+        public void CancelPendingFlush()
+        {
+            // TODO we may not want to support this.
+            _pipeWriter.CancelPendingFlush();
+        }
+
+        // This method is for chunked http responses
+        public ValueTask<FlushResult> WriteChunkAsync(ReadOnlySpan<byte> buffer, CancellationToken cancellationToken)
         {
             lock (_contextLock)
             {
                 if (_completed)
                 {
-                    return Task.CompletedTask;
+                    return default;
                 }
 
+                CommitChunkToPipe();
+
                 if (buffer.Length > 0)
                 {
                     var writer = new BufferWriter<PipeWriter>(_pipeWriter);
@@ -96,7 +172,7 @@ public Task WriteChunkAsync(ReadOnlySpan<byte> buffer, CancellationToken cancell
             return FlushAsync(cancellationToken);
         }
 
-        public void WriteResponseHeaders(int statusCode, string reasonPhrase, HttpResponseHeaders responseHeaders)
+        public void WriteResponseHeaders(int statusCode, string reasonPhrase, HttpResponseHeaders responseHeaders, bool autoChunk)
         {
             lock (_contextLock)
             {
@@ -117,6 +193,7 @@ public void WriteResponseHeaders(int statusCode, string reasonPhrase, HttpRespon
                 writer.Commit();
 
                 _unflushedBytes += writer.BytesCommitted;
+                _autoChunk = autoChunk;
             }
         }
 
@@ -139,7 +216,6 @@ public void Abort(ConnectionAbortedException error)
         {
             // Abort can be called after Dispose if there's a flush timeout.
             // It's important to still call _lifetimeFeature.Abort() in this case.
-
             lock (_contextLock)
             {
                 if (_aborted)
@@ -153,20 +229,33 @@ public void Abort(ConnectionAbortedException error)
             }
         }
 
-        public Task Write100ContinueAsync()
+        public ValueTask<FlushResult> Write100ContinueAsync()
         {
             return WriteAsync(_continueBytes.Span);
         }
 
-        private Task WriteAsync(
+        public void Complete(Exception exception = null)
+        {
+            // TODO What to do with exception.
+            // and how to handle writing to response here.
+        }
+
+        private ValueTask<FlushResult> WriteAsync(
             ReadOnlySpan<byte> buffer,
             CancellationToken cancellationToken = default)
         {
             lock (_contextLock)
             {
                 if (_completed)
                 {
-                    return Task.CompletedTask;
+                    return default;
+                }
+
+                if (_autoChunk)
+                {
+                    // If there is data that was chunked before writing (ex someone did GetMemory->Advance->WriteAsync)
+                    // make sure to write whatever was advanced first
+                    CommitChunkToPipe();
                 }
 
                 var writer = new BufferWriter<PipeWriter>(_pipeWriter);
@@ -188,5 +277,56 @@ private Task WriteAsync(
                     cancellationToken);
             }
         }
+
+        private Memory<byte> GetChunkedMemory(int sizeHint)
+        {
+            // The max size of a chunk will be the size of memory returned from the PipeWriter (today 4096)
+            // minus 5 for the max chunked prefix size and minus 2 for the chunked ending, leaving a total of
+            // 4089.
+
+            if (_currentChunkMemory.Length == 0)
+            {
+                // First time calling GetMemory
+                _currentChunkMemory = _pipeWriter.GetMemory(sizeHint);
+            }
+
+            var memoryMaxLength = _currentChunkMemory.Length - BeginChunkLengthMax - EndChunkLength;
+            if (_advancedBytesForChunk == memoryMaxLength)
+            {
+                // Chunk is completely written, commit it to the pipe so GetMemory will return a new chunk of memory.
+                CommitChunkToPipe();
+                _currentChunkMemory = _pipeWriter.GetMemory(sizeHint);
+            }
+
+            var actualMemory = _currentChunkMemory.Slice(
+                BeginChunkLengthMax + _advancedBytesForChunk,
+                memoryMaxLength - _advancedBytesForChunk);
+
+            return actualMemory;
+        }
+
+        private void CommitChunkToPipe()
+        {
+            var writer = new BufferWriter<PipeWriter>(_pipeWriter);
+
+            Debug.Assert(_advancedBytesForChunk <= _currentChunkMemory.Length);
+
+            if (_advancedBytesForChunk > 0)
+            {
+                var bytesWritten = writer.WriteBeginChunkBytes(_advancedBytesForChunk);
+                if (bytesWritten < BeginChunkLengthMax)
+                {
+                    // If the current chunk of memory isn't completely utilized, we need to copy the contents forwards.
+                    // This occurs if someone uses less than 255 bytes of the current Memory segment.
+                    // Therefore, we need to copy it forwards by either 1 or 2 bytes (depending on number of bytes)
+                    _currentChunkMemory.Slice(BeginChunkLengthMax, _advancedBytesForChunk).CopyTo(_currentChunkMemory.Slice(bytesWritten));
+                }
+
+                writer.Write(_currentChunkMemory.Slice(bytesWritten, _advancedBytesForChunk).Span);
+                writer.WriteEndChunkBytes();
+                writer.Commit();
+                _advancedBytesForChunk = 0;
+            }
+        }
     }
 }