Creating a pooled lifetime manager for Unity

My application has a child container per request, this is because there are a lot of services used in each request which depend upon an IUnitOfWork.  I thought it would be nice if I could define a pool of these IUnitOfWork instances so that any cost involved in creating them is reduced, they can just be reused in a round-robin fashion.  Well, more accurately, the object space (EcoSpace) on which they depend can anyway.

A pool can now be registered like so…

//Must be called once, when the container is created
container.AddNewExtension<PooledLifetimeExtension>();
//To register a pooled item
container.RegisterType
  <
    ISomeItemThatIsExpensiveToCreate, 
    SomeItemThatIsExpensiveToCreate
  >(new PooledLifetimeManager());

The number 10 specifies how many instances at once may be stored in the pool.  If an instance requests an item from the pool when it is empty it will get a new instance, it’s only at the point where an instance is returned to the pool that the MaxPoolSize is honoured – if the pool is not full then the instance goes back into the pool, if the pool is already full then Dispose is called on it instead.  Of course you should be careful when pooling that your pooled items don’t hold onto other injected references which are also pooled as that will cause problems.

I have also added the following interface so that an instance can (optionally) specify whether or not its current state should permit it to go back into the pool, and two methods which indicate to the instance when it has been retrieved from a pool or is about to be returned to the pool – in my case I deactivated the object space upon returning in order to flush any object cache, and reactivated it again when retrieved from the pool.

public interface IPooledResource
{
  bool CanReturnToPool { get; }
  void RetrievedFromPool();
  void ReturningToPool();
}

It’s important to note that this only works with child containers.  You should create your main controller as a kind of template by registering your request specific services (like IUnitOfWork) with HierarchicalLifetimeManager.  The PooledLifetimeManager will work in the same way as HierarchicalLifetimeManager except for the addition of pooling, and it will automatically attempt to return items to the pool once the child container is disposed of.

   //PooledLifetimeExtension.cs
   using Microsoft.Practices.Unity;
   using Microsoft.Practices.Unity.ObjectBuilder;
   
   namespace YourNameHere.Infrastructure.Unity
   {
     public class PooledLifetimeExtension : UnityContainerExtension
     {
       protected override void Initialize()
      {
        // Add to type mapping stage so that it runs before the lifetime stage
        Context.Strategies.AddNew<PooledLifetimeStrategy>
          (UnityBuildStage.TypeMapping);
      }
    }
  }
  
  
  
  //PooledLifetimeManager.cs
  using System;
  using System.Collections.Generic;
  using Microsoft.Practices.Unity;
  using Microsoft.Practices.ObjectBuilder;
  
  namespace YourNameHere.Infrastructure.Unity
  {
    public class PooledLifetimeManager : LifetimeManager
    {
      readonly PooledLifetimeManagerPool Pool;
  
      public PooledLifetimeManager(int maxPoolSize)
      {
        Pool = new PooledLifetimeManagerPool(maxPoolSize);
      }
  
      internal PooledLifetimeManagerPoolConsumer CreatePoolUserForChildContainer()
      {
        return new PooledLifetimeManagerPoolConsumer(Pool);
      }
  
      #region GetValue/SetValue/RemoveValue all redundant
      public override object GetValue()
      {
        throw new InvalidOperationException();
      }
  
      public override void SetValue(object newValue)
      {
        throw new InvalidOperationException();
      }
  
      public override void RemoveValue()
      {
        throw new InvalidOperationException();
      }
      #endregion
    }
  
    internal class PooledLifetimeManagerPool
    {
      readonly int MaxPoolSize;
      readonly object SyncRoot = new object();
      readonly Queue<object> Queue = new Queue<object>();
  
      public PooledLifetimeManagerPool(int maxPoolSize)
      {
        if (maxPoolSize < )
          throw new ArgumentOutOfRangeException(
            "MaxPoolSize", "Cannot be less than zero");
        MaxPoolSize = maxPoolSize;
      }
  
      internal object RetrieveFromPool()
      {
        object result = null;
        lock (SyncRoot)
          if (Queue.Count > )
            result = Queue.Dequeue();
        IPooledResource resultAsIPooledResource = result as IPooledResource;
        if (resultAsIPooledResource != null)
          resultAsIPooledResource.RetrievedFromPool();
        return result;
      }
  
      internal void ReturnToPool(object instance)
      {
        if (instance == null)
          throw new ArgumentNullException("Instance");
        IPooledResource instanceAsIPooledResource = instance as IPooledResource;
        if (instanceAsIPooledResource != null)
        {
          if (!instanceAsIPooledResource.CanReturnToPool)
            return;
          instanceAsIPooledResource.ReturningToPool();
        }
        bool wentIntoPool = false;
        lock (SyncRoot)
          if (Queue.Count < MaxPoolSize)
         {
           Queue.Enqueue(instance);
           wentIntoPool = true;
         }
 
       //Dispose any instances which did not go back into the pool
       if (!wentIntoPool && instance is IDisposable)
         ((IDisposable)instance).Dispose();
     }
   }
 
   public class PooledLifetimeManagerPoolConsumer :
     ContainerControlledLifetimeManager, IDisposable
   {
     readonly PooledLifetimeManagerPool SharedPool;
     object Instance;
 
     internal PooledLifetimeManagerPoolConsumer(
       PooledLifetimeManagerPool sharedPool)
     {
       if (sharedPool == null)
         throw new ArgumentNullException("SharedPool");
       SharedPool = sharedPool;
     }
 
     protected override object SynchronizedGetValue()
     {
       if (Instance == null)
         Instance = SharedPool.RetrieveFromPool();
       return Instance;
     }
 
     protected override void SynchronizedSetValue(object newValue)
     {
       if (Instance != null && newValue == null)
         SharedPool.ReturnToPool(Instance);
       Instance = newValue;
     }
 
     public override void RemoveValue()
     {
       throw new InvalidOperationException();
     }
 
     void IDisposable.Dispose()
     {
       SetValue(null);
       base.Dispose();
     }
   }
 }
 
 
 
 //PooledLifetimeStrategy.cs
 using Microsoft.Practices.ObjectBuilder;
 
 namespace YourNameHere.Infrastructure.Unity
 {
   public class PooledLifetimeStrategy : BuilderStrategy
   {
     public override void PreBuildUp(IBuilderContext context)
     {
       PooledLifetimeManager activeLifetime = 
         context.PersistentPolicies.Get<ILifetimePolicy>(context.BuildKey)
           as PooledLifetimeManager;
       if (activeLifetime != null)
       {
         // did this come from the local container or the parent?
         var localLifetime = 
           context.PersistentPolicies.Get<ILifetimePolicy>(
             context.BuildKey, true);
         if (localLifetime == null)
         {
           // came from parent, add a ContainerControlledLifetime here
           var newLifeTime = activeLifetime.CreatePoolUserForChildContainer();
           context.PersistentPolicies.Set<ILifetimePolicy>(
             newLifeTime, context.BuildKey);
           //Make sure it gets disposed
           //so that it can return its reference to the pool
           context.Lifetime.Add(newLifeTime);
         }
       }
     }
   }
 }
 
 
 
 //PooledResourceIntf.cs
 namespace YourNameHere.Infrastructure.Unity
 {
   public interface IPooledResource
   {
     bool CanReturnToPool { get; }
     void RetrievedFromPool();
     void ReturningToPool();
   }
 }