Expose response PipeWriter in Kestrel

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

@@ -2,6 +2,7 @@
 // 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.Text;
 using System.Threading;
 using System.Threading.Tasks;
@@ -60,8 +61,73 @@ public static class HttpResponseWritingExtensions
                 throw new ArgumentNullException(nameof(encoding));
             }
 
-            byte[] data = encoding.GetBytes(text);
-            return response.Body.WriteAsync(data, 0, data.Length, cancellationToken);
+            // Need to call StartAsync before GetMemory/GetSpan
+            if (!response.HasStarted)
+            {
+                var startAsyncTask = response.StartAsync(cancellationToken);
+                if (!startAsyncTask.IsCompletedSuccessfully)
+                {
+                    return StartAndWriteAsyncAwaited(response, text, encoding, cancellationToken, startAsyncTask);
+                }
+            }
+
+            Write(response, text, encoding);
+
+            var flushAsyncTask = response.BodyPipe.FlushAsync(cancellationToken);
+            if (flushAsyncTask.IsCompletedSuccessfully)
+            {
+                // Most implementations of ValueTask reset state in GetResult, so call it before returning a completed task.
+                flushAsyncTask.GetAwaiter().GetResult();
+                return Task.CompletedTask;
+            }
+
+            return flushAsyncTask.AsTask();
+        }
+
+        private static async Task StartAndWriteAsyncAwaited(this HttpResponse response, string text, Encoding encoding, CancellationToken cancellationToken, Task startAsyncTask)
+        {
+            await startAsyncTask;
+            Write(response, text, encoding);
+            await response.BodyPipe.FlushAsync(cancellationToken);
+        }
+
+        private static void Write(this HttpResponse response, string text, Encoding encoding)
+        {
+            var pipeWriter = response.BodyPipe;
+            var encodedLength = encoding.GetByteCount(text);
+            var destination = pipeWriter.GetSpan(encodedLength);
+
+            if (encodedLength <= destination.Length)
+            {
+                // Just call Encoding.GetBytes if everything will fit into a single segment.
+                var bytesWritten = encoding.GetBytes(text, destination);
+                pipeWriter.Advance(bytesWritten);
+            }
+            else
+            {
+                WriteMutliSegmentEncoded(pipeWriter, text, encoding, destination, encodedLength);
+            }
+        }
+
+        private static void WriteMutliSegmentEncoded(PipeWriter writer, string text, Encoding encoding, Span<byte> destination, int encodedLength)
+        {
+            var encoder = encoding.GetEncoder();
+            var source = text.AsSpan();
+            var completed = false;
+            var totalBytesUsed = 0;
+
+            // This may be a bug, but encoder.Convert returns completed = true for UTF7 too early.
+            // Therefore, we check encodedLength - totalBytesUsed too.
+            while (!completed || encodedLength - totalBytesUsed != 0)
+            {
+                encoder.Convert(source, destination, flush: source.Length == 0, out var charsUsed, out var bytesUsed, out completed);
+                totalBytesUsed += bytesUsed;
+
+                writer.Advance(bytesUsed);
+                source = source.Slice(charsUsed);
+
+                destination = writer.GetSpan(encodedLength - totalBytesUsed);
+            }
         }
     }
-}
+}

src/Http/Http.Abstractions/test/HttpResponseWritingExtensionsTests.cs

@@ -1,7 +1,11 @@
 // 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;
+using System.IO.Pipelines;
+using System.IO.Pipelines.Tests;
+using System.Text;
 using System.Threading.Tasks;
 using Xunit;
 
@@ -28,6 +32,65 @@ public async Task WritingText_MultipleWrites()
             Assert.Equal(22, context.Response.Body.Length);
         }
 
+        [Theory]
+        [MemberData(nameof(Encodings))]
+        public async Task WritingTextThatRequiresMultipleSegmentsWorks(Encoding encoding)
+        {
+            // Need to change the StreamPipeWriter with a capped MemoryPool
+            var memoryPool = new TestMemoryPool(maxBufferSize: 16);
+            var outputStream = new MemoryStream();
+            var streamPipeWriter = new StreamPipeWriter(outputStream, minimumSegmentSize: 0, memoryPool);
+
+            HttpContext context = new DefaultHttpContext();
+            context.Response.BodyPipe = streamPipeWriter;
+
+            var inputString = "昨日すき焼きを食べました";
+            var expected = encoding.GetBytes(inputString);
+            await context.Response.WriteAsync(inputString, encoding);
+
+            outputStream.Position = 0;
+            var actual = new byte[expected.Length];
+            var length = outputStream.Read(actual);
+
+            var res1 = encoding.GetString(actual);
+            var res2 = encoding.GetString(expected);
+            Assert.Equal(expected.Length, length);
+            Assert.Equal(expected, actual);
+            streamPipeWriter.Complete();
+        }
+
+        [Theory]
+        [MemberData(nameof(Encodings))]
+        public async Task WritingTextWithPassedInEncodingWorks(Encoding encoding)
+        {
+            HttpContext context = CreateRequest();
+
+            var inputString = "昨日すき焼きを食べました";
+            var expected = encoding.GetBytes(inputString);
+            await context.Response.WriteAsync(inputString, encoding);
+
+            context.Response.Body.Position = 0;
+            var actual = new byte[expected.Length * 2];
+            var length = context.Response.Body.Read(actual);
+
+            var actualShortened = new byte[length];
+            Array.Copy(actual, actualShortened, length);
+
+            Assert.Equal(expected.Length, length);
+            Assert.Equal(expected, actualShortened);
+        }
+
+        public static TheoryData<Encoding> Encodings =>
+            new TheoryData<Encoding>
+            {
+                        { Encoding.ASCII },
+                        { Encoding.BigEndianUnicode },
+                        { Encoding.Unicode },
+                        { Encoding.UTF32 },
+                        { Encoding.UTF7 },
+                        { Encoding.UTF8 }
+            };
+
         private HttpContext CreateRequest()
         {
             HttpContext context = new DefaultHttpContext();

src/Http/Http.Abstractions/test/Microsoft.AspNetCore.Http.Abstractions.Tests.csproj

@@ -4,6 +4,10 @@
     <TargetFramework>netcoreapp3.0</TargetFramework>
   </PropertyGroup>
 
+  <ItemGroup>
+    <ProjectReference Include="..\..\Http\test\Microsoft.AspNetCore.Http.Tests.csproj" />
+  </ItemGroup>
+
   <ItemGroup>
     <Reference Include="Microsoft.AspNetCore.Http" />
   </ItemGroup>

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/Internal/DefaultHttpResponse.cs

@@ -151,7 +151,7 @@ public override Task StartAsync(CancellationToken cancellationToken = default)
                 return HttpResponseFeature.Body.FlushAsync(cancellationToken);
             }
 
-            return HttpResponseStartFeature.StartAsync();
+            return HttpResponseStartFeature.StartAsync(cancellationToken);
         }
 
         struct FeatureInterfaces

src/Http/Http/src/ReadOnlyPipeStream.cs

@@ -13,7 +13,6 @@ namespace System.IO.Pipelines
     /// </summary>
     public class ReadOnlyPipeStream : Stream
     {
-        private readonly PipeReader _pipeReader;
         private bool _allowSynchronousIO = true;
 
         /// <summary>
@@ -33,7 +32,7 @@ public ReadOnlyPipeStream(PipeReader pipeReader) :
         public ReadOnlyPipeStream(PipeReader pipeReader, bool allowSynchronousIO)
         {
             _allowSynchronousIO = allowSynchronousIO;
-            _pipeReader = pipeReader;
+            InnerPipeReader = pipeReader;
         }
 
         /// <inheritdoc />
@@ -62,6 +61,8 @@ public override int WriteTimeout
             set => throw new NotSupportedException();
         }
 
+        public PipeReader InnerPipeReader { get; }
+
         /// <inheritdoc />
         public override void Write(byte[] buffer, int offset, int count)
             => throw new NotSupportedException();
@@ -160,7 +161,7 @@ private async ValueTask<int> ReadAsyncInternal(Memory<byte> buffer, Cancellation
         {
             while (true)
             {
-                var result = await _pipeReader.ReadAsync(cancellationToken);
+                var result = await InnerPipeReader.ReadAsync(cancellationToken);
                 var readableBuffer = result.Buffer;
                 var readableBufferLength = readableBuffer.Length;
 
@@ -186,7 +187,7 @@ private async ValueTask<int> ReadAsyncInternal(Memory<byte> buffer, Cancellation
                 }
                 finally
                 {
-                    _pipeReader.AdvanceTo(consumed);
+                    InnerPipeReader.AdvanceTo(consumed);
                 }
             }
         }
@@ -211,7 +212,7 @@ private async Task CopyToAsyncInternal(Stream destination, CancellationToken can
         {
             while (true)
             {
-                var result = await _pipeReader.ReadAsync(cancellationToken);
+                var result = await InnerPipeReader.ReadAsync(cancellationToken);
                 var readableBuffer = result.Buffer;
                 var readableBufferLength = readableBuffer.Length;
 
@@ -232,7 +233,7 @@ private async Task CopyToAsyncInternal(Stream destination, CancellationToken can
                 }
                 finally
                 {
-                    _pipeReader.AdvanceTo(readableBuffer.End);
+                    InnerPipeReader.AdvanceTo(readableBuffer.End);
                 }
             }
         }

src/Http/Http/src/StreamPipeReader.cs

@@ -17,7 +17,6 @@ public class StreamPipeReader : PipeReader, IDisposable
     {
         private readonly int _minimumSegmentSize;
         private readonly int _minimumReadThreshold;
-        private readonly Stream _readingStream;
         private readonly MemoryPool<byte> _pool;
 
         private CancellationTokenSource _internalTokenSource;
@@ -42,15 +41,14 @@ public StreamPipeReader(Stream readingStream)
         {
         }
 
-
         /// <summary>
         /// Creates a new StreamPipeReader.
         /// </summary>
         /// <param name="readingStream">The stream to read from.</param>
         /// <param name="options">The options to use.</param>
         public StreamPipeReader(Stream readingStream, StreamPipeReaderOptions options)
         {
-            _readingStream = readingStream ?? throw new ArgumentNullException(nameof(readingStream));
+            InnerStream = readingStream ?? throw new ArgumentNullException(nameof(readingStream));
 
             if (options == null)
             {
@@ -70,7 +68,7 @@ public StreamPipeReader(Stream readingStream, StreamPipeReaderOptions options)
         /// <summary>
         /// Gets the inner stream that is being read from.
         /// </summary>
-        public Stream InnerStream => _readingStream;
+        public Stream InnerStream { get; }
 
         /// <inheritdoc />
         public override void AdvanceTo(SequencePosition consumed)
@@ -235,7 +233,7 @@ public override async ValueTask<ReadResult> ReadAsync(CancellationToken cancella
                 {
                     AllocateReadTail();
 #if NETCOREAPP3_0
-                    var length = await _readingStream.ReadAsync(_readTail.AvailableMemory.Slice(_readTail.End), tokenSource.Token);
+                    var length = await InnerStream.ReadAsync(_readTail.AvailableMemory.Slice(_readTail.End), tokenSource.Token);
 #elif NETSTANDARD2_0
                     if (!MemoryMarshal.TryGetArray<byte>(_readTail.AvailableMemory.Slice(_readTail.End), out var arraySegment))
                     {

src/Http/Http/src/StreamPipeWriter.cs

@@ -14,7 +14,6 @@ namespace System.IO.Pipelines
     public class StreamPipeWriter : PipeWriter, IDisposable
     {
         private readonly int _minimumSegmentSize;
-        private readonly Stream _writingStream;
         private int _bytesWritten;
 
         private List<CompletedBuffer> _completedSegments;
@@ -53,14 +52,14 @@ public StreamPipeWriter(Stream writingStream) : this(writingStream, 4096)
         public StreamPipeWriter(Stream writingStream, int minimumSegmentSize, MemoryPool<byte> pool = null)
         {
             _minimumSegmentSize = minimumSegmentSize;
-            _writingStream = writingStream;
+            InnerStream = writingStream;
             _pool = pool ?? MemoryPool<byte>.Shared;
         }
 
         /// <summary>
         /// Gets the inner stream that is being written to.
         /// </summary>
-        public Stream InnerStream => _writingStream;
+        public Stream InnerStream { get; }
 
         /// <inheritdoc />
         public override void Advance(int count)
@@ -180,10 +179,10 @@ private async ValueTask<FlushResult> FlushAsyncInternal(CancellationToken cancel
                         {
                             var segment = _completedSegments[0];
 #if NETCOREAPP3_0
-                            await _writingStream.WriteAsync(segment.Buffer.Slice(0, segment.Length), localToken);
+                            await InnerStream.WriteAsync(segment.Buffer.Slice(0, segment.Length), localToken);
 #elif NETSTANDARD2_0
                             MemoryMarshal.TryGetArray<byte>(segment.Buffer, out var arraySegment);
-                            await _writingStream.WriteAsync(arraySegment.Array, 0, segment.Length, localToken);
+                            await InnerStream.WriteAsync(arraySegment.Array, 0, segment.Length, localToken);
 #else
 #error Target frameworks need to be updated.
 #endif
@@ -196,18 +195,18 @@ private async ValueTask<FlushResult> FlushAsyncInternal(CancellationToken cancel
                     if (!_currentSegment.IsEmpty)
                     {
 #if NETCOREAPP3_0
-                        await _writingStream.WriteAsync(_currentSegment.Slice(0, _position), localToken);
+                        await InnerStream.WriteAsync(_currentSegment.Slice(0, _position), localToken);
 #elif NETSTANDARD2_0
                         MemoryMarshal.TryGetArray<byte>(_currentSegment, out var arraySegment);
-                        await _writingStream.WriteAsync(arraySegment.Array, 0, _position, localToken);
+                        await InnerStream.WriteAsync(arraySegment.Array, 0, _position, localToken);
 #else
 #error Target frameworks need to be updated.
 #endif
                         _bytesWritten -= _position;
                         _position = 0;
                     }
 
-                    await _writingStream.FlushAsync(localToken);
+                    await InnerStream.FlushAsync(localToken);
 
                     return new FlushResult(isCanceled: false, _isCompleted);
                 }