Merge pull request #921 from dotnet/fix918

Auto-flush support for SimplexStream at 4KB threshold

Author: AArnott

src/Nerdbank.Streams/SimplexStream.cs

@@ -17,8 +17,26 @@ namespace Nerdbank.Streams
     /// A <see cref="Stream"/> that acts as a queue for bytes, in that what gets written to it
     /// can then be read from it, in order.
     /// </summary>
+    /// <remarks>
+    /// <para>
+    /// This stream buffers written content, as any other .NET stream might do.
+    /// To ensure bytes written are available to be read, callers should call <see cref="FlushAsync"/> or <see cref="Flush"/> after writing, and before reading.
+    /// Flushing automatically happens when the buffer is filled, so writes will block until the reader reads enough data to make room for the new data.
+    /// No flushing occurs when using the <see cref="IBufferWriter{T}"/> interface, so callers must call <see cref="FlushAsync"/> or <see cref="Flush"/> after calling <see cref="IBufferWriter{T}.Advance(int)"/>.
+    /// </para>
+    /// <para>
+    /// This class is thread safe for one concurrent reader and writer.
+    /// It is <em>not</em> thread safe for multiple concurrent readers or writers.
+    /// Disposal is not thread safe and must be executed exclusively of any concurrent reader or writer.
+    /// </para>
+    /// </remarks>
     public class SimplexStream : Stream, IBufferWriter<byte>, IDisposableObservable
     {
+        /// <summary>
+        /// The number of bytes to write before automatically flushing.
+        /// </summary>
+        private const int AutoFlushThreshold = 4096;
+
         /// <summary>
         /// The pipe that does all the hard work.
         /// </summary>
@@ -34,9 +52,17 @@ public class SimplexStream : Stream, IBufferWriter<byte>, IDisposableObservable
         /// </summary>
         private bool completed;
 
+        /// <summary>
+        /// The number of bytes written since the last flush.
+        /// </summary>
+        private int bytesSinceLastFlush;
+
         /// <summary>
         /// Initializes a new instance of the <see cref="SimplexStream"/> class.
         /// </summary>
+        /// <remarks>
+        /// The default thresholds for pausing and resuming the writer are 32KB and 16KB, respectively.
+        /// </remarks>
         public SimplexStream()
             : this(16 * 1024, 32 * 1024)
         {
@@ -78,6 +104,8 @@ public override long Position
             set => throw this.ThrowDisposedOr(new NotSupportedException());
         }
 
+        private long UnflushedBytes => this.pipe.Writer.CanGetUnflushedBytes ? this.pipe.Writer.UnflushedBytes : this.bytesSinceLastFlush;
+
         /// <summary>
         /// Signals that no more writing will take place, causing readers to receive 0 bytes when asking for any more data.
         /// </summary>
@@ -141,15 +169,33 @@ public override async Task<int> ReadAsync(byte[] buffer, int offset, int count,
         }
 
         /// <inheritdoc />
-        public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+        public override async Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
         {
+            Requires.NotNull(buffer, nameof(buffer));
+            Requires.Range(offset + count <= buffer.Length, nameof(count));
+            Requires.Range(offset >= 0, nameof(offset));
+            Requires.Range(count >= 0, nameof(count));
+            Verify.NotDisposed(this);
+
             cancellationToken.ThrowIfCancellationRequested();
-            this.Write(buffer, offset, count);
-            return Task.CompletedTask;
+            Memory<byte> memory = this.pipe.Writer.GetMemory(count);
+            buffer.AsMemory(offset, count).CopyTo(memory);
+            this.pipe.Writer.Advance(count);
+            this.RecordBytesWritten(count);
+
+            // Auto-flush if we've written enough data
+            if (this.UnflushedBytes >= AutoFlushThreshold)
+            {
+                await this.FlushAsync(cancellationToken).ConfigureAwait(false);
+            }
         }
 
         /// <inheritdoc />
-        void IBufferWriter<byte>.Advance(int count) => this.pipe.Writer.Advance(count);
+        void IBufferWriter<byte>.Advance(int count)
+        {
+            this.pipe.Writer.Advance(count);
+            this.RecordBytesWritten(count);
+        }
 
         /// <inheritdoc />
         Memory<byte> IBufferWriter<byte>.GetMemory(int sizeHint) => this.pipe.Writer.GetMemory(sizeHint);
@@ -174,6 +220,13 @@ public override void Write(byte[] buffer, int offset, int count)
             Memory<byte> memory = this.pipe.Writer.GetMemory(count);
             buffer.AsMemory(offset, count).CopyTo(memory);
             this.pipe.Writer.Advance(count);
+            this.RecordBytesWritten(count);
+
+            // Auto-flush if we've written enough data
+            if (this.UnflushedBytes >= AutoFlushThreshold)
+            {
+                this.Flush();
+            }
         }
 
         /// <inheritdoc />
@@ -195,5 +248,18 @@ private Exception ThrowDisposedOr(Exception ex)
             Verify.NotDisposed(this);
             throw ex;
         }
+
+        private void RecordBytesWritten(int count)
+        {
+            if (this.pipe.Writer.CanGetUnflushedBytes)
+            {
+                // The PipeWriter is tracking unflushed bytes for us, so we don't need to.
+                return;
+            }
+
+            this.bytesSinceLastFlush += count;
+        }
+
+        private void ResetBytesSinceLastFlush() => this.bytesSinceLastFlush = 0;
     }
 }

test/Nerdbank.Streams.Tests/SimplexStreamTests.cs

@@ -13,6 +13,9 @@ public class SimplexStreamTests : TestBase
 
     private const int PauseThreshold = 40;
 
+    // Test-local constant for auto-flush threshold. Keep in sync with production value in SimplexStream.
+    private const int AutoFlushThreshold = 4096;
+
     private readonly Random random = new Random();
 
     private SimplexStream stream = new SimplexStream(ResumeThreshold, PauseThreshold);
@@ -344,6 +347,203 @@ public void CompleteWriting_CompleteSuccessfullyThenWithError()
         Assert.Equal(0, this.stream.Read(buffer, 0, buffer.Length));
     }
 
+    [Theory]
+    [CombinatorialData]
+    public async Task AutoFlush_OccursAfter4KB(bool useAsync)
+    {
+        // Use a stream with larger thresholds to avoid blocking
+        using var largeStream = new SimplexStream(8192, 16384);
+
+        // Write exactly 4KB (4096 bytes) which should trigger auto-flush
+        byte[] sendBuffer = this.GetRandomBuffer(AutoFlushThreshold);
+        if (useAsync)
+        {
+            await largeStream.WriteAsync(sendBuffer, 0, sendBuffer.Length, this.TimeoutToken);
+        }
+        else
+        {
+            largeStream.Write(sendBuffer, 0, sendBuffer.Length);
+        }
+
+        // Data should be available for reading without explicit flush
+        byte[] recvBuffer = new byte[sendBuffer.Length];
+        await this.ReadAsync(largeStream, recvBuffer, isAsync: useAsync);
+        Assert.Equal(sendBuffer, recvBuffer);
+    }
+
+    [Theory]
+    [CombinatorialData]
+    public async Task AutoFlush_DoesNotOccurBelow4KB(bool useAsync)
+    {
+        // Use a stream with larger thresholds
+        using var largeStream = new SimplexStream(8192, 16384);
+
+        // Write less than 4KB
+        byte[] sendBuffer = this.GetRandomBuffer(4095);
+        if (useAsync)
+        {
+            await largeStream.WriteAsync(sendBuffer, 0, sendBuffer.Length, this.TimeoutToken);
+        }
+        else
+        {
+            largeStream.Write(sendBuffer, 0, sendBuffer.Length);
+        }
+
+        // Data should NOT be available without explicit flush - read should timeout
+        byte[] recvBuffer = new byte[1];
+        Task<int> readTask = largeStream.ReadAsync(recvBuffer, 0, 1, ExpectedTimeoutToken);
+
+        // This should timeout because data hasn't been flushed yet
+        await Assert.ThrowsAsync<OperationCanceledException>(async () => await readTask);
+
+        // After explicit flush, data should be available
+        await largeStream.FlushAsync();
+        recvBuffer = new byte[sendBuffer.Length];
+        await this.ReadAsync(largeStream, recvBuffer, isAsync: useAsync);
+        Assert.Equal(sendBuffer, recvBuffer);
+    }
+
+    [Theory]
+    [CombinatorialData]
+    public async Task AutoFlush_AccumulatesAcrossMultipleWrites(bool useAsync)
+    {
+        // Use a stream with larger thresholds to avoid blocking
+        using var largeStream = new SimplexStream(8192, 16384);
+
+        // Write 2KB three times (total 6KB) - should auto-flush after second write
+        byte[] sendBuffer = this.GetRandomBuffer(6144);
+
+        // First write (2KB) - no flush yet
+        if (useAsync)
+        {
+            await largeStream.WriteAsync(sendBuffer, 0, 2048, this.TimeoutToken);
+        }
+        else
+        {
+            largeStream.Write(sendBuffer, 0, 2048);
+        }
+
+        // Second write (2KB, total 4KB) - should auto-flush
+        if (useAsync)
+        {
+            await largeStream.WriteAsync(sendBuffer, 2048, 2048, this.TimeoutToken);
+        }
+        else
+        {
+            largeStream.Write(sendBuffer, 2048, 2048);
+        }
+
+        // Data should be available for reading (4KB)
+        byte[] recvBuffer = new byte[AutoFlushThreshold];
+        await this.ReadAsync(largeStream, recvBuffer, isAsync: useAsync);
+        Assert.Equal(sendBuffer.Take(AutoFlushThreshold), recvBuffer);
+
+        // Third write (2KB) - not flushed yet
+        if (useAsync)
+        {
+            await largeStream.WriteAsync(sendBuffer, AutoFlushThreshold, 2048, this.TimeoutToken);
+        }
+        else
+        {
+            largeStream.Write(sendBuffer, AutoFlushThreshold, 2048);
+        }
+
+        // Explicitly flush to make remaining data available
+        await largeStream.FlushAsync();
+        recvBuffer = new byte[2048];
+        await this.ReadAsync(largeStream, recvBuffer, isAsync: useAsync);
+        Assert.Equal(sendBuffer.Skip(AutoFlushThreshold).Take(2048), recvBuffer);
+    }
+
+    [Fact]
+    public async Task BackpressureWorks_WithAutoFlush()
+    {
+        // This test verifies that the pauseWriterThreshold works correctly with auto-flush
+        // by having concurrent reading and writing
+        var simplex = new SimplexStream(2048, AutoFlushThreshold);
+
+        try
+        {
+            byte[] sendBuffer = this.GetRandomBuffer(8192);
+            byte[] recvBuffer = new byte[8192];
+            int bytesRead = 0;
+
+            // Start concurrent reader
+            var readTask = Task.Run(async () =>
+            {
+                while (bytesRead < 8192)
+                {
+                    int count = await simplex.ReadAsync(recvBuffer, bytesRead, 1024, this.TimeoutToken);
+                    if (count == 0)
+                    {
+                        break;
+                    }
+
+                    bytesRead += count;
+                    await Task.Delay(10); // Simulate slow reader
+                }
+            });
+
+            // Write 8KB in 1KB chunks (should auto-flush twice at 4KB and 8KB)
+            for (int i = 0; i < 8; i++)
+            {
+                await simplex.WriteAsync(sendBuffer, i * 1024, 1024, this.TimeoutToken);
+            }
+
+            simplex.CompleteWriting();
+            await readTask.WithCancellation(this.TimeoutToken);
+
+            Assert.Equal(8192, bytesRead);
+            Assert.Equal(sendBuffer, recvBuffer);
+        }
+        finally
+        {
+            simplex.Dispose();
+        }
+    }
+
+    [Fact]
+    public async Task Issue918_LargeWriteSmallReadWithDispose()
+    {
+        // This is the scenario from issue #918
+        var simplex = new SimplexStream(0, 4096);
+
+        try
+        {
+            int written = 0;
+            var writeTask = Task.Run(async () =>
+            {
+                byte[] buffer = new byte[1024];
+                int totalToWrite = 10 * 1024 * 1024; // 10 MB
+
+                while (written < totalToWrite)
+                {
+                    await simplex.WriteAsync(buffer, 0, buffer.Length, this.TimeoutToken);
+                    written += buffer.Length;
+                }
+
+                simplex.CompleteWriting();
+            });
+
+            // Read only 1KB
+            byte[] readBuffer = new byte[1024];
+            int bytesRead = await simplex.ReadAsync(readBuffer, 0, readBuffer.Length, this.TimeoutToken);
+            Assert.Equal(1024, bytesRead);
+
+            // Dispose the stream - this should cause the writer to fail
+            simplex.Dispose();
+
+            // Wait for writer to complete (it should fail with ObjectDisposedException or similar)
+            Exception ex = await Assert.ThrowsAnyAsync<Exception>(() => writeTask.WithCancellation(this.TimeoutToken));
+            this.Logger.WriteLine($"Writer stopped after {written} bytes with: {ex.GetType().Name}: {ex.Message}");
+            simplex.CompleteWriting(ex);
+        }
+        finally
+        {
+            simplex.Dispose();
+        }
+    }
+
     protected override void Dispose(bool disposing)
     {
         this.stream.Dispose();