From 722c207ededd0efe7cf49eb5c35e649ae191eb55 Mon Sep 17 00:00:00 2001
From: Ben Adams <thundercat@illyriad.co.uk>
Date: Fri, 8 May 2020 14:16:25 +0100
Subject: [PATCH] Use Pinned Object Heap for MemoryPool

---
 .../Buffers.MemoryPool/MemoryPoolBlock.cs     |  2 +-
 .../Buffers.MemoryPool/MemoryPoolSlab.cs      | 56 +++----------------
 .../Buffers.MemoryPool/SlabMemoryPool.cs      |  4 +-
 3 files changed, 13 insertions(+), 49 deletions(-)

diff --git a/src/Shared/Buffers.MemoryPool/MemoryPoolBlock.cs b/src/Shared/Buffers.MemoryPool/MemoryPoolBlock.cs
index e23b4679c43f..0efe3dbef290 100644
--- a/src/Shared/Buffers.MemoryPool/MemoryPoolBlock.cs
+++ b/src/Shared/Buffers.MemoryPool/MemoryPoolBlock.cs
@@ -25,7 +25,7 @@ internal MemoryPoolBlock(SlabMemoryPool pool, MemoryPoolSlab slab, int offset, i
             Pool = pool;
             Slab = slab;
 
-            Memory = MemoryMarshal.CreateFromPinnedArray(slab.Array, _offset, _length);
+            Memory = MemoryMarshal.CreateFromPinnedArray(slab.PinnedArray, _offset, _length);
         }
 
         /// <summary>
diff --git a/src/Shared/Buffers.MemoryPool/MemoryPoolSlab.cs b/src/Shared/Buffers.MemoryPool/MemoryPoolSlab.cs
index a9baa560f532..5ed5cbee7762 100644
--- a/src/Shared/Buffers.MemoryPool/MemoryPoolSlab.cs
+++ b/src/Shared/Buffers.MemoryPool/MemoryPoolSlab.cs
@@ -1,8 +1,6 @@
 // Copyright (c) Microsoft. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
-using System.Runtime.InteropServices;
-
 namespace System.Buffers
 {
     /// <summary>
@@ -11,18 +9,9 @@ namespace System.Buffers
     /// </summary>
     internal class MemoryPoolSlab : IDisposable
     {
-        /// <summary>
-        /// This handle pins the managed array in memory until the slab is disposed. This prevents it from being
-        /// relocated and enables any subsections of the array to be used as native memory pointers to P/Invoked API calls.
-        /// </summary>
-        private GCHandle _gcHandle;
-        private bool _isDisposed;
-
-        public MemoryPoolSlab(byte[] data)
+        private MemoryPoolSlab(byte[] pinnedData)
         {
-            Array = data;
-            _gcHandle = GCHandle.Alloc(data, GCHandleType.Pinned);
-            NativePointer = _gcHandle.AddrOfPinnedObject();
+            PinnedArray = pinnedData;
         }
 
         /// <summary>
@@ -30,51 +19,24 @@ public MemoryPoolSlab(byte[] data)
         /// memory pool size an entire slab must be removed. That is done by (1) setting IsActive to false and removing the
         /// slab from the pool's _slabs collection, (2) as each block currently in use is Return()ed to the pool it will
         /// be allowed to be garbage collected rather than re-pooled, and (3) when all block tracking objects are garbage
-        /// collected and the slab is no longer references the slab will be garbage collected and the memory unpinned will
-        /// be unpinned by the slab's Dispose.
+        /// collected and the slab is no longer references the slab will be garbage collected
         /// </summary>
-        public bool IsActive => !_isDisposed;
+        public bool IsActive => PinnedArray != null;
 
-        public IntPtr NativePointer { get; private set; }
-
-        public byte[] Array { get; private set; }
+        public byte[] PinnedArray { get; private set; }
 
         public static MemoryPoolSlab Create(int length)
         {
-            // allocate and pin requested memory length
-            var array = new byte[length];
+            // allocate requested memory length from the pinned memory heap
+            var pinnedArray = GC.AllocateUninitializedArray<byte>(length, pinned: true);
 
             // allocate and return slab tracking object
-            return new MemoryPoolSlab(array);
-        }
-
-        protected void Dispose(bool disposing)
-        {
-            if (_isDisposed)
-            {
-                return;
-            }
-
-            _isDisposed = true;
-
-            Array = null;
-            NativePointer = IntPtr.Zero;
-
-            if (_gcHandle.IsAllocated)
-            {
-                _gcHandle.Free();
-            }
-        }
-
-        ~MemoryPoolSlab()
-        {
-            Dispose(false);
+            return new MemoryPoolSlab(pinnedArray);
         }
 
         public void Dispose()
         {
-            Dispose(true);
-            GC.SuppressFinalize(this);
+            PinnedArray = null;
         }
     }
 }
diff --git a/src/Shared/Buffers.MemoryPool/SlabMemoryPool.cs b/src/Shared/Buffers.MemoryPool/SlabMemoryPool.cs
index e80862401073..f39a0d58ca8d 100644
--- a/src/Shared/Buffers.MemoryPool/SlabMemoryPool.cs
+++ b/src/Shared/Buffers.MemoryPool/SlabMemoryPool.cs
@@ -3,6 +3,7 @@
 
 using System.Collections.Concurrent;
 using System.Diagnostics;
+using System.Runtime.InteropServices;
 using System.Threading;
 
 namespace System.Buffers
@@ -110,7 +111,8 @@ private MemoryPoolBlock AllocateSlab()
             var slab = MemoryPoolSlab.Create(_slabLength);
             _slabs.Push(slab);
 
-            var basePtr = slab.NativePointer;
+            // Get the address for alignment
+            IntPtr basePtr = Marshal.UnsafeAddrOfPinnedArrayElement(slab.PinnedArray, 0);
             // Page align the blocks
             var offset = (int)((((ulong)basePtr + (uint)_blockSize - 1) & ~((uint)_blockSize - 1)) - (ulong)basePtr);
             // Ensure page aligned