Cancel downstream operation in TimeoutStrategy.Pessimistic (#1697)

Update to extend the lifetime of the combined `CancellationTokenSource` so that it will cancel downstream operations in `TimeoutStrategy.Pessimistic`.
App-vNext · Oct 17, 2023 · 940c628 · 940c628
1 parent a3ca9d0
commit 940c628
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 7 deletions.
diff --git a/src/Polly/Timeout/AsyncTimeoutEngine.cs b/src/Polly/Timeout/AsyncTimeoutEngine.cs
@@ -37,7 +37,6 @@ internal static async Task<TResult> ImplementationAsync<TResult>(
             actionTask = action(context, combinedToken);
 
             return await (await Task.WhenAny(actionTask, timeoutTask).ConfigureAwait(continueOnCapturedContext)).ConfigureAwait(continueOnCapturedContext);
-
         }
         catch (Exception ex)
         {
@@ -51,6 +50,16 @@ internal static async Task<TResult> ImplementationAsync<TResult>(
 
             throw;
         }
+        finally
+        {
+            // If timeoutCancellationTokenSource was canceled & our combined token hasn't been signaled, cancel it.
+            // This avoids the exception propagating before the linked token can signal the downstream to cancel.
+            // See https://github.com/App-vNext/Polly/issues/722.
+            if (!combinedTokenSource.IsCancellationRequested && timeoutCancellationTokenSource.IsCancellationRequested)
+            {
+                combinedTokenSource.Cancel();
+            }
+        }
     }
 
     private static Task<TResult> AsTask<TResult>(this CancellationToken cancellationToken)
@@ -60,11 +69,11 @@ private static Task<TResult> AsTask<TResult>(this CancellationToken cancellation
         // A generalised version of this method would include a hotpath returning a canceled task (rather than setting up a registration) if (cancellationToken.IsCancellationRequested) on entry.  This is omitted, since we only start the timeout countdown in the token _after calling this method.
 
         IDisposable registration = null;
-            registration = cancellationToken.Register(() =>
-            {
-                tcs.TrySetCanceled();
-                registration?.Dispose();
-            }, useSynchronizationContext: false);
+        registration = cancellationToken.Register(() =>
+        {
+            tcs.TrySetCanceled();
+            registration?.Dispose();
+        }, useSynchronizationContext: false);
 
         return tcs.Task;
     }

diff --git a/test/Polly.Specs/Timeout/TimeoutAsyncSpecs.cs b/test/Polly.Specs/Timeout/TimeoutAsyncSpecs.cs
@@ -1,4 +1,6 @@
-namespace Polly.Specs.Timeout;
+using System;
+
+namespace Polly.Specs.Timeout;
 
 [Collection(Constants.SystemClockDependentTestCollection)]
 public class TimeoutAsyncSpecs : TimeoutSpecsBase
@@ -248,6 +250,29 @@ public async Task Should_rethrow_exception_from_inside_delegate__pessimistic()
         await policy.Awaiting(p => p.ExecuteAsync(() => throw new NotImplementedException())).Should().ThrowAsync<NotImplementedException>();
     }
 
+    [Fact]
+    public async Task Should_cancel_downstream_token_on_timeout__pessimistic()
+    {
+        // It seems that there's a difference in the mocked clock vs. the real time clock.
+        // This test does not fail when using the mocked timer.
+        // In the TimeoutSpecsBase we actually cancel the combined token. Which hides
+        // the fact that it doesn't actually cancel irl.
+        SystemClock.Reset();
+
+        var policy = Policy.TimeoutAsync(TimeSpan.FromMilliseconds(200), TimeoutStrategy.Pessimistic);
+        bool isCancelled = false;
+
+        var act = () => policy.ExecuteAsync(async (combinedToken) =>
+        {
+            combinedToken.Register(() => isCancelled = true);
+            await SystemClock.SleepAsync(TimeSpan.FromMilliseconds(1000), combinedToken);
+        }, CancellationToken.None);
+
+        await act.Should().ThrowAsync<TimeoutRejectedException>();
+
+        isCancelled.Should().BeTrue();
+    }
+
     #endregion
 
     #region Timeout operation - optimistic