From 7f98d7954e219aeb09996d873aab526339497888 Mon Sep 17 00:00:00 2001
From: Martin Tomka <martintomka@microsoft.com>
Date: Thu, 22 Jun 2023 12:22:24 +0200
Subject: [PATCH] Fix hedging being cancelled too early

---
 .../Hedging/HedgingResilienceStrategy.cs      |  8 +++++--
 .../Hedging/HedgingResilienceStrategyTests.cs | 22 +++++++++++++++++++
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/src/Polly.Core/Hedging/HedgingResilienceStrategy.cs b/src/Polly.Core/Hedging/HedgingResilienceStrategy.cs
index d497271ac1..d96e0197d9 100644
--- a/src/Polly.Core/Hedging/HedgingResilienceStrategy.cs
+++ b/src/Polly.Core/Hedging/HedgingResilienceStrategy.cs
@@ -67,12 +67,16 @@ private async ValueTask<Outcome<T>> ExecuteCoreAsync<TState>(
         ResilienceContext context,
         TState state)
     {
+        // Capture the original cancellation token so it stays the same while hedging is executing.
+        // If we do not do this the inner strategy can replace the cancellation token and with the concurrent
+        // nature of hedging this can cause issues.
+        var cancellationToken = context.CancellationToken;
+        var continueOnCapturedContext = context.ContinueOnCapturedContext;
+
         var attempt = -1;
         while (true)
         {
             attempt++;
-            var continueOnCapturedContext = context.ContinueOnCapturedContext;
-            var cancellationToken = context.CancellationToken;
             var start = _timeProvider.GetTimestamp();
             if (cancellationToken.IsCancellationRequested)
             {
diff --git a/test/Polly.Core.Tests/Hedging/HedgingResilienceStrategyTests.cs b/test/Polly.Core.Tests/Hedging/HedgingResilienceStrategyTests.cs
index 3119ff42a0..663067abc7 100644
--- a/test/Polly.Core.Tests/Hedging/HedgingResilienceStrategyTests.cs
+++ b/test/Polly.Core.Tests/Hedging/HedgingResilienceStrategyTests.cs
@@ -104,6 +104,28 @@ public void ExecutePrimaryAndSecondary_EnsureAttemptReported()
         attempts[1].Attempt.Should().Be(1);
     }
 
+    [Fact]
+    public async Task ExecutePrimary_Cancelled_SecondaryShouldBeExecuted()
+    {
+        _options.MaxHedgedAttempts = 2;
+
+        ConfigureHedging(o => o.Result == "primary", args => () => Outcome.FromResultAsTask("secondary"));
+        var strategy = Create();
+
+        var result = await strategy.ExecuteAsync(
+            context =>
+            {
+                var source = new CancellationTokenSource();
+                source.Cancel();
+                context.CancellationToken = source.Token;
+
+                return new ValueTask<string>("primary");
+            },
+            ResilienceContext.Get());
+
+        result.Should().Be("secondary");
+    }
+
     [InlineData(-1)]
     [InlineData(-1000)]
     [InlineData(0)]