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 buffer) public void SetMemory(OwnedMemory ownedMemory, int start, int end, bool readOnly = false) { _ownedMemory = ownedMemory; - _ownedMemory.Retain(); AvailableMemory = _ownedMemory.Memory; @@ -75,6 +74,7 @@ public void SetMemory(OwnedMemory 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 public ArrayMemoryPoolBuffer(int size) { _array = ArrayPool.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 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.