Extending Unity Container for IDisposable Instances (2nd approach)

In my previous blog post I detailed an approach of making the Unity Dependency Injection Container ț̦o automatically call Dispose() on all the IDisposable instances it builds and injects. The implementation described there, makes use of custom lifetime managers and it works fine for most of the cases except for the PerResolveLifetimeManager. Unity was extended with the PerResolve functionality in a tricky way, which makes my previous IDisposable implementation not to work for it. Therefore, I’ve came up with another solution which works in all cases. I will go into details about it here. Both these blog posts are a continuation of the Who Disposes Your Repository post, I’ve written a while ago, where I describe a context where such an automatically disposing mechanism is desired.

For both implementations we consider that a child container is created and associated with each new scoped operation. A scoped operation may be a request or a session in a web application, or it may be a window or a view in a desktop one. During that operation the child container will be used to inject the dependencies. This is also known as using Container Hierarchies or Scoped Containers. When the operation ends, it will dispose the child container, and what we want is that this also calls Dispose() on all the IDisposable instances that were created within that operation.

Core Solution

In the previous implementation, we were recording all the IDisposable instances that are created on the lifetime manager instance associated with a type registration. Now, because this does not work for the PerResolve case (where new PerResolveLifetimeManager instances are constructed during build), we need to come with another place where to record these instances. This new place also needs to be signaled when the owning container gets disposed, so we can dispose them.

The next thing to look at is to use a custom Unity builder strategy. It is one of the most powerful mechanisms that can be used to extend the build mechanism in Unity. To create one, we need to inherit from the BuilderStrategy base class and to override one of its methods to add the code we want to be executed when any new instance gets built or torn down. Here is the code of the base class:

// Represents a strategy in the chain of responsibility. 
// Strategies are required to support both BuildUp and TearDown. 
public abstract class BuilderStrategy : IBuilderStrategy 
{ 
	// Called during the chain of responsibility for a build operation. The 
	// PreBuildUp method is called when the chain is being executed in the 
	// forward direction. 
	public virtual void PreBuildUp(IBuilderContext context) 
	{ 
	}
	
	/// Called during the chain of responsibility for a build operation. The 
	/// PostBuildUp method is called when the chain has finished the PreBuildUp 
	/// phase and executes in reverse order from the PreBuildUp calls. 
	public virtual void PostBuildUp(IBuilderContext context) 
	{ 
	}
	
	/// Called during the chain of responsibility for a teardown operation. The 
	/// PreTearDown method is called when the chain is being executed in the 
	/// forward direction.
	public virtual void PreTearDown(IBuilderContext context) 
	{ 
	}
	
	/// Called during the chain of responsibility for a teardown operation. The 
	/// PostTearDown method is called when the chain has finished the PreTearDown 
	/// phase and executes in reverse order from the PreTearDown calls.  
	public virtual void PostTearDown(IBuilderContext context) 
	{ 
	} 
} 

For our case, we can override the PostBuildUp() and record all the IDisposable instances that get constructed. In comparison with the previous implementation, here we would keep on the strategy object all the IDisposable instances of different types (all types that implement IDisposable). There, on the lifetime manager we were keeping references only to the instances of the type registered with that lifetime manager object.

The next step is to trigger the dispose, down from the container to the recorded IDisposable instances. In its Dispose() the container disposes two things: the lifetime managers, and the extensions. See below:

protected virtual void Dispose(bool disposing) 
{ 
	if (disposing) 
	{ 
		if (lifetimeContainer != null) 
		{ 
			lifetimeContainer.Dispose(); 
			lifetimeContainer = null;
	
			if (parent != null && parent.lifetimeContainer != null) 
			{ 
				parent.lifetimeContainer.Remove(this); 
			} 
		}

		// this will trigger the Dispose() into our strategy object 
		extensions.OfType().ForEach(ex => ex.Dispose()); 
		extensions.Clear(); 
	} 
} 

The extensions dispose is our trigger point. We hold the IDisposable instances on the builder strategy object, not on the extension. To have them disposed, we can make the strategy object to implement IDisposable. There we can dispose all the instances. Then, going upwards, we can also make the custom extension that adds the strategy to the build strategies chain, to be IDisposable. It may keep a reference to the strategy and when it gets disposed by the container, it can dispose the strategy, which in its turn will dispose all the IDisposable instances that it records on the PostBuildUp(). So, when we put all the things together the code is like below

public class DisposeExtension : UnityContainerExtension, IDisposable 
{ 
	private DisposeStrategy strategy = new DisposeStrategy();

	protected override void Initialize() 
	{ 
		Context.Strategies.Add(strategy, UnityBuildStage.TypeMapping); 
	}
	
	public void Dispose() 
	{ 
		strategy.Dispose(); 
		strategy = null; 
	}
	
	class DisposeStrategy : BuilderStrategy, IDisposable 
	{ 
		private DisposablesObjectList disposables = new DisposablesObjectList();
	
		public override void PostBuildUp(IBuilderContext context) 
		{ 
			if (context != null) 
			{ 
				IDisposable instance = context.Existing as IDisposable; 
				if (instance != null) 
				{ 
					disposables.Add(instance); 
				} 
			}
		
			base.PostBuildUp(context); 
		} 
		… 
		public void Dispose() 
		{ 
			disposables.Dispose(); 
			disposables = null; 
		} 
	} 
} 

Having this done, if we add the extension to the child container after it is created, things work. All the IDisposable instances get disposed when the child container is disposed. The code snippet below shows the usage case:

// code in the application start-up area 
UnityContainer container = new UnityContainer(); 
DisposeExtension disposeExtension = new DisposeExtension(); 
container.AddExtension(disposeExtension);

container.RegisterType(new PerResolveLifetimeManager()); 
container.RegisterType(new PerResolveLifetimeManager());

// some instance created out of the scoped operation 
outerScopeSrv = container.Resolve();

// code in the area that marks and new operation. It creates a child container 
using (IUnityContainer childContainer = container.CreateChildContainer()) 
{ 
	var childDisposeExt = new DisposeExtension(); 
	childContainer.AddExtension(childDisposeExt);
	
	// instances built within the operation 
	scopedSrv1 = childContainer.Resolve(); 
	scopedSrv2 = childContainer.Resolve(); 
	…
} // end of the operation -> childContainer.Dispose()

AssertIsDisposed(scopedSrv1); // -> Pass 
AssertIsDisposed(scopedSrv2); // -> Pass 
AssertNotDisposed(outerScopeSrv); // -> Pass

// at application end 
container.Dispose();

AssertIsDisposed(outerScopeSrv); // -> Pass

However, there are two more aspects that need to be taken care of.

Issues

The first is about singletons. What happens when a singleton which is IDisposable gets injected through the child container? Now, it will be disposed by the first child container that used it when it gets disposed. From now on, all the other parts of the application will use the already disposed singleton instance. This is certainly a behavior we wouldn’t want. If a singleton is IDisposable (it’s another discussion why would we do that in the first place) it should not be disposed in the first operation that uses it. The code snippet below shows this case.

UnityContainer container = NewContainer();

// register as singleton 
container.RegisterType(new ContainerControlledLifetimeManager());

IService1 singletonSrv = container.Resolve(); 
IService1 scopedSrv;

using (IUnityContainer childContainer = CreateChildContainer(container)) 
{ 
	scopedSrv = childContainer.Resolve(); 
	Assert.AreSame(s, scopedService); 
}

AssertNotDisposed(singletonSrv); // -> Fail

To fix this, inside the strategy, we need to verify if the instance that is being built is singleton. If yes, it should not be recorded for dispose, in the strategy that belongs to the child container. To do this check, we look if the lifetime manager is ContainerControlledLifetimeManager and if it comes from the parent container. Below is the code that does it

… 
private bool IsControlledByParrent(IBuilderContext context) 
{ 
	IPolicyList lifetimePolicySource; 
	ILifetimePolicy activeLifetime = context.PersistentPolicies 
										.Get(context.BuildKey, out lifetimePolicySource);
	
	return activeLifetime is ContainerControlledLifetimeManager 
						&& !ReferenceEquals(lifetimePolicySource, context.PersistentPolicies); 
} 

This code already starts to make our extension a bit ugly. It hardcodes that the singletons are being registered with this particular lifetime manager type. If one will create a new custom lifetime manager to implement singletons, but with some additional behavior our extension may not work. Another thing that is a bit ugly, is that we need to search each time an instance is built, for the lifetime manager recursively upwards. However, I don’t see a better way to make this work for this case.

The second thing that needs to be fixed is more difficult to observe. The issue appears when we add this extension both to the parent and to the child container. This may happen either because we want the automatic dispose for the parent container too, or when we have deeper container hierarchies, meaning that we create a child container from another child container because of nested operations (probably not the case in a web app, but it could be the case in a desktop app).

By design, when a child container is created it inherits (by referencing) all the build strategies of the parent container. In general this is a good idea, because the strategies the parent container was configured with are also used by its children, and this brings consistency. However, for our case this is problematic. When the child container builds an IDisposable instance, that instance will be recorded twice: once in the child container DisposeStrategy object, and one more time in the parent DisposeStrategy object (we need to have different DisposeStrategy objects in the parent and in the child, because on these objects we record the IDisposable instances). Recording the instances built by the child container into the parent container strategy object may lead to memory leaks. The child containers are disposed once their operation ends, but the parent container will live typically as long as the application does. Its strategy object will also live with it, and it will grow with WeakRefernce objects for each instance built by all the child containers in all operations. This is bad!

It is also a sign that this approach to implement the automatic dispose, by keeping the reference on a builder strategy may not be in line with the design of Unity. The strategies were not meant to keep state from build to build. They were rather meant to extend the behavior of building instances, by adding logic that operates on the data from the IBuilderContext. However, because PerResolveLifetimeManager is implemented in a questionable way, now we are working around it with another questionable extension :(.

The fix for this case is not nice. In the PostBuildUp() method we need to determine if this strategy object belongs to the container that initiated the current build (the child) or it is an inherited strategy. If doesn’t belong to current container, we should not record currently built instance. The IBuilderContext which we get into the PostBuildUp(), does not contain any information on the container that builds this instance. So we cannot use a container match to distinguish if currently built instance should be registered in current builder strategy object or not. The only way we could distinguish between these cases is to relay on the order in which the strategies are put in the strategy chain which is constructed before a build starts. Always the strategy chain is created by placing the inherited strategies first, and at the end the specific strategies of the current container. Therefore, if the current strategy object is not the last in the strategy chain, it means that it is inherited and we should not record the current instance. We are relying on the ordering done in the StagedStrategyChain class which implements the IStagedStrategyChain interface. It is not very clear if the ordering we are relying on is by design, meaning that it should be preserved in future versions or it is just an implementation detail. Therefore, this should be an attention point when new versions of Unity will be released. The code for this fix is shown in the snippet below:

private bool IsInheritedStrategy(IBuilderContext context) 
{ 
	// unity container puts the parent container strategies before child strategies when it builds the chain 
	IBuilderStrategy lastStrategy = context.Strategies 
									.LastOrDefault(s => s is DisposeStrategy);

	return !ReferenceEquals(this, lastStrategy); 
}

In the end, if we put everything in together, we get the desired dispose behavior both for PerResolveLifetimeManager and for TrainsientLifetimeManager with this approach. The entire source code that implements this approach can be downloaded from here.

Conclusion

Both of the two approaches that I’ve presented in the previous blog post and in this one, have pluses and minuses. There is a trade-off to be made when picking one over the other. The first one, which uses the lifetime managers to record the IDisposable instances has a cleaner design. It is in-line with the way Unity was meant to be extended for lifetime and dispose management. Its disadvantage is that it does not work with the PerResolveLifetimeManager.

The second approach, which uses a builder strategy to record the IDisposable instances works with all the lifetime managers, including the PerResolveLifetimeManager. Its disadvantage is that it is an ugly extension. It holds state on a builder strategy and it may rely on some of the current implementation details. These makes it vulnerable to future versions of Unity or when combining it with other extensions.

If I were to choose, I would use the first one if I do not need the PerResolveLifetimeManager. If I need it, then I would fall back to the second one and carefully test if it works with the version of Unity that I am using and with the other extensions that I need. The good part is that switching from one implementation to the other is done by changing the way the container is configured. Usually this code is separated from the rest of the application, so by doing this switch there should be little if any change in the code that implements the use cases. Therefore, a switch with small costs.

Many examples like the above are included in my Code Design training