Inject MessageContext in Handler constructor

src/Abc.Zebus.Tests/Dispatch/SyncMessageHandlerInvokerTests.cs

@@ -32,6 +32,20 @@ namespace Abc.Zebus.Tests.Dispatch
             invoker.InvokeMessageHandlerAsync(invocation).RunSynchronously();
 
             invocation.ApplyContextCalled.ShouldBeTrue();
+        }
+
+        [Test]
+        public void should_inject_context_in_handler_constructor()
+        {
+            var container = new Container(x => x.For<IBus>().Use(new Mock<IBus>().Object));
+            var invoker = new SyncMessageHandlerInvoker(container, typeof(CommandHandlerWithMessageContextInConstructor), typeof(CommandHandlerWithMessageContextInConstructorCommand));
+            var messageContext = MessageContext.CreateOverride(new PeerId("Abc.Testing.0"), null);
+            var command = new CommandHandlerWithMessageContextInConstructorCommand();
+            var invocation = command.ToInvocation(messageContext);
+
+            invoker.InvokeMessageHandler(invocation);
+
+            command.Context.ShouldEqual(messageContext);
         }
 
         [Test]
@@ -123,6 +137,26 @@ namespace Abc.Zebus.Tests.Dispatch
             invoker.CreateHandler(messageContext);
 
             Measure.Execution(500000, () => invoker.CreateHandler(messageContext));
+        }
+
+        private class CommandHandlerWithMessageContextInConstructorCommand : ICommand
+        {
+            public MessageContext Context { get; set; }
+        }
+
+        private class CommandHandlerWithMessageContextInConstructor : IMessageHandler<CommandHandlerWithMessageContextInConstructorCommand>
+        {
+            public MessageContext Context { get; private set; }
+
+            public CommandHandlerWithMessageContextInConstructor(MessageContext context)
+            {
+                Context = context;
+            }
+
+            public void Handle(CommandHandlerWithMessageContextInConstructorCommand message)
+            {
+                message.Context = Context;
+            }
         }
 
         private class MessageContextAwareCommandHandler : IMessageHandler<ScanCommand1>, IMessageContextAware

src/Abc.Zebus/Abc.Zebus.csproj

@@ -91,7 +91,7 @@
     <Compile Include="Directory\UpdatePeerSubscriptionsCommand.cs" />
     <Compile Include="Directory\UpdatePeerSubscriptionsForTypesCommand.cs" />
     <Compile Include="Dispatch\AsyncMessageHandlerInvoker.cs" />
-    <Compile Include="Dispatch\BusConstructorInstance.cs" />
+    <Compile Include="Dispatch\MessageHandlerConstructorInstance.cs" />
     <Compile Include="Dispatch\DispatcherTaskScheduler.cs" />
     <Compile Include="Dispatch\DispatcherTaskSchedulerFactory.cs" />
     <Compile Include="Dispatch\IDispatcherTaskSchedulerFactory.cs" />

src/Abc.Zebus/Dispatch/BusConstructorInstance.cs

@@ -1,36 +0,0 @@
-using System;
-using StructureMap;
-using StructureMap.Pipeline;
-
-namespace Abc.Zebus.Dispatch
-{
-    internal class BusConstructorInstance : ConstructorInstance, IConfiguredInstance
-    {
-        private readonly IBus _bus;
-
-        public BusConstructorInstance(Type pluggedType, IBus bus) : base(pluggedType)
-        {
-            _bus = bus;
-        }
-
-        object IConfiguredInstance.Get(string propertyName, Type pluginType, BuildSession session)
-        {
-            // changed from ConstructorInstance
-
-            if (pluginType == typeof(IBus))
-                return _bus;
-
-            return new DefaultInstance().Build(pluginType, session);
-        }
-
-        T IConfiguredInstance.Get<T>(string propertyName, BuildSession session)
-        {
-            // changed from ConstructorInstance
-
-            if (typeof(T) == typeof(IBus))
-                return (T)_bus;
-
-            return (T)new DefaultInstance().Build(typeof(T), session);
-        }
-    }
-}

src/Abc.Zebus/Dispatch/MessageHandlerConstructorInstance.cs

@@ -0,0 +1,38 @@
+using System;
+using StructureMap;
+using StructureMap.Pipeline;
+
+namespace Abc.Zebus.Dispatch
+{
+    internal class MessageHandlerConstructorInstance : ConstructorInstance, IConfiguredInstance
+    {
+        private readonly IBus _bus;
+        private readonly MessageContext _messageContext;
+
+        public MessageHandlerConstructorInstance(Type pluggedType, IBus bus, MessageContext messageContext) : base(pluggedType)
+        {
+            _bus = bus;
+            _messageContext = messageContext;
+        }
+
+        object IConfiguredInstance.Get(string propertyName, Type pluginType, BuildSession session)
+        {
+            return GetSpecialParameter(pluginType, session);
+        }
+
+        T IConfiguredInstance.Get<T>(string propertyName, BuildSession session)
+        {
+            return (T)GetSpecialParameter(typeof(T), session);
+        }
+
+        object GetSpecialParameter(Type pluginType, BuildSession session)
+        {
+            if (pluginType == typeof(IBus))
+                return _bus;
+            if (pluginType == typeof(MessageContext))
+                return _messageContext;
+
+            return new DefaultInstance().Build(pluginType, session);
+        }
+    }
+}

src/Abc.Zebus/Dispatch/MessageHandlerInvoker.cs

@@ -59,7 +59,7 @@ namespace Abc.Zebus.Dispatch
                 return container.GetInstance(MessageHandlerType);
 
             var busProxy = new MessageContextAwareBus(_bus, messageContext);
-            var messageHandlerInstance = new BusConstructorInstance(MessageHandlerType, busProxy);
+            var messageHandlerInstance = new MessageHandlerConstructorInstance(MessageHandlerType, busProxy, messageContext);
 
             return container.GetInstance(MessageHandlerType, messageHandlerInstance);
         }