Fix Context.PolicyKey issue 463

Fix Context.PolicyKey such that, as execution bubbles back outwards through a PolicyWrap, Context.PolicyKey correctly adopts the value of the policies in play during the outward journey through the wrap.
App-vNext · Jul 14, 2018 · 61e42f1 · 61e42f1
1 parent 328faa6
commit 61e42f1
Show file tree

Hide file tree

Showing 8 changed files with 200 additions and 17 deletions.
diff --git a/src/Polly.Shared/Policy.Async.ExecuteOverloads.cs b/src/Polly.Shared/Policy.Async.ExecuteOverloads.cs
@@ -117,9 +117,17 @@ public Task ExecuteAsync(Func<Context, CancellationToken, Task> action, IDiction
         public Task ExecuteAsync(Func<Context, CancellationToken, Task> action, Context context, CancellationToken cancellationToken, bool continueOnCapturedContext)
         {
             if (context == null) throw new ArgumentNullException(nameof(context));
-            SetPolicyContext(context);
 
-            return ExecuteAsyncInternal(action, context, cancellationToken, continueOnCapturedContext);
+            SetPolicyContext(context, out string priorPolicyWrapKey, out string priorPolicyKey);
+
+            try
+            {
+                return ExecuteAsyncInternal(action, context, cancellationToken, continueOnCapturedContext);
+            }
+            finally
+            {
+                RestorePolicyContext(context, priorPolicyWrapKey, priorPolicyKey);
+            }
         }
 
         #region Overloads method-generic in TResult
@@ -245,9 +253,17 @@ public Task<TResult> ExecuteAsync<TResult>(Func<Context, CancellationToken, Task
         public Task<TResult> ExecuteAsync<TResult>(Func<Context, CancellationToken, Task<TResult>> action, Context context, CancellationToken cancellationToken, bool continueOnCapturedContext)
         {
             if (context == null) throw new ArgumentNullException(nameof(context));
-            SetPolicyContext(context);
 
-            return ExecuteAsyncInternal(action, context, cancellationToken, continueOnCapturedContext);
+            SetPolicyContext(context, out string priorPolicyWrapKey, out string priorPolicyKey);
+
+            try
+            {
+                return ExecuteAsyncInternal(action, context, cancellationToken, continueOnCapturedContext);
+            }
+            finally
+            {
+                RestorePolicyContext(context, priorPolicyWrapKey, priorPolicyKey);
+            }
         }
 
         #endregion

diff --git a/src/Polly.Shared/Policy.Async.TResult.ExecuteOverloads.cs b/src/Polly.Shared/Policy.Async.TResult.ExecuteOverloads.cs
@@ -128,9 +128,17 @@ public Task<TResult> ExecuteAsync(Func<Context, CancellationToken, Task<TResult>
         public Task<TResult> ExecuteAsync(Func<Context, CancellationToken, Task<TResult>> action, Context context, CancellationToken cancellationToken, bool continueOnCapturedContext)
         {
             if (context == null) throw new ArgumentNullException(nameof(context));
-            SetPolicyContext(context);
 
-            return ExecuteAsyncInternal(action, context, cancellationToken, continueOnCapturedContext);
+            SetPolicyContext(context, out string priorPolicyWrapKey, out string priorPolicyKey);
+
+            try
+            {
+                return ExecuteAsyncInternal(action, context, cancellationToken, continueOnCapturedContext);
+            }
+            finally
+            {
+                RestorePolicyContext(context, priorPolicyWrapKey, priorPolicyKey);
+            }
         }
 
         #endregion

diff --git a/src/Polly.Shared/Policy.ContextAndKeys.cs b/src/Polly.Shared/Policy.ContextAndKeys.cs
@@ -21,10 +21,27 @@ public abstract partial class PolicyBase
         /// Updates the execution <see cref="Context"/> with context from the executing <see cref="Policy"/>.
         /// </summary>
         /// <param name="executionContext">The execution <see cref="Context"/>.</param>
-        internal virtual void SetPolicyContext(Context executionContext)
+        /// <param name="priorPolicyWrapKey">The <see cref="M:Context.PolicyWrapKey"/> prior to changes by this method.</param>
+        /// <param name="priorPolicyKey">The <see cref="M:Context.PolicyKey"/> prior to changes by this method.</param>
+        internal virtual void SetPolicyContext(Context executionContext, out string priorPolicyWrapKey, out string priorPolicyKey) // priorPolicyWrapKey and priorPolicyKey could be handled as a (string, string) System.ValueTuple return type instead of out parameters, when our minimum supported target supports this.
         {
+            priorPolicyWrapKey = executionContext.PolicyWrapKey;
+            priorPolicyKey = executionContext.PolicyKey;
+
             executionContext.PolicyKey = PolicyKey;
         }
+
+        /// <summary>
+        /// Restores the supplied keys to the execution <see cref="Context"/>.
+        /// </summary>
+        /// <param name="executionContext">The execution <see cref="Context"/>.</param>
+        /// <param name="priorPolicyWrapKey">The <see cref="M:Context.PolicyWrapKey"/> prior to execution through this policy.</param>
+        /// <param name="priorPolicyKey">The <see cref="M:Context.PolicyKey"/> prior to execution through this policy.</param>
+        internal void RestorePolicyContext(Context executionContext, string priorPolicyWrapKey, string priorPolicyKey)
+        {
+            executionContext.PolicyWrapKey = priorPolicyWrapKey;
+            executionContext.PolicyKey = priorPolicyKey;
+        }
     }
 
     public abstract partial class Policy

diff --git a/src/Polly.Shared/Policy.ExecuteOverloads.cs b/src/Polly.Shared/Policy.ExecuteOverloads.cs
@@ -77,9 +77,16 @@ public void Execute(Action<Context, CancellationToken> action, Context context,
         {
             if (context == null) throw new ArgumentNullException(nameof(context));
 
-            SetPolicyContext(context);
+            SetPolicyContext(context, out string priorPolicyWrapKey, out string priorPolicyKey);
 
-            ExecuteInternal(action, context, cancellationToken);
+            try
+            {
+                ExecuteInternal(action, context, cancellationToken);
+            }
+            finally
+            {
+                RestorePolicyContext(context, priorPolicyWrapKey, priorPolicyKey);
+            }
         }
 
         #region Overloads method-generic in TResult
@@ -171,9 +178,16 @@ public TResult Execute<TResult>(Func<Context, CancellationToken, TResult> action
         {
             if (context == null) throw new ArgumentNullException(nameof(context));
 
-            SetPolicyContext(context);
+            SetPolicyContext(context, out string priorPolicyWrapKey, out string priorPolicyKey);
 
-            return ExecuteInternal(action, context, cancellationToken);
+            try
+            {
+                return ExecuteInternal(action, context, cancellationToken);
+            }
+            finally
+            {
+                RestorePolicyContext(context, priorPolicyWrapKey, priorPolicyKey);
+            }
         }
 
         #endregion

diff --git a/src/Polly.Shared/Policy.TResult.ExecuteOverloads.cs b/src/Polly.Shared/Policy.TResult.ExecuteOverloads.cs
@@ -92,9 +92,17 @@ public TResult Execute(Func<Context, CancellationToken, TResult> action, IDictio
         public TResult Execute(Func<Context, CancellationToken, TResult> action, Context context, CancellationToken cancellationToken)
         {
             if (context == null) throw new ArgumentNullException(nameof(context));
-            SetPolicyContext(context);
 
-            return ExecuteInternal(action, context, cancellationToken);
+            SetPolicyContext(context, out string priorPolicyWrapKey, out string priorPolicyKey);
+
+            try
+            {
+                return ExecuteInternal(action, context, cancellationToken);
+            }
+            finally
+            {
+                RestorePolicyContext(context, priorPolicyWrapKey, priorPolicyKey);
+            }
         }
 
         #endregion

diff --git a/src/Polly.Shared/Wrap/PolicyWrap.ContextAndKeys.cs b/src/Polly.Shared/Wrap/PolicyWrap.ContextAndKeys.cs
@@ -10,11 +10,16 @@ public partial class PolicyWrap
         /// Updates the execution <see cref="Context"/> with context from the executing <see cref="PolicyWrap"/>.
         /// </summary>
         /// <param name="executionContext">The execution <see cref="Context"/>.</param>
-        internal override void SetPolicyContext(Context executionContext)
+        /// <param name="priorPolicyWrapKey">The <see cref="M:Context.PolicyWrapKey"/> prior to changes by this method.</param>
+        /// <param name="priorPolicyKey">The <see cref="M:Context.PolicyKey"/> prior to changes by this method.</param>
+        internal override void SetPolicyContext(Context executionContext, out string priorPolicyWrapKey, out string priorPolicyKey)
         {
+            priorPolicyWrapKey = executionContext.PolicyWrapKey;
+            priorPolicyKey = executionContext.PolicyKey;
+
             if (executionContext.PolicyWrapKey == null) executionContext.PolicyWrapKey = PolicyKey;
 
-            base.SetPolicyContext(executionContext);
+            base.SetPolicyContext(executionContext, out _, out _);
         }
     }
 
@@ -24,11 +29,16 @@ public partial class PolicyWrap<TResult>
         /// Updates the execution <see cref="Context"/> with context from the executing <see cref="PolicyWrap{TResult}"/>.
         /// </summary>
         /// <param name="executionContext">The execution <see cref="Context"/>.</param>
-        internal override void SetPolicyContext(Context executionContext)
+        /// <param name="priorPolicyWrapKey">The <see cref="M:Context.PolicyWrapKey"/> prior to changes by this method.</param>
+        /// <param name="priorPolicyKey">The <see cref="M:Context.PolicyKey"/> prior to changes by this method.</param>
+        internal override void SetPolicyContext(Context executionContext, out string priorPolicyWrapKey, out string priorPolicyKey)
         {
+            priorPolicyWrapKey = executionContext.PolicyWrapKey;
+            priorPolicyKey = executionContext.PolicyKey;
+
             if (executionContext.PolicyWrapKey == null) executionContext.PolicyWrapKey = PolicyKey;
 
-            base.SetPolicyContext(executionContext);
+            base.SetPolicyContext(executionContext, out _, out _);
         }
     }
 }
diff --git a/src/Polly.SharedSpecs/Wrap/PolicyWrapContextAndKeySpecs.cs b/src/Polly.SharedSpecs/Wrap/PolicyWrapContextAndKeySpecs.cs
@@ -59,6 +59,33 @@ public void Should_pass_PolicyKey_to_execution_context_of_inner_policy_as_Policy
             policyWrapKeySetOnExecutionContext.Should().Be(wrapKey);
         }
 
+        [Fact]
+        public void Should_restore_PolicyKey_of_outer_policy_to_execution_context_as_move_outwards_through_PolicyWrap()
+        {
+            ISyncPolicy fallback = Policy
+                .Handle<Exception>()
+                .Fallback(_ => {}, onFallback: (_, context) =>
+                {
+                    context.PolicyWrapKey.Should().Be("PolicyWrap");
+                    context.PolicyKey.Should().Be("FallbackPolicy"); 
+                })
+                .WithPolicyKey("FallbackPolicy");
+
+            ISyncPolicy retry = Policy
+                .Handle<Exception>()
+                .Retry(1, onRetry: (result, retryCount, context) =>
+                {
+                    context.PolicyWrapKey.Should().Be("PolicyWrap");
+                    context.PolicyKey.Should().Be("RetryPolicy");
+                })
+                .WithPolicyKey("RetryPolicy");
+
+            ISyncPolicy policyWrap = Policy.Wrap(fallback, retry)
+                .WithPolicyKey("PolicyWrap");
+
+            policyWrap.Execute(() => throw new Exception());
+        }
+
         [Fact]
         public void Should_pass_outmost_PolicyWrap_Key_as_PolicyWrapKey_ignoring_inner_PolicyWrap_keys_even_when_executing_policies_in_inner_wrap()
         {
@@ -189,6 +216,33 @@ public void Should_pass_PolicyKey_to_execution_context_of_inner_policy_as_Policy
             policyWrapKeySetOnExecutionContext.Should().Be(wrapKey);
         }
 
+        [Fact]
+        public void Should_restore_PolicyKey_of_outer_policy_to_execution_context_as_move_outwards_through_PolicyWrap()
+        {
+            ISyncPolicy<ResultPrimitive> fallback = Policy<ResultPrimitive>
+                .Handle<Exception>()
+                .Fallback<ResultPrimitive>(ResultPrimitive.Undefined, onFallback: (result, context) =>
+                {
+                    context.PolicyWrapKey.Should().Be("PolicyWrap");
+                    context.PolicyKey.Should().Be("FallbackPolicy");
+                })
+                .WithPolicyKey("FallbackPolicy");
+
+            ISyncPolicy<ResultPrimitive> retry = Policy<ResultPrimitive>
+                .Handle<Exception>()
+                .Retry(1, onRetry: (result, retryCount, context) =>
+                {
+                    context.PolicyWrapKey.Should().Be("PolicyWrap");
+                    context.PolicyKey.Should().Be("RetryPolicy");
+                })
+                .WithPolicyKey("RetryPolicy");
+
+            Policy<ResultPrimitive> policyWrap = Policy.Wrap(fallback, retry)
+                .WithPolicyKey("PolicyWrap");
+
+            policyWrap.Execute(() => throw new Exception());
+        }
+
         [Fact]
         public void Should_pass_outmost_PolicyWrap_Key_as_PolicyWrapKey_ignoring_inner_PolicyWrap_keys_even_when_executing_policies_in_inner_wrap()
         {

diff --git a/src/Polly.SharedSpecs/Wrap/PolicyWrapContextAndKeySpecsAsync.cs b/src/Polly.SharedSpecs/Wrap/PolicyWrapContextAndKeySpecsAsync.cs
@@ -61,6 +61,34 @@ public async Task Should_pass_PolicyKey_to_execution_context_of_inner_policy_as_
             policyWrapKeySetOnExecutionContext.Should().Be(wrapKey);
         }
 
+        [Fact]
+        public void Should_restore_PolicyKey_of_outer_policy_to_execution_context_as_move_outwards_through_PolicyWrap()
+        {
+            IAsyncPolicy fallback = Policy
+                .Handle<Exception>()
+                .FallbackAsync((_,__) => TaskHelper.EmptyTask, (_, context) =>
+                {
+                    context.PolicyWrapKey.Should().Be("PolicyWrap");
+                    context.PolicyKey.Should().Be("FallbackPolicy");
+                    return TaskHelper.EmptyTask;
+                })
+                .WithPolicyKey("FallbackPolicy");
+
+            IAsyncPolicy retry = Policy
+                .Handle<Exception>()
+                .RetryAsync(1, onRetry: (result, retryCount, context) =>
+                {
+                    context.PolicyWrapKey.Should().Be("PolicyWrap");
+                    context.PolicyKey.Should().Be("RetryPolicy");
+                })
+                .WithPolicyKey("RetryPolicy");
+
+            IAsyncPolicy policyWrap = Policy.WrapAsync(fallback, retry)
+                .WithPolicyKey("PolicyWrap");
+
+            policyWrap.ExecuteAsync(() => throw new Exception());
+        }
+
         [Fact]
         public async Task Should_pass_outmost_PolicyWrap_Key_as_PolicyWrapKey_ignoring_inner_PolicyWrap_keys_even_when_executing_policies_in_inner_WrapAsync()
         {
@@ -192,6 +220,34 @@ public async Task Should_pass_PolicyKey_to_execution_context_of_inner_policy_as_
             policyWrapKeySetOnExecutionContext.Should().Be(wrapKey);
         }
 
+        [Fact]
+        public void Should_restore_PolicyKey_of_outer_policy_to_execution_context_as_move_outwards_through_PolicyWrap()
+        {
+            IAsyncPolicy<ResultPrimitive> fallback = Policy<ResultPrimitive>
+                .Handle<Exception>()
+                .FallbackAsync((_, __) => Task.FromResult(ResultPrimitive.Undefined), (_, context) =>
+                {
+                    context.PolicyWrapKey.Should().Be("PolicyWrap");
+                    context.PolicyKey.Should().Be("FallbackPolicy");
+                    return TaskHelper.EmptyTask;
+                })
+                .WithPolicyKey("FallbackPolicy");
+
+            IAsyncPolicy<ResultPrimitive> retry = Policy<ResultPrimitive>
+                .Handle<Exception>()
+                .RetryAsync(1, onRetry: (result, retryCount, context) =>
+                {
+                    context.PolicyWrapKey.Should().Be("PolicyWrap");
+                    context.PolicyKey.Should().Be("RetryPolicy");
+                })
+                .WithPolicyKey("RetryPolicy");
+
+            IAsyncPolicy<ResultPrimitive> policyWrap = Policy.WrapAsync(fallback, retry)
+                .WithPolicyKey("PolicyWrap");
+
+            policyWrap.ExecuteAsync(() => throw new Exception());
+        }
+
         [Fact]
         public async Task Should_pass_outmost_PolicyWrap_Key_as_PolicyWrapKey_ignoring_inner_PolicyWrap_keys_even_when_executing_policies_in_inner_WrapAsync()
         {