Update to cancel downstream operation in TimeoutStrategy.Pessimistic

src/Polly/Timeout/AsyncTimeoutEngine.cs

@@ -15,6 +15,8 @@ internal static async Task<TResult> ImplementationAsync<TResult>(
         TimeSpan timeout = timeoutProvider(context);
 
         using var timeoutCancellationTokenSource = new CancellationTokenSource();
+        // Do not use a using here, the exception will exit the scope before we have time to
+        // notify the downstream of the cancellation.
         using var combinedTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutCancellationTokenSource.Token);
 
         Task<TResult> actionTask = null;
@@ -37,7 +39,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 +52,15 @@ internal static async Task<TResult> ImplementationAsync<TResult>(
 
             throw;
         }
+        finally
+        {
+            // If the timeoutCancellation was canceled & our combined token hasn't been signaled, signal it.
+            // This avoids the exception propagating before the linked token can signal the downstream to cancel.
+            if (!combinedTokenSource.IsCancellationRequested && timeoutCancellationTokenSource.IsCancellationRequested)
+            {
+                combinedTokenSource.Cancel();
+            }
+        }
     }
 
     private static Task<TResult> AsTask<TResult>(this CancellationToken cancellationToken)
@@ -60,11 +70,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;
     }

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,31 @@ 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);
+        var upstreamToken = new CancellationToken();
+        bool isCancelled = false;
+
+        var act = () => policy.ExecuteAsync(async (combinedToken) =>
+        {
+            combinedToken.Register(() => isCancelled = true);
+            await SystemClock.SleepAsync(TimeSpan.FromMilliseconds(300), combinedToken);
+        }, upstreamToken);
+
+        await act.Should().ThrowAsync<TimeoutRejectedException>();
+
+        await SystemClock.SleepAsync(TimeSpan.FromMilliseconds(400), default);
+        isCancelled.Should().BeTrue();
+    }
+
     #endregion
 
     #region Timeout operation - optimistic