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 121223cce2d8..4a7ebf56121e 100644 --- a/src/System.IO.Pipelines/src/System/IO/Pipelines/BufferSegment.cs +++ b/src/System.IO.Pipelines/src/System/IO/Pipelines/BufferSegment.cs @@ -74,7 +74,6 @@ 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.IO.Pipelines/tests/FlushAsyncCancellationTests.cs b/src/System.IO.Pipelines/tests/FlushAsyncCancellationTests.cs index 3a3b328828a6..94c25220ff91 100644 --- a/src/System.IO.Pipelines/tests/FlushAsyncCancellationTests.cs +++ b/src/System.IO.Pipelines/tests/FlushAsyncCancellationTests.cs @@ -312,7 +312,7 @@ public static class TestWriterExtensions { public static PipeWriter WriteEmpty(this PipeWriter writer, int count) { - writer.GetMemory(count); + writer.GetSpan(count).Slice(0, count).Fill(0); writer.Advance(count); return writer; } diff --git a/src/System.IO.Pipelines/tests/PipePoolTests.cs b/src/System.IO.Pipelines/tests/PipePoolTests.cs index 3f49ef11d03e..79e6674babbf 100644 --- a/src/System.IO.Pipelines/tests/PipePoolTests.cs +++ b/src/System.IO.Pipelines/tests/PipePoolTests.cs @@ -3,6 +3,7 @@ // See the LICENSE file in the project root for more information. using System.Buffers; +using System.Diagnostics; using System.Threading.Tasks; using Xunit; @@ -13,6 +14,7 @@ public class PipePoolTests private class DisposeTrackingBufferPool : TestMemoryPool { public int ReturnedBlocks { get; set; } + public int DisposedBlocks { get; set; } public int CurrentlyRentedBlocks { get; set; } public override OwnedMemory Rent(int size) @@ -26,10 +28,12 @@ protected override void Dispose(bool disposing) private class DisposeTrackingOwnedMemory : OwnedMemory { - private readonly byte[] _array; + private byte[] _array; private readonly DisposeTrackingBufferPool _bufferPool; + private int _refCount = 1; + public DisposeTrackingOwnedMemory(byte[] array, DisposeTrackingBufferPool bufferPool) { _array = array; @@ -49,9 +53,9 @@ public override Span Span } } - public override bool IsDisposed { get; } + public override bool IsDisposed => _array == null; - protected override bool IsRetained => true; + protected override bool IsRetained => _refCount > 0; public override MemoryHandle Pin(int byteOffset = 0) { @@ -68,18 +72,26 @@ protected override bool TryGetArray(out ArraySegment arraySegment) protected override void Dispose(bool disposing) { - throw new NotImplementedException(); + if (IsRetained) + { + throw new InvalidOperationException(); + } + _bufferPool.DisposedBlocks++; + + _array = null; } public override bool Release() { _bufferPool.ReturnedBlocks++; _bufferPool.CurrentlyRentedBlocks--; + _refCount--; return IsRetained; } public override void Retain() { + _refCount++; } } } @@ -102,6 +114,8 @@ public async Task AdvanceToEndReturnsAllBlocks() pipe.Reader.AdvanceTo(readResult.Buffer.End); Assert.Equal(0, pool.CurrentlyRentedBlocks); + Assert.Equal(0, pool.DisposedBlocks); + Assert.Equal(3, pool.ReturnedBlocks); } [Fact] @@ -128,6 +142,10 @@ public async Task CanWriteAfterReturningMultipleBlocks() // Try writing more await pipe.Writer.WriteAsync(new byte[writeSize]); + + Assert.Equal(1, pool.CurrentlyRentedBlocks); + Assert.Equal(0, pool.DisposedBlocks); + Assert.Equal(2, pool.ReturnedBlocks); } [Fact] @@ -141,10 +159,12 @@ public async Task MultipleCompleteReaderWriterCauseDisposeOnlyOnce() readerWriter.Writer.Complete(); readerWriter.Reader.Complete(); Assert.Equal(1, pool.ReturnedBlocks); + Assert.Equal(0, pool.DisposedBlocks); readerWriter.Writer.Complete(); readerWriter.Reader.Complete(); Assert.Equal(1, pool.ReturnedBlocks); + Assert.Equal(0, pool.DisposedBlocks); } [Fact] @@ -174,11 +194,13 @@ public void ReturnsWriteHeadOnComplete() { var pool = new DisposeTrackingBufferPool(); var pipe = new Pipe(new PipeOptions(pool, readerScheduler: PipeScheduler.Inline, writerScheduler: PipeScheduler.Inline)); - var memory = pipe.Writer.GetMemory(512); + pipe.Writer.GetMemory(512); pipe.Reader.Complete(); pipe.Writer.Complete(); Assert.Equal(0, pool.CurrentlyRentedBlocks); + Assert.Equal(1, pool.ReturnedBlocks); + Assert.Equal(0, pool.DisposedBlocks); } [Fact] @@ -186,12 +208,14 @@ public void ReturnsWriteHeadWhenRequestingLargerBlock() { var pool = new DisposeTrackingBufferPool(); var pipe = new Pipe(new PipeOptions(pool, readerScheduler: PipeScheduler.Inline, writerScheduler: PipeScheduler.Inline)); - var memory = pipe.Writer.GetMemory(512); + pipe.Writer.GetMemory(512); pipe.Writer.GetMemory(4096); pipe.Reader.Complete(); pipe.Writer.Complete(); Assert.Equal(0, pool.CurrentlyRentedBlocks); + Assert.Equal(2, pool.ReturnedBlocks); + Assert.Equal(0, pool.DisposedBlocks); } [Fact] diff --git a/src/System.IO.Pipelines/tests/TestMemoryPool.cs b/src/System.IO.Pipelines/tests/TestMemoryPool.cs index f53173244e32..5ca96dfe6846 100644 --- a/src/System.IO.Pipelines/tests/TestMemoryPool.cs +++ b/src/System.IO.Pipelines/tests/TestMemoryPool.cs @@ -52,6 +52,7 @@ public PooledMemory(OwnedMemory ownedMemory, TestMemoryPool pool) _ownedMemory = ownedMemory; _pool = pool; _leaser = Environment.StackTrace; + _referenceCount = 1; } ~PooledMemory() @@ -74,12 +75,15 @@ public override MemoryHandle Pin(int byteOffset = 0) public override void Retain() { _pool.CheckDisposed(); + _ownedMemory.Retain(); Interlocked.Increment(ref _referenceCount); } public override bool Release() { _pool.CheckDisposed(); + _ownedMemory.Release(); + int newRefCount = Interlocked.Decrement(ref _referenceCount); if (newRefCount < 0) diff --git a/src/System.Memory/src/System.Memory.csproj b/src/System.Memory/src/System.Memory.csproj index 226268054836..0dcf986ae881 100644 --- a/src/System.Memory/src/System.Memory.csproj +++ b/src/System.Memory/src/System.Memory.csproj @@ -147,6 +147,7 @@ +