From 4eecbab92a20ba79689f78b8853b73f657ca5f1a Mon Sep 17 00:00:00 2001
From: Pavel Krymets <pavel@krymets.com>
Date: Thu, 1 Mar 2018 14:21:20 -0800
Subject: [PATCH] More thread safety

---
 .../src/System/IO/Pipelines/BufferSegment.cs  |  2 +-
 .../ArrayMemoryPool.ArrayMemoryPoolBuffer.cs  | 35 ++++++++++---------
 .../src/System/Buffers/ArrayMemoryPool.cs     |  4 +--
 3 files changed, 21 insertions(+), 20 deletions(-)

diff --git a/src/System.IO.Pipelines/src/System/IO/Pipelines/BufferSegment.cs b/src/System.IO.Pipelines/src/System/IO/Pipelines/BufferSegment.cs
index 5ee8f1a55fc8..121223cce2d8 100644
--- a/src/System.IO.Pipelines/src/System/IO/Pipelines/BufferSegment.cs
+++ b/src/System.IO.Pipelines/src/System/IO/Pipelines/BufferSegment.cs
@@ -61,7 +61,6 @@ public void SetMemory(OwnedMemory<byte> buffer)
         public void SetMemory(OwnedMemory<byte> ownedMemory, int start, int end, bool readOnly = false)
         {
             _ownedMemory = ownedMemory;
-            _ownedMemory.Retain();
 
             AvailableMemory = _ownedMemory.Memory;
 
@@ -75,6 +74,7 @@ public void SetMemory(OwnedMemory<byte> ownedMemory, int start, int end, bool re
         public void ResetMemory()
         {
             _ownedMemory.Release();
+            _ownedMemory.Dispose();
             _ownedMemory = null;
             AvailableMemory = default;
         }
diff --git a/src/System.Memory/src/System/Buffers/ArrayMemoryPool.ArrayMemoryPoolBuffer.cs b/src/System.Memory/src/System/Buffers/ArrayMemoryPool.ArrayMemoryPoolBuffer.cs
index 5a93b92f95d4..7f20c50d3eb3 100644
--- a/src/System.Memory/src/System/Buffers/ArrayMemoryPool.ArrayMemoryPoolBuffer.cs
+++ b/src/System.Memory/src/System/Buffers/ArrayMemoryPool.ArrayMemoryPoolBuffer.cs
@@ -22,6 +22,7 @@ private sealed class ArrayMemoryPoolBuffer : OwnedMemory<T>
             public ArrayMemoryPoolBuffer(int size)
             {
                 _array = ArrayPool<T>.Shared.Rent(size);
+                _refCount = 1;
             }
 
             public sealed override int Length => _array.Length;
@@ -80,28 +81,30 @@ public sealed override MemoryHandle Pin(int byteOffset = 0)
 
             public sealed override void Retain()
             {
-                if (IsDisposed)
-                    ThrowHelper.ThrowObjectDisposedException_ArrayMemoryPoolBuffer();
-
-                Interlocked.Increment(ref _refCount);
+                while (true)
+                {
+                    int currentCount = Volatile.Read(ref _refCount);
+                    if (currentCount <= 0) ThrowHelper.ThrowObjectDisposedException_ArrayMemoryPoolBuffer();
+                    if (Interlocked.CompareExchange(ref _refCount, currentCount + 1, currentCount) == currentCount) break;
+                }
             }
 
             public sealed override bool Release()
             {
-                if (IsDisposed)
-                    ThrowHelper.ThrowObjectDisposedException_ArrayMemoryPoolBuffer();
-
-                int newRefCount = Interlocked.Decrement(ref _refCount);
-                if (newRefCount == 0)
+                while (true)
                 {
-                    Dispose();
+                    int currentCount = Volatile.Read(ref _refCount);
+                    if (currentCount <= 0) ThrowHelper.ThrowObjectDisposedException_ArrayMemoryPoolBuffer();
+                    if (Interlocked.CompareExchange(ref _refCount, currentCount - 1, currentCount) == currentCount)
+                    {
+                        if (currentCount == 1)
+                        {
+                            Dispose();
+                            return false;
+                        }
+                        return true;
+                    }
                 }
-
-                // Other thread already disposed
-                if (newRefCount < 0)
-                    ThrowHelper.ThrowObjectDisposedException_ArrayMemoryPoolBuffer();
-
-                return newRefCount != 0;
             }
         }
     }
diff --git a/src/System.Memory/src/System/Buffers/ArrayMemoryPool.cs b/src/System.Memory/src/System/Buffers/ArrayMemoryPool.cs
index a5b64f20da1f..5b5472951eed 100644
--- a/src/System.Memory/src/System/Buffers/ArrayMemoryPool.cs
+++ b/src/System.Memory/src/System/Buffers/ArrayMemoryPool.cs
@@ -22,9 +22,7 @@ public sealed override OwnedMemory<T> Rent(int minimumBufferSize = -1)
             else if (((uint)minimumBufferSize) > s_maxBufferSize)
                 ThrowHelper.ThrowArgumentOutOfRangeException(ExceptionArgument.minimumBufferSize);
 
-            var buffer = new ArrayMemoryPoolBuffer(minimumBufferSize);
-            buffer.Retain();
-            return buffer;
+            return new ArrayMemoryPoolBuffer(minimumBufferSize);
         }
 
         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.