Allow reusing CircuitBreakerManualControl across multiple CBs (#1375)

src/Polly.Core/CircuitBreaker/CircuitBreakerManualControl.cs

@@ -1,66 +1,63 @@
+using System.Collections.Generic;
+
 namespace Polly.CircuitBreaker;
 
 /// <summary>
 /// Allows manual control of the circuit-breaker.
 /// </summary>
+/// <remarks>
+/// The instance of this class can be reused across multiple circuit breakers.
+/// </remarks>
 public sealed class CircuitBreakerManualControl : IDisposable
 {
-    private Action? _onDispose;
-    private Func<ResilienceContext, Task>? _onIsolate;
-    private Func<ResilienceContext, Task>? _onReset;
+    private readonly HashSet<Action> _onDispose = new();
+    private readonly HashSet<Func<ResilienceContext, Task>> _onIsolate = new();
+    private readonly HashSet<Func<ResilienceContext, Task>> _onReset = new();
+    private bool _isolated;
 
     internal void Initialize(Func<ResilienceContext, Task> onIsolate, Func<ResilienceContext, Task> onReset, Action onDispose)
     {
-        if (_onIsolate != null)
+        _onDispose.Add(onDispose);
+        _onIsolate.Add(onIsolate);
+        _onReset.Add(onReset);
+
+        if (_isolated)
         {
-            throw new InvalidOperationException($"This instance of '{nameof(CircuitBreakerManualControl)}' is already initialized and cannot be used in a different circuit-breaker strategy.");
-        }
+            var context = ResilienceContext.Get().Initialize<VoidResult>(isSynchronous: true);
 
-        _onDispose = onDispose;
-        _onIsolate = onIsolate;
-        _onReset = onReset;
+            // if the control indicates that circuit breaker should be isolated, we isolate it right away
+            IsolateAsync(context).GetAwaiter().GetResult();
+        }
     }
 
-    /// <summary>
-    /// Gets a value indicating whether the manual control is initialized.
-    /// </summary>
-    /// <remarks>
-    /// The initialization happens when the circuit-breaker strategy is attached to this class.
-    /// This happens when the final strategy is created by the <see cref="ResilienceStrategyBuilder.Build"/> call.
-    /// </remarks>
-    internal bool IsInitialized => _onIsolate != null;
-
     /// <summary>
     /// Isolates (opens) the circuit manually, and holds it in this state until a call to <see cref="CloseAsync(CancellationToken)"/> is made.
     /// </summary>
     /// <param name="context">The resilience context.</param>
     /// <returns>The instance of <see cref="Task"/> that represents the asynchronous execution.</returns>
     /// <exception cref="ArgumentNullException">Thrown when <paramref name="context"/> is <see langword="null"/>.</exception>
-    /// <exception cref="InvalidOperationException">Thrown when manual control is not initialized.</exception>
     /// <exception cref="ObjectDisposedException">Thrown when calling this method after this object is disposed.</exception>
-    internal Task IsolateAsync(ResilienceContext context)
+    internal async Task IsolateAsync(ResilienceContext context)
     {
         Guard.NotNull(context);
 
-        if (_onIsolate == null)
+        _isolated = true;
+
+        foreach (var action in _onIsolate)
         {
-            throw new InvalidOperationException("The circuit-breaker manual control is not initialized");
+            await action(context).ConfigureAwait(context.ContinueOnCapturedContext);
         }
-
-        context.Initialize<VoidResult>(isSynchronous: false);
-        return _onIsolate(context);
     }
 
     /// <summary>
     /// Isolates (opens) the circuit manually, and holds it in this state until a call to <see cref="CloseAsync(CancellationToken)"/> is made.
     /// </summary>
     /// <param name="cancellationToken">The cancellation token.</param>
     /// <returns>The instance of <see cref="Task"/> that represents the asynchronous execution.</returns>
-    /// <exception cref="InvalidOperationException">Thrown if manual control is not initialized.</exception>
     /// <exception cref="ObjectDisposedException">Thrown when calling this method after this object is disposed.</exception>
     public async Task IsolateAsync(CancellationToken cancellationToken = default)
     {
-        var context = ResilienceContext.Get();
+        var context = ResilienceContext.Get().Initialize<VoidResult>(isSynchronous: false);
         context.CancellationToken = cancellationToken;
 
         try
@@ -79,27 +76,26 @@ public async Task IsolateAsync(CancellationToken cancellationToken = default)
     /// <param name="context">The resilience context.</param>
     /// <returns>The instance of <see cref="Task"/> that represents the asynchronous execution.</returns>
     /// <exception cref="ArgumentNullException">Thrown when <paramref name="context"/> is <see langword="null"/>.</exception>
-    /// <exception cref="InvalidOperationException">Thrown if manual control is not initialized.</exception>
     /// <exception cref="ObjectDisposedException">Thrown when calling this method after this object is disposed.</exception>
-    internal Task CloseAsync(ResilienceContext context)
+    internal async Task CloseAsync(ResilienceContext context)
     {
         Guard.NotNull(context);
 
-        if (_onReset == null)
-        {
-            throw new InvalidOperationException("The circuit-breaker manual control is not initialized");
-        }
+        _isolated = false;
 
         context.Initialize<VoidResult>(isSynchronous: false);
-        return _onReset(context);
+
+        foreach (var action in _onReset)
+        {
+            await action(context).ConfigureAwait(context.ContinueOnCapturedContext);
+        }
     }
 
     /// <summary>
     /// Closes the circuit, and resets any statistics controlling automated circuit-breaking.
     /// </summary>
     /// <param name="cancellationToken">The cancellation token.</param>
     /// <returns>The instance of <see cref="Task"/> that represents the asynchronous execution.</returns>
-    /// <exception cref="InvalidOperationException">Thrown if manual control is not initialized.</exception>
     /// <exception cref="ObjectDisposedException">Thrown when calling this method after this object is disposed.</exception>
     public async Task CloseAsync(CancellationToken cancellationToken = default)
     {
@@ -119,5 +115,15 @@ public async Task CloseAsync(CancellationToken cancellationToken = default)
     /// <summary>
     /// Disposes the current class.
     /// </summary>
-    public void Dispose() => _onDispose?.Invoke();
+    public void Dispose()
+    {
+        foreach (var action in _onDispose)
+        {
+            action();
+        }
+
+        _onDispose.Clear();
+        _onIsolate.Clear();
+        _onReset.Clear();
+    }
 }

test/Polly.Core.Tests/CircuitBreaker/CircuitBreakerManualControlTests.cs

@@ -4,46 +4,58 @@ namespace Polly.Core.Tests.CircuitBreaker;
 
 public class CircuitBreakerManualControlTests
 {
-    [Fact]
-    public void Ctor_Ok()
+    [InlineData(true)]
+    [InlineData(false)]
+    [Theory]
+    public async Task IsolateAsync_NotInitialized_Ok(bool closedAfter)
     {
         using var control = new CircuitBreakerManualControl();
+        await control.IsolateAsync();
+        if (closedAfter)
+        {
+            await control.CloseAsync();
+        }
 
-        control.IsInitialized.Should().BeFalse();
-    }
+        var isolated = false;
 
-    [Fact]
-    public async Task IsolateAsync_NotInitialized_Throws()
-    {
-        using var control = new CircuitBreakerManualControl();
+        control.Initialize(
+            c =>
+            {
+                c.IsSynchronous.Should().BeTrue();
+                isolated = true;
+                return Task.CompletedTask;
+            },
+            _ => Task.CompletedTask,
+            () => { });
 
-        await control
-            .Invoking(c => c.IsolateAsync(CancellationToken.None))
-            .Should()
-            .ThrowAsync<InvalidOperationException>();
+        isolated.Should().Be(!closedAfter);
     }
 
     [Fact]
-    public async Task ResetAsync_NotInitialized_Throws()
+    public async Task ResetAsync_NotInitialized_Ok()
     {
         using var control = new CircuitBreakerManualControl();
 
         await control
             .Invoking(c => c.CloseAsync(CancellationToken.None))
             .Should()
-            .ThrowAsync<InvalidOperationException>();
+            .NotThrowAsync();
     }
 
     [Fact]
-    public void Initialize_Twice_Throws()
+    public async Task Initialize_Twice_Ok()
     {
-        using var control = new CircuitBreakerManualControl();
+        int called = 0;
+        var control = new CircuitBreakerManualControl();
         control.Initialize(_ => Task.CompletedTask, _ => Task.CompletedTask, () => { });
+        control.Initialize(_ => { called++; return Task.CompletedTask; }, _ => { called++; return Task.CompletedTask; }, () => { called++; });
 
-        control
-            .Invoking(c => c.Initialize(_ => Task.CompletedTask, _ => Task.CompletedTask, () => { }))
-            .Should()
-            .Throw<InvalidOperationException>();
+        await control.IsolateAsync();
+        await control.CloseAsync();
+
+        control.Dispose();
+
+        called.Should().Be(3);
     }
 
     [Fact]

test/Polly.Core.Tests/CircuitBreaker/CircuitBreakerResilienceStrategyBuilderTests.cs

@@ -179,4 +179,27 @@ public void AddAdvancedCircuitBreaker_IntegrationTest()
         halfOpened.Should().Be(2);
         closed.Should().Be(1);
     }
+
+    [Fact]
+    public async Task AddCircuitBreakers_WithIsolatedManualControl_ShouldBeIsolated()
+    {
+        var manualControl = new CircuitBreakerManualControl();
+        await manualControl.IsolateAsync();
+
+        var strategy1 = new ResilienceStrategyBuilder()
+            .AddAdvancedCircuitBreaker(new() { ManualControl = manualControl })
+            .Build();
+
+        var strategy2 = new ResilienceStrategyBuilder()
+            .AddAdvancedCircuitBreaker(new() { ManualControl = manualControl })
+            .Build();
+
+        strategy1.Invoking(s => s.Execute(() => { })).Should().Throw<IsolatedCircuitException>();
+        strategy2.Invoking(s => s.Execute(() => { })).Should().Throw<IsolatedCircuitException>();
+
+        await manualControl.CloseAsync();
+
+        strategy1.Execute(() => { });
+        strategy2.Execute(() => { });
+    }
 }

test/Polly.Core.Tests/CircuitBreaker/CircuitBreakerResilienceStrategyTests.cs

@@ -54,8 +54,6 @@ public async Task Ctor_ManualControl_EnsureAttached()
         _options.ManualControl = new CircuitBreakerManualControl();
         var strategy = Create();
 
-        _options.ManualControl.IsInitialized.Should().BeTrue();
-
         await _options.ManualControl.IsolateAsync(CancellationToken.None);
         strategy.Invoking(s => s.Execute(_ => 0)).Should().Throw<IsolatedCircuitException>();
 

test/Polly.Core.Tests/CircuitBreaker/CircuitBreakerStateProviderTests.cs

@@ -29,7 +29,7 @@ public async Task ResetAsync_NotInitialized_Throws()
         await control
             .Invoking(c => c.CloseAsync(CancellationToken.None))
             .Should()
-            .ThrowAsync<InvalidOperationException>();
+            .NotThrowAsync<InvalidOperationException>();
     }
 
     [Fact]