diff --git a/src/Transports/RabbitMQ/Wolverine.RabbitMQ.Tests/drain_wait_for_prefetch.cs b/src/Transports/RabbitMQ/Wolverine.RabbitMQ.Tests/drain_wait_for_prefetch.cs new file mode 100644 index 000000000..93b852790 --- /dev/null +++ b/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().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 Handled = new(); +} + +public static class DrainPrefetchMessageHandler +{ + public static void Handle(DrainPrefetchMessage message, DrainPrefetchTracker tracker) + { + tracker.Handled.Add(message.Id); + } +} + +/// +/// 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. +/// +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()); + } +} diff --git a/src/Transports/RabbitMQ/Wolverine.RabbitMQ/Internal/RabbitMqListener.cs b/src/Transports/RabbitMQ/Wolverine.RabbitMQ/Internal/RabbitMqListener.cs index 96b354b6a..29e25da40 100644 --- a/src/Transports/RabbitMQ/Wolverine.RabbitMQ/Internal/RabbitMqListener.cs +++ b/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); } } } diff --git a/src/Transports/RabbitMQ/Wolverine.RabbitMQ/Internal/RabbitMqQueue.cs b/src/Transports/RabbitMQ/Wolverine.RabbitMQ/Internal/RabbitMqQueue.cs index a374bb8ce..582b8f1e3 100644 --- a/src/Transports/RabbitMQ/Wolverine.RabbitMQ/Internal/RabbitMqQueue.cs +++ b/src/Transports/RabbitMQ/Wolverine.RabbitMQ/Internal/RabbitMqQueue.cs @@ -99,6 +99,17 @@ public ushort PreFetchCount set => _preFetchCount = value; } + /// + /// 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 ConsumerDispatchConcurrency 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. + /// + public bool DrainWaitForPrefetch { get; set; } + /// /// Use to override the dead letter queue for this queue /// diff --git a/src/Transports/RabbitMQ/Wolverine.RabbitMQ/Internal/WorkerQueueMessageConsumer.cs b/src/Transports/RabbitMQ/Wolverine.RabbitMQ/Internal/WorkerQueueMessageConsumer.cs index fd29a6e10..2f3537fc2 100644 --- a/src/Transports/RabbitMQ/Wolverine.RabbitMQ/Internal/WorkerQueueMessageConsumer.cs +++ b/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(); } + /// + /// Wait for the broker's cancel-ok on 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). + /// + internal async Task WaitForCancelOksAsync(int count, CancellationToken token) + { + for (var i = 0; i < count; i++) + { + await _cancelOks.WaitAsync(token); + } + } + + /// + /// 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. + /// + 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 body, diff --git a/src/Transports/RabbitMQ/Wolverine.RabbitMQ/RabbitMqListenerConfiguration.cs b/src/Transports/RabbitMQ/Wolverine.RabbitMQ/RabbitMqListenerConfiguration.cs index b94fd4070..df04c6dc5 100644 --- a/src/Transports/RabbitMQ/Wolverine.RabbitMQ/RabbitMqListenerConfiguration.cs +++ b/src/Transports/RabbitMQ/Wolverine.RabbitMQ/RabbitMqListenerConfiguration.cs @@ -131,6 +131,20 @@ public RabbitMqListenerConfiguration ConsumerDispatchConcurrency(ushort concurre return this; } + /// + /// 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 ConsumerDispatchConcurrency of 1 (the default); at higher + /// concurrency it degrades to best-effort as deliveries and cancel-ok are handled in parallel. + /// + public RabbitMqListenerConfiguration DrainWaitForPrefetch() + { + add(e => e.DrainWaitForPrefetch = true); + return this; + } + /// /// 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