Check for in-progress write in ServerStreamWriter (#903)

Author: JamesNK

src/Grpc.AspNetCore.Server/Internal/HttpContextStreamWriter.cs

@@ -27,11 +27,14 @@ internal class HttpContextStreamWriter<TResponse> : IServerStreamWriter<TRespons
     {
         private readonly HttpContextServerCallContext _context;
         private readonly Action<TResponse, SerializationContext> _serializer;
+        private readonly object _writeLock;
+        private Task? _writeTask;
 
         public HttpContextStreamWriter(HttpContextServerCallContext context, Action<TResponse, SerializationContext> serializer)
         {
             _context = context;
             _serializer = serializer;
+            _writeLock = new object();
         }
 
         public WriteOptions WriteOptions
@@ -44,15 +47,40 @@ public Task WriteAsync(TResponse message)
         {
             if (message == null)
             {
-                throw new ArgumentNullException(nameof(message));
+                return Task.FromException(new ArgumentNullException(nameof(message)));
             }
 
             if (_context.CancellationToken.IsCancellationRequested)
             {
-                throw new InvalidOperationException("Cannot write message after request is complete.");
+                return Task.FromException(new InvalidOperationException("Cannot write message after request is complete."));
             }
 
-            return _context.HttpContext.Response.BodyWriter.WriteMessageAsync(message, _context, _serializer, canFlush: true);
+            lock (_writeLock)
+            {
+                // Pending writes need to be awaited first
+                if (IsWriteInProgressUnsynchronized)
+                {
+                    return Task.FromException(new InvalidOperationException("Can't write the message because the previous write is in progress."));
+                }
+
+                // Save write task to track whether it is complete. Must be set inside lock.
+                _writeTask = _context.HttpContext.Response.BodyWriter.WriteMessageAsync(message, _context, _serializer, canFlush: true);
+            }
+
+            return _writeTask;
+        }
+
+        /// <summary>
+        /// A value indicating whether there is an async write already in progress.
+        /// Should only check this property when holding the write lock.
+        /// </summary>
+        private bool IsWriteInProgressUnsynchronized
+        {
+            get
+            {
+                var writeTask = _writeTask;
+                return writeTask != null && !writeTask.IsCompleted;
+            }
         }
     }
 }

src/Grpc.Net.Client/Internal/HttpContentClientStreamWriter.cs

@@ -129,7 +129,7 @@ public Task WriteAsync(TRequest message)
                         return CreateErrorTask("Can't write the message because the previous write is in progress.");
                     }
 
-                    // Save write task to track whether it is complete
+                    // Save write task to track whether it is complete. Must be set inside lock.
                     _writeTask = WriteAsyncCore(message);
                 }
             }

test/Grpc.AspNetCore.Server.Tests/HttpContextStreamWriterTests.cs

@@ -16,6 +16,7 @@
 
 #endregion
 
+using System;
 using System.IO;
 using System.IO.Pipelines;
 using System.Threading.Tasks;
@@ -110,5 +111,35 @@ await writer.WriteAsync(new HelloReply
             var writtenMessage2 = await MessageHelpers.AssertReadStreamMessageAsync<HelloReply>(pipeReader);
             Assert.AreEqual("Hello world 2", writtenMessage2!.Message);
         }
+
+        [Test]
+        public async Task WriteAsync_WriteInProgress_Error()
+        {
+            // Arrange
+            var tcs = new TaskCompletionSource<object?>(TaskCreationOptions.RunContinuationsAsynchronously);
+
+            var httpContext = new DefaultHttpContext();
+            httpContext.Features.Set<IHttpResponseBodyFeature>(new TestResponseBodyFeature(PipeWriter.Create(new MemoryStream()), startAsyncTask: tcs.Task));
+            var serverCallContext = HttpContextServerCallContextHelper.CreateServerCallContext(httpContext);
+            var writer = new HttpContextStreamWriter<HelloReply>(serverCallContext, MessageHelpers.ServiceMethod.ResponseMarshaller.ContextualSerializer);
+
+            // Act
+            _ = writer.WriteAsync(new HelloReply
+            {
+                Message = "Hello world 1"
+            });
+
+            var ex = await ExceptionAssert.ThrowsAsync<InvalidOperationException>(() =>
+            {
+                return writer.WriteAsync(new HelloReply
+                {
+                    Message = "Hello world 2"
+                });
+            });
+
+            // Assert
+            Assert.AreEqual("Can't write the message because the previous write is in progress.", ex.Message);
+        }
+
     }
 }

test/Shared/TestResponseBodyFeature.cs

@@ -27,9 +27,12 @@ namespace Grpc.Tests.Shared
 {
     public class TestResponseBodyFeature : IHttpResponseBodyFeature
     {
-        public TestResponseBodyFeature(PipeWriter writer)
+        private readonly Task _startAsyncTask;
+
+        public TestResponseBodyFeature(PipeWriter writer, Task? startAsyncTask = null)
         {
             Writer = writer;
+            _startAsyncTask = startAsyncTask ?? Task.CompletedTask;
         }
 
         public PipeWriter Writer { get; }
@@ -52,7 +55,7 @@ public Task SendFileAsync(string path, long offset, long? count, CancellationTok
 
         public Task StartAsync(CancellationToken cancellationToken = default)
         {
-            return Task.CompletedTask;
+            return _startAsyncTask;
         }
     }
 }