Latch receivers immediately on ApplicationStopping to prevent processing during shutdown

Hook into IHostApplicationLifetime.ApplicationStopping so all message receivers
are latched the moment SIGTERM fires, rather than waiting for IHostedService.StopAsync
which may be delayed by other hosted services stopping first. This prevents messages
already in internal queues from being picked up after the shutdown signal.

Also reorders StopAsync to drain endpoints before releasing ownership, adds bounded
WaitForCompletionAsync calls in DurableReceiver and BufferedReceiver with configurable
DrainTimeout (default 30s), and includes GracefulShutdown/RollingRestart chaos tests.

Closes #2282

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: jeremydmiller

src/Transports/RabbitMQ/ChaosTesting/ChaosSpecifications.cs

@@ -143,4 +143,30 @@ public Task RabbitMqDurableListener_Marten_ReceiverGoesUpAndDown() =>
     [Fact]
     public Task RabbitMqFiveDurableListeners_Marten_ReceiverGoesUpAndDown() =>
         execute<MartenStorageStrategy, ReceiverGoesUpAndDown>(RabbitMqFiveDurableListeners);
+
+    // Graceful shutdown tests
+    [Fact]
+    public Task RabbitMqDurableListener_Marten_GracefulShutdown() =>
+        execute<MartenStorageStrategy, GracefulShutdown>(RabbitMqDurableListener);
+
+    [Fact]
+    public Task RabbitMqFiveDurableListeners_Marten_GracefulShutdown() =>
+        execute<MartenStorageStrategy, GracefulShutdown>(RabbitMqFiveDurableListeners);
+
+    [Fact]
+    public Task RabbitMqBufferedListener_Marten_GracefulShutdown() =>
+        execute<MartenStorageStrategy, GracefulShutdown>(RabbitMqBufferedListener);
+
+    [Fact]
+    public Task RabbitMqOneInlineListener_Marten_GracefulShutdown() =>
+        execute<MartenStorageStrategy, GracefulShutdown>(RabbitMqOneInlineListener);
+
+    // Rolling restart tests
+    [Fact]
+    public Task RabbitMqDurableListener_Marten_RollingRestart() =>
+        execute<MartenStorageStrategy, RollingRestart>(RabbitMqDurableListener);
+
+    [Fact]
+    public Task RabbitMqFiveDurableListeners_Marten_RollingRestart() =>
+        execute<MartenStorageStrategy, RollingRestart>(RabbitMqFiveDurableListeners);
 }

src/Transports/RabbitMQ/ChaosTesting/Scripts/GracefulShutdown.cs

@@ -0,0 +1,66 @@
+using JasperFx.Core;
+
+namespace ChaosTesting.Scripts;
+
+/// <summary>
+/// Tests that in-flight messages complete during graceful shutdown and no messages
+/// are left orphaned. Sends messages, then stops the receiver while messages are
+/// still being processed, starts a new receiver, and verifies all messages complete.
+/// </summary>
+public class GracefulShutdown : ChaosScript
+{
+    public GracefulShutdown()
+    {
+        TimeOut = 2.Minutes();
+    }
+
+    public override async Task Drive(ChaosDriver driver)
+    {
+        await driver.StartReceiver("one");
+        await driver.StartSender("one");
+
+        // Send a batch of messages
+        await driver.SendMessages("one", 200);
+
+        // Give some time for processing to start but not complete
+        await Task.Delay(500.Milliseconds());
+
+        // Gracefully stop the receiver while messages are in-flight
+        await driver.StopReceiver("one");
+
+        // Start a new receiver to pick up any remaining messages
+        await driver.StartReceiver("two");
+    }
+}
+
+/// <summary>
+/// Tests that multiple rapid shutdown/restart cycles don't lose messages.
+/// Simulates rolling deployment behavior where receivers are stopped and new
+/// ones started in quick succession.
+/// </summary>
+public class RollingRestart : ChaosScript
+{
+    public RollingRestart()
+    {
+        TimeOut = 3.Minutes();
+    }
+
+    public override async Task Drive(ChaosDriver driver)
+    {
+        await driver.StartReceiver("one");
+        await driver.StartSender("one");
+
+        // Send initial batch
+        await driver.SendMessages("one", 300);
+
+        // Simulate a rolling restart: stop old, start new, repeat
+        await Task.Delay(1.Seconds());
+        await driver.StopReceiver("one");
+
+        await driver.StartReceiver("two");
+        await Task.Delay(1.Seconds());
+        await driver.StopReceiver("two");
+
+        await driver.StartReceiver("three");
+    }
+}

src/Wolverine/Configuration/EndpointCollection.cs

@@ -382,6 +382,24 @@ private ISendingAgent buildSendingAgent(Uri uri, Action<Endpoint>? configureNewE
         return endpoint.StartSending(_runtime, transport.ReplyEndpoint()?.Uri);
     }
 
+    /// <summary>
+    /// Immediately latch all receivers to stop picking up new messages from their internal queues.
+    /// This is called as early as possible during shutdown (via IHostApplicationLifetime.ApplicationStopping)
+    /// so that messages already queued internally are not processed after the shutdown signal.
+    /// </summary>
+    public void LatchAllReceivers()
+    {
+        foreach (var listener in _listeners.Values)
+        {
+            listener.LatchReceiver();
+        }
+
+        foreach (var queue in _localSenders.Enumerate().Select(x => x.Value).OfType<DurableLocalQueue>())
+        {
+            queue.LatchReceiver();
+        }
+    }
+
     public async Task DrainAsync()
     {
         // Drain the listeners

src/Wolverine/DurabilitySettings.cs

@@ -247,6 +247,12 @@ internal set
     /// </summary>
     public TimeSpan? NodeAssignmentHealthCheckTraceSamplingPeriod { get; set; }
 
+    /// <summary>
+    /// Maximum time to wait for in-flight message handlers to complete during graceful
+    /// shutdown before proceeding with the shutdown sequence. Default is 30 seconds.
+    /// </summary>
+    public TimeSpan DrainTimeout { get; set; } = 30.Seconds();
+
     /// <summary>
     ///     Get or set the logical Wolverine service name. By default, this is
     ///     derived from the name of a custom WolverineOptions

src/Wolverine/Runtime/WolverineRuntime.HostService.cs

@@ -2,6 +2,7 @@
 using JasperFx.CodeGeneration;
 using JasperFx.Core;
 using JasperFx.Core.Reflection;
+using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Logging;
 using Wolverine.Configuration;
 using Wolverine.Persistence.Durability;
@@ -102,6 +103,19 @@ public async Task StartAsync(CancellationToken cancellationToken)
 
             await Observer.RuntimeIsFullyStarted();
             _hasStarted = true;
+
+            // Subscribe to the host shutdown signal so we can immediately latch all receivers
+            // the moment SIGTERM/ApplicationStopping fires, rather than waiting until our
+            // IHostedService.StopAsync is called (which may be delayed by other hosted services)
+            try
+            {
+                var lifetime = _container.Services.GetService(typeof(IHostApplicationLifetime)) as IHostApplicationLifetime;
+                lifetime?.ApplicationStopping.Register(OnApplicationStopping);
+            }
+            catch (Exception e)
+            {
+                Logger.LogDebug(e, "Could not subscribe to IHostApplicationLifetime.ApplicationStopping");
+            }
         }
         catch (Exception? e)
         {
@@ -110,6 +124,12 @@ public async Task StartAsync(CancellationToken cancellationToken)
         }
     }
 
+    internal void OnApplicationStopping()
+    {
+        Logger.LogInformation("Application stopping signal received, latching all message receivers");
+        _endpoints.LatchAllReceivers();
+    }
+
     private bool _hasMigratedStorage;
 
     private async Task tryMigrateStorage()
@@ -196,7 +216,21 @@ public async Task StopAsync(CancellationToken cancellationToken)
         DisableHealthChecks();
 
         _idleAgentCleanupLoop?.SafeDispose();
-        
+
+        if (StopMode == StopMode.Normal)
+        {
+            // Step 1: Drain endpoints first — stop listeners from accepting new messages
+            // and wait for in-flight handlers to complete before releasing ownership.
+            // Receivers were already latched via IHostApplicationLifetime.ApplicationStopping
+            // to prevent new messages from being picked up, so this just waits for completion.
+            await _endpoints.DrainAsync();
+
+            if (_accumulator.IsValueCreated)
+            {
+                await _accumulator.Value.DrainAsync();
+            }
+        }
+
         if (_stores.IsValueCreated && StopMode == StopMode.Normal)
         {
             try
@@ -210,7 +244,8 @@ public async Task StopAsync(CancellationToken cancellationToken)
 
             try
             {
-                // New to 3.0, try to release any ownership on the way out. Do this *after* the drain
+                // Release any ownership on the way out. Do this *after* draining endpoints
+                // so in-flight messages complete before their ownership is released.
                 await _stores.Value.ReleaseAllOwnershipAsync(DurabilitySettings.AssignedNodeNumber);
             }
             catch (ObjectDisposedException)
@@ -221,15 +256,8 @@ public async Task StopAsync(CancellationToken cancellationToken)
 
         if (StopMode == StopMode.Normal)
         {
-            // This MUST be called before draining the endpoints
+            // Step 2: Now teardown agents — safe after endpoints drained and ownership released
             await teardownAgentsAsync();
-            
-            await _endpoints.DrainAsync();
-
-            if (_accumulator.IsValueCreated)
-            {
-                await _accumulator.Value.DrainAsync();
-            }
         }
 
         DurabilitySettings.Cancel();

src/Wolverine/Runtime/WorkerQueues/BufferedReceiver.cs

@@ -114,21 +114,38 @@ async ValueTask IChannelCallback.DeferAsync(Envelope envelope)
 
     public int QueueCount => (int)_receivingBlock.Count;
 
+    /// <summary>
+    /// Immediately latch to stop processing new messages without draining.
+    /// </summary>
+    public void Latch()
+    {
+        _latched = true;
+    }
+
     public async ValueTask DrainAsync()
     {
         _latched = true;
         _receivingBlock.Complete();
 
+        // Wait for in-flight handler executions to complete, bounded by a timeout
+        // to prevent hanging during shutdown
+        try
+        {
+            var completion = _receivingBlock.WaitForCompletionAsync();
+            await Task.WhenAny(completion, Task.Delay(_settings.DrainTimeout));
+        }
+        catch (Exception e)
+        {
+            _logger.LogDebug(e, "Error waiting for in-flight message processing to complete at {Uri}", Uri);
+        }
+
         await _completeBlock.DrainAsync();
         await _deferBlock.DrainAsync();
 
         if (_moveToErrors != null)
         {
             await _moveToErrors.DrainAsync();
         }
-
-        // It hangs, nothing to be done about this I think
-        //await _receivingBlock.Completion;
     }
 
     public void Enqueue(Envelope envelope)

src/Wolverine/Runtime/WorkerQueues/DurableReceiver.cs

@@ -286,8 +286,17 @@ public async ValueTask DrainAsync()
         _latched = true;
         _receiver.Complete();
 
-        // Latching is the best you can do here, otherwise it can hang
-        //await _receiver.Completion;
+        // Wait for in-flight handler executions to complete, bounded by a timeout
+        // to prevent hanging during shutdown
+        try
+        {
+            var completion = _receiver.WaitForCompletionAsync();
+            await Task.WhenAny(completion, Task.Delay(_settings.DrainTimeout));
+        }
+        catch (Exception e)
+        {
+            _logger.LogDebug(e, "Error waiting for in-flight message processing to complete at {Uri}", Uri);
+        }
 
         await _incrementAttempts.DrainAsync();
         await _scheduleExecution.DrainAsync();

src/Wolverine/Transports/ListeningAgent.cs

@@ -139,6 +139,24 @@ public async Task EnqueueDirectlyAsync(IEnumerable<Envelope> envelopes)
 
     public ListeningStatus Status { get; private set; } = ListeningStatus.Stopped;
 
+    /// <summary>
+    /// Immediately latch the receiver to stop processing new messages from its internal queue.
+    /// Does not stop the listener or drain — just prevents the receiver from executing any more messages.
+    /// </summary>
+    public void LatchReceiver()
+    {
+        if (_receiver is DurableReceiver dr)
+        {
+            dr.Latch();
+        }
+        // BufferedReceiver latches via its _latched field in DrainAsync,
+        // but we need an immediate latch here too
+        else if (_receiver is BufferedReceiver br)
+        {
+            br.Latch();
+        }
+    }
+
     public async ValueTask StopAndDrainAsync()
     {
         if (Status == ListeningStatus.Stopped || Status == ListeningStatus.GloballyLatched)

src/Wolverine/Transports/Local/DurableLocalQueue.cs

@@ -63,6 +63,15 @@ public DurableLocalQueue(Endpoint endpoint, WolverineRuntime runtime)
 
     public CircuitBreaker? CircuitBreaker { get; }
 
+    /// <summary>
+    /// Immediately latch the receiver to stop processing new messages.
+    /// </summary>
+    public void LatchReceiver()
+    {
+        Latched = true;
+        _receiver?.Latch();
+    }
+
     int IListenerCircuit.QueueCount => _receiver?.QueueCount ?? 0;
 
     async Task IListenerCircuit.EnqueueDirectlyAsync(IEnumerable<Envelope> envelopes)