From 3e62e9211fd746955f68db828637247eca22ea06 Mon Sep 17 00:00:00 2001
From: David Fowler <davidfowl@gmail.com>
Date: Wed, 17 Apr 2019 23:13:07 -0700
Subject: [PATCH] Handle concurrent dispose and WriteAsync calls - lock the
 entire Write operation - Added a test that Completes both reader and writer
 in the middle of writing a large buffer

---
 .../src/Resources/Strings.resx                |  6 ----
 .../src/System/IO/Pipelines/Pipe.cs           | 19 +++++++------
 .../tests/PipeWriterTests.cs                  | 28 +++++++++++++++++++
 3 files changed, 39 insertions(+), 14 deletions(-)

diff --git a/src/System.IO.Pipelines/src/Resources/Strings.resx b/src/System.IO.Pipelines/src/Resources/Strings.resx
index 8dc450466b2b..1801eb4fbe37 100644
--- a/src/System.IO.Pipelines/src/Resources/Strings.resx
+++ b/src/System.IO.Pipelines/src/Resources/Strings.resx
@@ -120,12 +120,6 @@
   <data name="AdvanceToInvalidCursor" xml:space="preserve">
     <value>The PipeReader has already advanced past the provided position.</value>
   </data>
-  <data name="CannotCompleteWhileReading" xml:space="preserve">
-    <value>Can't complete reader while reading.</value>
-  </data>
-  <data name="CannotCompleteWhiteWriting" xml:space="preserve">
-    <value>Can't complete writer while writing.</value>
-  </data>
   <data name="ConcurrentOperationsNotSupported" xml:space="preserve">
     <value>Concurrent reads or writes are not supported.</value>
   </data>
diff --git a/src/System.IO.Pipelines/src/System/IO/Pipelines/Pipe.cs b/src/System.IO.Pipelines/src/System/IO/Pipelines/Pipe.cs
index 3678ad63f03b..69925aa971c3 100644
--- a/src/System.IO.Pipelines/src/System/IO/Pipelines/Pipe.cs
+++ b/src/System.IO.Pipelines/src/System/IO/Pipelines/Pipe.cs
@@ -934,16 +934,19 @@ internal ValueTask<FlushResult> WriteAsync(ReadOnlyMemory<byte> source, Cancella
             // state as writing
             AllocateWriteHeadIfNeeded(0);
 
-            if (source.Length <= _writingHeadMemory.Length)
+            lock (_sync)
             {
-                source.CopyTo(_writingHeadMemory);
+                if (source.Length <= _writingHeadMemory.Length)
+                {
+                    source.CopyTo(_writingHeadMemory);
 
-                AdvanceCore(source.Length);
-            }
-            else
-            {
-                // This is the multi segment copy
-                WriteMultiSegment(source.Span);
+                    AdvanceCore(source.Length);
+                }
+                else
+                {
+                    // This is the multi segment copy
+                    WriteMultiSegment(source.Span);
+                }
             }
 
             return FlushAsync(cancellationToken);
diff --git a/src/System.IO.Pipelines/tests/PipeWriterTests.cs b/src/System.IO.Pipelines/tests/PipeWriterTests.cs
index 187d8de8d40a..0ae5110673bb 100644
--- a/src/System.IO.Pipelines/tests/PipeWriterTests.cs
+++ b/src/System.IO.Pipelines/tests/PipeWriterTests.cs
@@ -232,5 +232,33 @@ public async Task WritesUsingGetMemoryWorks()
 
             pipe.Reader.Complete();
         }
+
+        [Fact]
+        public async Task CompleteWithLargeWriteThrows()
+        {
+            var pipe = new Pipe();
+            pipe.Reader.Complete();
+
+            var task = Task.Run(async () =>
+            {
+                await Task.Delay(10);
+                pipe.Writer.Complete();
+            });
+
+            try
+            {
+                for (int i = 0; i < 1000; i++)
+                {
+                    var buffer = new byte[10000000];
+                    await pipe.Writer.WriteAsync(buffer);
+                }
+            }
+            catch (InvalidOperationException)
+            {
+                // Complete while writing
+            }
+
+            await task;
+        }
     }
 }