From b84d77267384e480324f2cd4ed5f4a7ff879f2c6 Mon Sep 17 00:00:00 2001
From: James Newton-King <james@newtonking.com>
Date: Tue, 2 Jan 2024 14:14:07 +0800
Subject: [PATCH 1/2] Fix cancellation token reported when using retries

---
 .../Internal/Retry/RetryCallBase.cs           | 17 ++++++++--
 .../Internal/Retry/StatusGrpcCall.cs          | 13 ++++++--
 .../Client/CancellationTests.cs               | 33 +++++++++++++++++++
 3 files changed, 59 insertions(+), 4 deletions(-)

diff --git a/src/Grpc.Net.Client/Internal/Retry/RetryCallBase.cs b/src/Grpc.Net.Client/Internal/Retry/RetryCallBase.cs
index e481dfda5..db4ecdb81 100644
--- a/src/Grpc.Net.Client/Internal/Retry/RetryCallBase.cs
+++ b/src/Grpc.Net.Client/Internal/Retry/RetryCallBase.cs
@@ -529,7 +529,7 @@ protected void HandleUnexpectedError(Exception ex)
         CommitReason commitReason;
 
         // Cancellation token triggered by dispose could throw here.
-        if (ex is OperationCanceledException && CancellationTokenSource.IsCancellationRequested)
+        if (ex is OperationCanceledException operationCanceledException && CancellationTokenSource.IsCancellationRequested)
         {
             // Cancellation could have been caused by an exceeded deadline.
             if (IsDeadlineExceeded())
@@ -542,7 +542,20 @@ protected void HandleUnexpectedError(Exception ex)
             else
             {
                 commitReason = CommitReason.Canceled;
-                resolvedCall = CreateStatusCall(Disposed ? GrpcProtocolConstants.CreateDisposeCanceledStatus(ex) : GrpcProtocolConstants.CreateClientCanceledStatus(ex));
+                Status status;
+                if (Disposed)
+                {
+                    status = GrpcProtocolConstants.CreateDisposeCanceledStatus(exception: null);
+                }
+                else
+                {
+                    if (Options.CancellationToken.IsCancellationRequested && Options.CancellationToken != operationCanceledException.CancellationToken)
+                    {
+                        ex = new OperationCanceledException(Options.CancellationToken);
+                    }
+                    status = GrpcProtocolConstants.CreateClientCanceledStatus(ex);
+                }
+                resolvedCall = CreateStatusCall(status);
             }
         }
         else
diff --git a/src/Grpc.Net.Client/Internal/Retry/StatusGrpcCall.cs b/src/Grpc.Net.Client/Internal/Retry/StatusGrpcCall.cs
index 0b9426aaf..00885b9fc 100644
--- a/src/Grpc.Net.Client/Internal/Retry/StatusGrpcCall.cs
+++ b/src/Grpc.Net.Client/Internal/Retry/StatusGrpcCall.cs
@@ -56,12 +56,21 @@ public void Dispose()
 
     public Task<TResponse> GetResponseAsync()
     {
-        return Task.FromException<TResponse>(new RpcException(_status));
+        return CreateErrorTask<TResponse>();
     }
 
     public Task<Metadata> GetResponseHeadersAsync()
     {
-        return Task.FromException<Metadata>(new RpcException(_status));
+        return CreateErrorTask<Metadata>();
+    }
+
+    private Task<T> CreateErrorTask<T>()
+    {
+        if (_channel.ThrowOperationCanceledOnCancellation && _status.DebugException is OperationCanceledException ex)
+        {
+            return Task.FromException<T>(ex);
+        }
+        return Task.FromException<T>(new RpcException(_status));
     }
 
     public Status GetStatus()
diff --git a/test/FunctionalTests/Client/CancellationTests.cs b/test/FunctionalTests/Client/CancellationTests.cs
index 0b85d093f..20af72dbd 100644
--- a/test/FunctionalTests/Client/CancellationTests.cs
+++ b/test/FunctionalTests/Client/CancellationTests.cs
@@ -509,6 +509,39 @@ async Task<DataMessage> UnaryMethod(DataMessage request, ServerCallContext conte
         Assert.AreEqual(StatusCode.Cancelled, call.GetStatus().StatusCode);
     }
 
+    [Test]
+    public async Task Unary_Retry_CancellationImmediately_TokenMatchesSource()
+    {
+        var tcs = new TaskCompletionSource<object?>(TaskCreationOptions.RunContinuationsAsynchronously);
+        async Task<DataMessage> UnaryMethod(DataMessage request, ServerCallContext context)
+        {
+            await tcs.Task;
+            return new DataMessage();
+        }
+
+        SetExpectedErrorsFilter(writeContext =>
+        {
+            return true;
+        });
+
+        // Arrange
+        var method = Fixture.DynamicGrpc.AddUnaryMethod<DataMessage, DataMessage>(UnaryMethod);
+        var serviceConfig = ServiceConfigHelpers.CreateRetryServiceConfig();
+        var channel = CreateChannel(throwOperationCanceledOnCancellation: true, serviceConfig: serviceConfig);
+        var client = TestClientFactory.Create(channel, method);
+
+        // Act
+        var cts = new CancellationTokenSource();
+        cts.Cancel();
+
+        var call = client.UnaryCall(new DataMessage(), new CallOptions(cancellationToken: cts.Token));
+
+        // Assert
+        var ex = await ExceptionAssert.ThrowsAsync<OperationCanceledException>(() => call.ResponseAsync).DefaultTimeout();
+        Assert.AreEqual(cts.Token, ex.CancellationToken);
+        Assert.AreEqual(StatusCode.Cancelled, call.GetStatus().StatusCode);
+    }
+
     [Test]
     public async Task ServerStreaming_CancellationDuringCall_TokenMatchesSource()
     {

From 8a6b8b89b5d180e4900f38343849498c2618cf5f Mon Sep 17 00:00:00 2001
From: James Newton-King <james@newtonking.com>
Date: Tue, 2 Jan 2024 20:45:39 +0800
Subject: [PATCH 2/2] Comment

---
 src/Grpc.Net.Client/Internal/Retry/RetryCallBase.cs | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/Grpc.Net.Client/Internal/Retry/RetryCallBase.cs b/src/Grpc.Net.Client/Internal/Retry/RetryCallBase.cs
index db4ecdb81..7f54895f1 100644
--- a/src/Grpc.Net.Client/Internal/Retry/RetryCallBase.cs
+++ b/src/Grpc.Net.Client/Internal/Retry/RetryCallBase.cs
@@ -549,6 +549,7 @@ protected void HandleUnexpectedError(Exception ex)
                 }
                 else
                 {
+                    // Replace the OCE from CancellationTokenSource with an OCE that has the passed in cancellation token if it is canceled.
                     if (Options.CancellationToken.IsCancellationRequested && Options.CancellationToken != operationCanceledException.CancellationToken)
                     {
                         ex = new OperationCanceledException(Options.CancellationToken);