WriteAsync cancellation throws an error with the calls completed stat…

…us if possible (#2170)

src/Grpc.Net.Client/Internal/GrpcCall.NonGeneric.cs

@@ -1,4 +1,4 @@
-#region Copyright notice and license
+#region Copyright notice and license
 
 // Copyright 2019 The gRPC Authors
 //
@@ -53,6 +53,7 @@ public DefaultDeserializationContext DeserializationContext
 
  public string? RequestGrpcEncoding { get; internal set; }
 
+ public abstract Task<Status> CallTask { get; }
  public abstract CancellationToken CancellationToken { get; }
  public abstract Type RequestType { get; }
  public abstract Type ResponseType { get; }
@@ -64,6 +65,29 @@ protected GrpcCall(CallOptions options, GrpcChannel channel)
  Logger = channel.LoggerFactory.CreateLogger(LoggerName);
  }
 
+ public Exception CreateCanceledStatusException(Exception? ex = null)
+ {
+ var status = (CallTask.IsCompletedSuccessfully()) ? CallTask.Result : new Status(StatusCode.Cancelled, string.Empty, ex);
+ return CreateRpcException(status);
+ }
+
+ public CancellationToken GetCanceledToken(CancellationToken methodCancellationToken)
+ {
+ if (methodCancellationToken.IsCancellationRequested)
+ {
+ return methodCancellationToken;
+ }
+ else if (Options.CancellationToken.IsCancellationRequested)
+ {
+ return Options.CancellationToken;
+ }
+ else if (CancellationToken.IsCancellationRequested)
+ {
+ return CancellationToken;
+ }
+ return CancellationToken.None;
+ }
+
  internal RpcException CreateRpcException(Status status)
  {
  // This code can be called from a background task.
@@ -84,6 +108,23 @@ internal RpcException CreateRpcException(Status status)
  return new RpcException(status, trailers ?? Metadata.Empty);
  }
 
+ public Exception CreateFailureStatusException(Status status)
+ {
+ if (Channel.ThrowOperationCanceledOnCancellation &&
+ (status.StatusCode == StatusCode.DeadlineExceeded || status.StatusCode == StatusCode.Cancelled))
+ {
+ // Convert status response of DeadlineExceeded to OperationCanceledException when
+ // ThrowOperationCanceledOnCancellation is true.
+ // This avoids a race between the client-side timer and the server status throwing different
+ // errors on deadline exceeded.
+ return new OperationCanceledException();
+ }
+ else
+ {
+ return CreateRpcException(status);
+ }
+ }
+
  protected bool TryGetTrailers([NotNullWhen(true)] out Metadata? trailers)
  {
  if (Trailers == null)

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

@@ -90,7 +90,7 @@ private void ValidateDeadline(DateTime? deadline)
  }
  }
 
- public Task<Status> CallTask => _callTcs.Task;
+ public override Task<Status> CallTask => _callTcs.Task;
 
  public override CancellationToken CancellationToken => _callCts.Token;
 
@@ -248,12 +248,6 @@ public void EnsureNotDisposed()
  }
  }
 
- public Exception CreateCanceledStatusException(Exception? ex = null)
- {
- var status = (CallTask.IsCompletedSuccessfully()) ? CallTask.Result : new Status(StatusCode.Cancelled, string.Empty, ex);
- return CreateRpcException(status);
- }
-
  private void FinishResponseAndCleanUp(Status status)
  {
  ResponseFinished = true;
@@ -760,23 +754,6 @@ public Exception EnsureUserCancellationTokenReported(Exception ex, CancellationT
  return ex;
  }
 
- public CancellationToken GetCanceledToken(CancellationToken methodCancellationToken)
- {
- if (methodCancellationToken.IsCancellationRequested)
- {
- return methodCancellationToken;
- }
- else if (Options.CancellationToken.IsCancellationRequested)
- {
- return Options.CancellationToken;
- }
- else if (CancellationToken.IsCancellationRequested)
- {
- return CancellationToken;
- }
- return CancellationToken.None;
- }
-
  private void SetFailedResult(Status status)
  {
  CompatibilityHelpers.Assert(_responseTcs != null);
@@ -795,23 +772,6 @@ private void SetFailedResult(Status status)
  }
  }
 
- public Exception CreateFailureStatusException(Status status)
- {
- if (Channel.ThrowOperationCanceledOnCancellation &&
- (status.StatusCode == StatusCode.DeadlineExceeded || status.StatusCode == StatusCode.Cancelled))
- {
- // Convert status response of DeadlineExceeded to OperationCanceledException when
- // ThrowOperationCanceledOnCancellation is true.
- // This avoids a race between the client-side timer and the server status throwing different
- // errors on deadline exceeded.
- return new OperationCanceledException();
- }
- else
- {
- return CreateRpcException(status);
- }
- }
-
  private (bool diagnosticSourceEnabled, Activity? activity) InitializeCall(HttpRequestMessage request, TimeSpan? timeout)
  {
  GrpcCallLog.StartingCall(Logger, Method.Type, request.RequestUri!);

src/Grpc.Net.Client/Internal/Http/PushStreamContent.cs

@@ -1,4 +1,4 @@
-#region Copyright notice and license
+#region Copyright notice and license
 
 // Copyright 2019 The gRPC Authors
 //
@@ -72,4 +72,10 @@ protected override bool TryComputeLength(out long length)
  // Hacky. ReadAsStreamAsync does not complete until SerializeToStreamAsync finishes.
  // WARNING: Will run SerializeToStreamAsync again on .NET Framework.
  internal Task PushComplete => ReadAsStreamAsync();
-}
+
+ // Internal for testing.
+ internal Task SerializeToStreamAsync(Stream stream)
+ {
+ return SerializeToStreamAsync(stream, context: null);
+ }
+}

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

@@ -1,4 +1,4 @@
-#region Copyright notice and license
+#region Copyright notice and license
 
 // Copyright 2019 The gRPC Authors
 //
@@ -23,6 +23,7 @@
 using System.Runtime.InteropServices;
 using Grpc.Core;
 using Grpc.Net.Compression;
+using Grpc.Shared;
 using Microsoft.Extensions.Logging;
 
 #if NETSTANDARD2_0
@@ -318,10 +319,9 @@ public static async ValueTask WriteMessageAsync<TMessage>(
  {
  GrpcCallLog.ErrorSendingMessage(call.Logger, ex);
 
- // Cancellation from disposing response while waiting for WriteAsync can throw ObjectDisposedException.
- if (ex is ObjectDisposedException && call.CancellationToken.IsCancellationRequested)
+ if (TryCreateCallCompleteException(ex, call, out var statusException))
  {
- throw new OperationCanceledException();
+ throw statusException;
  }
 
  throw;
@@ -342,24 +342,41 @@ public static async ValueTask WriteMessageAsync(
  {
  GrpcCallLog.SendingMessage(call.Logger);
 
- try
- {
- // Sending the header+content in a single WriteAsync call has significant performance benefits
- // https://github.com/dotnet/runtime/issues/35184#issuecomment-626304981
- await stream.WriteAsync(data, cancellationToken).ConfigureAwait(false);
- }
- catch (ObjectDisposedException) when (call.CancellationToken.IsCancellationRequested)
- {
- // Cancellation from disposing response while waiting for WriteAsync can throw ObjectDisposedException.
- throw new OperationCanceledException();
- }
+ // Sending the header+content in a single WriteAsync call has significant performance benefits
+ // https://github.com/dotnet/runtime/issues/35184#issuecomment-626304981
+ await stream.WriteAsync(data, cancellationToken).ConfigureAwait(false);
 
  GrpcCallLog.MessageSent(call.Logger);
  }
  catch (Exception ex)
  {
  GrpcCallLog.ErrorSendingMessage(call.Logger, ex);
+
+ if (TryCreateCallCompleteException(ex, call, out var statusException))
+ {
+ throw statusException;
+ }
+
  throw;
  }
  }
+
+ private static bool TryCreateCallCompleteException(Exception originalException, GrpcCall call, [NotNullWhen(true)] out Exception? exception)
+ {
+ // The call may have been completed while WriteAsync was running and caused WriteAsync to throw.
+ // In this situation, report the call's completed status.
+ //
+ // Replace exception with the status error if:
+ // 1. The original exception is one Stream.WriteAsync throws if the call was completed during a write, and
+ // 2. The call has already been successfully completed.
+ if (originalException is OperationCanceledException or ObjectDisposedException &&
+ call.CallTask.IsCompletedSuccessfully())
+ {
+ exception = call.CreateFailureStatusException(call.CallTask.Result);
+ return true;
+ }
+
+ exception = null;
+ return false;
+ }
 }

test/FunctionalTests/Client/StreamingTests.cs

@@ -1,4 +1,4 @@
-#region Copyright notice and license
+#region Copyright notice and license
 
 // Copyright 2019 The gRPC Authors
 //

test/Grpc.Net.Client.Tests/AsyncClientStreamingCallTests.cs

@@ -1,4 +1,4 @@
-#region Copyright notice and license
+#region Copyright notice and license
 
 // Copyright 2019 The gRPC Authors
 //
@@ -107,7 +107,6 @@ public async Task AsyncClientStreamingCall_Success_RequestContentSent()
  var responseTask = call.ResponseAsync;
  Assert.IsFalse(responseTask.IsCompleted, "Response not returned until client stream is complete.");
 
-
  await call.RequestStream.WriteAsync(new HelloRequest { Name = "1" }).DefaultTimeout();
  await call.RequestStream.WriteAsync(new HelloRequest { Name = "2" }).DefaultTimeout();
 
@@ -268,6 +267,106 @@ public async Task ClientStreamWriter_WriteAfterResponseHasFinished_ErrorThrown()
  Assert.AreEqual("Hello world", result.Message);
  }
 
+ [Test]
+ public async Task AsyncClientStreamingCall_ErrorWhileWriting_StatusExceptionThrown()
+ {
+ // Arrange
+ PushStreamContent<HelloRequest, HelloReply>? content = null;
+
+ var responseTcs = new TaskCompletionSource<HttpResponseMessage>(TaskCreationOptions.RunContinuationsAsynchronously);
+ var httpClient = ClientTestHelpers.CreateTestClient(request =>
+ {
+ content = (PushStreamContent<HelloRequest, HelloReply>)request.Content!;
+ return responseTcs.Task;
+ });
+
+ var invoker = HttpClientCallInvokerFactory.Create(httpClient);
+
+ // Act
+
+ // Client starts call
+ var call = invoker.AsyncClientStreamingCall<HelloRequest, HelloReply>(ClientTestHelpers.ServiceMethod, string.Empty, new CallOptions());
+ // Client starts request stream write
+ var writeTask = call.RequestStream.WriteAsync(new HelloRequest());
+
+ // Simulate HttpClient starting to accept the write. Stream.WriteAsync is blocked.
+ var writeSyncPoint = new SyncPoint(runContinuationsAsynchronously: true);
+ var testStream = new TestStream(writeSyncPoint);
+ var serializeToStreamTask = content!.SerializeToStreamAsync(testStream);
+
+ // Server completes response.
+ await writeSyncPoint.WaitForSyncPoint().DefaultTimeout();
+ responseTcs.SetResult(ResponseUtils.CreateResponse(HttpStatusCode.OK, new ByteArrayContent(Array.Empty<byte>()), grpcStatusCode: StatusCode.InvalidArgument));
+
+ await ExceptionAssert.ThrowsAsync<RpcException>(() => call.ResponseAsync).DefaultTimeout();
+ Assert.AreEqual(StatusCode.InvalidArgument, call.GetStatus().StatusCode);
+
+ // Unblock Stream.WriteAsync
+ writeSyncPoint.Continue();
+
+ // Get error thrown from write task. It should have the status returned by the server.
+ var ex = await ExceptionAssert.ThrowsAsync<RpcException>(() => writeTask).DefaultTimeout();
+
+ // Assert
+ Assert.AreEqual(StatusCode.InvalidArgument, ex.StatusCode);
+ Assert.AreEqual(StatusCode.InvalidArgument, call.GetStatus().StatusCode);
+ Assert.AreEqual(string.Empty, call.GetStatus().Detail);
+ }
+
+ private sealed class TestStream : Stream
+ {
+ private readonly SyncPoint _writeSyncPoint;
+
+ public TestStream(SyncPoint writeSyncPoint)
+ {
+ _writeSyncPoint = writeSyncPoint;
+ }
+
+ public override bool CanRead { get; }
+ public override bool CanSeek { get; }
+ public override bool CanWrite { get; }
+ public override long Length { get; }
+ public override long Position { get; set; }
+
+ public override void Flush()
+ {
+ }
+
+ public override int Read(byte[] buffer, int offset, int count)
+ {
+ throw new NotImplementedException();
+ }
+
+ public override long Seek(long offset, SeekOrigin origin)
+ {
+ throw new NotImplementedException();
+ }
+
+ public override void SetLength(long value)
+ {
+ throw new NotImplementedException();
+ }
+
+ public override void Write(byte[] buffer, int offset, int count)
+ {
+ throw new NotImplementedException();
+ }
+
+#if !NET472_OR_GREATER
+ public override async ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
+ {
+ await _writeSyncPoint.WaitToContinue();
+ throw new OperationCanceledException();
+ }
+#else
+ public override async Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+ {
+ await _writeSyncPoint.WaitToContinue();
+ throw new OperationCanceledException();
+ }
+#endif
+ }
+
  [Test]
  public async Task ClientStreamWriter_CancelledBeforeCallStarts_ThrowsError()
  {

test/Grpc.Net.Client.Tests/GrpcCallSerializationContextTests.cs

@@ -1,4 +1,4 @@
-#region Copyright notice and license
+#region Copyright notice and license
 
 // Copyright 2019 The gRPC Authors
 //
@@ -318,6 +318,7 @@ public TestGrpcCall(CallOptions options, GrpcChannel channel) : base(options, ch
  public override Type RequestType { get; } = typeof(int);
  public override Type ResponseType { get; } = typeof(string);
  public override CancellationToken CancellationToken { get; }
+ public override Task<Status> CallTask => Task.FromResult(Status.DefaultCancelled);
  }
 
  private GrpcCallSerializationContext CreateSerializationContext(string? requestGrpcEncoding = null, int? maxSendMessageSize = null)

test/Grpc.Net.Client.Tests/Infrastructure/StreamSerializationHelper.cs

@@ -1,4 +1,4 @@
-#region Copyright notice and license
+#region Copyright notice and license
 
 // Copyright 2019 The gRPC Authors
 //
@@ -65,5 +65,6 @@ public TestGrpcCall(CallOptions options, GrpcChannel channel, Type type) : base(
  public override Type RequestType => _type;
  public override Type ResponseType => _type;
  public override CancellationToken CancellationToken { get; }
+ public override Task<Status> CallTask => Task.FromResult(Status.DefaultCancelled);
  }
 }