Castle Windsor NServiceBus lifestyle manager – PerMessageLifestyleManager

I found myself in a similar position as with my previous post…  I want some instances to be bound to the lifetime of processing a message in NServiceBus…  Within about half a day of posting my WCF solution someone pointed me at one that already existed in the WCF Facility for Windsor.  So – I wonder if this was re-inventing the wheel (I promise I do look for what’s there before building it myself!)…

I’ve made a bit of an assumption that a Message is always handled by a single thread.  In WCF / ASP.Net land you can’t guarantee this…  But looking at some code that’s already out there (for example the saga persister) I noticed that it was using thread static to store state…  So, Udi – or any other contributors to NSB (to whom I’m very grateful already because the project is awesome) – let me know if this is a wrong assumption…

I decided to create a MessageContext (which is ThreadStatic) where I could put the instances resolved from the container.  I was thinking it might be useful to have such a thing to stuff various things into in the future…  Couldn’t see anything obvious that already existed..?

Next thing to do is create a MessageModule which will ensure that objects are evicted from the container once a message has been handled.

Lastly, create a lifestyle which checks the MessageContext for an existing instance.  If no instance is found then a new one is resolved from the container and it’s registered for eviction with the message module…  (This is a little different to how I did it for the WCF version).

Here’s the code…

public class MessageContext

{

    [ThreadStatic]

    private static MessageContext current;

    public static MessageContext Current

    {

        get

        {

            if (current == null)

            {

                current = new MessageContext();

            }

            return current;

        }

    }

    private readonly IDictionary<string, object> items = new Dictionary<string, object>();

    public object this[string key]

    {

        get

        {

            object item;

            items.TryGetValue(key, out item);

            return item;

        }

    }

    public void Add(string key, object item)

    {

        items.Add(key, item);

    }

    public bool Remove(string key)

    {

        return items.Remove(key);

    }

}

public class PerMessageLifestyleManager : AbstractLifestyleManager

    {

        private bool evicting;

        private readonly string perMessageKey = "PerMessageKey_" + Guid.NewGuid();

        public override object Resolve(Castle.MicroKernel.CreationContext context)

        {

            var instance = MessageContext.Current[perMessageKey];

            if (instance == null)

            {

                instance = base.Resolve(context);

                MessageContext.Current.Add(perMessageKey, instance);

                WindsorLifestyleMessageModule.RegisterForEviction(this, instance);

            }

            return instance;

        }

        public override bool Release(object instance)

        {

            if (!evicting) return false;

            bool released = base.Release(instance);

            MessageContext.Current.Remove(perMessageKey);

            return released;

        }

        public void Evict(object instance)

        {

            using (new EvictionScope(this))

            {

                Release(instance);

            }

        }

        public override void Dispose()

        {

        }

        private class EvictionScope : IDisposable

        {

            private readonly PerMessageLifestyleManager owner;

            public EvictionScope(PerMessageLifestyleManager owner)

            {

                this.owner = owner;

                this.owner.evicting = true;

            }

            public void Dispose()

            {

                owner.evicting = false;

            }

        }

    }

public class WindsorLifestyleMessageModule : IMessageModule

    {

        /// <summary>

        /// NServiceBus has 1 thread per message handler.

        /// </summary>

        [ThreadStatic] 

        private static IDictionary<PerMessageLifestyleManager, object> perThreadEvict;

        public static void RegisterForEviction(PerMessageLifestyleManager manager, object instance)

        {

            if (perThreadEvict == null)

            {

                perThreadEvict = new Dictionary<PerMessageLifestyleManager, object>();

            }

            perThreadEvict.Add(manager, instance);

        }

        public void HandleBeginMessage()

        {

        }

        public void HandleEndMessage()

        {

            EvictInstancesCreatedDuringMessageHandling();

        }

        public void HandleError()

        {

            EvictInstancesCreatedDuringMessageHandling();

        }

        private void EvictInstancesCreatedDuringMessageHandling()

        {

            if (perThreadEvict == null)

                return;

            foreach (var itemToEvict in perThreadEvict)

            {

                var manager = itemToEvict.Key;

                manager.Evict(itemToEvict.Value);

            }

            perThreadEvict.Clear();

            perThreadEvict = null;

        }

    }