Fixes inner container disposing outer scope binding

gustavopsantos · Dec 12, 2022 · de8c737 · de8c737
1 parent 0867f1a
commit de8c737
Show file tree

Hide file tree

Showing 9 changed files with 59 additions and 18 deletions.
diff --git a/Assets/Reflex.Tests/ScopedContainerTests.cs b/Assets/Reflex.Tests/ScopedContainerTests.cs
@@ -109,7 +109,7 @@ public void DisposingInnerScopeShouldNotDisposeInstancesFromOuterScope()
                     inner.Resolve<DependsOnFoo>().Foo.IsDisposed.Should().BeFalse();
                 }
 
-                outer.Resolve<Foo>().IsDisposed.Should().BeTrue();
+                outer.Resolve<Foo>().IsDisposed.Should().BeFalse();
             }
         }
 
@@ -158,12 +158,11 @@ public void Dispose()
         public void InnerScopedContainerShouldNotDisposeOuterBindings()
         {
             var outer = new Container("Outer");
-            outer.BindSingleton<IManager, Manager>();
-            // outer.Resolve<IManager>(); // Workaround: resolving it in here avoids it being disposed at inner 
+            outer.BindSingleton<Foo, Foo>();
             var inner = outer.Scope("Inner");
-            inner.Resolve<IManager>();
+            inner.Resolve<Foo>();
             inner.Dispose();
-            outer.Resolve<IManager>().Disposed.Should().BeFalse();
+            outer.Resolve<Foo>().IsDisposed.Should().BeFalse();
         }
     }
 }
diff --git a/Assets/Reflex/Scripts/Core/Container.cs b/Assets/Reflex/Scripts/Core/Container.cs
@@ -3,7 +3,6 @@
 using Reflex.Injectors;
 using Reflex.Scripts.Utilities;
 using System.Collections.Generic;
-using System.Linq;
 using Reflex.Scripts.Enums;
 using Reflex.Scripts.Extensions;
 
@@ -74,7 +73,9 @@ public bool Contains(Type type)
 
         public void BindFunction<TContract>(Func<TContract> function)
         {
-            _resolvers.Add(typeof(TContract), new FunctionResolver(function as Func<object>));
+            var resolver = new FunctionResolver(function as Func<object>);
+            _disposables.Add(resolver);
+            _resolvers.Add(typeof(TContract), resolver);
         }
 
         public void BindInstance(object instance)
@@ -89,17 +90,23 @@ public void BindInstanceAs<TContract>(TContract instance)
 
         public void BindInstanceAs(object instance, Type asType)
         {
-            _resolvers[asType] = new InstanceResolver(instance);
+            var resolver = new InstanceResolver(instance);
+            _disposables.Add(resolver);
+            _resolvers[asType] = resolver;
         }
 
         public void BindTransient<TContract, TConcrete>() where TConcrete : TContract
         {
-            _resolvers[typeof(TContract)] = new TransientResolver(typeof(TConcrete));
+            var resolver = new TransientResolver(typeof(TConcrete));
+            _disposables.Add(resolver);
+            _resolvers[typeof(TContract)] = resolver;
         }
 
         public void BindSingleton<TContract, TConcrete>() where TConcrete : TContract
         {
-            _resolvers[typeof(TContract)] = new SingletonResolver(typeof(TConcrete));
+            var resolver = new SingletonResolver(typeof(TConcrete));
+            _disposables.Add(resolver);
+            _resolvers[typeof(TContract)] = resolver;
         }
 
         public T Construct<T>()

diff --git a/Assets/Reflex/Scripts/Injectors/ConstructorInjector.cs b/Assets/Reflex/Scripts/Injectors/ConstructorInjector.cs
@@ -19,10 +19,6 @@ internal static object ConstructAndInject(Type concrete, Container container)
 			try
 			{
 				var instance = info.Activator(objects);
-				if (instance is IDisposable disposable)
-				{
-					container.AddDisposable(disposable);
-				}
 				ArrayPool<object>.Shared.Return(objects);
 				return instance;
 			}

diff --git a/Assets/Reflex/Scripts/Resolvers/FunctionResolver.cs b/Assets/Reflex/Scripts/Resolvers/FunctionResolver.cs
@@ -18,5 +18,10 @@ public object Resolve(Container container)
             Resolutions++;
             return _function.Invoke();
         }
+
+        public void Dispose()
+        {
+            // Objects created by user functions, should be disposed by users
+        }
     }
 }
diff --git a/Assets/Reflex/Scripts/Resolvers/IResolver.cs b/Assets/Reflex/Scripts/Resolvers/IResolver.cs
@@ -2,7 +2,7 @@
 
 namespace Reflex
 {
-	internal interface IResolver
+	internal interface IResolver : IDisposable
 	{
 		Type Concrete { get; }
 		int Resolutions { get; }

diff --git a/Assets/Reflex/Scripts/Resolvers/InstanceResolver.cs b/Assets/Reflex/Scripts/Resolvers/InstanceResolver.cs
@@ -20,5 +20,10 @@ public object Resolve(Container container)
             Resolutions++;
             return _instance;
         }
+
+        public void Dispose()
+        {
+            // Objects created by user, should be disposed by users
+        }
     }
 }
diff --git a/Assets/Reflex/Scripts/Resolvers/SingletonResolver.cs b/Assets/Reflex/Scripts/Resolvers/SingletonResolver.cs
@@ -13,10 +13,18 @@ public SingletonResolver(Type concrete)
             Concrete = concrete;
         }
 
+        public void Dispose()
+        {
+            if (_instance is IDisposable disposable)
+            {
+                disposable.Dispose();
+            }
+        }
+
         public object Resolve(Container container)
         {
             Resolutions++;
-            
+
             if (_instance == null)
             {
                 _instance = container.Construct(Concrete);

diff --git a/Assets/Reflex/Scripts/Resolvers/TransientResolver.cs b/Assets/Reflex/Scripts/Resolvers/TransientResolver.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using Reflex.Injectors;
 
 namespace Reflex
@@ -7,6 +8,7 @@ internal class TransientResolver : IResolver
     {
         public Type Concrete { get; }
         public int Resolutions { get; private set; }
+        private readonly Stack<object> _instances = new Stack<object>();
 
         public TransientResolver(Type concrete)
         {
@@ -16,7 +18,21 @@ public TransientResolver(Type concrete)
         public object Resolve(Container container)
         {
             Resolutions++;
-            return ConstructorInjector.ConstructAndInject(Concrete, container);
+            var instance = ConstructorInjector.ConstructAndInject(Concrete, container);
+            _instances.Push(instance);
+            return instance;
+        }
+
+        public void Dispose()
+        {
+            while (_instances.Count > 0)
+            {
+                var instance = _instances.Pop();
+                if (instance is IDisposable disposable)
+                {
+                    disposable.Dispose();
+                }
+            }
         }
     }
 }
diff --git a/Assets/Reflex/Scripts/Utilities/CompositeDisposable.cs b/Assets/Reflex/Scripts/Utilities/CompositeDisposable.cs
@@ -7,11 +7,16 @@ public class CompositeDisposable : IDisposable
     {
         private readonly List<IDisposable> _registry = new List<IDisposable>();
 
+        public void Add(IDisposable disposable)
+        {
+            _registry.Add(disposable);
+        }
+
         public void TryAdd(object obj)
         {
             if (obj is IDisposable disposable)
             {
-                _registry.Add(disposable);
+                Add(disposable);
             }
         }