Skip to content

Commit 2be676a

Browse files
JamesNKJamesNK
authored
WriteAsync cancellation throws an error with the calls completed status if possible (#2170)
1 parent 73c726b commit 2be676a

File tree

8 files changed

+189
-64
lines changed

8 files changed

+189
-64
lines changed

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

Lines changed: 42 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
#region Copyright notice and license
1+
#region Copyright notice and license
22

33
// Copyright 2019 The gRPC Authors
44
//
@@ -53,6 +53,7 @@ public DefaultDeserializationContext DeserializationContext
5353

5454
public string? RequestGrpcEncoding { get; internal set; }
5555

56+
public abstract Task<Status> CallTask { get; }
5657
public abstract CancellationToken CancellationToken { get; }
5758
public abstract Type RequestType { get; }
5859
public abstract Type ResponseType { get; }
@@ -64,6 +65,29 @@ protected GrpcCall(CallOptions options, GrpcChannel channel)
6465
Logger = channel.LoggerFactory.CreateLogger(LoggerName);
6566
}
6667

68+
public Exception CreateCanceledStatusException(Exception? ex = null)
69+
{
70+
var status = (CallTask.IsCompletedSuccessfully()) ? CallTask.Result : new Status(StatusCode.Cancelled, string.Empty, ex);
71+
return CreateRpcException(status);
72+
}
73+
74+
public CancellationToken GetCanceledToken(CancellationToken methodCancellationToken)
75+
{
76+
if (methodCancellationToken.IsCancellationRequested)
77+
{
78+
return methodCancellationToken;
79+
}
80+
else if (Options.CancellationToken.IsCancellationRequested)
81+
{
82+
return Options.CancellationToken;
83+
}
84+
else if (CancellationToken.IsCancellationRequested)
85+
{
86+
return CancellationToken;
87+
}
88+
return CancellationToken.None;
89+
}
90+
6791
internal RpcException CreateRpcException(Status status)
6892
{
6993
// This code can be called from a background task.
@@ -84,6 +108,23 @@ internal RpcException CreateRpcException(Status status)
84108
return new RpcException(status, trailers ?? Metadata.Empty);
85109
}
86110

111+
public Exception CreateFailureStatusException(Status status)
112+
{
113+
if (Channel.ThrowOperationCanceledOnCancellation &&
114+
(status.StatusCode == StatusCode.DeadlineExceeded || status.StatusCode == StatusCode.Cancelled))
115+
{
116+
// Convert status response of DeadlineExceeded to OperationCanceledException when
117+
// ThrowOperationCanceledOnCancellation is true.
118+
// This avoids a race between the client-side timer and the server status throwing different
119+
// errors on deadline exceeded.
120+
return new OperationCanceledException();
121+
}
122+
else
123+
{
124+
return CreateRpcException(status);
125+
}
126+
}
127+
87128
protected bool TryGetTrailers([NotNullWhen(true)] out Metadata? trailers)
88129
{
89130
if (Trailers == null)

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

Lines changed: 1 addition & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,7 @@ private void ValidateDeadline(DateTime? deadline)
9090
}
9191
}
9292

93-
public Task<Status> CallTask => _callTcs.Task;
93+
public override Task<Status> CallTask => _callTcs.Task;
9494

9595
public override CancellationToken CancellationToken => _callCts.Token;
9696

@@ -248,12 +248,6 @@ public void EnsureNotDisposed()
248248
}
249249
}
250250

251-
public Exception CreateCanceledStatusException(Exception? ex = null)
252-
{
253-
var status = (CallTask.IsCompletedSuccessfully()) ? CallTask.Result : new Status(StatusCode.Cancelled, string.Empty, ex);
254-
return CreateRpcException(status);
255-
}
256-
257251
private void FinishResponseAndCleanUp(Status status)
258252
{
259253
ResponseFinished = true;
@@ -760,23 +754,6 @@ public Exception EnsureUserCancellationTokenReported(Exception ex, CancellationT
760754
return ex;
761755
}
762756

763-
public CancellationToken GetCanceledToken(CancellationToken methodCancellationToken)
764-
{
765-
if (methodCancellationToken.IsCancellationRequested)
766-
{
767-
return methodCancellationToken;
768-
}
769-
else if (Options.CancellationToken.IsCancellationRequested)
770-
{
771-
return Options.CancellationToken;
772-
}
773-
else if (CancellationToken.IsCancellationRequested)
774-
{
775-
return CancellationToken;
776-
}
777-
return CancellationToken.None;
778-
}
779-
780757
private void SetFailedResult(Status status)
781758
{
782759
CompatibilityHelpers.Assert(_responseTcs != null);
@@ -795,23 +772,6 @@ private void SetFailedResult(Status status)
795772
}
796773
}
797774

798-
public Exception CreateFailureStatusException(Status status)
799-
{
800-
if (Channel.ThrowOperationCanceledOnCancellation &&
801-
(status.StatusCode == StatusCode.DeadlineExceeded || status.StatusCode == StatusCode.Cancelled))
802-
{
803-
// Convert status response of DeadlineExceeded to OperationCanceledException when
804-
// ThrowOperationCanceledOnCancellation is true.
805-
// This avoids a race between the client-side timer and the server status throwing different
806-
// errors on deadline exceeded.
807-
return new OperationCanceledException();
808-
}
809-
else
810-
{
811-
return CreateRpcException(status);
812-
}
813-
}
814-
815775
private (bool diagnosticSourceEnabled, Activity? activity) InitializeCall(HttpRequestMessage request, TimeSpan? timeout)
816776
{
817777
GrpcCallLog.StartingCall(Logger, Method.Type, request.RequestUri!);

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

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
#region Copyright notice and license
1+
#region Copyright notice and license
22

33
// Copyright 2019 The gRPC Authors
44
//
@@ -72,4 +72,10 @@ protected override bool TryComputeLength(out long length)
7272
// Hacky. ReadAsStreamAsync does not complete until SerializeToStreamAsync finishes.
7373
// WARNING: Will run SerializeToStreamAsync again on .NET Framework.
7474
internal Task PushComplete => ReadAsStreamAsync();
75-
}
75+
76+
// Internal for testing.
77+
internal Task SerializeToStreamAsync(Stream stream)
78+
{
79+
return SerializeToStreamAsync(stream, context: null);
80+
}
81+
}

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

Lines changed: 32 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
#region Copyright notice and license
1+
#region Copyright notice and license
22

33
// Copyright 2019 The gRPC Authors
44
//
@@ -23,6 +23,7 @@
2323
using System.Runtime.InteropServices;
2424
using Grpc.Core;
2525
using Grpc.Net.Compression;
26+
using Grpc.Shared;
2627
using Microsoft.Extensions.Logging;
2728

2829
#if NETSTANDARD2_0
@@ -318,10 +319,9 @@ public static async ValueTask WriteMessageAsync<TMessage>(
318319
{
319320
GrpcCallLog.ErrorSendingMessage(call.Logger, ex);
320321

321-
// Cancellation from disposing response while waiting for WriteAsync can throw ObjectDisposedException.
322-
if (ex is ObjectDisposedException && call.CancellationToken.IsCancellationRequested)
322+
if (TryCreateCallCompleteException(ex, call, out var statusException))
323323
{
324-
throw new OperationCanceledException();
324+
throw statusException;
325325
}
326326

327327
throw;
@@ -342,24 +342,41 @@ public static async ValueTask WriteMessageAsync(
342342
{
343343
GrpcCallLog.SendingMessage(call.Logger);
344344

345-
try
346-
{
347-
// Sending the header+content in a single WriteAsync call has significant performance benefits
348-
// https://github.com/dotnet/runtime/issues/35184#issuecomment-626304981
349-
await stream.WriteAsync(data, cancellationToken).ConfigureAwait(false);
350-
}
351-
catch (ObjectDisposedException) when (call.CancellationToken.IsCancellationRequested)
352-
{
353-
// Cancellation from disposing response while waiting for WriteAsync can throw ObjectDisposedException.
354-
throw new OperationCanceledException();
355-
}
345+
// Sending the header+content in a single WriteAsync call has significant performance benefits
346+
// https://github.com/dotnet/runtime/issues/35184#issuecomment-626304981
347+
await stream.WriteAsync(data, cancellationToken).ConfigureAwait(false);
356348

357349
GrpcCallLog.MessageSent(call.Logger);
358350
}
359351
catch (Exception ex)
360352
{
361353
GrpcCallLog.ErrorSendingMessage(call.Logger, ex);
354+
355+
if (TryCreateCallCompleteException(ex, call, out var statusException))
356+
{
357+
throw statusException;
358+
}
359+
362360
throw;
363361
}
364362
}
363+
364+
private static bool TryCreateCallCompleteException(Exception originalException, GrpcCall call, [NotNullWhen(true)] out Exception? exception)
365+
{
366+
// The call may have been completed while WriteAsync was running and caused WriteAsync to throw.
367+
// In this situation, report the call's completed status.
368+
//
369+
// Replace exception with the status error if:
370+
// 1. The original exception is one Stream.WriteAsync throws if the call was completed during a write, and
371+
// 2. The call has already been successfully completed.
372+
if (originalException is OperationCanceledException or ObjectDisposedException &&
373+
call.CallTask.IsCompletedSuccessfully())
374+
{
375+
exception = call.CreateFailureStatusException(call.CallTask.Result);
376+
return true;
377+
}
378+
379+
exception = null;
380+
return false;
381+
}
365382
}

‎test/FunctionalTests/Client/StreamingTests.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
#region Copyright notice and license
1+
#region Copyright notice and license
22

33
// Copyright 2019 The gRPC Authors
44
//

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

Lines changed: 101 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
#region Copyright notice and license
1+
#region Copyright notice and license
22

33
// Copyright 2019 The gRPC Authors
44
//
@@ -107,7 +107,6 @@ public async Task AsyncClientStreamingCall_Success_RequestContentSent()
107107
var responseTask = call.ResponseAsync;
108108
Assert.IsFalse(responseTask.IsCompleted, "Response not returned until client stream is complete.");
109109

110-
111110
await call.RequestStream.WriteAsync(new HelloRequest { Name = "1" }).DefaultTimeout();
112111
await call.RequestStream.WriteAsync(new HelloRequest { Name = "2" }).DefaultTimeout();
113112

@@ -268,6 +267,106 @@ public async Task ClientStreamWriter_WriteAfterResponseHasFinished_ErrorThrown()
268267
Assert.AreEqual("Hello world", result.Message);
269268
}
270269

270+
[Test]
271+
public async Task AsyncClientStreamingCall_ErrorWhileWriting_StatusExceptionThrown()
272+
{
273+
// Arrange
274+
PushStreamContent<HelloRequest, HelloReply>? content = null;
275+
276+
var responseTcs = new TaskCompletionSource<HttpResponseMessage>(TaskCreationOptions.RunContinuationsAsynchronously);
277+
var httpClient = ClientTestHelpers.CreateTestClient(request =>
278+
{
279+
content = (PushStreamContent<HelloRequest, HelloReply>)request.Content!;
280+
return responseTcs.Task;
281+
});
282+
283+
var invoker = HttpClientCallInvokerFactory.Create(httpClient);
284+
285+
// Act
286+
287+
// Client starts call
288+
var call = invoker.AsyncClientStreamingCall<HelloRequest, HelloReply>(ClientTestHelpers.ServiceMethod, string.Empty, new CallOptions());
289+
// Client starts request stream write
290+
var writeTask = call.RequestStream.WriteAsync(new HelloRequest());
291+
292+
// Simulate HttpClient starting to accept the write. Stream.WriteAsync is blocked.
293+
var writeSyncPoint = new SyncPoint(runContinuationsAsynchronously: true);
294+
var testStream = new TestStream(writeSyncPoint);
295+
var serializeToStreamTask = content!.SerializeToStreamAsync(testStream);
296+
297+
// Server completes response.
298+
await writeSyncPoint.WaitForSyncPoint().DefaultTimeout();
299+
responseTcs.SetResult(ResponseUtils.CreateResponse(HttpStatusCode.OK, new ByteArrayContent(Array.Empty<byte>()), grpcStatusCode: StatusCode.InvalidArgument));
300+
301+
await ExceptionAssert.ThrowsAsync<RpcException>(() => call.ResponseAsync).DefaultTimeout();
302+
Assert.AreEqual(StatusCode.InvalidArgument, call.GetStatus().StatusCode);
303+
304+
// Unblock Stream.WriteAsync
305+
writeSyncPoint.Continue();
306+
307+
// Get error thrown from write task. It should have the status returned by the server.
308+
var ex = await ExceptionAssert.ThrowsAsync<RpcException>(() => writeTask).DefaultTimeout();
309+
310+
// Assert
311+
Assert.AreEqual(StatusCode.InvalidArgument, ex.StatusCode);
312+
Assert.AreEqual(StatusCode.InvalidArgument, call.GetStatus().StatusCode);
313+
Assert.AreEqual(string.Empty, call.GetStatus().Detail);
314+
}
315+
316+
private sealed class TestStream : Stream
317+
{
318+
private readonly SyncPoint _writeSyncPoint;
319+
320+
public TestStream(SyncPoint writeSyncPoint)
321+
{
322+
_writeSyncPoint = writeSyncPoint;
323+
}
324+
325+
public override bool CanRead { get; }
326+
public override bool CanSeek { get; }
327+
public override bool CanWrite { get; }
328+
public override long Length { get; }
329+
public override long Position { get; set; }
330+
331+
public override void Flush()
332+
{
333+
}
334+
335+
public override int Read(byte[] buffer, int offset, int count)
336+
{
337+
throw new NotImplementedException();
338+
}
339+
340+
public override long Seek(long offset, SeekOrigin origin)
341+
{
342+
throw new NotImplementedException();
343+
}
344+
345+
public override void SetLength(long value)
346+
{
347+
throw new NotImplementedException();
348+
}
349+
350+
public override void Write(byte[] buffer, int offset, int count)
351+
{
352+
throw new NotImplementedException();
353+
}
354+
355+
#if !NET472_OR_GREATER
356+
public override async ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
357+
{
358+
await _writeSyncPoint.WaitToContinue();
359+
throw new OperationCanceledException();
360+
}
361+
#else
362+
public override async Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
363+
{
364+
await _writeSyncPoint.WaitToContinue();
365+
throw new OperationCanceledException();
366+
}
367+
#endif
368+
}
369+
271370
[Test]
272371
public async Task ClientStreamWriter_CancelledBeforeCallStarts_ThrowsError()
273372
{

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
#region Copyright notice and license
1+
#region Copyright notice and license
22

33
// Copyright 2019 The gRPC Authors
44
//
@@ -318,6 +318,7 @@ public TestGrpcCall(CallOptions options, GrpcChannel channel) : base(options, ch
318318
public override Type RequestType { get; } = typeof(int);
319319
public override Type ResponseType { get; } = typeof(string);
320320
public override CancellationToken CancellationToken { get; }
321+
public override Task<Status> CallTask => Task.FromResult(Status.DefaultCancelled);
321322
}
322323

323324
private GrpcCallSerializationContext CreateSerializationContext(string? requestGrpcEncoding = null, int? maxSendMessageSize = null)

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
#region Copyright notice and license
1+
#region Copyright notice and license
22

33
// Copyright 2019 The gRPC Authors
44
//
@@ -65,5 +65,6 @@ public TestGrpcCall(CallOptions options, GrpcChannel channel, Type type) : base(
6565
public override Type RequestType => _type;
6666
public override Type ResponseType => _type;
6767
public override CancellationToken CancellationToken { get; }
68+
public override Task<Status> CallTask => Task.FromResult(Status.DefaultCancelled);
6869
}
6970
}

0 commit comments

Comments
 (0)