Enhance OnHedgingArguments

src/Polly.Core/Hedging/Controller/HedgingHandler.cs

@@ -16,6 +16,7 @@ internal sealed record class HedgingHandler<T>(
         if (IsGeneric)
         {
             var copiedArgs = new HedgingActionGeneratorArguments<T>(
+                args.PrimaryContext,
                 args.Context,
                 args.Attempt,
                 (Func<ResilienceContext, ValueTask<Outcome<T>>>)(object)args.Callback);
@@ -29,7 +30,7 @@ internal sealed record class HedgingHandler<T>(
     private Func<ValueTask<Outcome<TResult>>>? CreateNonGenericAction<TResult>(HedgingActionGeneratorArguments<TResult> args)
     {
         var generator = (Func<HedgingActionGeneratorArguments<object>, Func<ValueTask<Outcome<object>>>?>)(object)ActionGenerator;
-        var action = generator(new HedgingActionGeneratorArguments<object>(args.Context, args.Attempt, async context =>
+        var action = generator(new HedgingActionGeneratorArguments<object>(args.PrimaryContext, args.Context, args.Attempt, async context =>
         {
             var outcome = await args.Callback(context).ConfigureAwait(context.ContinueOnCapturedContext);
             return outcome.AsOutcome();

src/Polly.Core/Hedging/Controller/TaskExecution.cs

@@ -100,7 +100,7 @@ public async ValueTask<bool> InitializeAsync<TResult, TState>(
 
             try
             {
-                action = _handler.GenerateAction(CreateArguments(primaryCallback, state, attempt));
+                action = _handler.GenerateAction(CreateArguments(primaryCallback, snapshot.Context, state, attempt));
                 if (action == null)
                 {
                     await ResetAsync().ConfigureAwait(false);
@@ -125,8 +125,9 @@ public async ValueTask<bool> InitializeAsync<TResult, TState>(
 
     private HedgingActionGeneratorArguments<TResult> CreateArguments<TResult, TState>(
         Func<ResilienceContext, TState, ValueTask<Outcome<TResult>>> primaryCallback,
+        ResilienceContext primaryContext,
         TState state,
-        int attempt) => new(Context, attempt, (context) => primaryCallback(context, state));
+        int attempt) => new(primaryContext, Context, attempt, (context) => primaryCallback(context, state));
 
     public async ValueTask ResetAsync()
     {

src/Polly.Core/Hedging/HedgingActionGeneratorArguments.cs

@@ -4,10 +4,15 @@ namespace Polly.Hedging;
 /// Represents arguments used in the hedging resilience strategy.
 /// </summary>
 /// <typeparam name="TResult">The type of the result.</typeparam>
-/// <param name="Context">The context associated with the execution of a user-provided callback.</param>
+/// <param name="PrimaryContext">The primary resilience context.</param>
+/// <param name="Context">The context associated with the execution of a user-provided callback that is cloned from the primary context.</param>
 /// <param name="Attempt">The zero-based hedging attempt number.</param>
 /// <param name="Callback">The callback passed to hedging strategy.</param>
 /// <remarks>
 /// Always use the constructor when creating this struct, otherwise we do not guarantee binary compatibility.
 /// </remarks>
-public readonly record struct HedgingActionGeneratorArguments<TResult>(ResilienceContext Context, int Attempt, Func<ResilienceContext, ValueTask<Outcome<TResult>>> Callback);
+public readonly record struct HedgingActionGeneratorArguments<TResult>(
+    ResilienceContext PrimaryContext,
+    ResilienceContext Context,
+    int Attempt,
+    Func<ResilienceContext, ValueTask<Outcome<TResult>>> Callback);

src/Polly.Core/Hedging/HedgingResilienceStrategy.cs

@@ -84,7 +84,7 @@ protected internal override async ValueTask<Outcome<TResult>> ExecuteCoreAsync<T
                     return outcome;
                 }
 
-                var onHedgingArgs = new OutcomeArguments<TResult, OnHedgingArguments>(context, outcome, new OnHedgingArguments(hedgingContext.LoadedTasks - 1));
+                var onHedgingArgs = new OutcomeArguments<TResult, OnHedgingArguments>(context, outcome, new OnHedgingArguments(context, hedgingContext.LoadedTasks - 1));
                 _telemetry.Report(HedgingConstants.OnHedgingEventName, onHedgingArgs);
 
                 if (OnHedging is not null)

src/Polly.Core/Hedging/OnHedgingArguments.cs

@@ -3,8 +3,9 @@ namespace Polly.Hedging;
 /// <summary>
 /// Represents arguments used by the on-hedging event.
 /// </summary>
+/// <param name="Context">The context associated with the execution of a user-provided callback.</param>
 /// <param name="Attempt">The zero-based hedging attempt number.</param>
 /// <remarks>
 /// Always use the constructor when creating this struct, otherwise we do not guarantee binary compatibility.
 /// </remarks>
-public readonly record struct OnHedgingArguments(int Attempt);
+public record OnHedgingArguments(ResilienceContext Context, int Attempt);

test/Polly.Core.Tests/Hedging/HedgingActionGeneratorArgumentsTests.cs

@@ -7,8 +7,9 @@ public class HedgingActionGeneratorArgumentsTests
     [Fact]
     public void Ctor_Ok()
     {
-        var args = new HedgingActionGeneratorArguments<string>(ResilienceContext.Get(), 5, _ => "dummy".AsOutcomeAsync());
+        var args = new HedgingActionGeneratorArguments<string>(ResilienceContext.Get(), ResilienceContext.Get(), 5, _ => "dummy".AsOutcomeAsync());
 
+        args.PrimaryContext.Should().NotBeNull();
         args.Context.Should().NotBeNull();
         args.Attempt.Should().Be(5);
         args.Callback.Should().NotBeNull();

test/Polly.Core.Tests/Hedging/HedgingHandlerTests.cs

@@ -15,7 +15,7 @@ public async Task GenerateAction_Generic_Ok()
             args => () => "ok".AsOutcomeAsync(),
             true);
 
-        var action = handler.GenerateAction(new HedgingActionGeneratorArguments<string>(ResilienceContext.Get(), 0, _ => "primary".AsOutcomeAsync()))!;
+        var action = handler.GenerateAction(new HedgingActionGeneratorArguments<string>(ResilienceContext.Get(), ResilienceContext.Get(), 0, _ => "primary".AsOutcomeAsync()))!;
         var res = await action();
 
         res.Result.Should().Be("ok");
@@ -39,7 +39,7 @@ public async Task GenerateAction_NonGeneric_Ok(bool nullAction)
             },
             false);
 
-        var action = handler.GenerateAction(new HedgingActionGeneratorArguments<string>(ResilienceContext.Get(), 0, _ => "primary".AsOutcomeAsync()))!;
+        var action = handler.GenerateAction(new HedgingActionGeneratorArguments<string>(ResilienceContext.Get(), ResilienceContext.Get(), 0, _ => "primary".AsOutcomeAsync()))!;
         if (nullAction)
         {
             action.Should().BeNull();
@@ -59,7 +59,7 @@ public async Task GenerateAction_NonGeneric_FromCallback()
             args => () => args.Callback(args.Context),
             false);
 
-        var action = handler.GenerateAction(new HedgingActionGeneratorArguments<string>(ResilienceContext.Get(), 0, _ => "callback".AsOutcomeAsync()))!;
+        var action = handler.GenerateAction(new HedgingActionGeneratorArguments<string>(ResilienceContext.Get(), ResilienceContext.Get(), 0, _ => "callback".AsOutcomeAsync()))!;
         var res = await action();
         res.Result.Should().Be("callback");
     }

test/Polly.Core.Tests/Hedging/HedgingResilienceStrategyTests.cs

@@ -120,6 +120,42 @@ public async Task ExecuteAsync_ShouldReturnAnyPossibleResult()
         result.Should().Be("Oranges");
     }
 
+    [Fact]
+    public async Task ExecuteAsync_EnsurePrimaryContextFlows()
+    {
+        var primaryContext = ResilienceContext.Get();
+        var attempts = 0;
+        var key = new ResiliencePropertyKey<string>("primary-key");
+
+        _options.MaxHedgedAttempts = 4;
+        _options.OnHedging = args =>
+        {
+            args.Context.Should().Be(primaryContext);
+
+            if (args.Arguments.Attempt == 0)
+            {
+                args.Context.Properties.Set(key, "dummy");
+            }
+
+            attempts++;
+
+            return default;
+        };
+
+        ConfigureHedging(args =>
+        {
+            args.PrimaryContext.Properties.GetValue(key, string.Empty).Should().Be("dummy");
+            args.PrimaryContext.Should().Be(primaryContext);
+            return () => Failure.AsOutcomeAsync();
+        });
+
+        var strategy = Create();
+        var result = await strategy.ExecuteAsync(_ => new ValueTask<string>(Failure), primaryContext);
+
+        attempts.Should().Be(4);
+        primaryContext.Properties.GetValue(key, string.Empty).Should().Be("dummy");
+    }
+
     [Fact]
     public async void ExecuteAsync_EnsureHedgedTasksCancelled_Ok()
     {

test/Polly.Core.Tests/Hedging/HedgingStrategyOptionsTests.cs

@@ -18,7 +18,7 @@ public async Task Ctor_EnsureDefaults()
         options.MaxHedgedAttempts.Should().Be(2);
         options.OnHedging.Should().BeNull();
 
-        var action = options.HedgingActionGenerator(new HedgingActionGeneratorArguments<int>(ResilienceContext.Get(), 1, c => 99.AsOutcomeAsync()))!;
+        var action = options.HedgingActionGenerator(new HedgingActionGeneratorArguments<int>(ResilienceContext.Get(), ResilienceContext.Get(), 1, c => 99.AsOutcomeAsync()))!;
         action.Should().NotBeNull();
         (await action()).Result.Should().Be(99);
     }