Allow to dispose linked resources on pipeline disposal (#1511)

Author: martintmk

src/Polly.Core/PublicAPI.Unshipped.txt

@@ -159,6 +159,7 @@ Polly.PredicateBuilder<TResult>.PredicateBuilder() -> void
 Polly.PredicateResult
 Polly.Registry.ConfigureBuilderContext<TKey>
 Polly.Registry.ConfigureBuilderContext<TKey>.EnableReloads(System.Func<System.Func<System.Threading.CancellationToken>!>! tokenProducerFactory) -> void
+Polly.Registry.ConfigureBuilderContext<TKey>.OnPipelineDisposed(System.Action! callback) -> void
 Polly.Registry.ConfigureBuilderContext<TKey>.PipelineKey.get -> TKey
 Polly.Registry.ResiliencePipelineProvider<TKey>
 Polly.Registry.ResiliencePipelineProvider<TKey>.ResiliencePipelineProvider() -> void

src/Polly.Core/Registry/ConfigureBuilderContext.cs

@@ -33,6 +33,8 @@ internal ConfigureBuilderContext(TKey strategyKey, string builderName, string? b
 
     internal Func<Func<CancellationToken>>? ReloadTokenProducer { get; private set; }
 
+    internal List<Action> DisposeCallbacks { get; } = new();
+
     /// <summary>
     /// Enables dynamic reloading of the strategy retrieved from <see cref="ResiliencePipelineRegistry{TKey}"/>.
     /// </summary>
@@ -48,4 +50,15 @@ public void EnableReloads(Func<Func<CancellationToken>> tokenProducerFactory)
 
         ReloadTokenProducer = tokenProducerFactory;
     }
+
+    /// <summary>
+    /// Registers a callback that is called when the pipeline instance being configured is disposed.
+    /// </summary>
+    /// <param name="callback">The callback delegate.</param>
+    public void OnPipelineDisposed(Action callback)
+    {
+        Guard.NotNull(callback);
+
+        DisposeCallbacks.Add(callback);
+    }
 }

src/Polly.Core/Registry/RegistryPipelineComponentBuilder.cs

@@ -60,13 +60,17 @@ private Builder CreateBuilder()
         _configure(builder, context);
 
         return new(
-            builder.BuildPipelineComponent,
+            () => PipelineComponentFactory.WithDisposableCallbacks(
+                    builder.BuildPipelineComponent(),
+                    context.DisposeCallbacks),
             context.ReloadTokenProducer,
+            context.DisposeCallbacks,
             builder.TelemetryListener);
     }
 
     private record Builder(
         Func<PipelineComponent> ComponentFactory,
         Func<Func<CancellationToken>>? ReloadTokenProducer,
+        List<Action> DisposeCallbacks,
         TelemetryListener? Listener);
 }

src/Polly.Core/Utils/Pipeline/ComponentWithDisposeCallbacks.cs

@@ -0,0 +1,43 @@
+namespace Polly.Utils.Pipeline;
+
+internal class ComponentWithDisposeCallbacks : PipelineComponent
+{
+    private readonly List<Action> _callbacks;
+
+    public ComponentWithDisposeCallbacks(PipelineComponent component, List<Action> callbacks)
+    {
+        Component = component;
+        _callbacks = callbacks;
+    }
+
+    internal PipelineComponent Component { get; }
+
+    public override void Dispose()
+    {
+        ExecuteCallbacks();
+
+        Component.Dispose();
+    }
+
+    public override ValueTask DisposeAsync()
+    {
+        ExecuteCallbacks();
+
+        return Component.DisposeAsync();
+    }
+
+    internal override ValueTask<Outcome<TResult>> ExecuteCore<TResult, TState>(
+        Func<ResilienceContext, TState, ValueTask<Outcome<TResult>>> callback,
+        ResilienceContext context,
+        TState state) => Component.ExecuteCore(callback, context, state);
+
+    private void ExecuteCallbacks()
+    {
+        foreach (var callback in _callbacks)
+        {
+            callback();
+        }
+
+        _callbacks.Clear();
+    }
+}

src/Polly.Core/Utils/Pipeline/PipelineComponentFactory.cs

@@ -1,6 +1,7 @@
 using System;
 using System.Collections.Generic;
 using Polly.Telemetry;
+
 namespace Polly.Utils.Pipeline;
 
 internal static class PipelineComponentFactory
@@ -13,6 +14,16 @@ internal static class PipelineComponentFactory
 
     public static PipelineComponent FromStrategy<T>(ResilienceStrategy<T> strategy) => new BridgeComponent<T>(strategy);
 
+    public static PipelineComponent WithDisposableCallbacks(PipelineComponent component, IEnumerable<Action> callbacks)
+    {
+        if (!callbacks.Any())
+        {
+            return component;
+        }
+
+        return new ComponentWithDisposeCallbacks(component, callbacks.ToList());
+    }
+
     public static PipelineComponent CreateComposite(
         IReadOnlyList<PipelineComponent> components,
         ResilienceStrategyTelemetry telemetry,

src/Polly.Extensions/DependencyInjection/AddResiliencePipelineContext.cs

@@ -2,6 +2,7 @@
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Options;
 using Polly.Registry;
+using Polly.Utils;
 
 namespace Polly.DependencyInjection;
 
@@ -63,4 +64,15 @@ internal AddResiliencePipelineContext(ConfigureBuilderContext<TKey> registryCont
 
         return name == null ? monitor.CurrentValue : monitor.Get(name);
     }
+
+    /// <summary>
+    /// Registers a callback that is called when the pipeline instance being configured is disposed.
+    /// </summary>
+    /// <param name="callback">The callback delegate.</param>
+    public void OnPipelineDisposed(Action callback)
+    {
+        Guard.NotNull(callback);
+
+        RegistryContext.OnPipelineDisposed(callback);
+    }
 }

src/Polly.Extensions/PublicAPI.Unshipped.txt

@@ -3,6 +3,7 @@ abstract Polly.Telemetry.MeteringEnricher.Enrich<TResult, TArgs>(in Polly.Teleme
 Polly.DependencyInjection.AddResiliencePipelineContext<TKey>
 Polly.DependencyInjection.AddResiliencePipelineContext<TKey>.EnableReloads<TOptions>(string? name = null) -> void
 Polly.DependencyInjection.AddResiliencePipelineContext<TKey>.GetOptions<TOptions>(string? name = null) -> TOptions
+Polly.DependencyInjection.AddResiliencePipelineContext<TKey>.OnPipelineDisposed(System.Action! callback) -> void
 Polly.DependencyInjection.AddResiliencePipelineContext<TKey>.PipelineKey.get -> TKey
 Polly.DependencyInjection.AddResiliencePipelineContext<TKey>.ServiceProvider.get -> System.IServiceProvider!
 Polly.PollyServiceCollectionExtensions

src/Polly.Testing/ResiliencePipelineExtensions.cs

@@ -62,7 +62,7 @@ private static object GetStrategyInstance<T>(PipelineComponent component)
         return component;
     }
 
-    private static bool ShouldSkip(object instance) => instance is ReloadableComponent;
+    private static bool ShouldSkip(object instance) => instance is ReloadableComponent || instance is ComponentWithDisposeCallbacks;
 
     private static void ExpandComponents(this PipelineComponent component, List<PipelineComponent> components)
     {
@@ -78,6 +78,10 @@ private static void ExpandComponents(this PipelineComponent component, List<Pipe
             components.Add(reloadable);
             ExpandComponents(reloadable.Component, components);
         }
+        else if (component is ComponentWithDisposeCallbacks callbacks)
+        {
+            ExpandComponents(callbacks.Component, components);
+        }
         else
         {
             components.Add(component);

test/Polly.Core.Tests/Registry/ResiliencePipelineRegistryTests.cs

@@ -285,6 +285,37 @@ public void EnableReloads_Ok()
         tries.Should().Be(retryCount + 1);
     }
 
+    [Fact]
+    public void EnableReloads_EnsureDisposedCallbackCalled()
+    {
+        // arrange
+        var registry = new ResiliencePipelineRegistry<string>();
+        using var changeSource = new CancellationTokenSource();
+        var disposedCalls = 0;
+
+        registry.TryAddBuilder("dummy", (builder, context) =>
+        {
+            // this call enables dynamic reloads for the dummy strategy
+            context.EnableReloads(() => () => changeSource.Token);
+            context.OnPipelineDisposed(() => disposedCalls++);
+            builder.AddTimeout(TimeSpan.FromSeconds(1));
+        });
+
+        // act
+        var strategy = registry.GetPipeline("dummy");
+
+        // assert
+        disposedCalls.Should().Be(0);
+        strategy.Execute(() => { });
+
+        changeSource.Cancel();
+        disposedCalls.Should().Be(1);
+        strategy.Execute(() => { });
+
+        registry.Dispose();
+        disposedCalls.Should().Be(2);
+    }
+
     [Fact]
     public void EnableReloads_Generic_Ok()
     {

test/Polly.Core.Tests/Utils/Pipeline/ComponentWithDisposeCallbacksTests.cs

@@ -0,0 +1,46 @@
+using NSubstitute;
+using Polly.Utils.Pipeline;
+
+namespace Polly.Core.Tests.Utils.Pipeline;
+
+public class ComponentWithDisposeCallbacksTests
+{
+    [InlineData(true)]
+    [InlineData(false)]
+    [Theory]
+    public async Task Dispose_Ok(bool isAsync)
+    {
+        // Arrange
+        var called1 = 0;
+        var called2 = 0;
+
+        var callbacks = new List<Action>
+        {
+            () => called1++,
+            () => called2++
+        };
+        var component = Substitute.For<PipelineComponent>();
+        var sut = new ComponentWithDisposeCallbacks(component, callbacks);
+
+        // Act
+        if (isAsync)
+        {
+            await sut.DisposeAsync();
+            await sut.DisposeAsync();
+            await component.Received(2).DisposeAsync();
+        }
+        else
+        {
+            sut.Dispose();
+#pragma warning disable S3966 // Objects should not be disposed more than once
+            sut.Dispose();
+#pragma warning restore S3966 // Objects should not be disposed more than once
+            component.Received(2).Dispose();
+        }
+
+        // Assert
+        callbacks.Should().BeEmpty();
+        called1.Should().Be(1);
+        called2.Should().Be(1);
+    }
+}