From 1c096d3ca3dba74fc04624ee345d88d1c2bf2e7a Mon Sep 17 00:00:00 2001
From: hadashiA <dev@hadashikick.jp>
Date: Fri, 29 Nov 2024 18:33:32 +0900
Subject: [PATCH] Fix bug where DisposeCallback would not close to parent and
 child

---
 VContainer/Assets/Tests/ContainerTest.cs      | 25 +++++++++++++++++++
 .../Runtime/ContainerBuilderExtensions.cs     | 11 ++++++--
 .../Internal/BuilderCallbackDisposable.cs     |  9 +++----
 3 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/VContainer/Assets/Tests/ContainerTest.cs b/VContainer/Assets/Tests/ContainerTest.cs
index 635d6618..6e5e86b1 100644
--- 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
index e82d7ade..5484effa 100644
--- 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
index c7693c3b..7cf49079 100644
--- 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);
         }
     }
 }
\ No newline at end of file