More thread safety

dotnet · Mar 1, 2018 · 27e7b74 · 27e7b74
1 parent ebbb756
commit 27e7b74
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 18 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
@@ -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
@@ -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
@@ -23,7 +23,6 @@ public sealed override OwnedMemory<T> Rent(int minimumBufferSize = -1)
                 ThrowHelper.ThrowArgumentOutOfRangeException(ExceptionArgument.minimumBufferSize);
 
             var buffer = new ArrayMemoryPoolBuffer(minimumBufferSize);
-            buffer.Retain();
             return buffer;
         }