Make sure to complete copy before returning

sharwell · sharwell · commit ebb674f4b49d · 2020-09-29T22:19:13.000-07:00
diff --git a/src/Workspaces/Remote/Core/RemoteHostAssetSerialization.cs b/src/Workspaces/Remote/Core/RemoteHostAssetSerialization.cs
@@ -86,7 +86,7 @@ static void WriteAsset(ObjectWriter writer, ISerializerService serializer, Check
             }
         }
 
-        public static ValueTask<ImmutableArray<(Checksum, object)>> ReadDataAsync(PipeReader pipeReader, int scopeId, ISet<Checksum> checksums, ISerializerService serializerService, CancellationToken cancellationToken)
+        public static async ValueTask<ImmutableArray<(Checksum, object)>> ReadDataAsync(PipeReader pipeReader, int scopeId, ISet<Checksum> checksums, ISerializerService serializerService, CancellationToken cancellationToken)
         {
             // Workaround for ObjectReader not supporting async reading.
             // Unless we read from the RPC stream asynchronously and with cancallation support we might hang when the server cancels.
@@ -98,7 +98,7 @@ static void WriteAsset(ObjectWriter writer, ISerializerService serializer, Check
             Exception? exception = null;
 
             // start a task on a thread pool thread copying from the RPC pipe to a local pipe:
-            Task.Run(async () =>
+            var copyTask = Task.Run(async () =>
             {
                 try
                 {
@@ -113,20 +113,26 @@ static void WriteAsset(ObjectWriter writer, ISerializerService serializer, Check
                     await localPipe.Writer.CompleteAsync(exception).ConfigureAwait(false);
                     await pipeReader.CompleteAsync(exception).ConfigureAwait(false);
                 }
-            }, cancellationToken).Forget();
+            }, cancellationToken);
 
             // blocking read from the local pipe on the current thread:
             try
             {
                 using var stream = localPipe.Reader.AsStream(leaveOpen: false);
-                return new(ReadData(stream, scopeId, checksums, serializerService, cancellationToken));
+                return ReadData(stream, scopeId, checksums, serializerService, cancellationToken);
             }
             catch (EndOfStreamException)
             {
                 cancellationToken.ThrowIfCancellationRequested();
 
                 throw exception ?? ExceptionUtilities.Unreachable;
             }
+            finally
+            {
+                // Make sure to complete the copy and pipes before returning, otherwise the caller could complete the
+                // reader and/or writer while they are still in use.
+                await copyTask.ConfigureAwait(false);
+            }
         }
 
         public static ImmutableArray<(Checksum, object)> ReadData(Stream stream, int scopeId, ISet<Checksum> checksums, ISerializerService serializerService, CancellationToken cancellationToken)

Original file line number	Diff line number	Diff line change
`@@ -86,7 +86,7 @@ static void WriteAsset(ObjectWriter writer, ISerializerService serializer, Check`
`86`	`86`	`}`
`87`	`87`	`}`
`88`	`88`
`89`		`- public static ValueTask<ImmutableArray<(Checksum, object)>> ReadDataAsync(PipeReader pipeReader, int scopeId, ISet<Checksum> checksums, ISerializerService serializerService, CancellationToken cancellationToken)`
	`89`	`+ public static async ValueTask<ImmutableArray<(Checksum, object)>> ReadDataAsync(PipeReader pipeReader, int scopeId, ISet<Checksum> checksums, ISerializerService serializerService, CancellationToken cancellationToken)`
`90`	`90`	`{`
`91`	`91`	`// Workaround for ObjectReader not supporting async reading.`
`92`	`92`	`// Unless we read from the RPC stream asynchronously and with cancallation support we might hang when the server cancels.`
`@@ -98,7 +98,7 @@ static void WriteAsset(ObjectWriter writer, ISerializerService serializer, Check`
`98`	`98`	`Exception? exception = null;`
`99`	`99`
`100`	`100`	`// start a task on a thread pool thread copying from the RPC pipe to a local pipe:`
`101`		`- Task.Run(async () =>`
	`101`	`+ var copyTask = Task.Run(async () =>`
`102`	`102`	`{`
`103`	`103`	`try`
`104`	`104`	`{`
`@@ -113,20 +113,26 @@ static void WriteAsset(ObjectWriter writer, ISerializerService serializer, Check`
`113`	`113`	`await localPipe.Writer.CompleteAsync(exception).ConfigureAwait(false);`
`114`	`114`	`await pipeReader.CompleteAsync(exception).ConfigureAwait(false);`
`115`	`115`	`}`
`116`		`- }, cancellationToken).Forget();`
	`116`	`+ }, cancellationToken);`
`117`	`117`
`118`	`118`	`// blocking read from the local pipe on the current thread:`
`119`	`119`	`try`
`120`	`120`	`{`
`121`	`121`	`using var stream = localPipe.Reader.AsStream(leaveOpen: false);`
`122`		`- return new(ReadData(stream, scopeId, checksums, serializerService, cancellationToken));`
	`122`	`+ return ReadData(stream, scopeId, checksums, serializerService, cancellationToken);`
`123`	`123`	`}`
`124`	`124`	`catch (EndOfStreamException)`
`125`	`125`	`{`
`126`	`126`	`cancellationToken.ThrowIfCancellationRequested();`
`127`	`127`
`128`	`128`	`throw exception ?? ExceptionUtilities.Unreachable;`
`129`	`129`	`}`
	`130`	`+ finally`
	`131`	`+ {`
	`132`	`+ // Make sure to complete the copy and pipes before returning, otherwise the caller could complete the`
	`133`	`+ // reader and/or writer while they are still in use.`
	`134`	`+ await copyTask.ConfigureAwait(false);`
	`135`	`+ }`
`130`	`136`	`}`
`131`	`137`
`132`	`138`	`public static ImmutableArray<(Checksum, object)> ReadData(Stream stream, int scopeId, ISet<Checksum> checksums, ISerializerService serializerService, CancellationToken cancellationToken)`