Fix child container reference holding issue

z4kn4fein · Aug 16, 2023 · 5dddd3b · 5dddd3b
1 parent 7d2cf08
commit 5dddd3b
Show file tree

Hide file tree

Showing 3 changed files with 72 additions and 36 deletions.
diff --git a/.version b/.version
@@ -1 +1 @@
-5.11.0
+5.11.1
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,10 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [v5.11.1] - 2023-08-16
+### Fixed
+- [#142](https://github.com/z4kn4fein/stashbox/issues/142): Upon disposing child containers, their parents still held a strong reference to them.
+
 ## [v5.11.0] - 2023-06-21
 ### Changed
 - Moved several functions of `IDependencyResolver` to extension methods.
@@ -384,6 +388,7 @@ The validation was executed only at the expression tree building phase, so an al
 - Removed the legacy container extension functionality.
 - Removed the support of PCL v259.
 
+[v5.11.1]: https://github.com/z4kn4fein/stashbox/compare/5.11.0...5.11.1
 [v5.11.0]: https://github.com/z4kn4fein/stashbox/compare/5.10.2...5.11.0
 [v5.10.2]: https://github.com/z4kn4fein/stashbox/compare/5.10.1...5.10.2
 [v5.10.1]: https://github.com/z4kn4fein/stashbox/compare/5.10.0...5.10.1

diff --git a/src/StashboxContainer.cs b/src/StashboxContainer.cs
@@ -20,42 +20,49 @@ public sealed partial class StashboxContainer : IStashboxContainer
 {
     private sealed class ChildContainerStore
     {
-        public ImmutableTree<object, IStashboxContainer> ChildContainers = ImmutableTree<object, IStashboxContainer>.Empty;
+        public ImmutableTree<object, IStashboxContainer> ChildContainers =
+            ImmutableTree<object, IStashboxContainer>.Empty;
     }
-    
+
     private static int globalContainerId = int.MinValue;
-    
+
     private int disposed;
     private readonly ChildContainerStore childContainerStore;
     private readonly ChildContainerStore? directParentChildContainerStore;
     private readonly ContainerConfigurator containerConfigurator;
     private readonly ResolutionScope rootScope;
     private readonly object containerId;
+    private readonly bool shouldDisposeWithParent;
 
     /// <summary>
     /// Constructs a <see cref="StashboxContainer"/>.
     /// </summary>
     public StashboxContainer(Action<ContainerConfigurator>? config = null)
-        : this(null, new ResolutionStrategy(), new ContainerConfigurator(), ReserveContainerId(), config)
-    { }
+        : this(null, new ResolutionStrategy(), new ContainerConfigurator(), ReserveContainerId(), false, config)
+    {
+    }
 
     private StashboxContainer(StashboxContainer? parentContainer, IResolutionStrategy resolutionStrategy,
-        ContainerConfigurator containerConfigurator, object containerId, Action<ContainerConfigurator>? config = null)
+        ContainerConfigurator containerConfigurator, object containerId, bool shouldDisposeWithParent,
+        Action<ContainerConfigurator>? config = null)
     {
         this.childContainerStore = new ChildContainerStore();
         this.directParentChildContainerStore = parentContainer?.childContainerStore;
-        this.ContainerContext = new ContainerContext(parentContainer?.ContainerContext, resolutionStrategy, containerConfigurator.ContainerConfiguration);
+        this.ContainerContext = new ContainerContext(parentContainer?.ContainerContext, resolutionStrategy,
+            containerConfigurator.ContainerConfiguration);
         this.rootScope = (ResolutionScope)this.ContainerContext.RootScope;
         this.containerConfigurator = containerConfigurator;
         this.containerId = containerId;
+        this.shouldDisposeWithParent = shouldDisposeWithParent;
         config?.Invoke(this.containerConfigurator);
     }
 
     /// <inheritdoc />
     public IContainerContext ContainerContext { get; }
 
     /// <inheritdoc />
-    public IEnumerable<ReadOnlyKeyValue<object, IStashboxContainer>> ChildContainers => this.childContainerStore.ChildContainers.Walk();
+    public IEnumerable<ReadOnlyKeyValue<object, IStashboxContainer>> ChildContainers =>
+        this.childContainerStore.ChildContainers.Walk();
 
     /// <inheritdoc />
     public void RegisterResolver(IResolver resolver)
@@ -101,51 +108,55 @@ public void Validate()
                 exceptions.Add(ex);
             }
         }
-        
-        foreach (var child in ChildContainers)
+
+        foreach (var child in this.childContainerStore.ChildContainers.Walk())
         {
             try
             {
                 child.Value.Validate();
             }
             catch (AggregateException ex)
             {
-                exceptions.Add(new AggregateException($"Child container validation failed for '{child.Key}'. See the inner exceptions for details.", ex.InnerExceptions));
+                exceptions.Add(new AggregateException(
+                    $"Child container validation failed for '{child.Key}'. See the inner exceptions for details.",
+                    ex.InnerExceptions));
             }
         }
 
         if (exceptions.Length > 0)
-            throw new AggregateException("Container validation failed. See the inner exceptions for details.", exceptions);
+            throw new AggregateException("Container validation failed. See the inner exceptions for details.",
+                exceptions);
     }
 
     /// <inheritdoc />
-    public IStashboxContainer CreateChildContainer(Action<ContainerConfigurator>? config = null, bool attachToParent = true) =>
-        this.CreateChildContainer(ReserveContainerId(), config == null ? null : cont => cont.Configure(config), attachToParent);
-
+    public IStashboxContainer CreateChildContainer(Action<ContainerConfigurator>? config = null,
+        bool attachToParent = true) =>
+        this.CreateChildContainer(ReserveContainerId(), config == null ? null : cont => cont.Configure(config),
+            attachToParent);
+
     /// <inheritdoc />
-    public IStashboxContainer CreateChildContainer(object identifier, Action<IStashboxContainer>? config = null, bool attachToParent = true)
+    public IStashboxContainer CreateChildContainer(object identifier, Action<IStashboxContainer>? config = null,
+        bool attachToParent = true)
     {
         this.ThrowIfDisposed();
 
         var child = new StashboxContainer(this, this.ContainerContext.ResolutionStrategy,
-            new ContainerConfigurator(this.ContainerContext.ContainerConfiguration.Clone()), identifier);
+            new ContainerConfigurator(this.ContainerContext.ContainerConfiguration.Clone()), identifier,
+            attachToParent);
 
         config?.Invoke(child);
-        
-        Swap.SwapValue(ref this.childContainerStore.ChildContainers, (t1, t2, _, _, childStore) => 
+
+        Swap.SwapValue(ref this.childContainerStore.ChildContainers, (t1, t2, _, _, childStore) =>
                 childStore.AddOrUpdate(t1, t2, false, false),
             identifier, child, Constants.DelegatePlaceholder, Constants.DelegatePlaceholder);
-
-        if (attachToParent)
-            this.rootScope.AddDisposableTracking(child);
 
         return child;
     }
-    
+
     /// <inheritdoc />
     public IStashboxContainer? GetChildContainer(object identifier) =>
         this.childContainerStore.ChildContainers.GetOrDefaultByValue(identifier);
-    
+
     /// <inheritdoc />
     public void Configure(Action<ContainerConfigurator> config)
     {
@@ -188,31 +199,51 @@ public void Dispose()
         if (Interlocked.CompareExchange(ref this.disposed, 1, 0) != 0)
             return;
 
-        this.RemoveFromChildContainers();
-        this.ContainerContext.RootScope.Dispose();
+        this.RemoveSelfFromParentChildContainers();
+        this.DisposeChildContainers();
+        this.rootScope.Dispose();
     }
 
 #if HAS_ASYNC_DISPOSABLE
     /// <inheritdoc />
-    public ValueTask DisposeAsync()
+    public async ValueTask DisposeAsync()
     {
         if (Interlocked.CompareExchange(ref this.disposed, 1, 0) != 0)
-            return new ValueTask(Task.CompletedTask);
+            return;
 
-        this.RemoveFromChildContainers();
-        return this.ContainerContext.RootScope.DisposeAsync();
+        this.RemoveSelfFromParentChildContainers();
+        await this.DisposeChildContainersAsync().ConfigureAwait(false);
+        await this.rootScope.DisposeAsync().ConfigureAwait(false);
+    }
+
+    private async ValueTask DisposeChildContainersAsync()
+    {
+        foreach (var childContainer in this.childContainerStore.ChildContainers.Walk())
+        {
+            if (childContainer.Value is StashboxContainer { shouldDisposeWithParent: true } container)
+                await container.DisposeAsync().ConfigureAwait(false);
+        }
     }
-
 #endif
 
-    private void RemoveFromChildContainers()
+    private void DisposeChildContainers()
+    {
+        foreach (var childContainer in this.childContainerStore.ChildContainers.Walk())
+        {
+            if (childContainer.Value is StashboxContainer { shouldDisposeWithParent: true } container)
+                container.Dispose();
+        }
+    }
+
+    private void RemoveSelfFromParentChildContainers()
     {
         if (this.directParentChildContainerStore != null)
             Swap.SwapValue(ref this.directParentChildContainerStore.ChildContainers, (t1, _, _, _, childRepo) =>
-                    childRepo.Remove(t1, false), 
-                this.containerId, Constants.DelegatePlaceholder, Constants.DelegatePlaceholder, Constants.DelegatePlaceholder);
+                    childRepo.Remove(t1, false),
+                this.containerId, Constants.DelegatePlaceholder, Constants.DelegatePlaceholder,
+                Constants.DelegatePlaceholder);
     }
-    
+
     private static int ReserveContainerId() =>
         Interlocked.Increment(ref globalContainerId);
 }