Handle CancellationToken for Retry

App-vNext · Nov 18, 2024 · 1401b0e · 1401b0e
1 parent 543cd50
commit 1401b0e
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 5 deletions.
diff --git a/src/Polly.Core/Retry/RetryResilienceStrategy.cs b/src/Polly.Core/Retry/RetryResilienceStrategy.cs
@@ -53,6 +53,15 @@ protected internal override async ValueTask<Outcome<T>> ExecuteCore<TState>(Func
         {
             var startTimestamp = _timeProvider.GetTimestamp();
             var outcome = await StrategyHelper.ExecuteCallbackSafeAsync(callback, context, state).ConfigureAwait(context.ContinueOnCapturedContext);
+            try
+            {
+                context.CancellationToken.ThrowIfCancellationRequested();
+            }
+            catch (OperationCanceledException e)
+            {
+                outcome = Outcome.FromException<T>(e);
+            }
+
             var shouldRetryArgs = new RetryPredicateArguments<T>(context, outcome, attempt);
             var handle = await ShouldHandle(shouldRetryArgs).ConfigureAwait(context.ContinueOnCapturedContext);
             var executionTime = _timeProvider.GetElapsedTime(startTimestamp);
@@ -67,7 +76,7 @@ protected internal override async ValueTask<Outcome<T>> ExecuteCore<TState>(Func
                 TelemetryUtil.ReportExecutionAttempt(_telemetry, context, outcome, attempt, executionTime, handle);
             }
 
-            if (context.CancellationToken.IsCancellationRequested || isLastAttempt || !handle)
+            if (isLastAttempt || !handle)
             {
                 return outcome;
             }

diff --git a/test/Polly.Core.Tests/Retry/RetryResilienceStrategyTests.cs b/test/Polly.Core.Tests/Retry/RetryResilienceStrategyTests.cs
@@ -31,11 +31,13 @@ public void ExecuteAsync_EnsureResultNotDisposed()
     }
 
     [Fact]
-    public async Task ExecuteAsync_CancellationRequested_EnsureNotRetried()
+    public async Task ExecuteAsync_CancellationRequestedBeforeCallback_EnsureNoAttempt()
     {
         SetupNoDelay();
-        var sut = CreateSut();
         using var cancellationToken = new CancellationTokenSource();
+        _options.ShouldHandle = _ => PredicateResult.True();
+        var sut = CreateSut();
+
         cancellationToken.Cancel();
         var context = ResilienceContextPool.Shared.Get();
         context.CancellationToken = cancellationToken.Token;
@@ -47,10 +49,36 @@ public async Task ExecuteAsync_CancellationRequested_EnsureNotRetried()
     }
 
     [Fact]
-    public async Task ExecuteAsync_CancellationRequestedAfterCallback_EnsureNotRetried()
+    public async Task ExecuteAsync_CancellationRequestedDuringCallback_EnsureNotRetried()
     {
+        SetupNoDelay();
         using var cancellationToken = new CancellationTokenSource();
+        _options.ShouldHandle = _ => PredicateResult.True();
+        var sut = CreateSut();
 
+        var context = ResilienceContextPool.Shared.Get();
+        context.CancellationToken = cancellationToken.Token;
+        var executed = false;
+        var attemptCounter = 0;
+
+        var result = await sut.ExecuteOutcomeAsync((_, _) =>
+        {
+            executed = true;
+            ++attemptCounter;
+            cancellationToken.Cancel();
+            return Outcome.FromResultAsValueTask("dummy");
+        }, context, "state");
+
+        result.Exception.Should().BeOfType<OperationCanceledException>();
+        executed.Should().BeTrue();
+        attemptCounter.Should().Be(1);
+    }
+
+    [Fact]
+    public async Task ExecuteAsync_CancellationRequestedAfterCallback_EnsureNotRetried()
+    {
+        SetupNoDelay();
+        using var cancellationToken = new CancellationTokenSource();
         _options.ShouldHandle = _ => PredicateResult.True();
         _options.OnRetry = _ =>
         {
@@ -62,10 +90,18 @@ public async Task ExecuteAsync_CancellationRequestedAfterCallback_EnsureNotRetri
         var context = ResilienceContextPool.Shared.Get();
         context.CancellationToken = cancellationToken.Token;
         var executed = false;
+        var attemptCounter = 0;
+
+        var result = await sut.ExecuteOutcomeAsync((_, _) =>
+        {
+            executed = true;
+            ++attemptCounter;
+            return Outcome.FromResultAsValueTask("dummy");
+        }, context, "state");
 
-        var result = await sut.ExecuteOutcomeAsync((_, _) => { executed = true; return Outcome.FromResultAsValueTask("dummy"); }, context, "state");
         result.Exception.Should().BeOfType<OperationCanceledException>();
         executed.Should().BeTrue();
+        attemptCounter.Should().Be(1);
     }
 
     [Fact]