Cleanup memory pool code add some tests

shared/Microsoft.Extensions.Buffers.MemoryPool.Sources/DiagnosticMemoryPool.cs

@@ -0,0 +1,165 @@
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading.Tasks;
+
+namespace System.Buffers
+{
+    /// <summary>
+    /// Used to allocate and distribute re-usable blocks of memory.
+    /// </summary>
+    internal class DiagnosticMemoryPool : MemoryPool<byte>
+    {
+        private readonly MemoryPool<byte> _pool;
+
+        private readonly bool _allowLateReturn;
+
+        private readonly bool _rentTracking;
+
+        private readonly object _syncObj;
+
+        private readonly HashSet<DiagnosticPoolBlock> _blocks;
+
+        private readonly List<Exception> _blockAccessExceptions;
+
+        private readonly TaskCompletionSource<object> _allBlocksRetuned;
+
+        private int _totalBlocks;
+
+        /// <summary>
+        /// This default value passed in to Rent to use the default value for the pool.
+        /// </summary>
+        private const int AnySize = -1;
+
+        public DiagnosticMemoryPool(MemoryPool<byte> pool, bool allowLateReturn = false, bool rentTracking = false)
+        {
+            _pool = pool;
+            _allowLateReturn = allowLateReturn;
+            _rentTracking = rentTracking;
+            _blocks = new HashSet<DiagnosticPoolBlock>();
+            _syncObj = new object();
+            _allBlocksRetuned = new TaskCompletionSource<object>();
+            _blockAccessExceptions = new List<Exception>();
+        }
+
+        public bool IsDisposed { get; private set; }
+
+        public override IMemoryOwner<byte> Rent(int size = AnySize)
+        {
+            lock (_syncObj)
+            {
+                if (IsDisposed)
+                {
+                    MemoryPoolThrowHelper.ThrowObjectDisposedException(MemoryPoolThrowHelper.ExceptionArgument.MemoryPool);
+                }
+
+                var diagnosticPoolBlock = new DiagnosticPoolBlock(this, _pool.Rent(size));
+                if (_rentTracking)
+                {
+                    diagnosticPoolBlock.Track();
+                }
+                _totalBlocks++;
+                _blocks.Add(diagnosticPoolBlock);
+                return diagnosticPoolBlock;
+            }
+        }
+
+        public override int MaxBufferSize => _pool.MaxBufferSize;
+
+        internal void Return(DiagnosticPoolBlock block)
+        {
+            bool returnedAllBlocks;
+            lock (_syncObj)
+            {
+                _blocks.Remove(block);
+                returnedAllBlocks = _blocks.Count == 0;
+            }
+
+            if (IsDisposed)
+            {
+                if (!_allowLateReturn)
+                {
+                    MemoryPoolThrowHelper.ThrowInvalidOperationException_BlockReturnedToDisposedPool(block);
+                }
+
+                if (returnedAllBlocks)
+                {
+                    SetAllBlocksReturned();
+                }
+            }
+
+        }
+
+        internal void ReportException(Exception exception)
+        {
+            lock (_syncObj)
+            {
+                _blockAccessExceptions.Add(exception);
+            }
+        }
+
+        protected override void Dispose(bool disposing)
+        {
+            if (IsDisposed)
+            {
+                MemoryPoolThrowHelper.ThrowInvalidOperationException_DoubleDispose();
+            }
+
+            bool allBlocksReturned = false;
+            try
+            {
+                lock (_syncObj)
+                {
+                    IsDisposed = true;
+                    allBlocksReturned = _blocks.Count == 0;
+                    if (!allBlocksReturned && !_allowLateReturn)
+                    {
+                        MemoryPoolThrowHelper.ThrowInvalidOperationException_DisposingPoolWithActiveBlocks(_totalBlocks - _blocks.Count, _totalBlocks, _blocks.ToArray());
+                    }
+
+                    if (_blockAccessExceptions.Any())
+                    {
+                        throw CreateAccessExceptions();
+                    }
+                }
+            }
+            finally
+            {
+                if (allBlocksReturned)
+                {
+                    SetAllBlocksReturned();
+                }
+            }
+        }
+
+        private void SetAllBlocksReturned()
+        {
+            if (_blockAccessExceptions.Any())
+            {
+                _allBlocksRetuned.SetException(CreateAccessExceptions());
+            }
+            else
+            {
+                _allBlocksRetuned.SetResult(null);
+            }
+        }
+
+        private AggregateException CreateAccessExceptions()
+        {
+            return new AggregateException("Exceptions occurred while accessing blocks", _blockAccessExceptions.ToArray());
+        }
+
+        public async Task WhenAllBlocksReturnedAsync(TimeSpan timeout)
+        {
+            var task = await Task.WhenAny(_allBlocksRetuned.Task, Task.Delay(timeout));
+            if (task != _allBlocksRetuned.Task)
+            {
+                MemoryPoolThrowHelper.ThrowInvalidOperationException_BlocksWereNotReturnedInTime(_totalBlocks - _blocks.Count, _totalBlocks, _blocks.ToArray());
+            }
+
+            await task;
+        }
+    }
+}

shared/Microsoft.Extensions.Buffers.MemoryPool.Sources/DiagnosticPoolBlock.cs

@@ -0,0 +1,224 @@
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+using System.Threading;
+using System.Diagnostics;
+using System.Runtime.InteropServices;
+
+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.
+    /// </summary>
+    internal sealed class DiagnosticPoolBlock : MemoryManager<byte>
+    {
+        /// <summary>
+        /// Back-reference to the memory pool which this block was allocated from. It may only be returned to this pool.
+        /// </summary>
+        private readonly DiagnosticMemoryPool _pool;
+
+        private readonly IMemoryOwner<byte> _memoryOwner;
+        private MemoryHandle? _memoryHandle;
+        private Memory<byte> _memory;
+
+        private readonly object _syncObj = new object();
+        private bool _isDisposed;
+        private int _pinCount;
+
+
+        /// <summary>
+        /// This object cannot be instantiated outside of the static Create method
+        /// </summary>
+        internal DiagnosticPoolBlock(DiagnosticMemoryPool pool, IMemoryOwner<byte> memoryOwner)
+        {
+            _pool = pool;
+            _memoryOwner = memoryOwner;
+            _memory = memoryOwner.Memory;
+        }
+
+        public override Memory<byte> Memory
+        {
+            get
+            {
+                try
+                {
+                    lock (_syncObj)
+                    {
+                        if (_isDisposed)
+                        {
+                            MemoryPoolThrowHelper.ThrowObjectDisposedException(MemoryPoolThrowHelper.ExceptionArgument.MemoryPoolBlock);
+                        }
+
+                        if (_pool.IsDisposed)
+                        {
+                            MemoryPoolThrowHelper.ThrowInvalidOperationException_BlockIsBackedByDisposedSlab(this);
+                        }
+
+                        return CreateMemory(_memory.Length);
+                    }
+                }
+                catch (Exception exception)
+                {
+                    _pool.ReportException(exception);
+                    throw;
+                }
+            }
+        }
+
+        protected override void Dispose(bool disposing)
+        {
+            try
+            {
+                lock (_syncObj)
+                {
+                    if (Volatile.Read(ref _pinCount) > 0)
+                    {
+                        MemoryPoolThrowHelper.ThrowInvalidOperationException_ReturningPinnedBlock(this);
+                    }
+
+                    if (_isDisposed)
+                    {
+                        MemoryPoolThrowHelper.ThrowInvalidOperationException_BlockDoubleDispose(this);
+                    }
+
+                    _memoryOwner.Dispose();
+
+                    _pool.Return(this);
+
+                    _isDisposed = true;
+                }
+            }
+            catch (Exception exception)
+            {
+                _pool.ReportException(exception);
+                throw;
+            }
+        }
+
+        public override Span<byte> GetSpan()
+        {
+            try
+            {
+                lock (_syncObj)
+                {
+                    if (_isDisposed)
+                    {
+                        MemoryPoolThrowHelper.ThrowObjectDisposedException(MemoryPoolThrowHelper.ExceptionArgument.MemoryPoolBlock);
+                    }
+
+                    if (_pool.IsDisposed)
+                    {
+                        MemoryPoolThrowHelper.ThrowInvalidOperationException_BlockIsBackedByDisposedSlab(this);
+                    }
+
+                    return _memory.Span;
+                }
+            }
+            catch (Exception exception)
+            {
+                _pool.ReportException(exception);
+                throw;
+            }
+        }
+
+        public override MemoryHandle Pin(int byteOffset = 0)
+        {
+            try
+            {
+                lock (_syncObj)
+                {
+                    if (_isDisposed)
+                    {
+                        MemoryPoolThrowHelper.ThrowObjectDisposedException(MemoryPoolThrowHelper.ExceptionArgument.MemoryPoolBlock);
+                    }
+
+                    if (_pool.IsDisposed)
+                    {
+                        MemoryPoolThrowHelper.ThrowInvalidOperationException_BlockIsBackedByDisposedSlab(this);
+                    }
+
+                    if (byteOffset < 0 || byteOffset > _memory.Length)
+                    {
+                        MemoryPoolThrowHelper.ThrowArgumentOutOfRangeException(_memory.Length, byteOffset);
+                    }
+
+                    _pinCount++;
+
+                    _memoryHandle = _memoryHandle ?? _memory.Pin();
+
+                    unsafe
+                    {
+                        return new MemoryHandle(((IntPtr)_memoryHandle.Value.Pointer + byteOffset).ToPointer(), default, this);
+                    }
+                }
+            }
+            catch (Exception exception)
+            {
+                _pool.ReportException(exception);
+                throw;
+            }
+        }
+
+        protected override bool TryGetArray(out ArraySegment<byte> segment)
+        {
+            try
+            {
+                lock (_syncObj)
+                {
+                    if (_isDisposed)
+                    {
+                        MemoryPoolThrowHelper.ThrowObjectDisposedException(MemoryPoolThrowHelper.ExceptionArgument.MemoryPoolBlock);
+                    }
+
+                    if (_pool.IsDisposed)
+                    {
+                        MemoryPoolThrowHelper.ThrowInvalidOperationException_BlockIsBackedByDisposedSlab(this);
+                    }
+
+                    return MemoryMarshal.TryGetArray(_memory, out segment);
+                }
+            }
+            catch (Exception exception)
+            {
+                _pool.ReportException(exception);
+                throw;
+            }
+        }
+
+        public override void Unpin()
+        {
+            try
+            {
+                lock (_syncObj)
+                {
+                    if (_pinCount == 0)
+                    {
+                        MemoryPoolThrowHelper.ThrowInvalidOperationException_PinCountZero(this);
+                    }
+
+                    _pinCount--;
+
+                    if (_pinCount == 0)
+                    {
+                        Debug.Assert(_memoryHandle.HasValue);
+                        _memoryHandle.Value.Dispose();
+                        _memoryHandle = null;
+                    }
+                }
+            }
+            catch (Exception exception)
+            {
+                _pool.ReportException(exception);
+                throw;
+            }
+        }
+
+        public StackTrace Leaser { get; set; }
+
+        public void Track()
+        {
+            Leaser = new StackTrace(false);
+        }
+    }
+}