Remove slabs from the slab memory pool (#30732)

- Slabs add an extra complication when trying to free blocks (blocks are attached to slabs and would need to be ref counted). Instead, since we can now directly allocate on the POH, slabs are less meaningful.
- Use 4K arrays on the pinned heap
- Remove the finalizer from MemoryPoolBlock

Author: davidfowl

src/Shared/Buffers.MemoryPool/MemoryPoolBlock.cs

@@ -6,52 +6,29 @@
 namespace System.Buffers
 {
     /// <summary>
-    /// Block tracking object used by the byte buffer memory pool. A slab is a large allocation which is divided into smaller blocks. The
-    /// individual blocks are then treated as independent array segments.
+    /// Wraps an array allocated in the pinned object heap in a reusable block of managed memory
     /// </summary>
     internal sealed class MemoryPoolBlock : IMemoryOwner<byte>
     {
-        private readonly int _offset;
-        private readonly int _length;
-
-        /// <summary>
-        /// This object cannot be instantiated outside of the static Create method
-        /// </summary>
-        internal MemoryPoolBlock(SlabMemoryPool pool, MemoryPoolSlab slab, int offset, int length)
+        internal MemoryPoolBlock(SlabMemoryPool pool, int length)
         {
-            _offset = offset;
-            _length = length;
-
             Pool = pool;
-            Slab = slab;
 
-            Memory = MemoryMarshal.CreateFromPinnedArray(slab.PinnedArray, _offset, _length);
+            var pinnedArray = GC.AllocateUninitializedArray<byte>(length, pinned: true);
+
+            Memory = MemoryMarshal.CreateFromPinnedArray(pinnedArray, 0, pinnedArray.Length);
         }
 
         /// <summary>
         /// Back-reference to the memory pool which this block was allocated from. It may only be returned to this pool.
         /// </summary>
         public SlabMemoryPool Pool { get; }
 
-        /// <summary>
-        /// Back-reference to the slab from which this block was taken, or null if it is one-time-use memory.
-        /// </summary>
-        public MemoryPoolSlab Slab { get; }
-
         public Memory<byte> Memory { get; }
 
-        ~MemoryPoolBlock()
-        {
-            Pool.RefreshBlock(Slab, _offset, _length);
-        }
-
         public void Dispose()
         {
             Pool.Return(this);
         }
-
-        public void Lease()
-        {
-        }
     }
 }

src/Shared/Buffers.MemoryPool/MemoryPoolSlab.cs

@@ -2,9 +2,6 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using System.Collections.Concurrent;
-using System.Diagnostics;
-using System.Runtime.InteropServices;
-using System.Threading;
 
 #nullable enable
 
@@ -20,13 +17,6 @@ internal sealed class SlabMemoryPool : MemoryPool<byte>
         /// </summary>
         private const int _blockSize = 4096;
 
-        /// <summary>
-        /// Allocating 32 contiguous blocks per slab makes the slab size 128k. This is larger than the 85k size which will place the memory
-        /// in the large object heap. This means the GC will not try to relocate this array, so the fact it remains pinned does not negatively
-        /// affect memory management's compactification.
-        /// </summary>
-        private const int _blockCount = 32;
-
         /// <summary>
         /// Max allocation block size for pooled blocks,
         /// larger values can be leased but they will be disposed after use rather than returned to the pool.
@@ -38,30 +28,17 @@ internal sealed class SlabMemoryPool : MemoryPool<byte>
         /// </summary>
         public static int BlockSize => _blockSize;
 
-        /// <summary>
-        /// 4096 * 32 gives you a slabLength of 128k contiguous bytes allocated per slab
-        /// </summary>
-        private static readonly int _slabLength = _blockSize * _blockCount;
-
         /// <summary>
         /// Thread-safe collection of blocks which are currently in the pool. A slab will pre-allocate all of the block tracking objects
         /// and add them to this collection. When memory is requested it is taken from here first, and when it is returned it is re-added.
         /// </summary>
         private readonly ConcurrentQueue<MemoryPoolBlock> _blocks = new ConcurrentQueue<MemoryPoolBlock>();
 
-        /// <summary>
-        /// Thread-safe collection of slabs which have been allocated by this pool. As long as a slab is in this collection and slab.IsActive,
-        /// the blocks will be added to _blocks when returned.
-        /// </summary>
-        private readonly ConcurrentStack<MemoryPoolSlab> _slabs = new ConcurrentStack<MemoryPoolSlab>();
-
         /// <summary>
         /// This is part of implementing the IDisposable pattern.
         /// </summary>
         private bool _isDisposed; // To detect redundant calls
 
-        private int _totalAllocatedBlocks;
-
         private readonly object _disposeSync = new object();
 
         /// <summary>
@@ -76,16 +53,6 @@ public override IMemoryOwner<byte> Rent(int size = AnySize)
                 MemoryPoolThrowHelper.ThrowArgumentOutOfRangeException_BufferRequestTooLarge(_blockSize);
             }
 
-            var block = Lease();
-            return block;
-        }
-
-        /// <summary>
-        /// Called to take a block from the pool.
-        /// </summary>
-        /// <returns>The block that is reserved for the called. It must be passed to Return when it is no longer being used.</returns>
-        private MemoryPoolBlock Lease()
-        {
             if (_isDisposed)
             {
                 MemoryPoolThrowHelper.ThrowObjectDisposedException(MemoryPoolThrowHelper.ExceptionArgument.MemoryPool);
@@ -94,53 +61,9 @@ private MemoryPoolBlock Lease()
             if (_blocks.TryDequeue(out var block))
             {
                 // block successfully taken from the stack - return it
-
-                block.Lease();
                 return block;
             }
-            // no blocks available - grow the pool
-            block = AllocateSlab();
-            block.Lease();
-            return block;
-        }
-
-        /// <summary>
-        /// Internal method called when a block is requested and the pool is empty. It allocates one additional slab, creates all of the
-        /// block tracking objects, and adds them all to the pool.
-        /// </summary>
-        private MemoryPoolBlock AllocateSlab()
-        {
-            var slab = MemoryPoolSlab.Create(_slabLength);
-            _slabs.Push(slab);
-
-            // 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
-            Debug.Assert(((ulong)basePtr + (uint)offset) % _blockSize == 0);
-
-            var blockCount = (_slabLength - offset) / _blockSize;
-            Interlocked.Add(ref _totalAllocatedBlocks, blockCount);
-
-            MemoryPoolBlock? block = null;
-
-            for (int i = 0; i < blockCount; i++)
-            {
-                block = new MemoryPoolBlock(this, slab, offset, _blockSize);
-
-                if (i != blockCount - 1) // last block
-                {
-#if BLOCK_LEASE_TRACKING
-                    block.IsLeased = true;
-#endif
-                    Return(block);
-                }
-
-                offset += _blockSize;
-            }
-
-            return block!;
+            return new MemoryPoolBlock(this, BlockSize);
         }
 
         /// <summary>
@@ -163,25 +86,6 @@ internal void Return(MemoryPoolBlock block)
             {
                 _blocks.Enqueue(block);
             }
-            else
-            {
-                GC.SuppressFinalize(block);
-            }
-        }
-
-        // This method can ONLY be called from the finalizer of MemoryPoolBlock
-        internal void RefreshBlock(MemoryPoolSlab slab, int offset, int length)
-        {
-            lock (_disposeSync)
-            {
-                if (!_isDisposed && slab != null && slab.IsActive)
-                {
-                    // Need to make a new object because this one is being finalized
-                    // Note, this must be called within the _disposeSync lock because the block
-                    // could be disposed at the same time as the finalizer.
-                    Return(new MemoryPoolBlock(this, slab, offset, length));
-                }
-            }
         }
 
         protected override void Dispose(bool disposing)
@@ -197,17 +101,11 @@ protected override void Dispose(bool disposing)
 
                 if (disposing)
                 {
-                    while (_slabs.TryPop(out var slab))
+                    // Discard blocks in pool
+                    while (_blocks.TryDequeue(out _))
                     {
-                        // dispose managed state (managed objects).
-                        slab.Dispose();
-                    }
-                }
 
-                // Discard blocks in pool
-                while (_blocks.TryDequeue(out var block))
-                {
-                    GC.SuppressFinalize(block);
+                    }
                 }
             }
         }