Fix for retry cancellation (#2456)

Fixes #2375.
App-vNext · Feb 4, 2025 · a953ba8 · a953ba8
1 parent 4f207a1
commit a953ba8
Show file tree

Hide file tree

Showing 2 changed files with 122 additions and 28 deletions.
diff --git a/src/Polly.Core/Retry/RetryResilienceStrategy.cs b/src/Polly.Core/Retry/RetryResilienceStrategy.cs
@@ -67,7 +67,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;
             }
@@ -100,17 +100,20 @@ protected internal override async ValueTask<Outcome<T>> ExecuteCore<TState>(Func
                 await DisposeHelper.TryDisposeSafeAsync(resultValue, context.IsSynchronous).ConfigureAwait(context.ContinueOnCapturedContext);
             }
 
-            // stryker disable once all : no means to test this
-            if (delay > TimeSpan.Zero)
+            try
             {
-                try
+                context.CancellationToken.ThrowIfCancellationRequested();
+
+                // stryker disable once all : no means to test this
+                if (delay > TimeSpan.Zero)
                 {
                     await _timeProvider.DelayAsync(delay, context).ConfigureAwait(context.ContinueOnCapturedContext);
                 }
-                catch (OperationCanceledException e)
-                {
-                    return Outcome.FromException<T>(e);
-                }
+
+            }
+            catch (OperationCanceledException e)
+            {
+                return Outcome.FromException<T>(e);
             }
 
             if (incrementAttempts)

diff --git a/test/Polly.Core.Tests/Retry/RetryResilienceStrategyTests.cs b/test/Polly.Core.Tests/Retry/RetryResilienceStrategyTests.cs
@@ -31,39 +31,130 @@ public void ExecuteAsync_EnsureResultNotDisposed()
     }
 
     [Fact]
-    public async Task ExecuteAsync_CancellationRequested_EnsureNotRetried()
+    public async Task ExecuteAsync_CanceledBeforeExecution_EnsureNotExecuted()
     {
-        SetupNoDelay();
         var sut = CreateSut();
-        using var cts = new CancellationTokenSource();
-        cts.Cancel();
-        var context = ResilienceContextPool.Shared.Get(cts.Token);
         var executed = false;
 
-        var result = await sut.ExecuteOutcomeAsync((_, _) => { executed = true; return Outcome.FromResultAsValueTask("dummy"); }, context, "state");
-        result.Exception.ShouldBeOfType<OperationCanceledException>();
+        var result = await sut.ExecuteOutcomeAsync(
+            (_, _) =>
+            {
+                executed = true;
+                return Outcome.FromResultAsValueTask(new object());
+            },
+            ResilienceContextPool.Shared.Get(new CancellationToken(canceled: true)),
+            default(object));
+
+        result.Exception.ShouldBeAssignableTo<OperationCanceledException>();
         executed.ShouldBeFalse();
     }
 
     [Fact]
-    public async Task ExecuteAsync_CancellationRequestedAfterCallback_EnsureNotRetried()
+    public async Task ExecuteAsync_CanceledDuringExecution_EnsureResultReturned()
     {
-        using var cts = new CancellationTokenSource();
+        var sut = CreateSut();
+        using var cancellation = new CancellationTokenSource();
+        var executions = 0;
+
+        var result = await sut.ExecuteOutcomeAsync(
+            (_, _) =>
+            {
+                executions++;
+                cancellation.Cancel();
+                return Outcome.FromResultAsValueTask(new object());
+            },
+            ResilienceContextPool.Shared.Get(cancellation.Token),
+            default(object));
+
+        result.Exception.ShouldBeNull();
+        executions.ShouldBe(1);
+    }
+
+    [Fact]
+    public async Task ExecuteAsync_CanceledDuringExecution_EnsureNotExecutedAgain()
+    {
+        var reported = false;
 
         _options.ShouldHandle = _ => PredicateResult.True();
-        _options.OnRetry = _ =>
-        {
-            cts.Cancel();
-            return default;
-        };
+        _options.OnRetry =
+            args =>
+            {
+                reported = true;
+                return default;
+            };
 
-        var sut = CreateSut(TimeProvider.System);
-        var context = ResilienceContextPool.Shared.Get(cts.Token);
-        var executed = false;
+        var sut = CreateSut();
+        using var cancellation = new CancellationTokenSource();
+        var executions = 0;
+
+        var result = await sut.ExecuteOutcomeAsync(
+            (_, _) =>
+            {
+                executions++;
+                cancellation.Cancel();
+                return Outcome.FromResultAsValueTask(new object());
+            },
+            ResilienceContextPool.Shared.Get(cancellation.Token),
+            default(object));
+
+        result.Exception.ShouldBeAssignableTo<OperationCanceledException>();
+        executions.ShouldBe(1);
+        reported.ShouldBeTrue();
+    }
+
+    [Fact]
+    public async Task ExecuteAsync_CanceledAfterExecution_EnsureNotExecutedAgain()
+    {
+        using var cancellation = new CancellationTokenSource();
+
+        _options.ShouldHandle = _ => PredicateResult.True();
+        _options.OnRetry =
+            args =>
+            {
+                cancellation.Cancel();
+                return default;
+            };
+
+        var sut = CreateSut();
+        var executions = 0;
+
+        var result = await sut.ExecuteOutcomeAsync(
+            (_, _) =>
+            {
+                executions++;
+                return Outcome.FromResultAsValueTask(new object());
+            },
+            ResilienceContextPool.Shared.Get(cancellation.Token),
+            default(object));
+
+        result.Exception.ShouldBeAssignableTo<OperationCanceledException>();
+        executions.ShouldBe(1);
+    }
+
+    [Fact]
+    public async Task ExecuteAsync_CanceledDuringDelay_EnsureNotExecutedAgain()
+    {
+        _options.ShouldHandle = _ => PredicateResult.True();
+
+        using var cancellation = _timeProvider.CreateCancellationTokenSource(_options.Delay);
+
+        var sut = CreateSut();
+        var executions = 0;
+
+        var resultTask = sut.ExecuteOutcomeAsync(
+            (_, _) =>
+            {
+                executions++;
+                return Outcome.FromResultAsValueTask(new object());
+            },
+            ResilienceContextPool.Shared.Get(cancellation.Token),
+            default(object));
+
+        _timeProvider.Advance(_options.Delay);
+        var result = await resultTask;
 
-        var result = await sut.ExecuteOutcomeAsync((_, _) => { executed = true; return Outcome.FromResultAsValueTask("dummy"); }, context, "state");
-        result.Exception.ShouldBeOfType<OperationCanceledException>();
-        executed.ShouldBeTrue();
+        result.Exception.ShouldBeAssignableTo<OperationCanceledException>();
+        executions.ShouldBe(1);
     }
 
     [Fact]