Unity lifetime management for IDisposable, part 2

Contents:

Two New Lifetime Managers

Having explained the issue, how can we solve it?

Before starting, I should mention that my implementation is based on previous work by Rory Primrose, who wrote an extension that changes the behavior of the TransientLifetimeManager in order to dispose of objects. He does an excellent job of explaining some of the rationale behind his implementation, so be sure to read his article.

I've placed the complete source code for my article (including some quick unit tests) on GitHub.

My approach is a little different: I did not want to change the default behavior of Unity, so instead I defined two new lifetime managers to explicitly set the desired behavior with: DisposingTransientLifetimeManager and DisposingSharedLifetimeManager.

DisposingTransientLifetimeManager

This lifetime manager behaves pretty much like the built-in TransientLifetimeManager; every time a new object is created. The only difference is, that on Teardown of any object, we will call Dispose() on objects that were created as a result of the original build-up operation (provided they are using this lifetime manager, of course).

Consider these two classes:

public class View
{
    public View(ViewModel viewModel) { /* ... */ }
}

public class ViewModel : IDisposable
{
    // ...
}

In the example below, we will want to call Dispose() on the implictly created ViewModel object:

container.RegisterType<ViewModel>(new DisposingTransientLifetimeManager());

// Create the View and accompanying ViewModel
View view = container.Resolve<View>()
container.Teardown(view);

DisposingSharedLifetimeManager

In many cases it's not desirable to create a new instance of an object every time. This is what this lifetime manager is intended for.

// We want to reuse the same ViewModel for every View
container.RegisterType<ViewModel>(new DisposingSharedLifetimeManager());

View view1 = container.Resolve<View>()
View view2 = container.Resolve<View>()

// We cannot Dispose the ViewModel yet, view2 is still using it
container.Teardown(view1);

// Now we can dispose the ViewModel
container.Teardown(view2);

DisposingLifetimeStrategy

Creating a lifetime manager with custom behavior is a litte more involved than inheriting from LifetimeManager or ILifetimePolicy. Much of the logic is is actually in the various builder strategies. They implement the following interface:

public interface IBuilderStrategy
{
    void PreBuildUp(IBuilderContext context);
    void PostBuildUp(IBuilderContext context);
    void PreTearDown(IBuilderContext context);
    void PostTearDown(IBuilderContext context);
}

Builder strategies can be added at each of the various stages of the Unity strategy pipeline: Setup, TypeMapping, Lifetime, PreCreation, Creation, Initialization and PostInitialization. As you'll see later, we're going to (ab)use the TypeMapping stage for our purpose.

PreBuildUp and PostBuildUp

In order to know when a dependency can be disposed, we need to keep track of which objects are using that specific instance of the dependency. This is done by creating a build tree. This is where the PreBuildUp and PostBuildUp methods come into play. Imagine we're letting Unity resolve all the dependencies in the following example:

public class View
{
    public View(IViewModel viewModel) { /* ... */ }
}

public class ViewModel : IViewModel
{
    public ViewModel(IServiceDependency) { /* ... */ }
}

public class ServiceDependency : IServiceDependency { /* ... */ }

Upon calling container.Resolve<View>() Unity will execute the strategy chain for each of these in forward order by calling PreBuildUp, then in reverse order for PostBuildUp:

  • PreBuildUp View
    • PreBuildUp ViewModel
      • PreBuildUp ServiceDependency
      • PostBuildUp ServiceDependency
    • PostBuildUp ViewModel
  • PostBuildUp View

At the end of the build up, the strategy will store a build tree like so:

  • View
    • ViewModel
      • ServiceDependency

This tree structure is entirely generated during the PreBuildUp phase. At this point, no actual objects have been created yet, meaning context.Existin will be null, unless there is a pre-existing object which will be used:

public override void PreBuildUp(IBuilderContext context)
{
    base.PreBuildUp(context);

    bool nodeCreatedByContainer = context.Existing == null;
    var newTreeNode = new BuildTreeItemNode(context.BuildKey, nodeCreatedByContainer,
                                            _currentBuildNode);

    if (_currentBuildNode != null)
    {
        // This is a child node. Add it to the parent node.
        _currentBuildNode.Children.Add(newTreeNode);
    }

    _currentBuildNode = newTreeNode;
}

Because we need to be able to call Dispose() on the created objects, we need to keep a reference to the newly created objects. These references are stored during the PostBuildUp phase. Also, we check if object actually implement IDisposable and if they're registered using either of our new lifetime managers. If the object nor any of its children in the tree do so, there is no need for us to track them:

private static bool ShouldTrackObject(IBuilderContext context)
{
    return
        context.Existing is IDisposable &&
        context.Lifetime.OfType<DisposingLifetimeManager>().
        Any(ltm => ltm.AppliesTo(context.Existing));
}

Every time an object is used is used in a build up, we increment a "reference count" to that object.

Teardown

Upon Teardown of an object, we simply get the original build tree for that object. Recursing through that tree, we decrease the reference counts for the tracked objects. Once the reference count has reached zero, we no longer need to store the build tree:

public override void PostTearDown(IBuilderContext context)
{
    base.PostTearDown(context);

    lock (_lock)
    {
        BuildTreeItemNode buildTree = GetBuildTreeForInstance(context.Existing);
        if (buildTree != null)
        {
            if (DecreaseTreeRefCount(context, buildTree))
            {
                _buildTrees.Remove(buildTree.ItemReference);
            }
        }
    }

    RemoveDeadTrees(context);
}

The last line in the method makes sure that the reference count to child objects is also decremented when an object is left to be garbage collected. However, I should warn you that failing to call Teardown on previously built up objects will cause undesired behavior.

Adding to the Strategy Chain

In order to use the two new lifetime managers, the accompanying builder strategy must be registered into the container. This has to be done through adding an extension to the container, in this case the DisposableStrategyExtension. It doesn't really do much otherwise:

public class DisposableStrategyExtension : UnityContainerExtension, IDisposable
{
    private readonly DisposingLifetimeStrategy _buildStrategy =
        new DisposingLifetimeStrategy();

    // ...
 	
    protected override void Initialize()
    {
        Context.Strategies.Add(_buildStrategy, UnityBuildStage.TypeMapping);
    }
}

How to Use

Unity has built in support for configuring the container using a configuration file, but here I'll quickly demonstrate how it is done from code:

using (var container = new UnityContainer())
{
    container.AddNewExtension<DisposableStrategyExtension>();
    // Use the container as you would normally
}

That's it! Now you can start using the two new lifetime managers. Some examples:

container.RegisterType(typeof(IViewModel), typeof(ViewModel),
    new DisposingSharedLifetimeManager());

container.RegisterType(typeof(IServiceDependency), typeof(ServiceDependency),
    new DisposingTransientLifetimeManager());

Whenever I create a View from the code example earlier on this page, I would get a reference to the same ViewModel. If other components also depend on IServiceDependency, they would get a separate instance of the ServiceDependency class.

Continue reading...

Files and Links:

Add comment

Loading