Merge pull request #728 from hadashiA/ku/fix-dispose-callback-scope

Fix bug where DisposeCallback would not close to parent and child
hadashiA · Nov 30, 2024 · b84428c · b84428c
2 parents 22f73e0 + 1c096d3
commit b84428c
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 7 deletions.
diff --git a/VContainer/Assets/Tests/ContainerTest.cs b/VContainer/Assets/Tests/ContainerTest.cs
@@ -617,6 +617,31 @@ public void OnContainerDisposeCallback()
             Assert.That(resolvedJustBeforeDispose2, Is.Not.Null);
         }
 
+        [Test]
+        public void OnContainerDisposeCallback_ParentChild()
+        {
+            var parentDisposeCalled = false;
+            var childDisposeCalled = false;
+
+            var builder = new ContainerBuilder();
+            builder.Register<NoDependencyServiceA>(Lifetime.Scoped);
+
+            builder.RegisterDisposeCallback(x => parentDisposeCalled = true);
+
+            var container = builder.Build();
+
+            var childContainer = container.CreateScope(childBuilder =>
+            {
+                childBuilder.Register<NoDependencyServiceB>(Lifetime.Scoped);
+                childBuilder.RegisterDisposeCallback(x => childDisposeCalled = true);
+            });
+
+            childContainer.Dispose();
+
+            Assert.That(childDisposeCalled, Is.True);
+            Assert.That(parentDisposeCalled, Is.False);
+        }
+
         [Test]
         public void TryResolveTransient()
         {

diff --git a/VContainer/Assets/VContainer/Runtime/ContainerBuilderExtensions.cs b/VContainer/Assets/VContainer/Runtime/ContainerBuilderExtensions.cs
@@ -143,8 +143,15 @@ public static RegistrationBuilder RegisterFactory<TParam1, TParam2, TParam3, TPa
 
         public static void RegisterDisposeCallback(this IContainerBuilder builder, Action<IObjectResolver> callback)
         {
-            builder.Register(container => new BuilderCallbackDisposable(callback, container), Lifetime.Scoped);
-            builder.RegisterBuildCallback(container => container.Resolve<IReadOnlyList<BuilderCallbackDisposable>>());
+            if (!builder.Exists(typeof(BuilderCallbackDisposable)))
+            {
+                builder.Register<BuilderCallbackDisposable>(Lifetime.Singleton);
+            }
+            builder.RegisterBuildCallback(container =>
+            {
+                var disposable = container.Resolve<BuilderCallbackDisposable>();
+                disposable.Disposing += callback;
+            });
         }
 
         [Obsolete("IObjectResolver is registered by default. This method does nothing.")]

diff --git a/VContainer/Assets/VContainer/Runtime/Internal/BuilderCallbackDisposable.cs b/VContainer/Assets/VContainer/Runtime/Internal/BuilderCallbackDisposable.cs
@@ -2,20 +2,19 @@
 
 namespace VContainer.Internal
 {
-    readonly struct BuilderCallbackDisposable : IDisposable
+    class BuilderCallbackDisposable : IDisposable
     {
-        readonly Action<IObjectResolver> callback;
+        public event Action<IObjectResolver> Disposing;
         readonly IObjectResolver container;
 
-        public BuilderCallbackDisposable(Action<IObjectResolver> callback, IObjectResolver container)
+        public BuilderCallbackDisposable(IObjectResolver container)
         {
-            this.callback = callback;
             this.container = container;
         }
 
         public void Dispose()
         {
-            callback.Invoke(container);
+            if (Disposing != null) Disposing.Invoke(container);
         }
     }
 }