Update to cancel downstream operation in TimeoutStrategy.Pessimistic

src/Polly/Timeout/AsyncTimeoutEngine.cs

@@ -15,17 +15,28 @@
         TimeSpan timeout = timeoutProvider(context);
 
         using var timeoutCancellationTokenSource = new CancellationTokenSource();
-        using var combinedTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutCancellationTokenSource.Token);
+        // Do not use a using here, the exception will exit the scope before we have time to
+        // notify the downstream of the cancellation.
+        var combinedTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutCancellationTokenSource.Token);
 
         Task<TResult> actionTask = null;
         CancellationToken combinedToken = combinedTokenSource.Token;
 
         try
         {
+
+            var actionTaskFn = async () =>
+            {
+                var result = await action(context, combinedToken).ConfigureAwait(continueOnCapturedContext);
+                // dispose of the token source after we've waited on the task to finish.
+                combinedTokenSource.Dispose();
+                return result;
+            };
+
             if (timeoutStrategy == TimeoutStrategy.Optimistic)
             {
                 SystemClock.CancelTokenAfter(timeoutCancellationTokenSource, timeout);
-                return await action(context, combinedToken).ConfigureAwait(continueOnCapturedContext);
+                return await actionTaskFn().ConfigureAwait(continueOnCapturedContext);
             }
 
             // else: timeoutStrategy == TimeoutStrategy.Pessimistic
@@ -34,7 +45,7 @@
 
             SystemClock.CancelTokenAfter(timeoutCancellationTokenSource, timeout);
 
-            actionTask = action(context, combinedToken);
+            actionTask = actionTaskFn();
 
             return await (await Task.WhenAny(actionTask, timeoutTask).ConfigureAwait(continueOnCapturedContext)).ConfigureAwait(continueOnCapturedContext);
 

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