Ensure context keys are set when execute overridden

src/Polly.Shared/Caching/CachePolicy.cs

@@ -51,7 +51,7 @@ internal CachePolicy(
         /// <param name="context">Execution context that is passed to the exception policy; defines the cache key to use in cache lookup.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns>The value returned by the action, or the cache.</returns>
-        public override TResult Execute<TResult>(Func<Context, CancellationToken, TResult> action, Context context, CancellationToken cancellationToken)
+        public override TResult ExecuteInternal<TResult>(Func<Context, CancellationToken, TResult> action, Context context, CancellationToken cancellationToken)
         {
             return CacheEngine.Implementation<TResult>(
                 _syncCacheProvider.For<TResult>(), 

src/Polly.Shared/Caching/CachePolicyAsync.cs

@@ -43,7 +43,7 @@ internal CachePolicy(
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <param name="continueOnCapturedContext">Whether to continue on a captured synchronization context.</param>
         /// <returns>The value returned by the action, or the cache.</returns>
-        public override Task<TResult> ExecuteAsync<TResult>(Func<Context, CancellationToken, Task<TResult>> action, Context context, CancellationToken cancellationToken, bool continueOnCapturedContext)
+        public override Task<TResult> ExecuteAsyncInternal<TResult>(Func<Context, CancellationToken, Task<TResult>> action, Context context, CancellationToken cancellationToken, bool continueOnCapturedContext)
         {
             return CacheEngine.ImplementationAsync<TResult>(
                 _asyncCacheProvider.AsyncFor<TResult>(), 

src/Polly.Shared/Policy.cs

@@ -296,18 +296,34 @@ public TResult Execute<TResult>(Func<Context, CancellationToken, TResult> action
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns>The value returned by the action</returns>
         [DebuggerStepThrough]
-        public virtual TResult Execute<TResult>(Func<Context, CancellationToken, TResult> action, Context context, CancellationToken cancellationToken)
+        public TResult Execute<TResult>(Func<Context, CancellationToken, TResult> action, Context context, CancellationToken cancellationToken)
         {
-            if (_exceptionPolicy == null) throw new InvalidOperationException(
-                "Please use the synchronous-defined policies when calling the synchronous Execute (and similar) methods.");
             if (context == null) throw new ArgumentNullException(nameof(context));
 
             SetPolicyContext(context);
 
+            return ExecuteInternal(action, context, cancellationToken);
+        }
+
+        /// <summary>
+        /// Executes the specified action within the policy and returns the result.
+        /// </summary>
+        /// <typeparam name="TResult">The type of the result.</typeparam>
+        /// <param name="action">The action to perform.</param>
+        /// <param name="context">Context data that is passed to the exception policy.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <returns>The value returned by the action</returns>
+        [DebuggerStepThrough]
+        public virtual TResult ExecuteInternal<TResult>(Func<Context, CancellationToken, TResult> action, Context context, CancellationToken cancellationToken)
+        {
+            if (_exceptionPolicy == null) throw new InvalidOperationException(
+                "Please use the synchronous-defined policies when calling the synchronous Execute (and similar) methods.");
+
             var result = default(TResult);
             _exceptionPolicy((ctx, ct) => { result = action(ctx, ct); }, context, cancellationToken);
             return result;
         }
+
         #endregion
 
         #endregion

src/Polly.Shared/PolicyAsync.cs

@@ -540,14 +540,30 @@ public Task<TResult> ExecuteAsync<TResult>(Func<Context, CancellationToken, Task
         /// <returns>The value returned by the action</returns>
         /// <exception cref="System.InvalidOperationException">Please use asynchronous-defined policies when calling asynchronous ExecuteAsync (and similar) methods.</exception>
         [DebuggerStepThrough]
-        public virtual async Task<TResult> ExecuteAsync<TResult>(Func<Context, CancellationToken, Task<TResult>> action, Context context, CancellationToken cancellationToken, bool continueOnCapturedContext)
+        public Task<TResult> ExecuteAsync<TResult>(Func<Context, CancellationToken, Task<TResult>> action, Context context, CancellationToken cancellationToken, bool continueOnCapturedContext)
         {
-            if (_asyncExceptionPolicy == null) throw new InvalidOperationException(
-                "Please use asynchronous-defined policies when calling asynchronous ExecuteAsync (and similar) methods.");
             if (context == null) throw new ArgumentNullException(nameof(context));
-
             SetPolicyContext(context);
 
+            return ExecuteAsyncInternal(action, context, cancellationToken, continueOnCapturedContext);
+        }
+
+        /// <summary>
+        ///     Executes the specified asynchronous action within the policy and returns the result.
+        /// </summary>
+        /// <typeparam name="TResult">The type of the result.</typeparam>
+        /// <param name="action">The action to perform.</param>
+        /// <param name="context">Context data that is passed to the exception policy.</param>
+        /// <param name="continueOnCapturedContext">Whether to continue on a captured synchronization context.</param>
+        /// <param name="cancellationToken">A cancellation token which can be used to cancel the action.  When a retry policy is in use, also cancels any further retries.</param>
+        /// <returns>The value returned by the action</returns>
+        /// <exception cref="System.InvalidOperationException">Please use asynchronous-defined policies when calling asynchronous ExecuteAsync (and similar) methods.</exception>
+        [DebuggerStepThrough]
+        public virtual async Task<TResult> ExecuteAsyncInternal<TResult>(Func<Context, CancellationToken, Task<TResult>> action, Context context, CancellationToken cancellationToken, bool continueOnCapturedContext)
+        {
+            if (_asyncExceptionPolicy == null) throw new InvalidOperationException(
+                "Please use asynchronous-defined policies when calling asynchronous ExecuteAsync (and similar) methods.");
+
             var result = default(TResult);
             await _asyncExceptionPolicy(async (ctx, ct) =>
             {
@@ -556,6 +572,7 @@ await _asyncExceptionPolicy(async (ctx, ct) =>
                 .ConfigureAwait(continueOnCapturedContext);
             return result;
         }
+
         #endregion
 
         #endregion

src/Polly.Shared/Wrap/PolicyWrap.cs

@@ -37,7 +37,7 @@ internal PolicyWrap(Action<Action<Context, CancellationToken>, Context, Cancella
         /// <param name="context">Execution context that is passed to the exception policy; defines the cache key to use in cache lookup.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns>The value returned by the action, or the cache.</returns>
-        public override TResult Execute<TResult>(Func<Context, CancellationToken, TResult> action, Context context, CancellationToken cancellationToken)
+        public override TResult ExecuteInternal<TResult>(Func<Context, CancellationToken, TResult> action, Context context, CancellationToken cancellationToken)
         {
             return PolicyWrapEngine.Implementation<TResult>(
                    action,

src/Polly.Shared/Wrap/PolicyWrapAsync.cs

@@ -23,7 +23,7 @@ internal PolicyWrap(Func<Func<Context, CancellationToken, Task>, Context, Cancel
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <param name="continueOnCapturedContext">Whether to continue on a captured synchronization context.</param>
         /// <returns>The value returned by the action, or the cache.</returns>
-        public override Task<TResult> ExecuteAsync<TResult>(Func<Context, CancellationToken, Task<TResult>> action, Context context, CancellationToken cancellationToken, bool continueOnCapturedContext)
+        public override Task<TResult> ExecuteAsyncInternal<TResult>(Func<Context, CancellationToken, Task<TResult>> action, Context context, CancellationToken cancellationToken, bool continueOnCapturedContext)
         {
             return PolicyWrapEngine.ImplementationAsync<TResult>(
                 action,

src/Polly.SharedSpecs/Wrap/PolicyWrapContextAndKeySpecs.cs

@@ -90,6 +90,47 @@ public void Should_pass_outmost_PolicyWrap_Key_as_PolicyWrapKey_ignoring_inner_P
             policyWrapKeySetOnExecutionContext.Should().Be(outerWrapKey);
         }
 
+        [Fact]
+        public void Should_pass_outmost_PolicyWrap_Key_as_PolicyWrapKey_to_innermost_Policy_when_execute_method_generic()
+        {
+            string retryKey = Guid.NewGuid().ToString();
+            string breakerKey = Guid.NewGuid().ToString();
+            string fallbackKey = Guid.NewGuid().ToString();
+            string innerWrapKey = Guid.NewGuid().ToString();
+            string outerWrapKey = Guid.NewGuid().ToString();
+
+            string policyWrapKeySetOnExecutionContext = null;
+            Action<Exception, TimeSpan, Context> onBreak = (e, t, context) =>
+            {
+                policyWrapKeySetOnExecutionContext = context.PolicyWrapKey;
+            };
+            Action<Context> doNothingOnReset = _ => { };
+
+            var retry = Policy.Handle<Exception>().Retry(1).WithPolicyKey(retryKey);
+            var breaker = Policy.Handle<Exception>().CircuitBreaker(1, TimeSpan.Zero, onBreak, doNothingOnReset).WithPolicyKey(breakerKey);
+            var fallback = Policy.Handle<Exception>().Fallback(() => { }).WithPolicyKey(fallbackKey);
+
+            var innerWrap = retry.Wrap(breaker).WithPolicyKey(innerWrapKey);
+            var outerWrap = fallback.Wrap(innerWrap).WithPolicyKey(outerWrapKey);
+
+            bool doneOnceOny = false;
+            outerWrap.Execute<int>(() =>
+            {
+                if (!doneOnceOny)
+                {
+                    doneOnceOny = true;
+                    throw new Exception();
+                }
+                return 0;
+            });
+
+            policyWrapKeySetOnExecutionContext.Should().NotBe(retryKey);
+            policyWrapKeySetOnExecutionContext.Should().NotBe(breakerKey);
+            policyWrapKeySetOnExecutionContext.Should().NotBe(fallbackKey);
+            policyWrapKeySetOnExecutionContext.Should().NotBe(innerWrapKey);
+            policyWrapKeySetOnExecutionContext.Should().Be(outerWrapKey);
+        }
+
         #endregion
 
     }

src/Polly.SharedSpecs/Wrap/PolicyWrapContextAndKeySpecsAsync.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Threading.Tasks;
 using FluentAssertions;
 using Polly.Specs.Helpers;
 using Polly.Utilities;
@@ -91,6 +92,46 @@ public void Should_pass_outmost_PolicyWrap_Key_as_PolicyWrapKey_ignoring_inner_P
             policyWrapKeySetOnExecutionContext.Should().Be(outerWrapKey);
         }
 
+        [Fact]
+        public async Task Should_pass_outmost_PolicyWrap_Key_as_PolicyWrapKey_to_innermost_Policy_when_execute_method_generic()
+        {
+            string retryKey = Guid.NewGuid().ToString();
+            string breakerKey = Guid.NewGuid().ToString();
+            string fallbackKey = Guid.NewGuid().ToString();
+            string innerWrapKey = Guid.NewGuid().ToString();
+            string outerWrapKey = Guid.NewGuid().ToString();
+
+            string policyWrapKeySetOnExecutionContext = null;
+            Action<Exception, TimeSpan, Context> onBreak = (e, t, context) =>
+            {
+                policyWrapKeySetOnExecutionContext = context.PolicyWrapKey;
+            };
+            Action<Context> doNothingOnReset = _ => { };
+
+            var retry = Policy.Handle<Exception>().RetryAsync(1).WithPolicyKey(retryKey);
+            var breaker = Policy.Handle<Exception>().CircuitBreakerAsync(1, TimeSpan.Zero, onBreak, doNothingOnReset).WithPolicyKey(breakerKey);
+            var fallback = Policy.Handle<Exception>().FallbackAsync(_ => TaskHelper.EmptyTask).WithPolicyKey(fallbackKey);
+
+            var innerWrap = retry.WrapAsync(breaker).WithPolicyKey(innerWrapKey);
+            var outerWrap = fallback.WrapAsync(innerWrap).WithPolicyKey(outerWrapKey);
+
+            bool doneOnceOny = false;
+            await outerWrap.ExecuteAsync<int>(() =>
+            {
+                if (!doneOnceOny)
+                {
+                    doneOnceOny = true;
+                    throw new Exception();
+                }
+                return TaskHelper.FromResult(0);
+            });
+
+            policyWrapKeySetOnExecutionContext.Should().NotBe(retryKey);
+            policyWrapKeySetOnExecutionContext.Should().NotBe(breakerKey);
+            policyWrapKeySetOnExecutionContext.Should().NotBe(fallbackKey);
+            policyWrapKeySetOnExecutionContext.Should().NotBe(innerWrapKey);
+            policyWrapKeySetOnExecutionContext.Should().Be(outerWrapKey);
+        }
         #endregion
 
     }