From 4c08acd8a43ea4eb6ab33fae5b74e5dc758fddf0 Mon Sep 17 00:00:00 2001
From: David Fowler <davidfowl@gmail.com>
Date: Sat, 10 Feb 2018 20:52:00 -0800
Subject: [PATCH] Don't allocate all the memory up front when receiving
 incoming sends (#1433)

---
 .../HttpConnectionDispatcher.cs               | 19 +++---
 .../Internal/PipeWriterStream.cs              | 66 +++++++++++++++++++
 2 files changed, 74 insertions(+), 11 deletions(-)
 create mode 100644 src/Microsoft.AspNetCore.Sockets.Http/Internal/PipeWriterStream.cs

diff --git a/src/Microsoft.AspNetCore.Sockets.Http/HttpConnectionDispatcher.cs b/src/Microsoft.AspNetCore.Sockets.Http/HttpConnectionDispatcher.cs
index 72766cc7889b..78b669683fbe 100644
--- a/src/Microsoft.AspNetCore.Sockets.Http/HttpConnectionDispatcher.cs
+++ b/src/Microsoft.AspNetCore.Sockets.Http/HttpConnectionDispatcher.cs
@@ -12,6 +12,7 @@
 using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Protocols;
 using Microsoft.AspNetCore.Sockets.Features;
+using Microsoft.AspNetCore.Sockets.Http.Internal;
 using Microsoft.AspNetCore.Sockets.Internal;
 using Microsoft.AspNetCore.Sockets.Internal.Transports;
 using Microsoft.Extensions.Logging;
@@ -431,18 +432,14 @@ private async Task ProcessSend(HttpContext context)
                 return;
             }
 
-            // TODO: Use a pool here
+            // Until the parsers are incremental, we buffer the entire request body before
+            // flushing the buffer. Using CopyToAsync allows us to avoid allocating a single giant
+            // buffer before writing.
+            var pipeWriterStream = new PipeWriterStream(connection.Application.Output);
+            await context.Request.Body.CopyToAsync(pipeWriterStream);
 
-            byte[] buffer;
-            using (var stream = new MemoryStream())
-            {
-                await context.Request.Body.CopyToAsync(stream);
-                await stream.FlushAsync();
-                buffer = stream.ToArray();
-            }
-
-            _logger.ReceivedBytes(buffer.Length);
-            await connection.Application.Output.WriteAsync(buffer);
+            _logger.ReceivedBytes(pipeWriterStream.Length);
+            await connection.Application.Output.FlushAsync();
         }
 
         private async Task<bool> EnsureConnectionStateAsync(DefaultConnectionContext connection, HttpContext context, TransportType transportType, TransportType supportedTransports, ConnectionLogScope logScope, HttpSocketOptions options)
diff --git a/src/Microsoft.AspNetCore.Sockets.Http/Internal/PipeWriterStream.cs b/src/Microsoft.AspNetCore.Sockets.Http/Internal/PipeWriterStream.cs
new file mode 100644
index 000000000000..4987b06f49f0
--- /dev/null
+++ b/src/Microsoft.AspNetCore.Sockets.Http/Internal/PipeWriterStream.cs
@@ -0,0 +1,66 @@
+// 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.Buffers;
+using System.IO;
+using System.IO.Pipelines;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.AspNetCore.Sockets.Http.Internal
+{
+    // Write only stream implementation for efficiently writing bytes from the request body
+    internal class PipeWriterStream : Stream
+    {
+        private long _length;
+        private readonly PipeWriter _pipeWriter;
+
+        public PipeWriterStream(PipeWriter pipeWriter)
+        {
+            _pipeWriter = pipeWriter;
+        }
+
+        public override bool CanRead => false;
+
+        public override bool CanSeek => false;
+
+        public override bool CanWrite => true;
+
+        public override long Length => _length;
+
+        public override long Position { get => throw new NotSupportedException(); set => throw new NotSupportedException(); }
+
+        public override void Flush()
+        {
+            throw new NotSupportedException();
+        }
+
+        public override int Read(byte[] buffer, int offset, int count)
+        {
+            throw new NotSupportedException();
+        }
+
+        public override long Seek(long offset, SeekOrigin origin)
+        {
+            throw new NotSupportedException();
+        }
+
+        public override void SetLength(long value)
+        {
+            throw new NotSupportedException();
+        }
+
+        public override void Write(byte[] buffer, int offset, int count)
+        {
+            _pipeWriter.Write(new ReadOnlySpan<byte>(buffer, offset, count));
+            _length += count;
+        }
+
+        public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+        {
+            Write(buffer, offset, count);
+            return Task.CompletedTask;
+        }
+    }
+}