Improve streaming error messages when call is complete (#1121)

Author: JamesNK

src/Grpc.AspNetCore.Server/Internal/CallHandlers/ClientStreamingServerCallHandler.cs

@@ -44,8 +44,17 @@ protected override async Task HandleCallAsyncCore(HttpContext httpContext, HttpC
             // Disable request body data rate for client streaming
             DisableMinRequestBodyDataRateAndMaxRequestBodySize(httpContext);
 
+            TResponse? response;
+
             var streamReader = new HttpContextStreamReader<TRequest>(serverCallContext, MethodInvoker.Method.RequestMarshaller.ContextualDeserializer);
-            var response = await _invoker.Invoke(httpContext, serverCallContext, streamReader);
+            try
+            {
+                response = await _invoker.Invoke(httpContext, serverCallContext, streamReader);
+            }
+            finally
+            {
+                streamReader.Complete();
+            }
 
             if (response == null)
             {

src/Grpc.AspNetCore.Server/Internal/CallHandlers/DuplexStreamingServerCallHandler.cs

@@ -42,15 +42,22 @@ public DuplexStreamingServerCallHandler(
             _invoker = invoker;
         }
 
-        protected override Task HandleCallAsyncCore(HttpContext httpContext, HttpContextServerCallContext serverCallContext)
+        protected override async Task HandleCallAsyncCore(HttpContext httpContext, HttpContextServerCallContext serverCallContext)
         {
             // Disable request body data rate for client streaming
             DisableMinRequestBodyDataRateAndMaxRequestBodySize(httpContext);
 
             var streamReader = new HttpContextStreamReader<TRequest>(serverCallContext, MethodInvoker.Method.RequestMarshaller.ContextualDeserializer);
             var streamWriter = new HttpContextStreamWriter<TResponse>(serverCallContext, MethodInvoker.Method.ResponseMarshaller.ContextualSerializer);
-
-            return _invoker.Invoke(httpContext, serverCallContext, streamReader, streamWriter);
+            try
+            {
+                await _invoker.Invoke(httpContext, serverCallContext, streamReader, streamWriter);
+            }
+            finally
+            {
+                streamReader.Complete();
+                streamWriter.Complete();
+            }
         }
     }
 }

src/Grpc.AspNetCore.Server/Internal/CallHandlers/ServerStreamingServerCallHandler.cs

@@ -48,7 +48,14 @@ protected override async Task HandleCallAsyncCore(HttpContext httpContext, HttpC
             var request = await httpContext.Request.BodyReader.ReadSingleMessageAsync<TRequest>(serverCallContext, MethodInvoker.Method.RequestMarshaller.ContextualDeserializer);
 
             var streamWriter = new HttpContextStreamWriter<TResponse>(serverCallContext, MethodInvoker.Method.ResponseMarshaller.ContextualSerializer);
-            await _invoker.Invoke(httpContext, serverCallContext, request, streamWriter);
+            try
+            {
+                await _invoker.Invoke(httpContext, serverCallContext, request, streamWriter);
+            }
+            finally
+            {
+                streamWriter.Complete();
+            }
         }
     }
 }

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

@@ -30,6 +30,7 @@ internal class HttpContextStreamReader<TRequest> : IAsyncStreamReader<TRequest>
 
         private readonly HttpContextServerCallContext _serverCallContext;
         private readonly Func<DeserializationContext, TRequest> _deserializer;
+        private bool _completed;
 
         public HttpContextStreamReader(HttpContextServerCallContext serverCallContext, Func<DeserializationContext, TRequest> deserializer)
         {
@@ -57,6 +58,11 @@ async Task<bool> MoveNextAsync(ValueTask<TRequest?> readStreamTask)
                 return Task.FromCanceled<bool>(cancellationToken);
             }
 
+            if (_completed || _serverCallContext.CancellationToken.IsCancellationRequested)
+            {
+                return Task.FromException<bool>(new InvalidOperationException("Can't read messages after the request is complete."));
+            }
+
             var request = _serverCallContext.HttpContext.Request.BodyReader.ReadStreamMessageAsync(_serverCallContext, _deserializer, cancellationToken);
             if (!request.IsCompletedSuccessfully)
             {
@@ -78,5 +84,10 @@ private bool ProcessPayload(TRequest? request)
             Current = request;
             return true;
         }
+
+        public void Complete()
+        {
+            _completed = true;
+        }
     }
 }

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

@@ -29,6 +29,7 @@ internal class HttpContextStreamWriter<TResponse> : IServerStreamWriter<TRespons
         private readonly Action<TResponse, SerializationContext> _serializer;
         private readonly object _writeLock;
         private Task? _writeTask;
+        private bool _completed;
 
         public HttpContextStreamWriter(HttpContextServerCallContext context, Action<TResponse, SerializationContext> serializer)
         {
@@ -50,9 +51,9 @@ public Task WriteAsync(TResponse message)
                 return Task.FromException(new ArgumentNullException(nameof(message)));
             }
 
-            if (_context.CancellationToken.IsCancellationRequested)
+            if (_completed || _context.CancellationToken.IsCancellationRequested)
             {
-                return Task.FromException(new InvalidOperationException("Cannot write message after request is complete."));
+                return Task.FromException(new InvalidOperationException("Can't write the message because the request is complete."));
             }
 
             lock (_writeLock)
@@ -70,6 +71,11 @@ public Task WriteAsync(TResponse message)
             return _writeTask;
         }
 
+        public void Complete()
+        {
+            _completed = true;
+        }
+
         /// <summary>
         /// A value indicating whether there is an async write already in progress.
         /// Should only check this property when holding the write lock.

test/FunctionalTests/Client/CancellationTests.cs

@@ -65,7 +65,14 @@ private async Task CancelInParallel(int tasks, bool waitForHeaders, int interati
 
                     // Cancellation when service is receiving message
                     if (writeContext.Exception is InvalidOperationException &&
-                        writeContext.Exception.Message == "Cannot write message after request is complete.")
+                        writeContext.Exception.Message == "Can't read messages after the request is complete.")
+                    {
+                        return true;
+                    }
+
+                    // Cancellation when service is writing message
+                    if (writeContext.Exception is InvalidOperationException &&
+                        writeContext.Exception.Message == "Can't write the message because the request is complete.")
                     {
                         return true;
                     }