Scale out error ingestion for EF (1/3): ingestion owns its dispatcher

src/ServiceControl.AcceptanceTests/Recoverability/MessageFailures/When_error_forwarding_is_enabled.cs

@@ -0,0 +1,79 @@
+namespace ServiceControl.AcceptanceTests.Recoverability.MessageFailures
+{
+    using System;
+    using System.Threading.Tasks;
+    using AcceptanceTesting;
+    using AcceptanceTesting.EndpointTemplates;
+    using NServiceBus;
+    using NServiceBus.AcceptanceTesting;
+    using NUnit.Framework;
+
+    class When_error_forwarding_is_enabled : AcceptanceTest
+    {
+        [Test]
+        public async Task Should_forward_the_failed_message_to_the_error_log_queue()
+        {
+            SetSettings = settings =>
+            {
+                settings.ForwardErrorMessages = true;
+                settings.ErrorLogQueue = NServiceBus.AcceptanceTesting.Customization.Conventions.EndpointNamingConvention(typeof(Spy));
+            };
+
+            var context = await Define<MyContext>()
+                .WithEndpoint<Failing>(b => b
+                    .When(session => session.SendLocal(new ForwardedMessage()))
+                    .DoNotFailOnErrorMessages())
+                // The ingestor probes the forwarding address on startup with an empty message the
+                // spy cannot deserialize, so the spy has to tolerate its own failures.
+                .WithEndpoint<Spy>(b => b.DoNotFailOnErrorMessages())
+                .Done(c => c.ForwardedMessageId != null)
+                .Run();
+
+            using (Assert.EnterMultipleScope())
+            {
+                Assert.That(context.ForwardedMessageId, Is.EqualTo(context.FailedMessageId));
+                Assert.That(context.ForwardedWithFailureHeaders, Is.True);
+            }
+        }
+
+        public class Failing : EndpointConfigurationBuilder
+        {
+            public Failing() => EndpointSetup<DefaultServerWithoutAudit>(c => c.NoRetries());
+
+            [Handler]
+            public class ForwardedMessageHandler(MyContext testContext) : IHandleMessages<ForwardedMessage>
+            {
+                public Task Handle(ForwardedMessage message, IMessageHandlerContext context)
+                {
+                    testContext.FailedMessageId = context.MessageId;
+                    throw new Exception("Simulated exception");
+                }
+            }
+        }
+
+        public class Spy : EndpointConfigurationBuilder
+        {
+            public Spy() => EndpointSetup<DefaultServerWithoutAudit>(c => c.NoRetries());
+
+            [Handler]
+            public class ForwardedMessageHandler(MyContext testContext) : IHandleMessages<ForwardedMessage>
+            {
+                public Task Handle(ForwardedMessage message, IMessageHandlerContext context)
+                {
+                    testContext.ForwardedWithFailureHeaders = context.MessageHeaders.ContainsKey("NServiceBus.ExceptionInfo.Message");
+                    testContext.ForwardedMessageId = context.MessageId;
+                    return Task.CompletedTask;
+                }
+            }
+        }
+
+        public class MyContext : ScenarioContext
+        {
+            public string FailedMessageId { get; set; }
+            public string ForwardedMessageId { get; set; }
+            public bool ForwardedWithFailureHeaders { get; set; }
+        }
+    }
+
+    public class ForwardedMessage : ICommand;
+}

src/ServiceControl.Audit/Auditing/AuditIngestion.cs

@@ -142,8 +142,9 @@ async Task SetUpAndStartInfrastructure(CancellationToken cancellationToken)
                 );
 
                 messageReceiver = transportInfrastructure.Receivers[inputEndpoint];
+                messageDispatcher = transportInfrastructure.Dispatcher;
 
-                await auditIngestor.VerifyCanReachForwardingAddress(cancellationToken);
+                await auditIngestor.VerifyCanReachForwardingAddress(messageDispatcher, cancellationToken);
                 await messageReceiver.StartReceive(cancellationToken);
 
                 logger.LogInformation(LogMessages.StartedInfrastructure);
@@ -172,10 +173,7 @@ async Task StopAndTeardownInfrastructure(CancellationToken cancellationToken)
                 logger.LogInformation("Stopping infrastructure");
                 try
                 {
-                    if (messageReceiver != null)
-                    {
-                        await messageReceiver.StopReceive(cancellationToken);
-                    }
+                    await StopReceiving(cancellationToken);
                 }
                 finally
                 {
@@ -184,6 +182,7 @@ async Task StopAndTeardownInfrastructure(CancellationToken cancellationToken)
 
                 messageReceiver = null;
                 transportInfrastructure = null;
+                receiveStopped = false;
 
                 logger.LogInformation(LogMessages.StoppedInfrastructure);
             }
@@ -214,6 +213,33 @@ async Task EnsureStopped(CancellationToken cancellationToken)
             }
         }
 
+        async Task EnsureReceivingStopped(CancellationToken cancellationToken)
+        {
+            await startStopSemaphore.WaitAsync(cancellationToken);
+
+            try
+            {
+                await StopReceiving(cancellationToken);
+            }
+            finally
+            {
+                startStopSemaphore.Release();
+            }
+        }
+
+        // Stops the receiver on its own, leaving the infrastructure up. Idempotent because a
+        // shutdown stops receiving before draining and then tears down, so this runs twice.
+        async Task StopReceiving(CancellationToken cancellationToken)
+        {
+            if (messageReceiver == null || receiveStopped)
+            {
+                return;
+            }
+
+            await messageReceiver.StopReceive(cancellationToken);
+            receiveStopped = true;
+        }
+
         async Task OnMessage(MessageContext messageContext, CancellationToken cancellationToken)
         {
             using var messageIngestionMetrics = metrics.BeginIngestion(messageContext);
@@ -258,7 +284,7 @@ protected override async Task ExecuteAsync(CancellationToken cancellationToken =
                             contexts.Add(context);
                         }
 
-                        await auditIngestor.Ingest(contexts, cancellationToken);
+                        await auditIngestor.Ingest(contexts, messageDispatcher, cancellationToken);
 
                         batchMetrics.Complete(contexts.Count);
                     }
@@ -300,28 +326,28 @@ public override async Task StopAsync(CancellationToken cancellationToken = defau
         {
             try
             {
-                await watchdog.Stop(cancellationToken);
+                // Order matters. Receiving stops under the shutdown token rather than a cancelled
+                // one, so messages already being processed finish and their receives commit instead
+                // of being abandoned and redelivered after having been forwarded. Nothing new enters
+                // the channel after this, and the infrastructure stays up until the channel drains.
+                await EnsureReceivingStopped(cancellationToken);
                 channel.Writer.Complete();
                 await base.StopAsync(cancellationToken);
             }
             finally
             {
-                if (transportInfrastructure != null)
-                {
-                    try
-                    {
-                        await transportInfrastructure.Shutdown(cancellationToken);
-                    }
-                    catch (OperationCanceledException e) when (cancellationToken.IsCancellationRequested)
-                    {
-                        logger.LogInformation(e, "Shutdown cancelled");
-                    }
-                }
+                // Tears the infrastructure down, now that nothing is left to dispatch.
+                await watchdog.Stop(cancellationToken);
             }
         }
 
         TransportInfrastructure transportInfrastructure;
         IMessageReceiver messageReceiver;
+        bool receiveStopped;
+
+        // Left in place when the infrastructure is torn down. A shutdown drains before tearing down,
+        // so this is still usable there.
+        IMessageDispatcher messageDispatcher;
 
         readonly int MaxBatchSize;
         readonly SemaphoreSlim startStopSemaphore = new(1);

src/ServiceControl.Audit/Auditing/AuditIngestor.cs

@@ -24,13 +24,11 @@ public AuditIngestor(
             EndpointInstanceMonitoring endpointInstanceMonitoring,
             IEnumerable<IEnrichImportedAuditMessages> auditEnrichers, // allows extending message enrichers with custom enrichers registered in the DI container
             IMessageSession messageSession,
-            Lazy<IMessageDispatcher> messageDispatcher,
             ITransportCustomization transportCustomization,
             ILogger<AuditIngestor> logger
         )
         {
             this.settings = settings;
-            this.messageDispatcher = messageDispatcher;
             this.logger = logger;
             var enrichers = new IEnrichImportedAuditMessages[] { new MessageTypeEnricher(), new EnrichWithTrackingIds(), new ProcessingStatisticsEnricher(), new DetectNewEndpointsFromAuditImportsEnricher(endpointInstanceMonitoring), new DetectSuccessfulRetriesEnricher(), new SagaRelationshipsEnricher() }.Concat(auditEnrichers).ToArray();
 
@@ -40,20 +38,19 @@ ILogger<AuditIngestor> logger
                 unitOfWorkFactory,
                 enrichers,
                 messageSession,
-                messageDispatcher,
                 logger
             );
         }
 
-        public async Task Ingest(List<MessageContext> contexts, CancellationToken cancellationToken = default)
+        public async Task Ingest(List<MessageContext> contexts, IMessageDispatcher dispatcher, CancellationToken cancellationToken = default)
         {
-            var stored = await auditPersister.Persist(contexts, cancellationToken);
+            var stored = await auditPersister.Persist(contexts, dispatcher, cancellationToken);
 
             try
             {
                 if (settings.ForwardAuditMessages)
                 {
-                    await Forward(stored, logQueueAddress, cancellationToken);
+                    await Forward(stored, logQueueAddress, dispatcher, cancellationToken);
                 }
 
                 foreach (var context in contexts)
@@ -74,7 +71,7 @@ public async Task Ingest(List<MessageContext> contexts, CancellationToken cancel
             }
         }
 
-        Task Forward(IReadOnlyCollection<MessageContext> messageContexts, string forwardingAddress, CancellationToken cancellationToken)
+        Task Forward(IReadOnlyCollection<MessageContext> messageContexts, string forwardingAddress, IMessageDispatcher dispatcher, CancellationToken cancellationToken)
         {
             var transportOperations = new List<TransportOperation>(messageContexts.Count);
             MessageContext anyContext = null;
@@ -99,13 +96,13 @@ Task Forward(IReadOnlyCollection<MessageContext> messageContexts, string forward
             }
 
             return anyContext != null
-                ? messageDispatcher.Value.Dispatch(
+                ? dispatcher.Dispatch(
                     new TransportOperations([.. transportOperations]),
                     anyContext.TransportTransaction, cancellationToken)
                 : Task.CompletedTask;
         }
 
-        public async Task VerifyCanReachForwardingAddress(CancellationToken cancellationToken = default)
+        public async Task VerifyCanReachForwardingAddress(IMessageDispatcher dispatcher, CancellationToken cancellationToken = default)
         {
             if (!settings.ForwardAuditMessages)
             {
@@ -122,7 +119,7 @@ public async Task VerifyCanReachForwardingAddress(CancellationToken cancellation
                     )
                 );
 
-                await messageDispatcher.Value.Dispatch(transportOperations, new TransportTransaction(), cancellationToken);
+                await dispatcher.Dispatch(transportOperations, new TransportTransaction(), cancellationToken);
             }
             catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
             {
@@ -136,7 +133,6 @@ public async Task VerifyCanReachForwardingAddress(CancellationToken cancellation
 
         readonly AuditPersister auditPersister;
         readonly Settings settings;
-        readonly Lazy<IMessageDispatcher> messageDispatcher;
         readonly string logQueueAddress;
 
         readonly ILogger logger;

src/ServiceControl.Audit/Auditing/AuditPersister.cs

@@ -19,10 +19,9 @@
     class AuditPersister(IAuditIngestionUnitOfWorkFactory unitOfWorkFactory,
         IEnrichImportedAuditMessages[] enrichers,
         IMessageSession messageSession,
-        Lazy<IMessageDispatcher> messageDispatcher,
         ILogger logger)
     {
-        public async Task<IReadOnlyList<MessageContext>> Persist(IReadOnlyList<MessageContext> contexts, CancellationToken cancellationToken = default)
+        public async Task<IReadOnlyList<MessageContext>> Persist(IReadOnlyList<MessageContext> contexts, IMessageDispatcher dispatcher, CancellationToken cancellationToken = default)
         {
             var storedContexts = new List<MessageContext>(contexts.Count);
             IAuditIngestionUnitOfWork unitOfWork = null;
@@ -33,7 +32,7 @@ public async Task<IReadOnlyList<MessageContext>> Persist(IReadOnlyList<MessageCo
                 var inserts = new List<Task>(contexts.Count);
                 foreach (var context in contexts)
                 {
-                    inserts.Add(ProcessMessage(context, cancellationToken));
+                    inserts.Add(ProcessMessage(context, dispatcher, cancellationToken));
                 }
 
                 await Task.WhenAll(inserts);
@@ -93,7 +92,7 @@ public async Task<IReadOnlyList<MessageContext>> Persist(IReadOnlyList<MessageCo
             return storedContexts;
         }
 
-        async Task ProcessMessage(MessageContext context, CancellationToken cancellationToken)
+        async Task ProcessMessage(MessageContext context, IMessageDispatcher dispatcher, CancellationToken cancellationToken)
         {
             if (context.Headers.TryGetValue(Headers.EnclosedMessageTypes, out var messageType)
                 && messageType == typeof(SagaUpdatedMessage).FullName)
@@ -102,7 +101,7 @@ async Task ProcessMessage(MessageContext context, CancellationToken cancellation
             }
             else
             {
-                await ProcessAuditMessage(context, cancellationToken);
+                await ProcessAuditMessage(context, dispatcher, cancellationToken);
             }
         }
 
@@ -127,7 +126,7 @@ void ProcessSagaAuditMessage(MessageContext context)
             }
         }
 
-        async Task ProcessAuditMessage(MessageContext context, CancellationToken cancellationToken)
+        async Task ProcessAuditMessage(MessageContext context, IMessageDispatcher dispatcher, CancellationToken cancellationToken)
         {
             if (!context.Headers.TryGetValue(Headers.MessageId, out var messageId))
             {
@@ -160,7 +159,7 @@ async Task ProcessAuditMessage(MessageContext context, CancellationToken cancell
                     await messageSession.Send(commandToEmit, cancellationToken);
                 }
 
-                await messageDispatcher.Value.Dispatch(new TransportOperations(messagesToEmit.ToArray()),
+                await dispatcher.Dispatch(new TransportOperations(messagesToEmit.ToArray()),
                     new TransportTransaction(), cancellationToken); //Do not hook into the incoming transaction
 
                 logger.LogDebug("{CommandsToEmitCount} commands and {MessagesToEmitCount} control messages emitted", commandsToEmit.Count, messagesToEmit.Count);

src/ServiceControl.Audit/Auditing/ImportFailedAudits.cs

@@ -14,18 +14,20 @@ public class ImportFailedAudits
         public ImportFailedAudits(
             IFailedAuditStorage failedAuditStore,
             AuditIngestor auditIngestor,
+            Lazy<IMessageDispatcher> messageDispatcher,
             Settings settings,
             ILogger<ImportFailedAudits> logger)
         {
             this.settings = settings;
             this.failedAuditStore = failedAuditStore;
             this.auditIngestor = auditIngestor;
+            this.messageDispatcher = messageDispatcher;
             this.logger = logger;
         }
 
         public async Task Run(CancellationToken cancellationToken = default)
         {
-            await auditIngestor.VerifyCanReachForwardingAddress(cancellationToken);
+            await auditIngestor.VerifyCanReachForwardingAddress(messageDispatcher.Value, cancellationToken);
 
             var succeeded = 0;
             var failed = 0;
@@ -47,7 +49,7 @@ await failedAuditStore.ProcessFailedMessages(
                             var taskCompletionSource = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
                             messageContext.SetTaskCompletionSource(taskCompletionSource);
 
-                            await auditIngestor.Ingest([messageContext], cancellationToken);
+                            await auditIngestor.Ingest([messageContext], messageDispatcher.Value, cancellationToken);
 
                             await taskCompletionSource.Task;
 
@@ -78,6 +80,7 @@ await failedAuditStore.ProcessFailedMessages(
 
         readonly IFailedAuditStorage failedAuditStore;
         readonly AuditIngestor auditIngestor;
+        readonly Lazy<IMessageDispatcher> messageDispatcher;
         readonly Settings settings;
         readonly ILogger<ImportFailedAudits> logger;