Fix semantics of ArrayMemoryPool

src/System.Memory/src/System/Buffers/ArrayMemoryPool.ArrayMemoryPoolBuffer.cs

@@ -3,6 +3,7 @@
 // See the LICENSE file in the project root for more information.
 
 using System.Runtime.InteropServices;
+using System.Threading;
 #if !netstandard
 using Internal.Runtime.CompilerServices;
 #else
@@ -27,7 +28,7 @@ public ArrayMemoryPoolBuffer(int size)
 
             public sealed override bool IsDisposed => _array == null;
 
-            protected sealed override bool IsRetained => _refCount > 0;
+            protected sealed override bool IsRetained => Volatile.Read(ref _refCount) > 0;
 
             public sealed override Span<T> Span
             {
@@ -82,17 +83,23 @@ public sealed override void Retain()
                 if (IsDisposed)
                     ThrowHelper.ThrowObjectDisposedException_ArrayMemoryPoolBuffer();
 
-                _refCount++;
+                Interlocked.Increment(ref _refCount);
             }
 
             public sealed override bool Release()
             {
                 if (IsDisposed)
                     ThrowHelper.ThrowObjectDisposedException_ArrayMemoryPoolBuffer();
 
-                int newRefCount = --_refCount;
+                int newRefCount = Interlocked.Decrement(ref _refCount);
+                if (newRefCount == 0)
+                {
+                    Dispose();
+                }
+
+                // Other thread already disposed
                 if (newRefCount < 0)
-                    ThrowHelper.ThrowInvalidOperationException();
+                    ThrowHelper.ThrowObjectDisposedException_ArrayMemoryPoolBuffer();
 
                 return newRefCount != 0;
             }

src/System.Memory/src/System/Buffers/ArrayMemoryPool.cs

@@ -22,7 +22,9 @@ public sealed override OwnedMemory<T> Rent(int minimumBufferSize = -1)
             else if (((uint)minimumBufferSize) > s_maxBufferSize)
                 ThrowHelper.ThrowArgumentOutOfRangeException(ExceptionArgument.minimumBufferSize);
 
-            return new ArrayMemoryPoolBuffer(minimumBufferSize);
+            var buffer = new ArrayMemoryPoolBuffer(minimumBufferSize);
+            buffer.Retain();
+            return buffer;
         }
 
         protected sealed override void Dispose(bool disposing) {}  // ArrayMemoryPool is a shared pool so Dispose() would be a nop even if there were native resources to dispose.

src/System.Memory/tests/MemoryPool/MemoryPool.cs

@@ -29,6 +29,7 @@ public static void DisposingTheSharedPoolIsANop()
             using (OwnedMemory<int> block = mp.Rent(10))
             {
                 Assert.True(block.Length >= 10);
+                block.Release();
             }
         }
 
@@ -55,6 +56,7 @@ public static void MemoryPoolSpan()
                         Assert.Equal((IntPtr)newMemoryHandle.Pointer, (IntPtr)pSpan);
                     }
                 }
+                block.Release();
             }
         }
 
@@ -77,6 +79,7 @@ public static void MemoryPoolPin(int byteOffset)
                         Assert.Equal((IntPtr)pSpan, ((IntPtr)newMemoryHandle.Pointer) - byteOffset);
                     }
                 }
+                block.Release();
             }
         }
 
@@ -109,7 +112,7 @@ public static void MemoryPoolPinOffsetAtEnd()
             {
                 return; // The pool gave us a very large block - too big to compute the byteOffset needed to carry out this test. Skip.
             }
-            
+
             using (MemoryHandle newMemoryHandle = block.Pin(byteOffset: byteOffset))
             {
                 unsafe
@@ -177,6 +180,7 @@ public static void EachRentalIsUniqueUntilDisposed()
 
             foreach (OwnedMemory<int> prior in priorBlocks)
             {
+                prior.Release();
                 prior.Dispose();
             }
         }
@@ -187,6 +191,7 @@ public static void RentWithDefaultSize()
             using (OwnedMemory<int> block = MemoryPool<int>.Shared.Rent(minBufferSize: -1))
             {
                 Assert.True(block.Length >= 1);
+                block.Release();
             }
         }
 
@@ -224,6 +229,7 @@ public static void MemoryPoolTryGetArray()
                         Assert.Equal((IntPtr)pSpan, (IntPtr)pArray);
                     }
                 }
+                block.Release();
             }
         }
 
@@ -243,10 +249,13 @@ public static void RefCounting()
                 moreToGo = block.Release();
                 Assert.True(moreToGo);
 
+                moreToGo = block.Release();
+                Assert.True(moreToGo);
+
                 moreToGo = block.Release();
                 Assert.False(moreToGo);
 
-                Assert.Throws<InvalidOperationException>(() => block.Release());
+                Assert.Throws<ObjectDisposedException>(() => block.Release());
             }
         }
 
@@ -255,7 +264,7 @@ public static void IsDisposed()
         {
             OwnedMemory<int> block = MemoryPool<int>.Shared.Rent(42);
             Assert.False(block.IsDisposed);
-            block.Dispose();
+            block.Release();
             Assert.True(block.IsDisposed);
             block.Dispose();
             Assert.True(block.IsDisposed);
@@ -265,6 +274,7 @@ public static void IsDisposed()
         public static void ExtraDisposesAreIgnored()
         {
             OwnedMemory<int> block = MemoryPool<int>.Shared.Rent(42);
+            block.Release();
             block.Dispose();
             block.Dispose();
         }
@@ -273,6 +283,7 @@ public static void ExtraDisposesAreIgnored()
         public static void NoSpanAfterDispose()
         {
             OwnedMemory<int> block = MemoryPool<int>.Shared.Rent(42);
+            block.Release();
             block.Dispose();
             Assert.Throws<ObjectDisposedException>(() => block.Span.DontBox());
         }
@@ -281,6 +292,7 @@ public static void NoSpanAfterDispose()
         public static void NoRetainAfterDispose()
         {
             OwnedMemory<int> block = MemoryPool<int>.Shared.Rent(42);
+            block.Release();
             block.Dispose();
             Assert.Throws<ObjectDisposedException>(() => block.Retain());
         }
@@ -289,6 +301,7 @@ public static void NoRetainAfterDispose()
         public static void NoRelease_AfterDispose()
         {
             OwnedMemory<int> block = MemoryPool<int>.Shared.Rent(42);
+            block.Release();
             block.Dispose();
             Assert.Throws<ObjectDisposedException>(() => block.Release());
         }
@@ -297,6 +310,7 @@ public static void NoRelease_AfterDispose()
         public static void NoPinAfterDispose()
         {
             OwnedMemory<int> block = MemoryPool<int>.Shared.Rent(42);
+            block.Release();
             block.Dispose();
             Assert.Throws<ObjectDisposedException>(() => block.Pin());
         }
@@ -306,9 +320,26 @@ public static void NoTryGetArrayAfterDispose()
         {
             OwnedMemory<int> block = MemoryPool<int>.Shared.Rent(42);
             Memory<int> memory = block.Memory;
+            block.Release();
             block.Dispose();
             Assert.Throws<ObjectDisposedException>(() => memory.TryGetArray(out ArraySegment<int> arraySegment));
         }
+
+        [Fact]
+        public static void IsRetainedWhenReturned()
+        {
+            OwnedMemory<int> block = MemoryPool<int>.Shared.Rent(42);
+            Assert.False(block.Release());
+        }
+
+        [Fact]
+        public static void IsDisposedWhenReleased()
+        {
+            OwnedMemory<int> block = MemoryPool<int>.Shared.Rent(42);
+            block.Release();
+
+            Assert.True(block.IsDisposed);
+        }
     }
 }