From 9399f09b723f9316d94dc1257eee864667981a48 Mon Sep 17 00:00:00 2001
From: Ben Adams <thundercat@illyriad.co.uk>
Date: Fri, 30 Aug 2019 18:31:32 +0100
Subject: [PATCH] BufferSegment use ArrayPool for over-sized allocs (#13495)

---
 .../Infrastructure/PipeWriterHelpers/BufferSegment.cs  | 10 +++-------
 .../PipeWriterHelpers/ConcurrentPipeWriter.cs          |  4 ++--
 2 files changed, 5 insertions(+), 9 deletions(-)

diff --git a/src/Servers/Kestrel/Core/src/Internal/Infrastructure/PipeWriterHelpers/BufferSegment.cs b/src/Servers/Kestrel/Core/src/Internal/Infrastructure/PipeWriterHelpers/BufferSegment.cs
index fdd8ac7367d5..a9bf8d942404 100644
--- a/src/Servers/Kestrel/Core/src/Internal/Infrastructure/PipeWriterHelpers/BufferSegment.cs
+++ b/src/Servers/Kestrel/Core/src/Internal/Infrastructure/PipeWriterHelpers/BufferSegment.cs
@@ -59,20 +59,16 @@ public void SetOwnedMemory(byte[] 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)
+            else
             {
-                ArrayPool<byte>.Shared.Return(array);
+                byte[] poolArray = (byte[])_memoryOwner;
+                ArrayPool<byte>.Shared.Return(poolArray);
             }
 
             // Order of below field clears is significant as it clears in a sequential order
diff --git a/src/Servers/Kestrel/Core/src/Internal/Infrastructure/PipeWriterHelpers/ConcurrentPipeWriter.cs b/src/Servers/Kestrel/Core/src/Internal/Infrastructure/PipeWriterHelpers/ConcurrentPipeWriter.cs
index 3ab051c8c413..4e09b0a8a40a 100644
--- a/src/Servers/Kestrel/Core/src/Internal/Infrastructure/PipeWriterHelpers/ConcurrentPipeWriter.cs
+++ b/src/Servers/Kestrel/Core/src/Internal/Infrastructure/PipeWriterHelpers/ConcurrentPipeWriter.cs
@@ -341,8 +341,8 @@ private BufferSegment AllocateSegmentUnsynchronized(int sizeHint)
             }
             else
             {
-                // We can't use the pool so allocate an array
-                newSegment.SetUnownedMemory(new byte[sizeHint]);
+                // We can't use the recommended pool so use the ArrayPool
+                newSegment.SetOwnedMemory(ArrayPool<byte>.Shared.Rent(sizeHint));
             }
 
             _tailMemory = newSegment.AvailableMemory;