ResilienceContextPool improvements

src/Polly.Core/PublicAPI.Unshipped.txt

@@ -1,7 +1,7 @@
 #nullable enable
 abstract Polly.Registry.ResiliencePipelineProvider<TKey>.TryGetPipeline(TKey key, out Polly.ResiliencePipeline? pipeline) -> bool
 abstract Polly.Registry.ResiliencePipelineProvider<TKey>.TryGetPipeline<TResult>(TKey key, out Polly.ResiliencePipeline<TResult>? pipeline) -> bool
-abstract Polly.ResilienceContextPool.Get(string? operationKey, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> Polly.ResilienceContext!
+abstract Polly.ResilienceContextPool.Get(Polly.ResilienceContextCreationArguments arguments) -> Polly.ResilienceContext!
 abstract Polly.ResilienceContextPool.Return(Polly.ResilienceContext! context) -> void
 abstract Polly.ResilienceStrategy.ExecuteCore<TResult, TState>(System.Func<Polly.ResilienceContext!, TState, System.Threading.Tasks.ValueTask<Polly.Outcome<TResult>>>! callback, Polly.ResilienceContext! context, TState state) -> System.Threading.Tasks.ValueTask<Polly.Outcome<TResult>>
 abstract Polly.ResilienceStrategy<TResult>.ExecuteCore<TState>(System.Func<Polly.ResilienceContext!, TState, System.Threading.Tasks.ValueTask<Polly.Outcome<TResult>>>! callback, Polly.ResilienceContext! context, TState state) -> System.Threading.Tasks.ValueTask<Polly.Outcome<TResult>>
@@ -198,12 +198,20 @@ Polly.Registry.ResiliencePipelineRegistryOptions<TKey>.ResiliencePipelineRegistr
 Polly.ResilienceContext
 Polly.ResilienceContext.CancellationToken.get -> System.Threading.CancellationToken
 Polly.ResilienceContext.ContinueOnCapturedContext.get -> bool
-Polly.ResilienceContext.ContinueOnCapturedContext.set -> void
 Polly.ResilienceContext.IsSynchronous.get -> bool
 Polly.ResilienceContext.OperationKey.get -> string?
 Polly.ResilienceContext.Properties.get -> Polly.ResilienceProperties!
 Polly.ResilienceContext.ResilienceEvents.get -> System.Collections.Generic.IReadOnlyList<Polly.Telemetry.ResilienceEvent>!
+Polly.ResilienceContextCreationArguments
+Polly.ResilienceContextCreationArguments.CancellationToken.get -> System.Threading.CancellationToken
+Polly.ResilienceContextCreationArguments.ContinueOnCapturedContext.get -> bool?
+Polly.ResilienceContextCreationArguments.OperationKey.get -> string?
+Polly.ResilienceContextCreationArguments.ResilienceContextCreationArguments() -> void
+Polly.ResilienceContextCreationArguments.ResilienceContextCreationArguments(string? operationKey, bool? continueOnCapturedContext, System.Threading.CancellationToken cancellationToken) -> void
 Polly.ResilienceContextPool
+Polly.ResilienceContextPool.Get(bool continueOnCapturedContext, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> Polly.ResilienceContext!
+Polly.ResilienceContextPool.Get(string? operationKey, bool? continueOnCapturedContext, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> Polly.ResilienceContext!
+Polly.ResilienceContextPool.Get(string? operationKey, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> Polly.ResilienceContext!
 Polly.ResilienceContextPool.Get(System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> Polly.ResilienceContext!
 Polly.ResilienceContextPool.ResilienceContextPool() -> void
 Polly.ResiliencePipeline

src/Polly.Core/ResilienceContext.cs

@@ -14,8 +14,6 @@ namespace Polly;
 /// </remarks>
 public sealed class ResilienceContext
 {
-    private const bool ContinueOnCapturedContextDefault = false;
-
     private readonly List<ResilienceEvent> _resilienceEvents = new();
 
     internal ResilienceContext()
@@ -54,9 +52,9 @@ internal ResilienceContext()
     internal bool IsVoid => ResultType == typeof(VoidResult);
 
     /// <summary>
-    /// Gets or sets a value indicating whether the execution should continue on the captured context.
+    /// Gets a value indicating whether the execution should continue on the captured context.
     /// </summary>
-    public bool ContinueOnCapturedContext { get; set; }
+    public bool ContinueOnCapturedContext { get; internal set; }
 
     /// <summary>
     /// Gets a value indicating whether the context is initialized.
@@ -98,7 +96,6 @@ internal ResilienceContext Initialize<TResult>(bool isSynchronous)
     {
         IsSynchronous = isSynchronous;
         ResultType = typeof(TResult);
-        ContinueOnCapturedContext = ContinueOnCapturedContextDefault;
 
         return this;
     }

src/Polly.Core/ResilienceContextCreationArguments.cs

@@ -0,0 +1,37 @@
+namespace Polly;
+
+#pragma warning disable CA1815 // Override equals and operator equals on value types
+
+/// <summary>
+/// Arguments used by the <see cref="ResilienceContextPool"/> when creating <see cref="ResilienceContext"/>.
+/// </summary>
+public readonly struct ResilienceContextCreationArguments
+{
+    /// <summary>
+    /// Initializes a new instance of the <see cref="ResilienceContextCreationArguments"/> struct.
+    /// </summary>
+    /// <param name="operationKey">The operation key, if any.</param>
+    /// <param name="continueOnCapturedContext">Value indicating whether to continue on captured context.</param>
+    /// <param name="cancellationToken">The cancellation token.</param>
+    public ResilienceContextCreationArguments(string? operationKey, bool? continueOnCapturedContext, CancellationToken cancellationToken)
+    {
+        OperationKey = operationKey;
+        ContinueOnCapturedContext = continueOnCapturedContext;
+        CancellationToken = cancellationToken;
+    }
+
+    /// <summary>
+    /// Gets the operation key, if any.
+    /// </summary>
+    public string? OperationKey { get; }
+
+    /// <summary>
+    /// Gets the value indicating whether to continue on captured context, if any.
+    /// </summary>
+    public bool? ContinueOnCapturedContext { get; }
+
+    /// <summary>
+    /// Gets the cancellation token.
+    /// </summary>
+    public CancellationToken CancellationToken { get; }
+}

src/Polly.Core/ResilienceContextPool.Shared.cs

@@ -2,16 +2,19 @@
 
 public abstract partial class ResilienceContextPool
 {
+    private const bool ContinueOnCapturedContextDefault = false;
+
     private sealed class SharedPool : ResilienceContextPool
     {
         private readonly ObjectPool<ResilienceContext> _pool = new(static () => new ResilienceContext(), static c => c.Reset());
 
-        public override ResilienceContext Get(string? operationKey, CancellationToken cancellationToken = default)
+        public override ResilienceContext Get(ResilienceContextCreationArguments arguments)
         {
             var context = _pool.Get();
 
-            context.OperationKey = operationKey;
-            context.CancellationToken = cancellationToken;
+            context.OperationKey = arguments.OperationKey;
+            context.CancellationToken = arguments.CancellationToken;
+            context.ContinueOnCapturedContext = arguments.ContinueOnCapturedContext ?? ContinueOnCapturedContextDefault;
 
             return context;
         }

src/Polly.Core/ResilienceContextPool.cs

@@ -1,7 +1,7 @@
 namespace Polly;
 
-#pragma warning disable CA1716 // Identifiers should not match keywords
 #pragma warning disable RS0026 // Do not add multiple public overloads with optional parameters
+#pragma warning disable CA1716 // Identifiers should not match keywords
 
 /// <summary>
 /// The pool of <see cref="ResilienceContext"/> instances.
@@ -34,7 +34,45 @@ public abstract partial class ResilienceContextPool
     /// After the execution is finished you should return the <see cref="ResilienceContext"/> back to the pool
     /// by calling <see cref="Return(ResilienceContext)"/> method.
     /// </remarks>
-    public abstract ResilienceContext Get(string? operationKey, CancellationToken cancellationToken = default);
+    public ResilienceContext Get(string? operationKey, CancellationToken cancellationToken = default) => Get(operationKey, null, cancellationToken);
+
+    /// <summary>
+    /// Gets a <see cref="ResilienceContext"/> instance from the pool.
+    /// </summary>
+    /// <param name="operationKey">An operation key associated with the context.</param>
+    /// <param name="continueOnCapturedContext">Value indicating whether to continue on captured context.</param>
+    /// <param name="cancellationToken">The cancellation token.</param>
+    /// <returns>An instance of <see cref="ResilienceContext"/>.</returns>
+    /// <remarks>
+    /// After the execution is finished you should return the <see cref="ResilienceContext"/> back to the pool
+    /// by calling <see cref="Return(ResilienceContext)"/> method.
+    /// </remarks>
+    public ResilienceContext Get(string? operationKey, bool? continueOnCapturedContext, CancellationToken cancellationToken = default)
+        => Get(new ResilienceContextCreationArguments(operationKey, continueOnCapturedContext, cancellationToken));
+
+    /// <summary>
+    /// Gets a <see cref="ResilienceContext"/> instance from the pool.
+    /// </summary>
+    /// <param name="continueOnCapturedContext">Value indicating whether to continue on captured context.</param>
+    /// <param name="cancellationToken">The cancellation token.</param>
+    /// <returns>An instance of <see cref="ResilienceContext"/>.</returns>
+    /// <remarks>
+    /// After the execution is finished you should return the <see cref="ResilienceContext"/> back to the pool
+    /// by calling <see cref="Return(ResilienceContext)"/> method.
+    /// </remarks>
+    public ResilienceContext Get(bool continueOnCapturedContext, CancellationToken cancellationToken = default)
+        => Get(new ResilienceContextCreationArguments(null, continueOnCapturedContext, cancellationToken));
+
+    /// <summary>
+    /// Gets a <see cref="ResilienceContext"/> instance from the pool.
+    /// </summary>
+    /// <param name="arguments">The creation arguments.</param>
+    /// <returns>An instance of <see cref="ResilienceContext"/>.</returns>
+    /// <remarks>
+    /// After the execution is finished you should return the <see cref="ResilienceContext"/> back to the pool
+    /// by calling <see cref="Return(ResilienceContext)"/> method.
+    /// </remarks>
+    public abstract ResilienceContext Get(ResilienceContextCreationArguments arguments);
 
     /// <summary>
     /// Returns a <paramref name="context"/> back to the pool.

src/Polly/Utilities/Wrappers/ResilienceContextFactory.cs

@@ -8,8 +8,7 @@ public static ResilienceContext Create(
         bool continueOnCapturedContext,
         out IDictionary<string, object> oldProperties)
     {
-        var resilienceContext = ResilienceContextPool.Shared.Get(context.OperationKey, cancellationToken);
-        resilienceContext.ContinueOnCapturedContext = continueOnCapturedContext;
+        var resilienceContext = ResilienceContextPool.Shared.Get(context.OperationKey, continueOnCapturedContext, cancellationToken);
         resilienceContext.Properties.SetProperties(context, out oldProperties);
 
         return resilienceContext;

test/Polly.Core.Tests/ResilienceContextPoolTests.cs

@@ -40,6 +40,36 @@ public void Get_CancellationToken_Ok()
         context.CancellationToken.Should().Be(token.Token);
     }
 
+    [Fact]
+    public void Get_ContinueOnCapturedContextDefault_ShouldBeFalse()
+    {
+        using var token = new CancellationTokenSource();
+
+        var context = ResilienceContextPool.Shared.Get();
+
+        context.ContinueOnCapturedContext.Should().BeFalse();
+    }
+
+    [Fact]
+    public void Get_ContinueOnCapturedContext_Ok()
+    {
+        var context = ResilienceContextPool.Shared.Get(true);
+
+        context.ContinueOnCapturedContext.Should().Be(true);
+    }
+
+    [Fact]
+    public void Get_OperationKeyContinueOnCapturedContext_Ok()
+    {
+        using var token = new CancellationTokenSource();
+
+        var context = ResilienceContextPool.Shared.Get("dummy", true, token.Token);
+
+        context.ContinueOnCapturedContext.Should().Be(true);
+        context.OperationKey.Should().Be("dummy");
+        context.CancellationToken.Should().Be(token.Token);
+    }
+
     [InlineData(null)]
     [InlineData("")]
     [InlineData("some-key")]