feat: Allow waiting for prefetched messages on drain

src/Transports/RabbitMQ/Wolverine.RabbitMQ.Tests/drain_wait_for_prefetch.cs

@@ -0,0 +1,211 @@
+using System.Collections.Concurrent;
+using System.Diagnostics;
+using IntegrationTests;
+using JasperFx.Core;
+using JasperFx.Resources;
+using Marten;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Hosting;
+using RabbitMQ.Client;
+using Shouldly;
+using Wolverine.ComplianceTests;
+using Wolverine.Marten;
+using Wolverine.RabbitMQ.Internal;
+using Wolverine.Tracking;
+using Wolverine.Transports;
+using Wolverine.Util;
+using Xunit;
+
+namespace Wolverine.RabbitMQ.Tests;
+
+// With DrainWaitForPrefetch on, StopAsync cancels WITHOUT nowait and awaits cancel-ok (plus a batch
+// flush in durable micro-batching mode), so prefetched deliveries land in the inbox instead of being
+// abandoned to broker redelivery. This asserts durable survival, not synchronous handling: the
+// receiver latches mid-drain and defers late arrivals to the durability agent (existing
+// DurableReceiver behavior), so the guarantee is "every prefetched delivery persists", not "every
+// handler ran before stop returned".
+public class drain_wait_for_prefetch : IAsyncLifetime
+{
+    private IHost _host = null!;
+    private string _queueName = null!;
+    private GateFirstDeliveryMapper _mapper = null!;
+    private DrainPrefetchTracker _tracker = null!;
+
+    public async ValueTask InitializeAsync()
+    {
+        _queueName = RabbitTesting.NextQueueName();
+        var schemaName = $"drain_prefetch_{Guid.NewGuid():N}";
+        _tracker = new DrainPrefetchTracker();
+
+        _host = await Host.CreateDefaultBuilder()
+            .UseWolverine(opts =>
+            {
+                opts.LocalRoutingConventionDisabled = true;
+
+                opts.Services.AddSingleton(_tracker);
+
+                opts.Services.AddMarten(m =>
+                {
+                    m.Connection(Servers.PostgresConnectionString);
+                    m.DatabaseSchemaName = schemaName;
+                    m.DisableNpgsqlLogging = true;
+                }).IntegrateWithWolverine(x => x.MessageStorageSchemaName = schemaName);
+
+                opts.UseRabbitMq().AutoProvision().AutoPurgeOnStartup();
+
+                opts.PublishMessage<DrainPrefetchMessage>().ToRabbitQueue(_queueName);
+
+                // Durable + the default MaximumMessagesToReceive (100) turns on the micro-batching
+                // channel, so this exercises StopAsync's DrainBatchedDeliveriesAsync path, not just
+                // the cancel-ok wait.
+                opts.ListenToRabbitQueue(_queueName)
+                    .UseDurableInbox()
+                    .DrainWaitForPrefetch()
+                    .PreFetchCount(100)
+                    .ConfigureQueue(q =>
+                    {
+                        var queue = (RabbitMqQueue)q;
+                        _mapper = new GateFirstDeliveryMapper(new RabbitMqEnvelopeMapper(queue, null!));
+                        queue.EnvelopeMapper = _mapper;
+                    });
+
+                opts.Services.AddResourceSetupOnStartup();
+            }).StartAsync();
+
+        await _host.ResetResourceState();
+    }
+
+    public async ValueTask DisposeAsync()
+    {
+        await _host.StopAsync();
+        _host.Dispose();
+    }
+
+    // Reverted to a fire-and-forget nowait cancel, this fails: the gated mapper keeps 24 of the 25
+    // messages undispatched in the client, and the channel would tear down before they reached
+    // HandleBasicDeliverAsync. Persisting all 25 is only possible because StopAsync awaited cancel-ok
+    // and drained the batching channel first.
+    [Fact]
+    public async Task prefetched_durable_batch_is_drained_through_on_stop()
+    {
+        const int messageCount = 25;
+
+        var bus = _host.MessageBus();
+        for (var i = 0; i < messageCount; i++)
+        {
+            await bus.PublishAsync(new DrainPrefetchMessage(Guid.NewGuid()));
+        }
+
+        // Block on the first delivery in the client's single dispatch thread (ConsumerDispatchConcurrency
+        // 1), so nothing dispatches past it. The broker still pushes the other 24 onto the wire while it
+        // sits blocked; give them a moment to land so they're prefetched, not still in flight.
+        _mapper.Entered.WaitOne(10.Seconds())
+            .ShouldBeTrue("the first prefetched delivery never reached the gated mapper");
+        await Task.Delay(1.Seconds(), TestContext.Current.CancellationToken);
+
+        var runtime = _host.GetRuntime();
+        var agent = runtime.Endpoints.ActiveListeners()
+            .Single(x => x.Uri == new Uri($"rabbitmq://queue/{_queueName}"));
+
+        var stopTask = agent.StopAndDrainAsync().AsTask();
+
+        _mapper.Release();
+
+        await stopTask.WaitAsync(30.Seconds(), TestContext.Current.CancellationToken);
+
+        // Every prefetched delivery is captured in the inbox rather than left to broker redelivery.
+        // Some may still be "Incoming" rather than "Handled": the receiver latches partway through the
+        // drain and defers later arrivals to the recovery sweep (pre-existing shutdown behavior).
+        var incoming = await runtime.Storage.Admin.AllIncomingAsync();
+        var persistedCount = incoming.Count(x => x.MessageType == typeof(DrainPrefetchMessage).ToMessageTypeName());
+
+        persistedCount.ShouldBe(messageCount);
+
+        // And the ones that beat the shutdown latch ran through the handler, proving drained
+        // deliveries reach the pipeline, not just the inbox table.
+        _tracker.Handled.Count.ShouldBeGreaterThan(0);
+    }
+}
+
+public record DrainPrefetchMessage(Guid Id);
+
+public class DrainPrefetchTracker
+{
+    public readonly ConcurrentBag<Guid> Handled = new();
+}
+
+public static class DrainPrefetchMessageHandler
+{
+    public static void Handle(DrainPrefetchMessage message, DrainPrefetchTracker tracker)
+    {
+        tracker.Handled.Add(message.Id);
+    }
+}
+
+/// <summary>
+/// Blocks the first incoming delivery's envelope mapping until released. Mapping runs synchronously in
+/// HandleBasicDeliverAsync, so at the default ConsumerDispatchConcurrency of 1 this blocks the client's
+/// single dispatch thread, keeping every other prefetched delivery undispatched.
+/// </summary>
+internal class GateFirstDeliveryMapper : IRabbitMqEnvelopeMapper
+{
+    private readonly IRabbitMqEnvelopeMapper _inner;
+    private readonly ManualResetEventSlim _entered = new(false);
+    private readonly ManualResetEventSlim _release = new(false);
+    private int _hits;
+
+    public GateFirstDeliveryMapper(IRabbitMqEnvelopeMapper inner)
+    {
+        _inner = inner;
+    }
+
+    public WaitHandle Entered => _entered.WaitHandle;
+
+    public void Release()
+    {
+        _release.Set();
+    }
+
+    public void MapIncomingToEnvelope(Envelope envelope, IReadOnlyBasicProperties incoming)
+    {
+        if (Interlocked.Increment(ref _hits) == 1)
+        {
+            _entered.Set();
+            _release.Wait(TimeSpan.FromSeconds(20));
+        }
+
+        _inner.MapIncomingToEnvelope(envelope, incoming);
+    }
+
+    public void MapEnvelopeToOutgoing(Envelope envelope, IBasicProperties outgoing)
+    {
+        _inner.MapEnvelopeToOutgoing(envelope, outgoing);
+    }
+}
+
+// Parity check: a listener that never opts in keeps the old fire-and-forget behavior and stops
+// promptly.
+public class drain_wait_for_prefetch_default_off
+{
+    [Fact]
+    public async Task listener_without_the_flag_still_stops_cleanly_and_promptly()
+    {
+        var queue = RabbitTesting.NextQueueName();
+        using var host = await WolverineHost.ForAsync(opts =>
+        {
+            opts.UseRabbitMq().AutoProvision().AutoPurgeOnStartup();
+            opts.ListenToRabbitQueue(queue); // DrainWaitForPrefetch() not called -- default off
+        });
+
+        var runtime = host.GetRuntime();
+        var agent = runtime.Endpoints.ActiveListeners()
+            .Single(x => x.Uri == new Uri($"rabbitmq://queue/{queue}"));
+
+        var sw = Stopwatch.StartNew();
+        await agent.StopAndDrainAsync();
+        sw.Stop();
+
+        agent.Status.ShouldBe(ListeningStatus.Stopped);
+        sw.Elapsed.ShouldBeLessThan(5.Seconds());
+    }
+}

src/Transports/RabbitMQ/Wolverine.RabbitMQ/Internal/RabbitMqListener.cs

@@ -150,9 +150,42 @@ public async ValueTask StopAsync()
         var channel = Channel;
         if (channel != null)
         {
-            foreach (var consumerTag in consumer.ConsumerTags)
+            if (!Queue.DrainWaitForPrefetch)
             {
-                await channel.BasicCancelAsync(consumerTag, true, default);
+                // nowait cancel: no cancel-ok, and still-prefetched deliveries are requeued by the
+                // broker on channel close.
+                foreach (var consumerTag in consumer.ConsumerTags)
+                {
+                    await channel.BasicCancelAsync(consumerTag, true, default);
+                }
+
+                return;
+            }
+
+            // Cancel WITHOUT nowait so the broker replies cancel-ok, which the client dispatches only
+            // after every prefetched delivery ahead of it in FIFO order. In durable micro-batching mode
+            // that only means the deliveries reached the batching channel, so drain that batch too --
+            // otherwise the caller latches the receiver first and the batch is redelivered. Bound the
+            // wait on the shared drain budget and log-and-continue so an unreachable broker can't abort
+            // the caller's stop-and-drain.
+            using var cts = new CancellationTokenSource(_runtime.DurabilitySettings.DrainTimeout);
+            try
+            {
+                var cancelled = 0;
+                foreach (var consumerTag in consumer.ConsumerTags)
+                {
+                    await channel.BasicCancelAsync(consumerTag, false, cts.Token);
+                    cancelled++;
+                }
+
+                await consumer.WaitForCancelOksAsync(cancelled, cts.Token);
+                await consumer.DrainBatchedDeliveriesAsync().WaitAsync(cts.Token);
+            }
+            catch (Exception e)
+            {
+                Logger.LogWarning(e,
+                    "Timed out or errored waiting for prefetched messages to drain at {Uri} during listener stop; continuing shutdown",
+                    Address);
             }
         }
     }

src/Transports/RabbitMQ/Wolverine.RabbitMQ/Internal/RabbitMqQueue.cs

@@ -99,6 +99,17 @@ public ushort PreFetchCount
         set => _preFetchCount = value;
     }
 
+    /// <summary>
+    /// When true, listener shutdown waits for prefetched messages in the RabbitMQ client's dispatch
+    /// buffer to reach the consumer before closing the channel, preventing silent redeliveries of
+    /// messages that were prefetched but not yet handled. Default is false.
+    ///
+    /// Only a hard guarantee at <c>ConsumerDispatchConcurrency</c> of 1 (the default). At higher
+    /// concurrency the client handles deliveries and cancel-ok in parallel, so it degrades to
+    /// best-effort -- some messages may still be redelivered, but fewer than without waiting.
+    /// </summary>
+    public bool DrainWaitForPrefetch { get; set; }
+
     /// <summary>
     ///     Use to override the dead letter queue for this queue
     /// </summary>

src/Transports/RabbitMQ/Wolverine.RabbitMQ/Internal/WorkerQueueMessageConsumer.cs

@@ -15,6 +15,7 @@ internal class WorkerQueueMessageConsumer : AsyncDefaultBasicConsumer, IDisposab
     private readonly IRabbitMqEnvelopeMapper _mapper;
     private readonly IReceiver _workerQueue;
     private bool _latched;
+    private readonly SemaphoreSlim _cancelOks = new(0);
 
     // GH-3492: durable endpoints coalesce prefetched deliveries into Envelope[] batches so the
     // inbox persists them with one multi-VALUES insert instead of gating on one INSERT round
@@ -76,6 +77,42 @@ public void Dispose()
         _batching?.TriggerBatch();
     }
 
+    /// <summary>
+    /// Wait for the broker's cancel-ok on <paramref name="count"/> cancelled consumer tags. The client
+    /// dispatches each cancel-ok only after every prefetched delivery ahead of it in FIFO order, so once
+    /// all arrive the prefetch backlog has drained (to the batching channel, in micro-batching mode).
+    /// </summary>
+    internal async Task WaitForCancelOksAsync(int count, CancellationToken token)
+    {
+        for (var i = 0; i < count; i++)
+        {
+            await _cancelOks.WaitAsync(token);
+        }
+    }
+
+    /// <summary>
+    /// Flush any batched-but-undelivered envelopes to the receiver and await that flush. cancel-ok only
+    /// guarantees deliveries reached the batching channel, so without this the receiver latches first and
+    /// the batch is redelivered.
+    /// </summary>
+    internal async Task DrainBatchedDeliveriesAsync()
+    {
+        if (_batching == null)
+        {
+            return;
+        }
+
+        _batching.TriggerBatch();
+        _batching.Complete();
+        await _batching.WaitForCompletionAsync();
+    }
+
+    public override Task HandleBasicCancelOkAsync(string consumerTag, CancellationToken cancellationToken = default)
+    {
+        _cancelOks.Release();
+        return base.HandleBasicCancelOkAsync(consumerTag, cancellationToken);
+    }
+
     //TODO do something with the token passed in here
     public override async Task HandleBasicDeliverAsync(string consumerTag, ulong deliveryTag, bool redelivered, string exchange,
         string routingKey, IReadOnlyBasicProperties properties, ReadOnlyMemory<byte> body,

src/Transports/RabbitMQ/Wolverine.RabbitMQ/RabbitMqListenerConfiguration.cs

@@ -131,6 +131,20 @@ public RabbitMqListenerConfiguration ConsumerDispatchConcurrency(ushort concurre
         return this;
     }
 
+    /// <summary>
+    /// When enabled, listener shutdown waits for prefetched messages in the RabbitMQ client's dispatch
+    /// buffer to be delivered before closing the channel, preventing silent redeliveries of
+    /// prefetched-but-unprocessed messages. Default is false.
+    ///
+    /// Only a hard guarantee at <c>ConsumerDispatchConcurrency</c> of 1 (the default); at higher
+    /// concurrency it degrades to best-effort as deliveries and cancel-ok are handled in parallel.
+    /// </summary>
+    public RabbitMqListenerConfiguration DrainWaitForPrefetch()
+    {
+        add(e => e.DrainWaitForPrefetch = true);
+        return this;
+    }
+
     /// <summary>
     /// For durable (inbox-backed) listeners, the maximum number of prefetched deliveries the
     /// consumer coalesces into one batched inbox insert (5ms max accumulation age). 1 reverts