WCF Unit of work Behavior

UnitOfWorkBehaviorAttribute.cs

public class UnitOfWorkBehaviorAttribute : Attribute, IContractBehavior, IContractBehaviorAttribute
    {
        public void Validate(ContractDescription contractDescription, ServiceEndpoint endpoint)
        { }

        public void ApplyDispatchBehavior(ContractDescription contractDescription, ServiceEndpoint endpoint, DispatchRuntime dispatchRuntime)
        {
            foreach(var operation in dispatchRuntime.EndpointDispatcher.DispatchRuntime.Operations)
            {
                operation.CallContextInitializers.Add(new PCMUnitOfWorkCallContextInitializer());
            }
        }

        public void ApplyClientBehavior(ContractDescription contractDescription, ServiceEndpoint endpoint, ClientRuntime clientRuntime)
        {}

        public void AddBindingParameters(ContractDescription contractDescription, ServiceEndpoint endpoint, BindingParameterCollection bindingParameters)
        {}

        public Type TargetContract
        {
            get { return null; }
        }
    }

    public class PCMUnitOfWorkCallContextInitializer : ICallContextInitializer
    {
        public object BeforeInvoke(InstanceContext instanceContext, IClientChannel channel, Message message)
        {
            var unitOfWork = ObjectFactory.Container.Resolve<IUnitOfWork>();

            if(unitOfWork.IsNotNull())
                unitOfWork.Start();

            return unitOfWork;
        }

        public void AfterInvoke(object correlationState)
        {
            var unitOfWork = correlationState as IUnitOfWork;

            if(unitOfWork.IsNotNull())
                unitOfWork.End();

        }
    }