From eca633f012df7eb317b825b231c31c108d65e9f1 Mon Sep 17 00:00:00 2001
From: hadashiA <dev@hadashikick.jp>
Date: Mon, 25 Dec 2023 11:29:30 +0900
Subject: [PATCH] FIx

---
 VContainer/Assets/Tests/ContainerTest.cs          | 15 ++++++++++-----
 .../Runtime/ContainerBuilderExtensions.cs         |  4 +++-
 ...Disposable.cs => BuilderCallbackDisposable.cs} |  4 ++--
 ....cs.meta => BuilderCallbackDisposable.cs.meta} |  0
 4 files changed, 15 insertions(+), 8 deletions(-)
 rename VContainer/Assets/VContainer/Runtime/Internal/{ActionDisposable.cs => BuilderCallbackDisposable.cs} (68%)
 rename VContainer/Assets/VContainer/Runtime/Internal/{ActionDisposable.cs.meta => BuilderCallbackDisposable.cs.meta} (100%)

diff --git a/VContainer/Assets/Tests/ContainerTest.cs b/VContainer/Assets/Tests/ContainerTest.cs
index da549744..57b2da5e 100644
--- a/VContainer/Assets/Tests/ContainerTest.cs
+++ b/VContainer/Assets/Tests/ContainerTest.cs
@@ -571,19 +571,24 @@ public void Inject()
         public void OnContainerDisposeCallback()
         {
             NoDependencyServiceA resolvedJustBeforeDispose = null;
-            
+            NoDependencyServiceB resolvedJustBeforeDispose2 = null;
+
             var builder = new ContainerBuilder();
-            
+
             builder.Register<NoDependencyServiceA>(Lifetime.Scoped);
+            builder.Register<NoDependencyServiceB>(Lifetime.Scoped);
             builder.RegisterDisposeCallback(resolver => resolvedJustBeforeDispose = resolver.Resolve<NoDependencyServiceA>());
+            builder.RegisterDisposeCallback(resolver => resolvedJustBeforeDispose2 = resolver.Resolve<NoDependencyServiceB>());
 
             var container = builder.Build();
-            
+
             Assert.That(resolvedJustBeforeDispose, Is.Null);
-            
+            Assert.That(resolvedJustBeforeDispose2, Is.Null);
+
             container.Dispose();
-            
+
             Assert.That(resolvedJustBeforeDispose, Is.Not.Null);
+            Assert.That(resolvedJustBeforeDispose2, Is.Not.Null);
         }
     }
 }
diff --git a/VContainer/Assets/VContainer/Runtime/ContainerBuilderExtensions.cs b/VContainer/Assets/VContainer/Runtime/ContainerBuilderExtensions.cs
index 5c707b0c..4dd8b8be 100644
--- a/VContainer/Assets/VContainer/Runtime/ContainerBuilderExtensions.cs
+++ b/VContainer/Assets/VContainer/Runtime/ContainerBuilderExtensions.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using System.Runtime.CompilerServices;
 using VContainer.Internal;
 
@@ -147,7 +148,8 @@ public static RegistrationBuilder RegisterFactory<TParam1, TParam2, TParam3, TPa
 
         public static void RegisterDisposeCallback(this IContainerBuilder builder, Action<IObjectResolver> callback)
         {
-            builder.Register(container => new ActionDisposable(callback, container), Lifetime.Scoped);
+            builder.Register(container => new BuilderCallbackDisposable(callback, container), Lifetime.Scoped);
+            builder.RegisterBuildCallback(container => container.Resolve<IReadOnlyList<BuilderCallbackDisposable>>());
         }
 
         [Obsolete("IObjectResolver is registered by default. This method does nothing.")]
diff --git a/VContainer/Assets/VContainer/Runtime/Internal/ActionDisposable.cs b/VContainer/Assets/VContainer/Runtime/Internal/BuilderCallbackDisposable.cs
similarity index 68%
rename from VContainer/Assets/VContainer/Runtime/Internal/ActionDisposable.cs
rename to VContainer/Assets/VContainer/Runtime/Internal/BuilderCallbackDisposable.cs
index 1b58b30c..c7693c3b 100644
--- a/VContainer/Assets/VContainer/Runtime/Internal/ActionDisposable.cs
+++ b/VContainer/Assets/VContainer/Runtime/Internal/BuilderCallbackDisposable.cs
@@ -2,12 +2,12 @@
 
 namespace VContainer.Internal
 {
-    struct ActionDisposable : IDisposable
+    readonly struct BuilderCallbackDisposable : IDisposable
     {
         readonly Action<IObjectResolver> callback;
         readonly IObjectResolver container;
 
-        public ActionDisposable(Action<IObjectResolver> callback, IObjectResolver container)
+        public BuilderCallbackDisposable(Action<IObjectResolver> callback, IObjectResolver container)
         {
             this.callback = callback;
             this.container = container;
diff --git a/VContainer/Assets/VContainer/Runtime/Internal/ActionDisposable.cs.meta b/VContainer/Assets/VContainer/Runtime/Internal/BuilderCallbackDisposable.cs.meta
similarity index 100%
rename from VContainer/Assets/VContainer/Runtime/Internal/ActionDisposable.cs.meta
rename to VContainer/Assets/VContainer/Runtime/Internal/BuilderCallbackDisposable.cs.meta