No more NHibernate LazyInitializationExceptions

For the last few months I have been working on a web application using S#arp Architecture with NHibernate.  Now S#arp Architecture give us a nice session-per-request model for the NHibernate session but we needed to keep state over web requests and preferred to do it without changing to the long running conversation model.  Therefore we stored our application service in the ASP .NET Cache to be picked up by the next web request for that session which could then carry on using the entities therein.  Great, until we tried to lazy load a property when we would get the NHibernate.LazyInitializationException because, of course, the session the entities were originally loaded into no longer existed.  A quick search on the web told us that our detached entities needed re-attaching to the session using either the Update or Lock methods depending on whether we knew they had changed or not - we had no way of knowing whether changes had been made so used the Update method.  This was not ideal as everything would then get written to the database when a transaction was committed - not what we needed when the majority of it was lookup data.


So, what to do?  Some investigation into the NHibernate source and I found that lazy initialization (of non-collections anyway) was performed by an implementation of the BasicLazyInitializer abstract class, the solution seemed quite simple create a subclass of our current LazyIntializer (in this case NHibernate.ByteCode.Castle.LazyInitializer) override the method that does the initialization to check if a session is available and if it's not get one before carrying on the call the base method.  It turned out this could be done in the Intercept method and the code for our lazy initializer looks like this:

namespaceProxy
{
    using System;
    using System.Collections.Generic;
    using System.Linq;
    using System.Reflection;
    using System.Text;
    using Castle.Core.Interceptor;
    using NHibernate.ByteCode.Castle;
    using NHibernate.Engine;
    using NHibernate.Type;

    [Serializable]
    public class LazyInitializer : LazyInitializer, IInterceptor
    {
        public LazyInitializer(string entityName, Type persistentClass, object id, MethodInfo getIdentifierMethod, MethodInfo setIdentifierMethod, IAbstractComponentType componentIdType, ISessionImplementor session)
            : base(entityName, persistentClass, id, getIdentifierMethod, setIdentifierMethod, componentIdType, session)
        {
        }

        public override void Intercept(Castle.Core.Interceptor.IInvocation invocation)
        {
            // If we don't have a session at this point then we will get a lazy initialization
            // error if we carry on.  So before we do we can set the session to our current default
            if (this.Session == null)
            {
                this.Session = Helper.SessionHelper.DefaultSession.GetSessionImplementation();
            }

            base.Intercept(invocation);
        }
    }
}

Then all that was required was to create a new ProxyFactory to return the new LazyInitializer as part of it's GetProxy method like so:
namespace Proxy
{
    using System;
    using Castle.DynamicProxy;
    using NHibernate;
    using NHibernate.ByteCode.Castle;
    using NHibernate.Proxy;

    public class NewProxyFactory : ProxyFactory
    {
        private static readonly ProxyGenerator proxyGenerator = new ProxyGenerator();

        public override INHibernateProxy GetProxy(object id, NHibernate.Engine.ISessionImplementor session)
        {
            INHibernateProxy proxy;
            try
            {
               LazyInitializer initializer = new LazyInitializer(this.EntityName, this.PersistentClass, id, this.GetIdentifierMethod, this.SetIdentifierMethod, this.ComponentIdType, session);
                object obj2 = this.IsClassProxy ? proxyGenerator.CreateClassProxy(this.PersistentClass, this.Interfaces, new Castle.Core.Interceptor.IInterceptor[] { initializer }) : proxyGenerator.CreateInterfaceProxyWithoutTarget(this.Interfaces[0], this.Interfaces, new Castle.Core.Interceptor.IInterceptor[] { initializer });
                initializer._constructed = true;
                proxy = (INHibernateProxy)obj2;
            }
            catch (Exception exception)
            {
                log.Error("Creating a proxy instance failed", exception);
                throw new HibernateException("Creating a proxy instance failed", exception);
            }

            return proxy;
        }
    }
}
And a new ProxyFactoryFactory which is the class referenced in the NHibernate config file:

namespace Proxy
{
    using System;
    using System.Collections.Generic;
    using System.Linq;
    using System.Text;
    using NHibernate.Bytecode;
    using NHibernate.Proxy;

    public class ProxyFactoryFactory : IProxyFactoryFactory
    {
        public NHibernate.Proxy.IProxyValidator ProxyValidator
        {
            get 
            {
                return new DynProxyTypeValidator();
            }
        }

        public NHibernate.Proxy.IProxyFactory BuildProxyFactory()
        {
            return new NewProxyFactory();
        }
    }
}
Unfortunately the above doesn't deal with collections so a separate set of classes was required to correctly lazy initialize a collection.  Working on the same technique as above I created a subclass of the PersistentGenericBag class with overrides for members that will require a session - in our case this turned out to be the Count property and the GetEnumerator method.  In addition to checking if a session was available it was also necessary to check if the collection had previously been initialized, using the WasInitialized property, so that we didn't accidentally add a session to a collection that shouldn't have one yet.  The class ended up like the following code:

namespace Proxy
{
    using System;
    using System.Collections.Generic;
    using NHibernate;
    using NHibernate.Collection.Generic;
    using NHibernate.Engine;
    using NHibernate.Persister.Collection;

    [Serializable]
    public class NewPersistentGenericBag<T> : PersistentGenericBag<T>, IList<T>
    {
        public NewPersistentGenericBag() 
            : base()
        { 
        }

        public NewPersistentGenericBag(ISessionImplementor session) 
            : base(session) 
        {
        }

        public NewPersistentGenericBag(ISessionImplementor session, ICollection<T> coll)
            : base(session, coll)
        {
        }

        int ICollection<T>.Count
        {
            get
            {
                this.CheckSession();

                return this.Count;
            }
        }

        IEnumerator<T> IEnumerable<T>.GetEnumerator()
        {
            this.CheckSession();

            return this.GetEnumerator() as IEnumerator<T>;
        }

        private void CheckSession()
        {
            if (this.Session == null && !this.WasInitialized)
            {
this.SetCurrentSession(Helper.SessionHelper.DefaultSession.GetSessionImplementation());
                ICollectionPersister persister = this.Session.Factory.GetCollectionPersister(this.Role);
                this.Session.PersistenceContext.AddUninitializedCollection(persister, this, this.Owner.Id);
                EntityKey ownerKey = new EntityKey(this.Owner.Id, persister.OwnerEntityPersister, this.Session.EntityMode);
                if (!this.Session.PersistenceContext.ContainsEntity(ownerKey))
                {
                    this.Session.PersistenceContext.AddEntity(this.Owner, Status.Loaded, null, ownerKey, persister.OwnerEntityPersister.GetVersion(this.Owner, this.Session.EntityMode), LockMode.None, true, persister.OwnerEntityPersister, false, true);
                }
            }
        }
    }
}
As well as the new collection type a CollectionType implementation was required to return the new persistent bag type:

namespace Proxy
{
    using System;
    using System.Collections.Generic;
    using NHibernate.Collection;
    using NHibernate.Engine;
    using NHibernate.Persister.Collection;
    using NHibernate.Type;

    public class PersistentGenericBagType<T> : BagType
    {
        public NewPersistentGenericBagType(string role, string propertyRef)
            : base(role, propertyRef, false)
        {
        }

        public override Type ReturnedClass
        {
            get 
            { 
                return typeof(IList<T>); 
            }
        }

        public override IPersistentCollection Wrap(ISessionImplementor session, object collection)
        {
            return new NewPersistentGenericBag<T>(session, collection as IList<T>);
        }

        public override object Instantiate(int anticipatedSize)
        {
            return anticipatedSize <= 0 ? new List<T>() : new List<T>(anticipatedSize + 1);
        }

        public override IPersistentCollection Instantiate(ISessionImplementor session, ICollectionPersister persister, object key)
        {
            return new NewPersistentGenericBag<T>(session);
        }
    }
}
And finally an implementation of ICollectionTypeFactory to return the new collection type when asked for a collection - this factory was then added to the NHibernate config file so that the new collection became the default collection type.  The simple code for this:

namespace Proxy
{
    using NHibernate.Type;

    public class CollectionTypeFactory : DefaultCollectionTypeFactory
    {
        public override CollectionType Bag<T>(string role, string propertyRef, bool embedded)
        {
            return new PersistentGenericBagType<T>(role, propertyRef);
        }
    }
}
All that was left was to add/change the entries in the config file for the proxyfactory.factory_class and collectiontype.factory_class properties to point to the new factory classes and we have a data layer in which we can use lazy loading without having to re-attach objects for each web request which increases performance and reduces code complexity.