From 1ca67695ed7f9517e23c2ffb69bf45b150c7cb81 Mon Sep 17 00:00:00 2001
From: Ben Adams <thundercat@illyriad.co.uk>
Date: Thu, 3 Dec 2015 19:17:40 +0000
Subject: [PATCH] Use pooled memory for filtered stream

---
 .../Filter/FilteredStreamAdapter.cs           |  8 +++++--
 .../Filter/StreamExtensions.cs                | 23 +++++++++++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)
 create mode 100644 src/Microsoft.AspNet.Server.Kestrel/Filter/StreamExtensions.cs

diff --git a/src/Microsoft.AspNet.Server.Kestrel/Filter/FilteredStreamAdapter.cs b/src/Microsoft.AspNet.Server.Kestrel/Filter/FilteredStreamAdapter.cs
index 598620d1e6..beb5c3b803 100644
--- a/src/Microsoft.AspNet.Server.Kestrel/Filter/FilteredStreamAdapter.cs
+++ b/src/Microsoft.AspNet.Server.Kestrel/Filter/FilteredStreamAdapter.cs
@@ -28,9 +28,13 @@ namespace Microsoft.AspNet.Server.Kestrel.Filter
             _filteredStream = filteredStream;
             _socketInputStream = new SocketInputStream(SocketInput);
 
-            // Don't use 81920 byte buffer
-            _filteredStream.CopyToAsync(_socketInputStream, 4096).ContinueWith((task, state) =>
+            var block = memory.Lease();
+            // Use pooled block for copy
+            _filteredStream.CopyToAsync(_socketInputStream, block).ContinueWith((task, state) =>
             {
+                var returnedBlock = task.Result;
+                returnedBlock.Pool?.Return(returnedBlock);
+
                 ((FilteredStreamAdapter)state).OnStreamClose(task);
             }, this);
         }
diff --git a/src/Microsoft.AspNet.Server.Kestrel/Filter/StreamExtensions.cs b/src/Microsoft.AspNet.Server.Kestrel/Filter/StreamExtensions.cs
new file mode 100644
index 0000000000..af3e77d15f
--- /dev/null
+++ b/src/Microsoft.AspNet.Server.Kestrel/Filter/StreamExtensions.cs
@@ -0,0 +1,23 @@
+// 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.IO;
+using System.Threading.Tasks;
+using Microsoft.AspNet.Server.Kestrel.Infrastructure;
+
+namespace Microsoft.AspNet.Server.Kestrel.Filter
+{
+    public static class StreamExtensions
+    {
+        public static async Task<MemoryPoolBlock2> CopyToAsync(this Stream source, Stream destination, MemoryPoolBlock2 block)
+        {
+            int bytesRead;
+            while ((bytesRead = await source.ReadAsync(block.Array, block.Data.Offset, block.Data.Count)) != 0)
+            {
+                await destination.WriteAsync(block.Array, block.Data.Offset, bytesRead);
+            }
+
+            return block;
+        }
+    }
+}