StreamPipeWriter: don't flush when the writer is completed with an ex…

…ception. (#40874) This provides a way to discard buffered data without pushing it to the inner Stream.
dotnet · Oct 21, 2019 · 03453d9 · 03453d9
1 parent db17568
commit 03453d9
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 9 deletions.
diff --git a/src/System.IO.Pipelines/src/System/IO/Pipelines/StreamPipeWriter.cs b/src/System.IO.Pipelines/src/System/IO/Pipelines/StreamPipeWriter.cs
@@ -212,7 +212,7 @@ public override void Complete(Exception? exception = null)
 
             _isCompleted = true;
 
-            FlushInternal();
+            FlushInternal(writeToStream: exception == null);
 
             _internalTokenSource?.Dispose();
 
@@ -231,7 +231,7 @@ public override async ValueTask CompleteAsync(Exception? exception = null)
 
             _isCompleted = true;
 
-            await FlushAsyncInternal().ConfigureAwait(false);
+            await FlushAsyncInternal(writeToStream: exception == null).ConfigureAwait(false);
 
             _internalTokenSource?.Dispose();
 
@@ -253,15 +253,15 @@ public override ValueTask<FlushResult> FlushAsync(CancellationToken cancellation
                 return new ValueTask<FlushResult>(new FlushResult(isCanceled: false, isCompleted: false));
             }
 
-            return FlushAsyncInternal(cancellationToken);
+            return FlushAsyncInternal(writeToStream: true, cancellationToken);
         }
 
         private void Cancel()
         {
             InternalTokenSource.Cancel();
         }
 
-        private async ValueTask<FlushResult> FlushAsyncInternal(CancellationToken cancellationToken = default)
+        private async ValueTask<FlushResult> FlushAsyncInternal(bool writeToStream, CancellationToken cancellationToken = default)
         {
             // Write all completed segments and whatever remains in the current segment
             // and flush the result.
@@ -291,7 +291,7 @@ private async ValueTask<FlushResult> FlushAsyncInternal(CancellationToken cancel
                         BufferSegment returnSegment = segment;
                         segment = segment.NextSegment;
 
-                        if (returnSegment.Length > 0)
+                        if (returnSegment.Length > 0 && writeToStream)
                         {
                             await InnerStream.WriteAsync(returnSegment.Memory, localToken).ConfigureAwait(false);
                         }
@@ -303,7 +303,7 @@ private async ValueTask<FlushResult> FlushAsyncInternal(CancellationToken cancel
                         _head = segment;
                     }
 
-                    if (_bytesBuffered > 0)
+                    if (_bytesBuffered > 0 && writeToStream)
                     {
                         await InnerStream.FlushAsync(localToken).ConfigureAwait(false);
                     }
@@ -335,7 +335,7 @@ private async ValueTask<FlushResult> FlushAsyncInternal(CancellationToken cancel
             }
         }
 
-        private void FlushInternal()
+        private void FlushInternal(bool writeToStream)
         {
             // Write all completed segments and whatever remains in the current segment
             // and flush the result.
@@ -354,7 +354,7 @@ private void FlushInternal()
                 BufferSegment returnSegment = segment;
                 segment = segment.NextSegment;
 
-                if (returnSegment.Length > 0)
+                if (returnSegment.Length > 0 && writeToStream)
                 {
 #if !NETSTANDARD2_0
                     InnerStream.Write(returnSegment.Memory.Span);
@@ -370,7 +370,7 @@ private void FlushInternal()
                 _head = segment;
             }
 
-            if (_bytesBuffered > 0)
+            if (_bytesBuffered > 0 && writeToStream)
             {
                 InnerStream.Flush();
             }

diff --git a/src/System.IO.Pipelines/tests/StreamPipeWriterTests.cs b/src/System.IO.Pipelines/tests/StreamPipeWriterTests.cs
@@ -61,6 +61,40 @@ public async Task DataFlushedOnCompleteAsync()
             Assert.Equal("Hello World", Encoding.ASCII.GetString(stream.ToArray()));
         }
 
+        [Fact]
+        public void DataNotFlushedOnCompleteWithException()
+        {
+            byte[] bytes = Encoding.ASCII.GetBytes("Hello World");
+            var stream = new MemoryStream();
+            PipeWriter writer = PipeWriter.Create(stream, new StreamPipeWriterOptions(leaveOpen: true));
+
+            bytes.AsSpan().CopyTo(writer.GetSpan(bytes.Length));
+            writer.Advance(bytes.Length);
+
+            Assert.Equal(0, stream.Length);
+
+            writer.Complete(new Exception());
+
+            Assert.Equal(0, stream.Length);
+        }
+
+        [Fact]
+        public async Task DataNotFlushedOnCompleteAsyncWithException()
+        {
+            byte[] bytes = Encoding.ASCII.GetBytes("Hello World");
+            var stream = new MemoryStream();
+            PipeWriter writer = PipeWriter.Create(stream, new StreamPipeWriterOptions(leaveOpen: true));
+
+            bytes.AsSpan().CopyTo(writer.GetSpan(bytes.Length));
+            writer.Advance(bytes.Length);
+
+            Assert.Equal(0, stream.Length);
+
+            await writer.CompleteAsync(new Exception());
+
+            Assert.Equal(0, stream.Length);
+        }
+
         [Fact]
         public async Task CompleteAsyncDoesNotThrowObjectDisposedException()
         {