Use ArrayPool for oversized allocations (dotnet/corefx#39643)

* Use ArrayPool for oversized allocations

Commit migrated from https://github.com/dotnet/corefx/commit/e7134de2b56bd1c4af0b5c58fe6e2cfdb8b72e5a

diff --git a/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/BufferSegment.cs b/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/BufferSegment.cs
index 756a03b..7d0a297 100644 (file)
--- a/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/BufferSegment.cs
+++ b/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/BufferSegment.cs
@@ -59,20 +59,16 @@ namespace System.IO.Pipelines
             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/libraries/System.IO.Pipelines/src/System/IO/Pipelines/Pipe.cs b/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/Pipe.cs
index 5842715..6523e48 100644 (file)
--- a/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/Pipe.cs
+++ b/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/Pipe.cs
@@ -223,20 +223,16 @@ namespace System.IO.Pipelines
         {
             BufferSegment newSegment = CreateSegmentUnsynchronized();
 
-            if (_pool is null)
+            if (_pool is null || sizeHint > _pool.MaxBufferSize)
             {
                 // Use the array pool
-                newSegment.SetOwnedMemory(ArrayPool<byte>.Shared.Rent(GetSegmentSize(sizeHint)));
-            }
-            else if (sizeHint <= _pool.MaxBufferSize)
-            {
-                // Use the specified pool if it fits
-                newSegment.SetOwnedMemory(_pool.Rent(GetSegmentSize(sizeHint, _pool.MaxBufferSize)));
+                int sizeToRequest = GetSegmentSize(sizeHint);
+                newSegment.SetOwnedMemory(ArrayPool<byte>.Shared.Rent(sizeToRequest));
             }
             else
             {
-                // We can't use the pool so allocate an array
-                newSegment.SetUnownedMemory(new byte[sizeHint]);
+                // Use the specified pool as it fits
+                newSegment.SetOwnedMemory(_pool.Rent(GetSegmentSize(sizeHint, _pool.MaxBufferSize)));
             }
 
             _writingHeadMemory = newSegment.AvailableMemory;

diff --git a/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/StreamPipeWriter.cs b/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/StreamPipeWriter.cs
index 3aa78de..d252f0b 100644 (file)
--- a/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/StreamPipeWriter.cs
+++ b/src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/StreamPipeWriter.cs
@@ -21,14 +21,14 @@ namespace System.IO.Pipelines
         private int _tailBytesBuffered;
         private int _bytesBuffered;
 
-        private MemoryPool<byte> _pool;
+        private readonly MemoryPool<byte> _pool;
 
         private CancellationTokenSource _internalTokenSource;
         private bool _isCompleted;
-        private object _lockObject = new object();
+        private readonly object _lockObject = new object();
 
         private BufferSegmentStack _bufferSegmentPool;
-        private bool _leaveOpen;
+        private readonly bool _leaveOpen;
 
         private CancellationTokenSource InternalTokenSource
         {
@@ -150,20 +150,16 @@ namespace System.IO.Pipelines
         {
             BufferSegment newSegment = CreateSegmentUnsynchronized();
 
-            if (_pool is null)
+            if (_pool is null || sizeHint > _pool.MaxBufferSize)
             {
                 // Use the array pool
-                newSegment.SetOwnedMemory(ArrayPool<byte>.Shared.Rent(GetSegmentSize(sizeHint)));
-            }
-            else if (sizeHint <= _pool.MaxBufferSize)
-            {
-                // Use the specified pool if it fits
-                newSegment.SetOwnedMemory(_pool.Rent(GetSegmentSize(sizeHint, _pool.MaxBufferSize)));
+                int sizeToRequest = GetSegmentSize(sizeHint);
+                newSegment.SetOwnedMemory(ArrayPool<byte>.Shared.Rent(sizeToRequest));
             }
             else
             {
-                // We can't use the pool so allocate an array
-                newSegment.SetUnownedMemory(new byte[sizeHint]);
+                // Use the specified pool as it fits
+                newSegment.SetOwnedMemory(_pool.Rent(GetSegmentSize(sizeHint, _pool.MaxBufferSize)));
             }
 
             _tailMemory = newSegment.AvailableMemory;

diff --git a/src/libraries/System.IO.Pipelines/tests/PipePoolTests.nonnetstandard.cs b/src/libraries/System.IO.Pipelines/tests/PipePoolTests.nonnetstandard.cs
index 946bfc5..fc45842 100644 (file)
--- a/src/libraries/System.IO.Pipelines/tests/PipePoolTests.nonnetstandard.cs
+++ b/src/libraries/System.IO.Pipelines/tests/PipePoolTests.nonnetstandard.cs
@@ -71,8 +71,8 @@ namespace System.IO.Pipelines.Tests
 
                 Assert.Same(startSegment, endSegment);
 
-                // Null owner implies that the buffer is allocated and wasn't rented from the pool
-                Assert.Null(startSegment.MemoryOwner);
+                // Check owner is array rather than a different type
+                Assert.IsType<byte[]>(startSegment.MemoryOwner);
 
                 pipe.Reader.AdvanceTo(result.Buffer.End);
                 pipe.Reader.Complete();

diff --git a/src/libraries/System.IO.Pipelines/tests/StreamPipeWriterTests.cs b/src/libraries/System.IO.Pipelines/tests/StreamPipeWriterTests.cs
index 80703cd..e113d5e 100644 (file)
--- a/src/libraries/System.IO.Pipelines/tests/StreamPipeWriterTests.cs
+++ b/src/libraries/System.IO.Pipelines/tests/StreamPipeWriterTests.cs
@@ -432,7 +432,7 @@ namespace System.IO.Pipelines.Tests
         }
 
         [Fact]
-        public void GetMemoryBiggerThanPoolSizeAllocatesUnpooledArray()
+        public void GetMemoryBiggerThanPoolSizeAllocatesArrayPoolArray()
         {
             using (var pool = new DisposeTrackingBufferPool())
             {
@@ -440,8 +440,7 @@ namespace System.IO.Pipelines.Tests
                 var options = new StreamPipeWriterOptions(pool);
                 PipeWriter writer = PipeWriter.Create(stream, options);
                 Memory<byte> memory = writer.GetMemory(pool.MaxBufferSize + 1);
-
-                Assert.Equal(pool.MaxBufferSize + 1, memory.Length);
+                Assert.True(memory.Length > pool.MaxBufferSize + 1);
                 Assert.Equal(0, pool.CurrentlyRentedBlocks);
                 Assert.Equal(0, pool.DisposedBlocks);