NHibernateSessionMiddleware.cs

using AstonishingLab.Core.NHibernate;
using Microsoft.Owin;
using NHibernate;
using NHibernate.Context;
using Owin;
using System;
using System.Threading.Tasks;

namespace AstonishingLab.Core.Owin
{
    internal class NHibernateSessionMiddleware : OwinMiddleware
    {
        protected readonly ISessionFactory SessionFactory;

        public NHibernateSessionMiddleware(OwinMiddleware next, ISessionFactory sessionFactory)
            : base(next)
        {
            SessionFactory = sessionFactory;
        }

        public async override Task Invoke(IOwinContext context)
        {
            var session = SessionFactory.OpenSession();
            session.BeginTransaction();
            CurrentSessionContext.Bind(session);

            var exceptions = false;
            try
            {
                await Next.Invoke(context);
            }
            catch (Exception)
            {
                exceptions = true;
                throw;
            }
            finally
            {
                CommitOrRollbackAndCloseSession(exceptions);
            }
        }

        protected virtual bool ShouldCommitOrRollbackTransaction(ITransaction transaction)
        {
            return transaction != null && transaction.IsActive && !transaction.WasRolledBack;
        }

        protected virtual void CommitOrRollbackAndCloseSession(bool exceptions)
        {
            var session = CurrentSessionContext.Unbind(SessionFactory);
            try
            {
                // If the transaction has been already commited/rolled back
                // We do nothing.
                var transaction = session.Transaction;
                if (ShouldCommitOrRollbackTransaction(transaction))
                {
                    // If there were exceptions, rollback.
                    if (exceptions)
                    {
                        transaction.Rollback();
                    }
                    else
                    {
                        // If there were no exceptions, commit.
                        transaction.Commit();
                    }
                    transaction.Dispose();
                }
            }
            catch (Exception e)
            {
                throw;
            }
            finally
            {
                session.Close();
            }
        }

    }

    public static class NHibernateSessionMiddlewareAppExtensions
    {
        public static IAppBuilder UseNHibernateSessionMiddleware(this IAppBuilder app)
        {
            return app.Use<NHibernateSessionMiddleware>();
        }
    }
}

OwinSessionContext.cs

using NHibernate.Engine;
using System;
using System.Collections;
using System.Runtime.Remoting.Messaging;

namespace NHibernate.Context
{
    [Serializable]
    public class OwinSessionContext : MapBasedSessionContext
    {
        private const string SessionFactoryMapKey = "NHibernate.Context.OwinSessionContext.SessionFactoryMapKey";

        public OwinSessionContext(ISessionFactoryImplementor factory) : base(factory) { }

        protected override IDictionary GetMap()
        {
            return CallContext.LogicalGetData(SessionFactoryMapKey) as IDictionary;
        }

        protected override void SetMap(IDictionary value)
        {
            CallContext.LogicalSetData(SessionFactoryMapKey, value);
        }
    }
}