Processing a control message causes the outbox to throw a null reference exception (#549) (#551)

* Processing a control message causes the outbox to throw a null reference exception

* Update src/NServiceBus.Persistence.CosmosDB/Outbox/OutboxBehavior.cs

Co-authored-by: Andreas Öhlund <andreas.ohlund@particular.net>

---------

Co-authored-by: Andreas Öhlund <andreas.ohlund@particular.net>
# Conflicts:
#	src/NServiceBus.Persistence.CosmosDB/Outbox/OutboxBehavior.cs

Author: danielmarbach

src/NServiceBus.Persistence.CosmosDB.PhysicalOutbox.AcceptanceTests/When_using_outbox_control_message.cs

@@ -0,0 +1,91 @@
+namespace NServiceBus.AcceptanceTests
+{
+    using System;
+    using System.Threading.Tasks;
+    using AcceptanceTesting;
+    using EndpointTemplates;
+    using Extensibility;
+    using NServiceBus.AcceptanceTesting.Support;
+    using NServiceBus.Features;
+    using NServiceBus.Pipeline;
+    using NServiceBus.Routing;
+    using NServiceBus.Transport;
+    using NServiceBus.Unicast.Transport;
+    using NUnit.Framework;
+
+    [TestFixture]
+    public class When_using_outbox_control_message : NServiceBusAcceptanceTest
+    {
+        [Test]
+        public async Task Should_work()
+        {
+            var runSettings = new RunSettings();
+            runSettings.DoNotRegisterDefaultPartitionKeyProvider();
+
+            var context = await Scenario.Define<Context>()
+                .WithEndpoint<Endpoint>()
+                .Done(c => c.ProcessedControlMessage)
+                .Run(runSettings)
+                .ConfigureAwait(false);
+
+            Assert.True(context.ProcessedControlMessage);
+        }
+
+        public class Context : ScenarioContext
+        {
+            public bool ProcessedControlMessage { get; set; }
+        }
+
+        public class Endpoint : EndpointConfigurationBuilder
+        {
+            public Endpoint() =>
+                EndpointSetup<DefaultServer>((config, runDescriptor) =>
+                {
+                    config.EnableOutbox();
+                    config.ConfigureTransport().Transactions(TransportTransactionMode.ReceiveOnly);
+                    config.EnableFeature<ControlMessageSender>();
+                    config.Pipeline.Register(new ControlMessageBehavior(runDescriptor.ScenarioContext as Context), "Checks that the control message was processed successfully");
+                });
+
+            class ControlMessageSender : Feature
+            {
+                protected override void Setup(FeatureConfigurationContext context)
+                {
+                    context.Container.ConfigureComponent<StartupTask>(DependencyLifecycle.InstancePerCall);
+                    context.RegisterStartupTask(b => b.Build<StartupTask>());
+                }
+
+                class StartupTask : FeatureStartupTask
+                {
+                    public StartupTask(IDispatchMessages dispatcher) => this.dispatcher = dispatcher;
+
+                    protected override Task OnStart(IMessageSession session)
+                    {
+                        var controlMessage = ControlMessageFactory.Create(MessageIntentEnum.Subscribe);
+                        var messageOperation = new TransportOperation(controlMessage, new UnicastAddressTag(AcceptanceTesting.Customization.Conventions.EndpointNamingConvention(typeof(Endpoint))));
+
+                        return dispatcher.Dispatch(new TransportOperations(messageOperation), new TransportTransaction(), new ContextBag());
+                    }
+
+                    protected override Task OnStop(IMessageSession sessio) => Task.CompletedTask;
+
+                    readonly IDispatchMessages dispatcher;
+                }
+            }
+
+            class ControlMessageBehavior : Behavior<IIncomingPhysicalMessageContext>
+            {
+                public ControlMessageBehavior(Context testContext) => this.testContext = testContext;
+
+                public override async Task Invoke(IIncomingPhysicalMessageContext context, Func<Task> next)
+                {
+                    await next();
+
+                    testContext.ProcessedControlMessage = true;
+                }
+
+                readonly Context testContext;
+            }
+        }
+    }
+}

src/NServiceBus.Persistence.CosmosDB/Outbox/OutboxBehavior.cs

@@ -65,7 +65,7 @@ public async Task Invoke(IIncomingLogicalMessageContext context, Func<IIncomingL
             // Outbox operating at the logical stage
             if (!context.Extensions.TryGet<PartitionKey>(out var partitionKey))
             {
-                throw new Exception("For the outbox to work the following information must be provided at latest up to the incoming physical or logical message stage. A partition key via `context.Extensions.Set<PartitionKey>(yourPartitionKey)`.");
+                throw new Exception("For the outbox to work a partition key must be provided at latest up to the incoming physical or logical message stage. Set one via '{nameof(CosmosPersistenceConfig.TransactionInformation)}'.");
             }
 
             var containerHolder = containerHolderResolver.ResolveAndSetIfAvailable(context.Extensions);
@@ -77,6 +77,8 @@ public async Task Invoke(IIncomingLogicalMessageContext context, Func<IIncomingL
             outboxTransaction.PartitionKey = partitionKey;
             outboxTransaction.StorageSession.ContainerHolder = containerHolder;
 
+            setAsDispatchedHolder.ThrowIfContainerIsNotSet();
+
             var outboxRecord = await containerHolder.Container.ReadOutboxRecord(context.MessageId, outboxTransaction.PartitionKey.Value, serializer, context.Extensions)
                 .ConfigureAwait(false);
 

src/NServiceBus.Persistence.CosmosDB/Outbox/OutboxPersister.cs

@@ -5,6 +5,8 @@
     using Microsoft.Azure.Cosmos;
     using Newtonsoft.Json;
     using Outbox;
+    using Transport;
+    using Headers = NServiceBus.Headers;
 
     class OutboxPersister : IOutboxStorage
     {
@@ -37,10 +39,19 @@ public async Task<OutboxMessage> Get(string messageId, ContextBag context)
 
             if (!context.TryGet<PartitionKey>(out var partitionKey))
             {
-                // we return null here to enable outbox work at logical stage
-                return null;
+                // because of the transactional session we cannot assume the incoming message is always present
+                if (!context.TryGet<IncomingMessage>(out var incomingMessage) ||
+                    !incomingMessage.Headers.ContainsKey(Headers.ControlMessageHeader))
+                {
+                    // we return null here to enable outbox work at logical stage
+                    return null;
+                }
+
+                partitionKey = new PartitionKey(messageId);
+                context.Set(partitionKey);
             }
 
+            setAsDispatchedHolder.ThrowIfContainerIsNotSet();
             setAsDispatchedHolder.PartitionKey = partitionKey;
 
             var outboxRecord = await setAsDispatchedHolder.ContainerHolder.Container.ReadOutboxRecord(messageId, partitionKey, serializer, context)
@@ -72,6 +83,7 @@ public Task Store(OutboxMessage message, OutboxTransaction transaction, ContextB
         public async Task SetAsDispatched(string messageId, ContextBag context)
         {
             var setAsDispatchedHolder = context.Get<SetAsDispatchedHolder>();
+            setAsDispatchedHolder.ThrowIfContainerIsNotSet();
 
             var partitionKey = setAsDispatchedHolder.PartitionKey;
             var containerHolder = setAsDispatchedHolder.ContainerHolder;

src/NServiceBus.Persistence.CosmosDB/Outbox/SetAsDispatchedHolderExtensions.cs

@@ -0,0 +1,19 @@
+#nullable enable
+
+namespace NServiceBus.Persistence.CosmosDB
+{
+    using System;
+
+    static class SetAsDispatchedHolderExtensions
+    {
+        public static void ThrowIfContainerIsNotSet(this SetAsDispatchedHolder setAsDispatchedHolder)
+        {
+            if (setAsDispatchedHolder.ContainerHolder != null && setAsDispatchedHolder.ContainerHolder.Container != null)
+            {
+                return;
+            }
+
+            throw new Exception($"For the outbox to work a container must be configured. Either configure a default one using '{nameof(CosmosPersistenceConfig.DefaultContainer)}' or set one via '{nameof(CosmosPersistenceConfig.TransactionInformation)}'.");
+        }
+    }
+}

src/SharedAcceptanceTests.RequirePartitionKey/When_no_container_information_is_configured.cs

@@ -0,0 +1,59 @@
+namespace NServiceBus.AcceptanceTests
+{
+    using System.Linq;
+    using System.Threading.Tasks;
+    using AcceptanceTesting;
+    using EndpointTemplates;
+    using NServiceBus.AcceptanceTesting.Support;
+    using NUnit.Framework;
+
+    [TestFixture]
+    public class When_no_container_information_is_configured : NServiceBusAcceptanceTest
+    {
+        [Test]
+        public async Task Should_throw_meaningful_exception()
+        {
+            var runSettings = new RunSettings();
+            runSettings.DoNotRegisterDefaultContainerInformationProvider();
+
+            var context = await Scenario.Define<Context>()
+                .WithEndpoint<Endpoint>(b =>
+                {
+                    b.DoNotFailOnErrorMessages();
+                    b.When(s => s.SendLocal(new MyMessage()));
+                })
+                .Done(c => c.FailedMessages.Any())
+                .Run(runSettings);
+
+            var failure = context.FailedMessages.FirstOrDefault()
+                .Value.First();
+
+            Assert.That(failure.Exception.Message, Does.Contain("container"));
+        }
+
+        class Context : ScenarioContext
+        {
+        }
+
+        class Endpoint : EndpointConfigurationBuilder
+        {
+            public Endpoint() =>
+                EndpointSetup<DefaultServer>((config, runDescriptor) =>
+                {
+                    config.EnableOutbox();
+                    config.ConfigureTransport().Transactions(TransportTransactionMode.ReceiveOnly);
+                });
+
+            class MyMessageHandler : IHandleMessages<MyMessage>
+            {
+                public Task Handle(MyMessage message, IMessageHandlerContext context)
+                {
+                    Assert.Fail("Should not be called");
+                    return Task.CompletedTask;
+                }
+            }
+        }
+
+        class MyMessage : IMessage { }
+    }
+}

src/SharedAcceptanceTests.RequirePartitionKey/When_no_partition_key_is_configured.cs

@@ -0,0 +1,59 @@
+namespace NServiceBus.AcceptanceTests
+{
+    using System.Linq;
+    using System.Threading.Tasks;
+    using AcceptanceTesting;
+    using EndpointTemplates;
+    using NServiceBus.AcceptanceTesting.Support;
+    using NUnit.Framework;
+
+    [TestFixture]
+    public class When_no_partition_key_is_configured : NServiceBusAcceptanceTest
+    {
+        [Test]
+        public async Task Should_throw_meaningful_exception()
+        {
+            var runSettings = new RunSettings();
+            runSettings.DoNotRegisterDefaultPartitionKeyProvider();
+
+            var context = await Scenario.Define<Context>()
+                .WithEndpoint<Endpoint>(b =>
+                {
+                    b.DoNotFailOnErrorMessages();
+                    b.When(s => s.SendLocal(new MyMessage()));
+                })
+                .Done(c => c.FailedMessages.Any())
+                .Run(runSettings);
+
+            var failure = context.FailedMessages.FirstOrDefault()
+                .Value.First();
+
+            Assert.That(failure.Exception.Message, Does.Contain("partition key"));
+        }
+
+        class Context : ScenarioContext
+        {
+        }
+
+        class Endpoint : EndpointConfigurationBuilder
+        {
+            public Endpoint() =>
+                EndpointSetup<DefaultServer>((config, runDescriptor) =>
+                {
+                    config.EnableOutbox();
+                    config.ConfigureTransport().Transactions(TransportTransactionMode.ReceiveOnly);
+                });
+
+            class MyMessageHandler : IHandleMessages<MyMessage>
+            {
+                public Task Handle(MyMessage message, IMessageHandlerContext context)
+                {
+                    Assert.Fail("Should not be called");
+                    return Task.CompletedTask;
+                }
+            }
+        }
+
+        class MyMessage : IMessage { }
+    }
+}