More thread safety

dotnet · Mar 1, 2018 · 4eecbab · 4eecbab
1 parent ebbb756
commit 4eecbab
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 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
@@ -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
@@ -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.