clientside writes should finish with TaskCanceledException if cancel was previously requested

4 files changed, 29 insertions, 19 deletions

diff --git a/src/csharp/Grpc.Core.Tests/Internal/AsyncCallTest.cs b/src/csharp/Grpc.Core.Tests/Internal/AsyncCallTest.cs
index 303a878742..81897f8c77 100644
--- a/src/csharp/Grpc.Core.Tests/Internal/AsyncCallTest.cs
+++ b/src/csharp/Grpc.Core.Tests/Internal/AsyncCallTest.cs
@@ -270,7 +270,7 @@ namespace Grpc.Core.Internal.Tests
         }
 
         [Test]
-        public void ClientStreaming_WriteAfterCancellationRequestThrowsOperationCancelledException()
+        public void ClientStreaming_WriteAfterCancellationRequestThrowsTaskCanceledException()
         {
             var resultTask = asyncCall.ClientStreamingCallAsync();
             var requestStream = new ClientRequestStream<string, string>(asyncCall);
@@ -278,8 +278,8 @@ namespace Grpc.Core.Internal.Tests
             asyncCall.Cancel();
             Assert.IsTrue(fakeCall.IsCancelled);
 
-            // TODO: awaiting the writeTask should throw TaskCanceledException
-            Assert.Throws(typeof(OperationCanceledException), () => requestStream.WriteAsync("request1"));
+            var writeTask = requestStream.WriteAsync("request1");
+            Assert.ThrowsAsync(typeof(TaskCanceledException), async () => await writeTask);
 
             fakeCall.UnaryResponseClientHandler(true,
                 CreateClientSideStatus(StatusCode.Cancelled),
@@ -416,7 +416,7 @@ namespace Grpc.Core.Internal.Tests
         }
 
         [Test]
-        public void DuplexStreaming_WriteAfterCancellationRequestThrowsOperationCancelledException()
+        public void DuplexStreaming_WriteAfterCancellationRequestThrowsTaskCanceledException()
         {
             asyncCall.StartDuplexStreamingCall();
             var requestStream = new ClientRequestStream<string, string>(asyncCall);
@@ -424,7 +424,9 @@ namespace Grpc.Core.Internal.Tests
 
             asyncCall.Cancel();
             Assert.IsTrue(fakeCall.IsCancelled);
-            Assert.Throws(typeof(OperationCanceledException), () => requestStream.WriteAsync("request1"));
+
+            var writeTask = requestStream.WriteAsync("request1");
+            Assert.ThrowsAsync(typeof(TaskCanceledException), async () => await writeTask);
 
             var readTask = responseStream.MoveNext();
             fakeCall.ReceivedMessageHandler(true, null);
diff --git a/src/csharp/Grpc.Core/Internal/AsyncCall.cs b/src/csharp/Grpc.Core/Internal/AsyncCall.cs
index ff70efc8b3..8652b297c8 100644
--- a/src/csharp/Grpc.Core/Internal/AsyncCall.cs
+++ b/src/csharp/Grpc.Core/Internal/AsyncCall.cs
@@ -252,7 +252,12 @@ namespace Grpc.Core.Internal
             lock (myLock)
             {
                 GrpcPreconditions.CheckState(started);
-                CheckSendPreconditionsClientSide();
+
+                var earlyResult = CheckSendPreconditionsClientSide();
+                if (earlyResult != null)
+                {
+                    return earlyResult;
+                }
 
                 if (disposed || finished)
                 {
@@ -338,7 +343,11 @@ namespace Grpc.Core.Internal
 
         protected override Task CheckSendAllowedOrEarlyResult()
         {
-            CheckSendPreconditionsClientSide();
+            var earlyResult = CheckSendPreconditionsClientSide();
+            if (earlyResult != null)
+            {
+                return earlyResult;
+            }
 
             if (finishedStatus.HasValue)
             {
@@ -355,11 +364,20 @@ namespace Grpc.Core.Internal
             return null;
         }
 
-        private void CheckSendPreconditionsClientSide()
+        private Task CheckSendPreconditionsClientSide()
         {
-            CheckNotCancelled();
+            if (cancelRequested)
+            {
+                // Return a cancelled task.
+                var tcs = new TaskCompletionSource<object>();
+                tcs.SetCanceled();
+                return tcs.Task;
+            }
+
             GrpcPreconditions.CheckState(!halfcloseRequested, "Request stream has already been completed.");
             GrpcPreconditions.CheckState(streamingWriteTcs == null, "Only one write can be pending at a time.");
+
+            return null;
         }
 
         private void Initialize(CompletionQueueSafeHandle cq)
diff --git a/src/csharp/Grpc.Core/Internal/AsyncCallBase.cs b/src/csharp/Grpc.Core/Internal/AsyncCallBase.cs
index d60876ddf3..5f561daedd 100644
--- a/src/csharp/Grpc.Core/Internal/AsyncCallBase.cs
+++ b/src/csharp/Grpc.Core/Internal/AsyncCallBase.cs
@@ -222,14 +222,6 @@ namespace Grpc.Core.Internal
         /// </summary>
         protected abstract Task CheckSendAllowedOrEarlyResult();
 
-        protected void CheckNotCancelled()
-        {
-            if (cancelRequested)
-            {
-                throw new OperationCanceledException("Remote call has been cancelled.");
-            }
-        }
-
         protected byte[] UnsafeSerialize(TWrite msg)
         {
             using (Profilers.ForCurrentThread().NewScope("AsyncCallBase.UnsafeSerialize"))
diff --git a/src/csharp/Grpc.Core/Internal/AsyncCallServer.cs b/src/csharp/Grpc.Core/Internal/AsyncCallServer.cs
index 109c207c75..d1bb80762e 100644
--- a/src/csharp/Grpc.Core/Internal/AsyncCallServer.cs
+++ b/src/csharp/Grpc.Core/Internal/AsyncCallServer.cs
@@ -200,8 +200,6 @@ namespace Grpc.Core.Internal
 
         protected override Task CheckSendAllowedOrEarlyResult()
         {
-            CheckNotCancelled();
-
             GrpcPreconditions.CheckState(!halfcloseRequested, "Response stream has already been completed.");
             GrpcPreconditions.CheckState(!finished, "Already finished.");
             GrpcPreconditions.CheckState(streamingWriteTcs == null, "Only one write can be pending at a time");