ResiliencePipelineRegistry is now disposable (#1496)

Author: martintmk

src/Polly.Core/CircuitBreaker/CircuitBreakerManualControl.cs

@@ -6,9 +6,9 @@ namespace Polly.CircuitBreaker;
 /// <remarks>
 /// The instance of this class can be reused across multiple circuit breakers.
 /// </remarks>
-public sealed class CircuitBreakerManualControl : IDisposable
+public sealed class CircuitBreakerManualControl
 {
-    private readonly HashSet<Action> _onDispose = new();
+    private readonly object _lock = new();
     private readonly HashSet<Func<ResilienceContext, Task>> _onIsolate = new();
     private readonly HashSet<Func<ResilienceContext, Task>> _onReset = new();
     private bool _isolated;
@@ -23,21 +23,34 @@ public CircuitBreakerManualControl()
     /// <summary>
     /// Initializes a new instance of the <see cref="CircuitBreakerManualControl"/> class.
     /// </summary>
-    /// <param name="isIsolated">Determines whether the circit breaker is isolated immediately after construction.</param>
+    /// <param name="isIsolated">Determines whether the circuit breaker is isolated immediately after construction.</param>
     public CircuitBreakerManualControl(bool isIsolated) => _isolated = isIsolated;
 
-    internal void Initialize(Func<ResilienceContext, Task> onIsolate, Func<ResilienceContext, Task> onReset, Action onDispose)
-    {
-        _onDispose.Add(onDispose);
-        _onIsolate.Add(onIsolate);
-        _onReset.Add(onReset);
+    internal bool IsEmpty => _onIsolate.Count == 0;
 
-        if (_isolated)
+    internal IDisposable Initialize(Func<ResilienceContext, Task> onIsolate, Func<ResilienceContext, Task> onReset)
+    {
+        lock (_lock)
         {
-            var context = ResilienceContextPool.Shared.Get().Initialize<VoidResult>(isSynchronous: true);
-
-            // if the control indicates that circuit breaker should be isolated, we isolate it right away
-            IsolateAsync(context).GetAwaiter().GetResult();
+            _onIsolate.Add(onIsolate);
+            _onReset.Add(onReset);
+
+            if (_isolated)
+            {
+                var context = ResilienceContextPool.Shared.Get().Initialize<VoidResult>(isSynchronous: true);
+
+                // if the control indicates that circuit breaker should be isolated, we isolate it right away
+                IsolateAsync(context).GetAwaiter().GetResult();
+            }
+
+            return new RegistrationDisposable(() =>
+            {
+                lock (_lock)
+                {
+                    _onIsolate.Remove(onIsolate);
+                    _onReset.Remove(onReset);
+                }
+            });
         }
     }
 
@@ -54,7 +67,14 @@ internal async Task IsolateAsync(ResilienceContext context)
 
         _isolated = true;
 
-        foreach (var action in _onIsolate)
+        Func<ResilienceContext, Task>[] callbacks;
+
+        lock (_lock)
+        {
+            callbacks = _onIsolate.ToArray();
+        }
+
+        foreach (var action in callbacks)
         {
             await action(context).ConfigureAwait(context.ContinueOnCapturedContext);
         }
@@ -95,7 +115,14 @@ internal async Task CloseAsync(ResilienceContext context)
 
         context.Initialize<VoidResult>(isSynchronous: false);
 
-        foreach (var action in _onReset)
+        Func<ResilienceContext, Task>[] callbacks;
+
+        lock (_lock)
+        {
+            callbacks = _onReset.ToArray();
+        }
+
+        foreach (var action in callbacks)
         {
             await action(context).ConfigureAwait(context.ContinueOnCapturedContext);
         }
@@ -121,18 +148,12 @@ public async Task CloseAsync(CancellationToken cancellationToken = default)
         }
     }
 
-    /// <summary>
-    /// Disposes the current class.
-    /// </summary>
-    public void Dispose()
+    private class RegistrationDisposable : IDisposable
     {
-        foreach (var action in _onDispose)
-        {
-            action();
-        }
+        private readonly Action _disposeAction;
+
+        public RegistrationDisposable(Action disposeAction) => _disposeAction = disposeAction;
 
-        _onDispose.Clear();
-        _onIsolate.Clear();
-        _onReset.Clear();
+        public void Dispose() => _disposeAction();
     }
 }

src/Polly.Core/CircuitBreaker/CircuitBreakerResilienceStrategy.cs

@@ -1,9 +1,10 @@
 namespace Polly.CircuitBreaker;
 
-internal sealed class CircuitBreakerResilienceStrategy<T> : ResilienceStrategy<T>
+internal sealed class CircuitBreakerResilienceStrategy<T> : ResilienceStrategy<T>, IDisposable
 {
     private readonly Func<OutcomeArguments<T, CircuitBreakerPredicateArguments>, ValueTask<bool>> _handler;
     private readonly CircuitStateController<T> _controller;
+    private readonly IDisposable? _manualControlRegistration;
 
     public CircuitBreakerResilienceStrategy(
         Func<OutcomeArguments<T, CircuitBreakerPredicateArguments>, ValueTask<bool>> handler,
@@ -15,10 +16,15 @@ public CircuitBreakerResilienceStrategy(
         _controller = controller;
 
         stateProvider?.Initialize(() => _controller.CircuitState, () => _controller.LastHandledOutcome);
-        manualControl?.Initialize(
+        _manualControlRegistration = manualControl?.Initialize(
             async c => await _controller.IsolateCircuitAsync(c).ConfigureAwait(c.ContinueOnCapturedContext),
-            async c => await _controller.CloseCircuitAsync(c).ConfigureAwait(c.ContinueOnCapturedContext),
-            _controller.Dispose);
+            async c => await _controller.CloseCircuitAsync(c).ConfigureAwait(c.ContinueOnCapturedContext));
+    }
+
+    public void Dispose()
+    {
+        _manualControlRegistration?.Dispose();
+        _controller.Dispose();
     }
 
     protected internal override async ValueTask<Outcome<T>> ExecuteCore<TState>(Func<ResilienceContext, TState, ValueTask<Outcome<T>>> callback, ResilienceContext context, TState state)

src/Polly.Core/PublicAPI.Unshipped.txt

@@ -24,7 +24,6 @@ Polly.CircuitBreaker.CircuitBreakerManualControl
 Polly.CircuitBreaker.CircuitBreakerManualControl.CircuitBreakerManualControl() -> void
 Polly.CircuitBreaker.CircuitBreakerManualControl.CircuitBreakerManualControl(bool isIsolated) -> void
 Polly.CircuitBreaker.CircuitBreakerManualControl.CloseAsync(System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task!
-Polly.CircuitBreaker.CircuitBreakerManualControl.Dispose() -> void
 Polly.CircuitBreaker.CircuitBreakerManualControl.IsolateAsync(System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task!
 Polly.CircuitBreaker.CircuitBreakerPredicateArguments
 Polly.CircuitBreaker.CircuitBreakerPredicateArguments.CircuitBreakerPredicateArguments() -> void
@@ -167,6 +166,8 @@ Polly.Registry.ConfigureBuilderContext<TKey>.PipelineKey.get -> TKey
 Polly.Registry.ResiliencePipelineProvider<TKey>
 Polly.Registry.ResiliencePipelineProvider<TKey>.ResiliencePipelineProvider() -> void
 Polly.Registry.ResiliencePipelineRegistry<TKey>
+Polly.Registry.ResiliencePipelineRegistry<TKey>.Dispose() -> void
+Polly.Registry.ResiliencePipelineRegistry<TKey>.DisposeAsync() -> System.Threading.Tasks.ValueTask
 Polly.Registry.ResiliencePipelineRegistry<TKey>.GetOrAddPipeline(TKey key, System.Action<Polly.ResiliencePipelineBuilder!, Polly.Registry.ConfigureBuilderContext<TKey>!>! configure) -> Polly.ResiliencePipeline!
 Polly.Registry.ResiliencePipelineRegistry<TKey>.GetOrAddPipeline(TKey key, System.Action<Polly.ResiliencePipelineBuilder!>! configure) -> Polly.ResiliencePipeline!
 Polly.Registry.ResiliencePipelineRegistry<TKey>.GetOrAddPipeline<TResult>(TKey key, System.Action<Polly.ResiliencePipelineBuilder<TResult>!, Polly.Registry.ConfigureBuilderContext<TKey>!>! configure) -> Polly.ResiliencePipeline<TResult>!

src/Polly.Core/Registry/ResiliencePipelineRegistry.TResult.cs

@@ -5,7 +5,7 @@ namespace Polly.Registry;
 public sealed partial class ResiliencePipelineRegistry<TKey> : ResiliencePipelineProvider<TKey>
     where TKey : notnull
 {
-    private sealed class GenericRegistry<TResult>
+    private sealed class GenericRegistry<TResult> : IDisposable, IAsyncDisposable
     {
         private readonly Func<ResiliencePipelineBuilder<TResult>> _activator;
         private readonly ConcurrentDictionary<TKey, Action<ResiliencePipelineBuilder<TResult>, ConfigureBuilderContext<TKey>>> _builders;
@@ -52,14 +52,34 @@ public ResiliencePipeline<TResult> GetOrAdd(TKey key, Action<ResiliencePipelineB
 #if NETCOREAPP3_0_OR_GREATER
             return _strategies.GetOrAdd(key, static (_, factory) =>
             {
-                return new ResiliencePipeline<TResult>(CreatePipelineComponent(factory.instance._activator, factory.context, factory.configure));
+                return new ResiliencePipeline<TResult>(CreatePipelineComponent(factory.instance._activator, factory.context, factory.configure), DisposeBehavior.Reject);
             },
             (instance: this, context, configure));
 #else
-            return _strategies.GetOrAdd(key, _ => new ResiliencePipeline<TResult>(CreatePipelineComponent(_activator, context, configure)));
+            return _strategies.GetOrAdd(key, _ => new ResiliencePipeline<TResult>(CreatePipelineComponent(_activator, context, configure), DisposeBehavior.Reject));
 #endif
         }
 
         public bool TryAddBuilder(TKey key, Action<ResiliencePipelineBuilder<TResult>, ConfigureBuilderContext<TKey>> configure) => _builders.TryAdd(key, configure);
+
+        public void Dispose()
+        {
+            foreach (var strategy in _strategies.Values)
+            {
+                strategy.DisposeHelper.ForceDispose();
+            }
+
+            _strategies.Clear();
+        }
+
+        public async ValueTask DisposeAsync()
+        {
+            foreach (var strategy in _strategies.Values)
+            {
+                await strategy.DisposeHelper.ForceDisposeAsync().ConfigureAwait(false);
+            }
+
+            _strategies.Clear();
+        }
     }
 }