Enhance OnHedgingArguments (#1314)

App-vNext · Jun 16, 2023 · 488bc10 · 488bc10
1 parent 4524822
commit 488bc10
Show file tree

Hide file tree

Showing 14 changed files with 84 additions and 40 deletions.
diff --git a/src/Polly.Core/Hedging/Controller/HedgingHandler.cs b/src/Polly.Core/Hedging/Controller/HedgingHandler.cs
@@ -16,7 +16,8 @@ internal sealed record class HedgingHandler<T>(
         if (IsGeneric)
         {
             var copiedArgs = new HedgingActionGeneratorArguments<T>(
-                args.Context,
+                args.PrimaryContext,
+                args.ActionContext,
                 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.ActionContext, args.Attempt, async context =>
         {
             var outcome = await args.Callback(context).ConfigureAwait(context.ContinueOnCapturedContext);
             return outcome.AsOutcome();
@@ -42,7 +43,7 @@ internal sealed record class HedgingHandler<T>(
 
         return async () =>
         {
-            var outcome = await action().ConfigureAwait(args.Context.ContinueOnCapturedContext);
+            var outcome = await action().ConfigureAwait(args.ActionContext.ContinueOnCapturedContext);
             return outcome.AsOutcome<TResult>();
         };
     }

diff --git a/src/Polly.Core/Hedging/Controller/TaskExecution.cs b/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()
     {

diff --git a/src/Polly.Core/Hedging/HedgingActionGeneratorArguments.cs b/src/Polly.Core/Hedging/HedgingActionGeneratorArguments.cs
@@ -4,10 +4,17 @@ 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="ActionContext">
+/// The context that will be passed to action generated by <see cref="HedgingStrategyOptions{TResult}.HedgingActionGenerator"/>.
+/// This context is cloned from <paramref name="PrimaryContext"/>.</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 ActionContext,
+    int Attempt,
+    Func<ResilienceContext, ValueTask<Outcome<TResult>>> Callback);
diff --git a/src/Polly.Core/Hedging/HedgingResilienceStrategy.cs b/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)

diff --git a/src/Polly.Core/Hedging/HedgingStrategyOptions.TResult.cs b/src/Polly.Core/Hedging/HedgingStrategyOptions.TResult.cs
@@ -61,7 +61,7 @@ public class HedgingStrategyOptions<TResult> : ResilienceStrategyOptions
     {
         return async () =>
         {
-            return await args.Callback(args.Context).ConfigureAwait(args.Context.ContinueOnCapturedContext);
+            return await args.Callback(args.ActionContext).ConfigureAwait(args.ActionContext.ContinueOnCapturedContext);
         };
     };
 

diff --git a/src/Polly.Core/Hedging/OnHedgingArguments.cs b/src/Polly.Core/Hedging/OnHedgingArguments.cs
@@ -3,8 +3,6 @@ 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);
diff --git a/test/Polly.Core.Tests/Hedging/Controller/HedgingExecutionContextTests.cs b/test/Polly.Core.Tests/Hedging/Controller/HedgingExecutionContextTests.cs
@@ -456,12 +456,12 @@ private void ConfigureSecondaryTasks(params TimeSpan[] delays)
                 return null;
             }
 
-            args.Context.AddResilienceEvent(new ResilienceEvent("dummy-event"));
+            args.ActionContext.AddResilienceEvent(new ResilienceEvent("dummy-event"));
 
             return async () =>
             {
-                args.Context.Properties.Set(new ResiliencePropertyKey<int>(attempt.ToString(CultureInfo.InvariantCulture)), attempt);
-                await _timeProvider.Delay(delays[attempt], args.Context.CancellationToken);
+                args.ActionContext.Properties.Set(new ResiliencePropertyKey<int>(attempt.ToString(CultureInfo.InvariantCulture)), attempt);
+                await _timeProvider.Delay(delays[attempt], args.ActionContext.CancellationToken);
                 return new DisposableResult(delays[attempt].ToString()).AsOutcome();
             };
         };

diff --git a/test/Polly.Core.Tests/Hedging/Controller/TaskExecutionTests.cs b/test/Polly.Core.Tests/Hedging/Controller/TaskExecutionTests.cs
@@ -74,7 +74,7 @@ public async Task Initialize_Secondary_Ok(string value, bool handled)
         var execution = Create();
         Generator = args =>
         {
-            AssertSecondaryContext(args.Context, execution);
+            AssertSecondaryContext(args.ActionContext, execution);
             args.Attempt.Should().Be(4);
             return () => new DisposableResult { Name = value }.AsOutcomeAsync();
         };
@@ -158,7 +158,7 @@ public async Task Initialize_Cancelled_EnsureRespected(bool primary)
         {
             return async () =>
             {
-                await _timeProvider.Delay(TimeSpan.FromDays(1), args.Context.CancellationToken);
+                await _timeProvider.Delay(TimeSpan.FromDays(1), args.ActionContext.CancellationToken);
                 return new DisposableResult { Name = Handled }.AsOutcome();
             };
         };

diff --git a/test/Polly.Core.Tests/Hedging/HedgingActionGeneratorArgumentsTests.cs b/test/Polly.Core.Tests/Hedging/HedgingActionGeneratorArgumentsTests.cs
@@ -7,9 +7,10 @@ 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.Context.Should().NotBeNull();
+        args.PrimaryContext.Should().NotBeNull();
+        args.ActionContext.Should().NotBeNull();
         args.Attempt.Should().Be(5);
         args.Callback.Should().NotBeNull();
     }

diff --git a/test/Polly.Core.Tests/Hedging/HedgingActions.cs b/test/Polly.Core.Tests/Hedging/HedgingActions.cs
@@ -24,7 +24,7 @@ public HedgingActions(TimeProvider timeProvider)
             {
                 return async () =>
                 {
-                    return await Functions[args.Attempt - 1]!(args.Context);
+                    return await Functions[args.Attempt - 1]!(args.ActionContext);
                 };
             }
 
@@ -58,7 +58,7 @@ private async ValueTask<Outcome<string>> GetOranges(ResilienceContext context)
 
     public static Func<HedgingActionGeneratorArguments<string>, Func<ValueTask<Outcome<string>>>?> GetGenerator(Func<ResilienceContext, ValueTask<Outcome<string>>> task)
     {
-        return args => () => task(args.Context);
+        return args => () => task(args.ActionContext);
     }
 
     public int MaxHedgedTasks => Functions.Count + 1;

diff --git a/test/Polly.Core.Tests/Hedging/HedgingHandlerTests.cs b/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();
@@ -56,10 +56,10 @@ public async Task GenerateAction_NonGeneric_FromCallback()
     {
         var handler = new HedgingHandler<object>(
             PredicateInvoker<HandleHedgingArguments>.Create<object>(args => PredicateResult.True, false)!,
-            args => () => args.Callback(args.Context),
+            args => () => args.Callback(args.ActionContext),
             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");
     }

diff --git a/test/Polly.Core.Tests/Hedging/HedgingResilienceStrategyBuilderExtensionsTests.cs b/test/Polly.Core.Tests/Hedging/HedgingResilienceStrategyBuilderExtensionsTests.cs
@@ -68,7 +68,7 @@ public async Task AddHedging_IntegrationTest()
                 {
                     return async () =>
                     {
-                        await Task.Delay(25, args.Context.CancellationToken);
+                        await Task.Delay(25, args.ActionContext.CancellationToken);
 
                         if (args.Attempt == 3)
                         {

diff --git a/test/Polly.Core.Tests/Hedging/HedgingResilienceStrategyTests.cs b/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()
     {
@@ -256,12 +292,12 @@ public async Task ExecuteAsync_EveryHedgedTaskShouldHaveDifferentContexts()
         {
             return async () =>
             {
-                tokenHashCodes.Add(args.Context.CancellationToken.GetHashCode());
-                args.Context.CancellationToken.CanBeCanceled.Should().BeTrue();
-                args.Context.Properties.GetValue(beforeKey, "wrong").Should().Be("before");
-                contexts.Add(args.Context);
+                tokenHashCodes.Add(args.ActionContext.CancellationToken.GetHashCode());
+                args.ActionContext.CancellationToken.CanBeCanceled.Should().BeTrue();
+                args.ActionContext.Properties.GetValue(beforeKey, "wrong").Should().Be("before");
+                contexts.Add(args.ActionContext);
                 await Task.Yield();
-                args.Context.Properties.Set(afterKey, "after");
+                args.ActionContext.Properties.Set(afterKey, "after");
                 return "secondary".AsOutcome();
             };
         });
@@ -327,10 +363,10 @@ public async Task ExecuteAsync_EnsurePropertiesConsistency(bool primaryFails)
         {
             return async () =>
             {
-                contexts.Add(args.Context);
-                args.Context.Properties.GetValue(primaryKey, string.Empty).Should().Be("primary");
-                args.Context.Properties.Set(secondaryKey, "secondary");
-                await _timeProvider.Delay(TimeSpan.FromHours(1), args.Context.CancellationToken);
+                contexts.Add(args.ActionContext);
+                args.ActionContext.Properties.GetValue(primaryKey, string.Empty).Should().Be("primary");
+                args.ActionContext.Properties.Set(secondaryKey, "secondary");
+                await _timeProvider.Delay(TimeSpan.FromHours(1), args.ActionContext.CancellationToken);
                 return (primaryFails ? Success : Failure).AsOutcome();
             };
         });
@@ -415,9 +451,9 @@ public async Task ExecuteAsync_Secondary_CustomPropertiesAvailable()
         {
             return () =>
             {
-                args.Context.Properties.TryGetValue(key2, out var val).Should().BeTrue();
+                args.ActionContext.Properties.TryGetValue(key2, out var val).Should().BeTrue();
                 val.Should().Be("my-value-2");
-                args.Context.Properties.Set(key, "my-value");
+                args.ActionContext.Properties.Set(key, "my-value");
                 return Success.AsOutcomeAsync();
             };
         });
@@ -820,7 +856,7 @@ private void ConfigureHedging()
 
     private void ConfigureHedging(Func<ResilienceContext, ValueTask<Outcome<string>>> background)
     {
-        ConfigureHedging(args => () => background(args.Context));
+        ConfigureHedging(args => () => background(args.ActionContext));
     }
 
     private void ConfigureHedging(Func<HedgingActionGeneratorArguments<string>, Func<ValueTask<Outcome<string>>>?> generator)

diff --git a/test/Polly.Core.Tests/Hedging/HedgingStrategyOptionsTests.cs b/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);
     }