diff --git a/src/Mocha/src/Mocha.Transport.AzureServiceBus/Middlewares/Receive/AzureServiceBusAcknowledgementMiddleware.cs b/src/Mocha/src/Mocha.Transport.AzureServiceBus/Middlewares/Receive/AzureServiceBusAcknowledgementMiddleware.cs index f2560916d37..8aa053243b5 100644 --- a/src/Mocha/src/Mocha.Transport.AzureServiceBus/Middlewares/Receive/AzureServiceBusAcknowledgementMiddleware.cs +++ b/src/Mocha/src/Mocha.Transport.AzureServiceBus/Middlewares/Receive/AzureServiceBusAcknowledgementMiddleware.cs @@ -25,19 +25,18 @@ public async ValueTask InvokeAsync(IReceiveContext context, ReceiveDelegate next feature.ProcessMessageEventArgs?.EntityPath ?? feature.ProcessSessionMessageEventArgs?.EntityPath ?? string.Empty; - var cancellationToken = context.CancellationToken; try { await next(context); - await CompleteAsync(feature.Actions, context.Services, feature.Message, entityPath, cancellationToken); + await CompleteAsync(feature.Actions, context.Services, feature.Message, entityPath, CancellationToken.None); } catch { try { - await AbandonAsync(feature.Actions, context.Services, feature.Message, entityPath, cancellationToken); + await AbandonAsync(feature.Actions, context.Services, feature.Message, entityPath, CancellationToken.None); } catch { diff --git a/src/Mocha/src/Mocha.Transport.InMemory/InMemoryReceiveEndpoint.cs b/src/Mocha/src/Mocha.Transport.InMemory/InMemoryReceiveEndpoint.cs index 9aff1043560..c72fdfa9155 100644 --- a/src/Mocha/src/Mocha.Transport.InMemory/InMemoryReceiveEndpoint.cs +++ b/src/Mocha/src/Mocha.Transport.InMemory/InMemoryReceiveEndpoint.cs @@ -78,6 +78,10 @@ await ExecuteAsync( item.Envelope, cancellationToken); } + catch (Exception) when (cancellationToken.IsCancellationRequested) + { + // Message processing was cancelled during shutdown. + } catch (Exception ex) { logger.LogCritical(ex, "Error processing message"); diff --git a/src/Mocha/src/Mocha.Transport.Postgres/Connection/PostgresMessageStore.cs b/src/Mocha/src/Mocha.Transport.Postgres/Connection/PostgresMessageStore.cs index fbf34a15eca..c64fc68da23 100644 --- a/src/Mocha/src/Mocha.Transport.Postgres/Connection/PostgresMessageStore.cs +++ b/src/Mocha/src/Mocha.Transport.Postgres/Connection/PostgresMessageStore.cs @@ -415,6 +415,40 @@ FROM updated_message public Task ReleaseMessageAsync(Guid transportMessageId, CancellationToken cancellationToken) => ReleaseMessageAsync(transportMessageId, errorInfo: null, cancellationToken); + /// + /// Releases the messages that are still leased by the specified consumer back to the queue. + /// Messages that were deleted or already released are not affected. + /// + public async Task ReleaseLeasedMessagesAsync( + Guid[] transportMessageIds, + Guid consumerId, + CancellationToken cancellationToken) + { + await using var connection = await _connectionManager.OpenConnectionAsync(cancellationToken); + await using var command = connection.CreateCommand(); + + command.CommandText = $""" + WITH released AS ( + UPDATE {_schemaOptions.MessageTable} + SET consumer_id = NULL + WHERE transport_message_id = ANY(@ids) + AND consumer_id = @consumer_id + RETURNING queue_id + ) + SELECT pg_notify( + '{_schemaOptions.NotificationChannel}', + q.name::text + ) + FROM (SELECT DISTINCT queue_id FROM released) r + JOIN {_schemaOptions.QueueTable} q ON r.queue_id = q.id; + """; + + command.Parameters.Add( + new NpgsqlParameter("ids", NpgsqlDbType.Array | NpgsqlDbType.Uuid) { Value = transportMessageIds }); + command.Parameters.Add(new NpgsqlParameter("consumer_id", NpgsqlDbType.Uuid) { Value = consumerId }); + await command.ExecuteNonQueryAsync(cancellationToken); + } + /// /// Appends error information to the error_reason JSONB array on a message. /// Called when a message processing attempt fails but has not yet exceeded the retry limit. diff --git a/src/Mocha/src/Mocha.Transport.Postgres/PostgresReceiveEndpoint.cs b/src/Mocha/src/Mocha.Transport.Postgres/PostgresReceiveEndpoint.cs index 9090afc3561..beb448b3ac5 100644 --- a/src/Mocha/src/Mocha.Transport.Postgres/PostgresReceiveEndpoint.cs +++ b/src/Mocha/src/Mocha.Transport.Postgres/PostgresReceiveEndpoint.cs @@ -117,14 +117,24 @@ private async Task PollMessagesAsync(ILogger logger, CancellationToken cancellat continue; } - await Parallel.ForEachAsync( - batch.Messages, - new ParallelOptions + try + { + await Parallel.ForEachAsync( + batch.Messages, + new ParallelOptions + { + MaxDegreeOfParallelism = _maxConcurrency, + CancellationToken = cancellationToken + }, + (message, ct) => new ValueTask(ProcessMessageAsync(message, logger, ct))); + } + finally + { + if (cancellationToken.IsCancellationRequested) { - MaxDegreeOfParallelism = _maxConcurrency, - CancellationToken = cancellationToken - }, - (message, ct) => new ValueTask(ProcessMessageAsync(message, logger, ct))); + await ReleaseLeasedMessagesAsync(batch, logger); + } + } // If we got a full batch, there may be more messages hasMore = batch.Count >= _maxBatchSize; @@ -190,9 +200,14 @@ await ExecuteAsync( message, cancellationToken); + // The message was handled, so it is deleted even when the endpoint is stopping. await transport.MessageStore.DeleteMessageAsync( message.TransportMessageId, - cancellationToken); + CancellationToken.None); + } + catch (Exception) when (cancellationToken.IsCancellationRequested) + { + // The interrupted message is released with the rest of its batch. } catch (Exception ex) { @@ -204,7 +219,7 @@ await transport.MessageStore.DeleteMessageAsync( await transport.MessageStore.ReleaseMessageAsync( message.TransportMessageId, errorInfo, - cancellationToken); + CancellationToken.None); } catch (Exception releaseEx) { @@ -213,6 +228,27 @@ await transport.MessageStore.ReleaseMessageAsync( } } + private async Task ReleaseLeasedMessagesAsync(PostgresMessageBatch batch, ILogger logger) + { + var transportMessageIds = new Guid[batch.Count]; + for (var i = 0; i < transportMessageIds.Length; i++) + { + transportMessageIds[i] = batch.Messages[i].TransportMessageId; + } + + try + { + await transport.MessageStore.ReleaseLeasedMessagesAsync( + transportMessageIds, + _consumerId, + CancellationToken.None); + } + catch (Exception ex) + { + logger.LeasedMessagesReleaseFailed(ex, Queue.Name); + } + } + private async Task UpdateScheduledTriggerAsync(CancellationToken cancellationToken) { var scheduledAt = await transport.MessageStore.GetNextScheduledTimeAsync( @@ -301,6 +337,9 @@ internal static partial class Logs [LoggerMessage(LogLevel.Error, "Error releasing message {TransportMessageId}.")] public static partial void MessageReleaseFailed(this ILogger logger, Exception exception, Guid transportMessageId); + [LoggerMessage(LogLevel.Error, "Error releasing the leased messages of queue {QueueName} while stopping.")] + public static partial void LeasedMessagesReleaseFailed(this ILogger logger, Exception exception, string queueName); + [LoggerMessage(LogLevel.Warning, "Database is unreachable for queue {QueueName}, waiting for connectivity to resume.")] public static partial void WaitingForDatabase(this ILogger logger, string queueName); } diff --git a/src/Mocha/src/Mocha.Transport.RabbitMQ/Connection/RabbitMQConsumerManager.cs b/src/Mocha/src/Mocha.Transport.RabbitMQ/Connection/RabbitMQConsumerManager.cs index 8754d58d010..f99ee4edc6f 100644 --- a/src/Mocha/src/Mocha.Transport.RabbitMQ/Connection/RabbitMQConsumerManager.cs +++ b/src/Mocha/src/Mocha.Transport.RabbitMQ/Connection/RabbitMQConsumerManager.cs @@ -374,9 +374,9 @@ private async Task DisconnectAsync(CancellationToken cancellationToken) /// public async ValueTask DisposeAsync() { - await DisconnectAsync(CancellationToken.None); Manager.RemoveConsumer(this); await _consumerCts.CancelAsync(); + await DisconnectAsync(CancellationToken.None); _consumerCts.Dispose(); } } diff --git a/src/Mocha/src/Mocha.Transport.RabbitMQ/Middlewares/Receive/RabbitMQAcknowledgementMiddleware.cs b/src/Mocha/src/Mocha.Transport.RabbitMQ/Middlewares/Receive/RabbitMQAcknowledgementMiddleware.cs index 1c83c1c7af2..8b90db26617 100644 --- a/src/Mocha/src/Mocha.Transport.RabbitMQ/Middlewares/Receive/RabbitMQAcknowledgementMiddleware.cs +++ b/src/Mocha/src/Mocha.Transport.RabbitMQ/Middlewares/Receive/RabbitMQAcknowledgementMiddleware.cs @@ -19,7 +19,6 @@ public async ValueTask InvokeAsync(IReceiveContext context, ReceiveDelegate next var feature = context.Features.GetOrSet(); var channel = feature.Channel; var eventArgs = feature.EventArgs; - var cancellationToken = context.CancellationToken; try { @@ -27,14 +26,14 @@ public async ValueTask InvokeAsync(IReceiveContext context, ReceiveDelegate next if (channel.IsOpen) { - await channel.BasicAckAsync(eventArgs.DeliveryTag, false, cancellationToken); + await channel.BasicAckAsync(eventArgs.DeliveryTag, false, CancellationToken.None); } } catch { if (channel.IsOpen) { - await channel.BasicNackAsync(eventArgs.DeliveryTag, false, true, cancellationToken); + await channel.BasicNackAsync(eventArgs.DeliveryTag, false, true, CancellationToken.None); } throw; diff --git a/src/Mocha/src/Mocha/Consumers/Batching/BatchCollector.cs b/src/Mocha/src/Mocha/Consumers/Batching/BatchCollector.cs index c5d9bcfe0c1..c59407b3c11 100644 --- a/src/Mocha/src/Mocha/Consumers/Batching/BatchCollector.cs +++ b/src/Mocha/src/Mocha/Consumers/Batching/BatchCollector.cs @@ -61,6 +61,33 @@ public async ValueTask> Add(IConsumeContext contex return entry; } + /// + /// Removes an entry that has not been emitted in a batch yet. + /// + /// true if the entry was removed; false if it was already emitted. + public bool TryRemove(BufferedEntry entry) + { + lock (_sync) + { + for (var i = 0; i < _buffer.Count; i++) + { + if (_buffer[i].Task == entry.Task) + { + _buffer.RemoveAt(i); + + if (_buffer.Count == 0) + { + _delay.Cancel(); + } + + return true; + } + } + } + + return false; + } + private async ValueTask OnDelayElapsed() { MessageBatch? batch; diff --git a/src/Mocha/src/Mocha/Consumers/Batching/BatchOptions.cs b/src/Mocha/src/Mocha/Consumers/Batching/BatchOptions.cs index 04976530ec6..c2a78f1d0c7 100644 --- a/src/Mocha/src/Mocha/Consumers/Batching/BatchOptions.cs +++ b/src/Mocha/src/Mocha/Consumers/Batching/BatchOptions.cs @@ -16,7 +16,7 @@ public sealed class BatchOptions public TimeSpan BatchTimeout { get; set; } = TimeSpan.FromSeconds(1); /// - /// Gets or sets the maximum number of batches that can be processed concurrently. + /// Gets or sets the maximum number of batches that each receive endpoint processes concurrently. /// Higher values improve throughput when batch processing is slow relative to message /// arrival rate, at the cost of losing ordering guarantees between batches. Default: 1. /// diff --git a/src/Mocha/src/Mocha/Consumers/Implementations/BatchConsumer.cs b/src/Mocha/src/Mocha/Consumers/Implementations/BatchConsumer.cs index 596a0e9b243..af95f7fa0a0 100644 --- a/src/Mocha/src/Mocha/Consumers/Implementations/BatchConsumer.cs +++ b/src/Mocha/src/Mocha/Consumers/Implementations/BatchConsumer.cs @@ -1,3 +1,4 @@ +using System.Collections.Concurrent; using System.Threading.Channels; using Microsoft.Extensions.DependencyInjection; using Microsoft.Extensions.Logging; @@ -13,18 +14,25 @@ namespace Mocha; /// Uses a TCS-based pattern to hold each per-message pipeline open until the batch handler /// completes. This preserves existing middleware semantics (ACK, fault, circuit breaker) /// without any modifications to the middleware chain. +/// Each receive endpoint collects and processes its own batches. /// internal sealed class BatchConsumer : Consumer where THandler : class, IBatchEventHandler { public BatchConsumer() : base(typeof(THandler)) { } - private BatchCollector _collector = null!; - private Channel> _channel = null!; - private ChannelProcessor> _processor = null!; +#if NET9_0_OR_GREATER + private readonly Lock _sync = new(); +#else + private readonly object _sync = new(); +#endif + private readonly ConcurrentDictionary _endpoints = new(); + private BatchOptions _options = null!; + private TimeProvider _timeProvider = null!; private IServiceProvider _applicationServices = null!; private ILogger _logger = null!; private MessageType? _itemMessageType; + private bool _disposed; protected override void Configure(IConsumerDescriptor descriptor) { @@ -39,25 +47,11 @@ protected override void OnAfterInitialize(IMessagingSetupContext context) var options = Configuration!.Features.Get() ?? new BatchOptions(); options.Validate(); + _options = options; _applicationServices = context.Services.GetRequiredService().ServiceProvider; _logger = context.Services.GetRequiredService>>(); - - var timeProvider = context.Services.GetRequiredService(); - - _channel = - Channel.CreateBounded>( - new BoundedChannelOptions(options.MaxConcurrentBatches) - { - SingleReader = options.MaxConcurrentBatches == 1 - }); - - _processor = new ChannelProcessor>( - _channel.Reader.ReadAllAsync, - ProcessBatchAsync, - options.MaxConcurrentBatches); - - _collector = new BatchCollector(options, batch => _channel.Writer.WriteAsync(batch), timeProvider); + _timeProvider = context.Services.GetRequiredService(); _itemMessageType = context.Messages.GetMessageType(typeof(TEvent)); } @@ -83,14 +77,56 @@ public override async ValueTask ProcessAsync(IReceiveContext context) // the message can be deserialized before adding to the batch _ = batchContext.Message; - var entry = await _collector.Add(batchContext); - await entry.Task; + var cancellationToken = context.CancellationToken; + var collector = GetEndpointBatches(context.Endpoint).Collector; + var entry = await collector.Add(batchContext); + + try + { + await entry.Task.WaitAsync(cancellationToken); + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + if (collector.TryRemove(entry)) + { + throw; + } + + // A dispatched batch observes the same cancellation and uses the context until it completes. + await entry.Task; + } + } + + private EndpointBatches GetEndpointBatches(ReceiveEndpoint endpoint) + { + if (_endpoints.TryGetValue(endpoint, out var batches)) + { + return batches; + } + + lock (_sync) + { + ObjectDisposedException.ThrowIf(_disposed, this); + + if (!_endpoints.TryGetValue(endpoint, out batches)) + { + batches = new EndpointBatches(_options, _timeProvider, ProcessBatchAsync); + _endpoints[endpoint] = batches; + } + + return batches; + } } private async Task ProcessBatchAsync(MessageBatch batch, CancellationToken cancellationToken) { + using var cancellation = LinkCancellation(batch, cancellationToken); + var batchToken = cancellation.Token; + try { + batchToken.ThrowIfCancellationRequested(); + _logger.DispatchingBatch(batch.Count, batch.CompletionMode); // The batch has no receive scope of its own, so it gets one here, mirroring the scope @@ -103,7 +139,7 @@ private async Task ProcessBatchAsync(MessageBatch batch, CancellationTok batch.GetContext(0), Guid.NewGuid().ToString(), _itemMessageType, - cancellationToken); + batchToken); await Pipeline(batchContext); @@ -112,10 +148,8 @@ private async Task ProcessBatchAsync(MessageBatch batch, CancellationTok entry.Complete(); } } - catch (OperationCanceledException) + catch (Exception) when (batchToken.IsCancellationRequested) { - // Handler observed cancellation - cancel all entries so per-message pipelines - // unblock for NACK/redelivery foreach (var entry in batch.Entries) { entry.Cancel(); @@ -140,6 +174,32 @@ private async Task ProcessBatchAsync(MessageBatch batch, CancellationTok } } + private static CancellationTokenSource LinkCancellation( + MessageBatch batch, + CancellationToken processorToken) + { + var first = batch.Entries[0].Context.CancellationToken; + List? others = null; + + for (var i = 1; i < batch.Count; i++) + { + var token = batch.Entries[i].Context.CancellationToken; + if (token != first && others?.Contains(token) != true) + { + (others ??= []).Add(token); + } + } + + if (others is null) + { + return CancellationTokenSource.CreateLinkedTokenSource(processorToken, first); + } + + others.Add(processorToken); + others.Add(first); + return CancellationTokenSource.CreateLinkedTokenSource([.. others]); + } + public override ConsumerDescription Describe() { return new ConsumerDescription( @@ -154,22 +214,66 @@ public override ConsumerDescription Describe() public override async ValueTask DisposeAsync() { - await _collector.DisposeAsync(); + EndpointBatches[] endpoints; - _channel.Writer.Complete(); - await _processor.DisposeAsync(); + lock (_sync) + { + _disposed = true; + endpoints = [.. _endpoints.Values]; + } - while (_channel.Reader.TryRead(out var batch)) + foreach (var batches in endpoints) { - foreach (var entry in batch.Entries) - { - try + await batches.DisposeAsync(); + } + } + + /// + /// The batch collector, channel and processor of one receive endpoint. + /// + private sealed class EndpointBatches : IAsyncDisposable + { + private readonly Channel> _channel; + private readonly ChannelProcessor> _processor; + + public EndpointBatches( + BatchOptions options, + TimeProvider timeProvider, + Func, CancellationToken, Task> processBatch) + { + _channel = Channel.CreateBounded>( + new BoundedChannelOptions(options.MaxConcurrentBatches) { - entry.Cancel(); - } - catch + SingleReader = options.MaxConcurrentBatches == 1 + }); + _processor = new ChannelProcessor>( + _channel.Reader.ReadAllAsync, + processBatch, + options.MaxConcurrentBatches); + Collector = new BatchCollector(options, batch => _channel.Writer.WriteAsync(batch), timeProvider); + } + + public BatchCollector Collector { get; } + + public async ValueTask DisposeAsync() + { + await Collector.DisposeAsync(); + + _channel.Writer.Complete(); + await _processor.DisposeAsync(); + + while (_channel.Reader.TryRead(out var batch)) + { + foreach (var entry in batch.Entries) { - // Best-effort cancellation + try + { + entry.Cancel(); + } + catch + { + // Best-effort cancellation + } } } } diff --git a/src/Mocha/src/Mocha/Endpoints/ReceiveEndpoint.cs b/src/Mocha/src/Mocha/Endpoints/ReceiveEndpoint.cs index 188e4ff37f8..030efbffe76 100644 --- a/src/Mocha/src/Mocha/Endpoints/ReceiveEndpoint.cs +++ b/src/Mocha/src/Mocha/Endpoints/ReceiveEndpoint.cs @@ -90,7 +90,7 @@ public abstract class ReceiveEndpoint(MessagingTransport transport) : IReceiveEn /// Allocates a scoped , retrieves a pooled /// , configures it via , and /// executes the compiled middleware pipeline. Exceptions that escape the pipeline are - /// caught and logged at the Critical level to prevent transport-level crashes. + /// logged at the Critical level unless message processing was cancelled. /// /// The type of caller-provided state passed to the configure action. /// @@ -104,6 +104,8 @@ public async ValueTask ExecuteAsync( TState state, CancellationToken cancellationToken) { + cancellationToken.ThrowIfCancellationRequested(); + var logger = _runtimeState!.Logger; var services = _runtimeState!.ServiceProvider; var pools = _runtimeState.Pools; @@ -123,7 +125,7 @@ public async ValueTask ExecuteAsync( await _pipeline(context); } - catch (Exception ex) + catch (Exception ex) when (!cancellationToken.IsCancellationRequested) { // exceptions should technically never bubble up here. logger.LogCritical(ex, "Error processing message"); diff --git a/src/Mocha/src/Mocha/Middlewares/Consume/Retry/RetryExecutor.cs b/src/Mocha/src/Mocha/Middlewares/Consume/Retry/RetryExecutor.cs index 6e2afa3cf8d..620ee4cfac1 100644 --- a/src/Mocha/src/Mocha/Middlewares/Consume/Retry/RetryExecutor.cs +++ b/src/Mocha/src/Mocha/Middlewares/Consume/Retry/RetryExecutor.cs @@ -30,12 +30,14 @@ public static async ValueTask ExecuteAsync( while (true) { + cancellationToken.ThrowIfCancellationRequested(); + try { await action(state); return; } - catch (Exception ex) + catch (Exception ex) when (!cancellationToken.IsCancellationRequested) { // Match exception against policy rules. var rule = ExceptionPolicyMatcher.Match(rules, ex); diff --git a/src/Mocha/src/Mocha/Middlewares/Receive/ReceiveDeadLetterMiddleware.cs b/src/Mocha/src/Mocha/Middlewares/Receive/ReceiveDeadLetterMiddleware.cs index e434a07aac7..40ced9bd7b3 100644 --- a/src/Mocha/src/Mocha/Middlewares/Receive/ReceiveDeadLetterMiddleware.cs +++ b/src/Mocha/src/Mocha/Middlewares/Receive/ReceiveDeadLetterMiddleware.cs @@ -27,7 +27,7 @@ public async ValueTask InvokeAsync(IReceiveContext context, ReceiveDelegate next { await next(context); } - catch (Exception ex) + catch (Exception ex) when (!context.CancellationToken.IsCancellationRequested) { logger.ExceptionOccurred(ex); } diff --git a/src/Mocha/src/Mocha/Middlewares/Receive/ReceiveFaultMiddleware.cs b/src/Mocha/src/Mocha/Middlewares/Receive/ReceiveFaultMiddleware.cs index 1250452c496..2524c0254df 100644 --- a/src/Mocha/src/Mocha/Middlewares/Receive/ReceiveFaultMiddleware.cs +++ b/src/Mocha/src/Mocha/Middlewares/Receive/ReceiveFaultMiddleware.cs @@ -26,7 +26,7 @@ public async ValueTask InvokeAsync(IReceiveContext context, ReceiveDelegate next { await next(context); } - catch (Exception ex) + catch (Exception ex) when (!context.CancellationToken.IsCancellationRequested) { var fault = FaultInfo.From(Guid.NewGuid(), provider.GetUtcNow(), ex); diff --git a/src/Mocha/src/Mocha/Middlewares/Receive/Redelivery/ReceiveRedeliveryMiddleware.cs b/src/Mocha/src/Mocha/Middlewares/Receive/Redelivery/ReceiveRedeliveryMiddleware.cs index e52e675c8e3..22cec645686 100644 --- a/src/Mocha/src/Mocha/Middlewares/Receive/Redelivery/ReceiveRedeliveryMiddleware.cs +++ b/src/Mocha/src/Mocha/Middlewares/Receive/Redelivery/ReceiveRedeliveryMiddleware.cs @@ -35,7 +35,7 @@ public async ValueTask InvokeAsync(IReceiveContext context, ReceiveDelegate next { await next(context); } - catch (Exception ex) + catch (Exception ex) when (!context.CancellationToken.IsCancellationRequested) { // Request/reply messages must not be redelivered - the caller is waiting. if (context.Envelope?.ResponseAddress is not null) diff --git a/src/Mocha/src/Mocha/Middlewares/ReceiveContext.cs b/src/Mocha/src/Mocha/Middlewares/ReceiveContext.cs index ce4b2f69036..ab10cd64a84 100644 --- a/src/Mocha/src/Mocha/Middlewares/ReceiveContext.cs +++ b/src/Mocha/src/Mocha/Middlewares/ReceiveContext.cs @@ -261,24 +261,24 @@ public void SetEnvelope(MessageEnvelope envelope) /// /// /// Sets the , , , - /// and properties and initializes the internal feature collection. + /// and properties and initializes the internal + /// feature collection. /// /// The scoped service provider for this receive operation. /// The receive endpoint that owns this operation. /// The messaging runtime providing host info and global options. - /// Reserved for future use; not currently consumed by this method. + /// The token that cancels the receive operation. public void Initialize( IServiceProvider services, ReceiveEndpoint endpoint, IMessagingRuntime runtime, -#pragma warning disable RCS1163 // Unused parameter CancellationToken cancellationToken) -#pragma warning restore RCS1163 // Unused parameter { Services = services; Endpoint = endpoint; Transport = endpoint.Transport; Runtime = runtime; + CancellationToken = cancellationToken; _features.Initialize(); } diff --git a/src/Mocha/src/Mocha/Runtime/MessagingRuntime.cs b/src/Mocha/src/Mocha/Runtime/MessagingRuntime.cs index 46e8c63121a..3456c1103dd 100644 --- a/src/Mocha/src/Mocha/Runtime/MessagingRuntime.cs +++ b/src/Mocha/src/Mocha/Runtime/MessagingRuntime.cs @@ -1,4 +1,5 @@ using System.Collections.Immutable; +using System.Runtime.ExceptionServices; using Microsoft.Extensions.Primitives; using Mocha.Middlewares; @@ -148,22 +149,34 @@ public async ValueTask StartAsync(CancellationToken cancellationToken) /// /// Stops the started transports so their receive endpoints stop consuming messages. A stopped - /// runtime cannot be started again. + /// runtime cannot be started again, and a transport that fails to stop does not prevent the + /// others from stopping; the first failure is rethrown. /// /// A token to cancel the shutdown operations. public async ValueTask StopAsync(CancellationToken cancellationToken) { _stopped = true; + ExceptionDispatchInfo? failure = null; foreach (var transport in transports) { - if (transport.IsStarted) + if (!transport.IsStarted) + { + continue; + } + + try { await transport.StopAsync(this, cancellationToken); } + catch (Exception ex) + { + failure ??= ExceptionDispatchInfo.Capture(ex); + } } IsStarted = false; + failure?.Throw(); } /// diff --git a/src/Mocha/src/Mocha/Transport/MessagingTransport.cs b/src/Mocha/src/Mocha/Transport/MessagingTransport.cs index 9b355e3cea0..c364e0bca50 100644 --- a/src/Mocha/src/Mocha/Transport/MessagingTransport.cs +++ b/src/Mocha/src/Mocha/Transport/MessagingTransport.cs @@ -1,5 +1,6 @@ using System.Collections.Immutable; using System.Diagnostics.CodeAnalysis; +using System.Runtime.ExceptionServices; using Microsoft.Extensions.Primitives; namespace Mocha; @@ -227,7 +228,9 @@ public async ValueTask StartAsync(IMessagingRuntimeContext context, Cancellation } /// - /// Stops the transport by invoking pre-stop hooks and deactivating all receive endpoints. + /// Stops the transport by invoking pre-stop hooks and deactivating all receive endpoints. A failing + /// step does not prevent the remaining steps from running; the first failure is rethrown and the + /// transport stays started. /// /// The runtime context providing access to services and configuration. /// A token to cancel the shutdown sequence. @@ -239,13 +242,39 @@ public async ValueTask StopAsync(IMessagingRuntimeContext context, CancellationT throw ThrowHelper.TransportNotStarted(); } - await OnBeforeStopAsync(cancellationToken); + ExceptionDispatchInfo? failure = null; + + try + { + await OnBeforeStopAsync(cancellationToken); + } + catch (Exception ex) + { + failure = ExceptionDispatchInfo.Capture(ex); + } + foreach (var endpoint in ReceiveEndpoints) { - await endpoint.StopAsync(context, cancellationToken); + try + { + await endpoint.StopAsync(context, cancellationToken); + } + catch (Exception ex) + { + failure ??= ExceptionDispatchInfo.Capture(ex); + } + } + + try + { + await OnAfterStopAsync(cancellationToken); + } + catch (Exception ex) + { + failure ??= ExceptionDispatchInfo.Capture(ex); } - await OnAfterStopAsync(cancellationToken); + failure?.Throw(); IsStarted = false; } diff --git a/src/Mocha/test/Mocha.EntityFrameworkCore.Postgres.Tests/PostgresOutboxIntegrationTests.cs b/src/Mocha/test/Mocha.EntityFrameworkCore.Postgres.Tests/PostgresOutboxIntegrationTests.cs index 2369c8c973b..b71ecde6910 100644 --- a/src/Mocha/test/Mocha.EntityFrameworkCore.Postgres.Tests/PostgresOutboxIntegrationTests.cs +++ b/src/Mocha/test/Mocha.EntityFrameworkCore.Postgres.Tests/PostgresOutboxIntegrationTests.cs @@ -196,7 +196,10 @@ public async Task Outbox_Should_ResumeProcessing_When_WorkerRestartedAfterInterr await db.Database.EnsureCreatedAsync(TestContext.Current.CancellationToken); } - var hostedServices = provider.GetServices().ToList(); + // the runtime is started above and stays running, only the outbox worker is restarted + var hostedServices = provider.GetServices() + .Where(s => s is not MessagingRuntimeHostedService) + .ToList(); foreach (var svc in hostedServices) { await svc.StartAsync(TestContext.Current.CancellationToken); diff --git a/src/Mocha/test/Mocha.TestHelpers/LogRecorder.cs b/src/Mocha/test/Mocha.TestHelpers/LogRecorder.cs new file mode 100644 index 00000000000..5cd80b206dc --- /dev/null +++ b/src/Mocha/test/Mocha.TestHelpers/LogRecorder.cs @@ -0,0 +1,33 @@ +using System.Collections.Concurrent; +using Microsoft.Extensions.Logging; + +namespace Mocha.TestHelpers; + +public sealed class LogRecorder : ILoggerProvider +{ + public ConcurrentQueue Entries { get; } = new(); + + public ILogger CreateLogger(string categoryName) => new Logger(Entries); + + public void Dispose() { } + + private sealed class Logger(ConcurrentQueue entries) : ILogger + { + public IDisposable? BeginScope(TState state) where TState : notnull => null; + + public bool IsEnabled(LogLevel logLevel) => logLevel >= LogLevel.Warning; + + public void Log( + LogLevel logLevel, + EventId eventId, + TState state, + Exception? exception, + Func formatter) + { + if (IsEnabled(logLevel)) + { + entries.Enqueue($"{logLevel}: {formatter(state, exception)}"); + } + } + } +} diff --git a/src/Mocha/test/Mocha.Tests/Consumers/Batching/BatchCollectorTests.cs b/src/Mocha/test/Mocha.Tests/Consumers/Batching/BatchCollectorTests.cs index 6af189dc9b4..23e27c3547a 100644 --- a/src/Mocha/test/Mocha.Tests/Consumers/Batching/BatchCollectorTests.cs +++ b/src/Mocha/test/Mocha.Tests/Consumers/Batching/BatchCollectorTests.cs @@ -54,6 +54,40 @@ public async Task Add_Should_DispatchBatch_When_TimerFires() Assert.Equal(BatchCompletionMode.Time, batch.CompletionMode); } + [Fact] + public async Task TryRemove_Should_ExcludeEntryFromBatch_When_EntryIsBuffered() + { + // arrange + var dispatched = new BatchRecorder(); + await using var collector = CreateCollector(dispatched, opts => opts.MaxBatchSize = 2); + var removed = await collector.Add(CreateContext("removed")); + + // act + var wasRemoved = collector.TryRemove(removed); + await AddEntries(collector, 2); + + // assert + Assert.True(wasRemoved); + Assert.True(await dispatched.WaitAsync(s_timeout), "Batch was not dispatched when MaxBatchSize reached"); + Assert.Equal(["msg-0", "msg-1"], dispatched.Single().Select(e => e.Id)); + } + + [Fact] + public async Task TryRemove_Should_ReturnFalse_When_EntryWasDispatched() + { + // arrange + var dispatched = new BatchRecorder(); + await using var collector = CreateCollector(dispatched, opts => opts.MaxBatchSize = 1); + var entry = await collector.Add(CreateContext("dispatched")); + + // act + var wasRemoved = collector.TryRemove(entry); + + // assert + Assert.False(wasRemoved); + Assert.Single(dispatched.Single()); + } + [Fact] public async Task DisposeAsync_Should_FlushRemaining_When_BufferHasMessages() { diff --git a/src/Mocha/test/Mocha.Tests/Middlewares/Consume/Retry/RetryExecutorTests.cs b/src/Mocha/test/Mocha.Tests/Middlewares/Consume/Retry/RetryExecutorTests.cs index 49d20c177ef..b2d11e1b42b 100644 --- a/src/Mocha/test/Mocha.Tests/Middlewares/Consume/Retry/RetryExecutorTests.cs +++ b/src/Mocha/test/Mocha.Tests/Middlewares/Consume/Retry/RetryExecutorTests.cs @@ -643,30 +643,65 @@ await Assert.ThrowsAsync(() => [Fact] public async Task ExecuteAsync_Should_ThrowOperationCanceled_When_CancellationRequestedDuringDelay() { - // arrange - use a short real delay with pre-cancelled token to verify cancellation propagation - var rules = BuildRules(p => p.On().Retry(3, TimeSpan.FromSeconds(1), RetryBackoffType.Constant)); + // arrange + var rules = BuildRules(p => p.On().Retry(3, TimeSpan.FromSeconds(10), RetryBackoffType.Constant)); using var cts = new CancellationTokenSource(); var counter = new Counter(); - // Pre-cancel so the delay will throw OperationCanceledException immediately - cts.Cancel(); + // act + var execution = RetryExecutor.ExecuteAsync( + rules, + counter, + static s => + { + s.Increment(); + throw new InvalidOperationException("fail"); + }, + onRetry: null, + cts.Token).AsTask(); + await cts.CancelAsync(); - // act & assert - await Assert.ThrowsAnyAsync(() => - RetryExecutor - .ExecuteAsync( - rules, - counter, - static (s) => - { - s.Increment(); - throw new InvalidOperationException("fail"); - }, - onRetry: null, - cts.Token) - .AsTask() - ); + // assert + await Assert.ThrowsAnyAsync(() => execution); + Assert.Equal(1, counter.Count); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task ExecuteAsync_Should_PropagateFailure_When_ActionCancelsProcessing(bool discard) + { + // arrange + var rules = BuildRules(p => + { + if (discard) + { + p.On().Discard(); + } + else + { + p.On().Retry(3, TimeSpan.Zero, RetryBackoffType.Constant); + } + }); + using var cts = new CancellationTokenSource(); + var counter = new Counter(); + var failure = new InvalidOperationException("Cancelled attempt"); + // act + var exception = await Assert.ThrowsAsync(() => RetryExecutor.ExecuteAsync( + rules, + counter, + s => + { + s.Increment(); + cts.Cancel(); + throw failure; + }, + onRetry: null, + cts.Token).AsTask()); + + // assert + Assert.Same(failure, exception); Assert.Equal(1, counter.Count); } diff --git a/src/Mocha/test/Mocha.Tests/Runtime/MessagingRuntimeShutdownTests.cs b/src/Mocha/test/Mocha.Tests/Runtime/MessagingRuntimeShutdownTests.cs new file mode 100644 index 00000000000..9c97f8d736d --- /dev/null +++ b/src/Mocha/test/Mocha.Tests/Runtime/MessagingRuntimeShutdownTests.cs @@ -0,0 +1,155 @@ +using System.Collections.Concurrent; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Hosting; +using Microsoft.Extensions.Logging; +using Mocha.Transport.InMemory; + +namespace Mocha.Tests; + +public sealed class MessagingRuntimeShutdownTests +{ + [Fact] + public async Task StopAsync_Should_CancelEventHandler_When_HostStops() + { + // arrange + var probe = new HandlerProbe(); + var logs = new LogRecorder(); + await using var provider = new ServiceCollection() + .AddSingleton(probe) + .AddLogging(b => b.AddProvider(logs)) + .AddMessageBus() + .AddEventHandler() + .AddInMemory(t => t.Queue("events").Handler()) + .Services.BuildServiceProvider(); + var hostedService = await StartAsync(provider); + await PublishAsync(provider); + await WaitAsync(probe.Started); + + // act + await hostedService.StopAsync(TestContext.Current.CancellationToken); + await WaitAsync(probe.Completed); + + // assert + Assert.True(Assert.Single(probe.Tokens).IsCancellationRequested); + Assert.Empty(logs.Entries); + } + + [Fact] + public async Task StopAsync_Should_CancelBatchHandler_When_HostStops() + { + // arrange + var probe = new HandlerProbe(); + var logs = new LogRecorder(); + await using var provider = new ServiceCollection() + .AddSingleton(probe) + .AddLogging(b => b.AddProvider(logs)) + .AddMessageBus() + .AddBatchHandler(o => o.MaxBatchSize = 1) + .AddInMemory(t => t.Queue("batches").Handler()) + .Services.BuildServiceProvider(); + var hostedService = await StartAsync(provider); + await PublishAsync(provider); + await WaitAsync(probe.Started); + + // act + await hostedService.StopAsync(TestContext.Current.CancellationToken); + await WaitAsync(probe.Completed); + + // assert + Assert.True(Assert.Single(probe.Tokens).IsCancellationRequested); + Assert.Empty(logs.Entries); + } + + [Fact] + public async Task StopAsync_Should_CancelOnlyStoppedEndpointBatch_When_EndpointsShareHandler() + { + // arrange + var probe = new HandlerProbe(); + await using var provider = new ServiceCollection() + .AddSingleton(probe) + .AddMessageBus() + .AddBatchHandler(o => o.MaxBatchSize = 1) + .AddInMemory(t => + { + t.Queue("first").Handler(); + t.Queue("second").Handler(); + }) + .Services.BuildServiceProvider(); + await StartAsync(provider); + await PublishAsync(provider); + await WaitAsync(probe.Started); + await WaitAsync(probe.Started); + var runtime = (MessagingRuntime)provider.GetRequiredService(); + var endpoints = runtime.Transports.Single().ReceiveEndpoints; + var first = endpoints.Single(e => e.Name == "first"); + var second = endpoints.Single(e => e.Name == "second"); + + // act + await first.StopAsync(runtime, TestContext.Current.CancellationToken); + await WaitAsync(probe.Completed); + + // assert + Assert.Single(probe.Tokens, token => token.IsCancellationRequested); + Assert.True(second.IsStarted); + } + + private static async Task StartAsync(IServiceProvider provider) + { + var hostedService = provider.GetServices().OfType().Single(); + await hostedService.StartAsync(TestContext.Current.CancellationToken); + return hostedService; + } + + private static async Task PublishAsync(IServiceProvider provider) + { + using var scope = provider.CreateScope(); + await scope.ServiceProvider.GetRequiredService() + .PublishAsync(new TestEvent(), TestContext.Current.CancellationToken); + } + + private static async Task WaitAsync(SemaphoreSlim signal) + { + if (!await signal.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken)) + { + throw new TimeoutException("The handler did not signal in time."); + } + } + + public sealed class TestEvent; + + public sealed class HandlerProbe + { + public ConcurrentQueue Tokens { get; } = new(); + + public SemaphoreSlim Started { get; } = new(0); + + public SemaphoreSlim Completed { get; } = new(0); + + public async ValueTask RunAsync(CancellationToken cancellationToken) + { + Tokens.Enqueue(cancellationToken); + Started.Release(); + + try + { + await Task.Delay(Timeout.InfiniteTimeSpan, cancellationToken); + } + finally + { + Completed.Release(); + } + } + } + + public sealed class EventHandler(HandlerProbe probe) : IEventHandler + { + public ValueTask HandleAsync(TestEvent message, CancellationToken cancellationToken) + => probe.RunAsync(cancellationToken); + } + + public sealed class BatchHandler(HandlerProbe probe) : IBatchEventHandler + { + public ValueTask HandleAsync(IMessageBatch messages, CancellationToken cancellationToken) + => probe.RunAsync(cancellationToken); + } +} diff --git a/src/Mocha/test/Mocha.Transport.AzureServiceBus.Tests/AzureServiceBusAcknowledgementMiddlewareTests.cs b/src/Mocha/test/Mocha.Transport.AzureServiceBus.Tests/AzureServiceBusAcknowledgementMiddlewareTests.cs index 21cc317d8b8..b8908a371f1 100644 --- a/src/Mocha/test/Mocha.Transport.AzureServiceBus.Tests/AzureServiceBusAcknowledgementMiddlewareTests.cs +++ b/src/Mocha/test/Mocha.Transport.AzureServiceBus.Tests/AzureServiceBusAcknowledgementMiddlewareTests.cs @@ -26,7 +26,7 @@ public async Task InvokeAsync_Should_Complete_When_NextSucceeds(bool isSession) // assert Assert.Equal(1, settlement.CompleteCallCount); Assert.Equal(0, settlement.AbandonCallCount); - Assert.Equal(cts.Token, settlement.LastCancellationToken); + Assert.Equal(CancellationToken.None, settlement.LastCancellationToken); Assert.Equal(entityPath, settlement.ObservedEntityPath); } @@ -50,7 +50,7 @@ public async Task InvokeAsync_Should_AbandonAndRethrow_When_NextFails(bool isSes Assert.Same(handlerException, thrown); Assert.Equal(0, settlement.CompleteCallCount); Assert.Equal(1, settlement.AbandonCallCount); - Assert.Equal(cts.Token, settlement.LastCancellationToken); + Assert.Equal(CancellationToken.None, settlement.LastCancellationToken); Assert.Equal(entityPath, settlement.ObservedEntityPath); } diff --git a/src/Mocha/test/Mocha.Transport.AzureServiceBus.Tests/ReceiveEndpointLifecycleUnitTests.cs b/src/Mocha/test/Mocha.Transport.AzureServiceBus.Tests/ReceiveEndpointLifecycleUnitTests.cs index 3e2c7d7a109..666c852dc37 100644 --- a/src/Mocha/test/Mocha.Transport.AzureServiceBus.Tests/ReceiveEndpointLifecycleUnitTests.cs +++ b/src/Mocha/test/Mocha.Transport.AzureServiceBus.Tests/ReceiveEndpointLifecycleUnitTests.cs @@ -384,6 +384,44 @@ await processor.RaiseProcessErrorAsync( Assert.False(endpoint.IsStarted); } + [Fact] + public async Task StopAsync_Should_StopRemainingEndpoints_When_AnEndpointFailsToStop() + { + // arrange + var client = new FakeServiceBusClient(_ => null); + var services = new ServiceCollection(); + services.AddSingleton(client); + services.AddSingleton(new FakeServiceBusAdministrationClient()); + var builder = services + .AddMessageBus() + .AddConsumer() + .AddAzureServiceBus(t => + { + t.AutoProvision(false); + t.Endpoint("first-ep").Consumer().Queue("first"); + t.Endpoint("second-ep").Consumer().Queue("second"); + }); + await using var bus = await builder.BuildTestBusAsync(); + + var runtime = (MessagingRuntime)bus.Provider.GetRequiredService(); + var transport = runtime.Transports.OfType().Single(); + // the endpoint that is stopped first fails, so every later endpoint depends on the stop continuing + var failing = (AzureServiceBusReceiveEndpoint)transport.ReceiveEndpoints.First(); + var failure = new TimeoutException("stop processing timed out"); + client.CreatedProcessors.Single(p => p.QueueName == failing.Queue.Name).Processor.OnStopProcessing = + () => throw failure; + + // act + var exception = await Assert.ThrowsAsync( + () => runtime.StopAsync(Xunit.TestContext.Current.CancellationToken).AsTask()); + + // assert + Assert.Same(failure, exception); + Assert.All(transport.ReceiveEndpoints.Where(e => e != failing), e => Assert.False(e.IsStarted)); + Assert.True(transport.IsStarted); + Assert.False(runtime.IsStarted); + } + public sealed class NoOpConsumer : IConsumer { public ValueTask ConsumeAsync(IConsumeContext context) => default; diff --git a/src/Mocha/test/Mocha.Transport.Postgres.Tests/Behaviors/ShutdownTests.cs b/src/Mocha/test/Mocha.Transport.Postgres.Tests/Behaviors/ShutdownTests.cs index cf41f30e66e..ec80b0d2e68 100644 --- a/src/Mocha/test/Mocha.Transport.Postgres.Tests/Behaviors/ShutdownTests.cs +++ b/src/Mocha/test/Mocha.Transport.Postgres.Tests/Behaviors/ShutdownTests.cs @@ -1,5 +1,6 @@ using Microsoft.Extensions.DependencyInjection; using Microsoft.Extensions.Hosting; +using Microsoft.Extensions.Logging; using Mocha.Transport.Postgres.Tests.Helpers; using Npgsql; @@ -13,7 +14,7 @@ public async Task StopAsync_Should_UnregisterConsumer_When_HostStops() { // arrange await using var db = await fixture.CreateDatabaseAsync("shutdown_unregister"); - await using var provider = CreateProvider(db.ConnectionString); + await using var provider = CreateProvider(db.ConnectionString, new HandlerGate(), new LogRecorder()); var services = provider.GetServices().ToArray(); await StartAsync(services); var consumersBeforeStop = await CountConsumersAsync(db.ConnectionString); @@ -26,14 +27,65 @@ public async Task StopAsync_Should_UnregisterConsumer_When_HostStops() Assert.Equal(0, await CountConsumersAsync(db.ConnectionString)); } - private static ServiceProvider CreateProvider(string connectionString) + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task StopAsync_Should_ReleaseInterruptedMessages_When_HostStops(bool failAfterCancellation) + { + // arrange + await using var db = await fixture.CreateDatabaseAsync($"shutdown_release_{failAfterCancellation}"); + var gate = new HandlerGate { FailAfterCancellation = failAfterCancellation }; + var logs = new LogRecorder(); + await using var provider = CreateProvider(db.ConnectionString, gate, logs); + var services = provider.GetServices().ToArray(); + await StartAsync(services); + await PublishAsync(provider, count: 2); + await gate.Started.Task.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken); + + // act + await StopAsync(services); + + // assert + Assert.True(gate.Cancelled); + Assert.Equal([(false, null), (false, null)], await ReadMessagesAsync(db.ConnectionString)); + Assert.Empty(logs.Entries); + } + + [Fact] + public async Task StopAsync_Should_DeleteMessage_When_HandlerCompletesAfterStopIsRequested() + { + // arrange + await using var db = await fixture.CreateDatabaseAsync("shutdown_complete"); + var gate = new HandlerGate(); + var logs = new LogRecorder(); + await using var provider = CreateProvider(db.ConnectionString, gate, logs); + var services = provider.GetServices().ToArray(); + await StartAsync(services); + await PublishAsync(provider, count: 1); + await gate.Started.Task.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken); + + // act + await StopAsync(services); + + // assert + Assert.Empty(await ReadMessagesAsync(db.ConnectionString)); + Assert.Empty(logs.Entries); + } + + private static ServiceProvider CreateProvider( + string connectionString, + HandlerGate gate, + LogRecorder logs) + where THandler : class, IEventHandler => new ServiceCollection() + .AddSingleton(gate) + .AddLogging(b => b.AddProvider(logs)) .AddMessageBus() - .AddEventHandler() + .AddEventHandler() .AddPostgres(t => { t.ConnectionString(connectionString); - t.Endpoint("shutdown").Handler(); + t.Endpoint("shutdown").Handler().MaxConcurrency(1); }) .Services.BuildServiceProvider(); @@ -49,7 +101,18 @@ private static async Task StopAsync(IHostedService[] services) { foreach (var service in services.Reverse()) { - await service.StopAsync(TestContext.Current.CancellationToken); + await service.StopAsync(CancellationToken.None) + .WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken); + } + } + + private static async Task PublishAsync(IServiceProvider provider, int count) + { + using var scope = provider.CreateScope(); + var bus = scope.ServiceProvider.GetRequiredService(); + for (var i = 0; i < count; i++) + { + await bus.PublishAsync(new TestEvent(), TestContext.Current.CancellationToken); } } @@ -62,11 +125,68 @@ private static async Task CountConsumersAsync(string connectionString) return (long)(await command.ExecuteScalarAsync(TestContext.Current.CancellationToken))!; } + private static async Task> ReadMessagesAsync(string connectionString) + { + await using var connection = new NpgsqlConnection(connectionString); + await connection.OpenAsync(TestContext.Current.CancellationToken); + await using var command = connection.CreateCommand(); + command.CommandText = "SELECT consumer_id IS NOT NULL, error_reason FROM mocha_message"; + await using var reader = await command.ExecuteReaderAsync(TestContext.Current.CancellationToken); + var messages = new List<(bool Leased, string? ErrorReason)>(); + while (await reader.ReadAsync(TestContext.Current.CancellationToken)) + { + messages.Add((reader.GetBoolean(0), reader.IsDBNull(1) ? null : reader.GetString(1))); + } + + return messages; + } + public sealed class TestEvent; + public sealed class HandlerGate + { + public TaskCompletionSource Started { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously); + public bool Cancelled { get; set; } + public bool FailAfterCancellation { get; init; } + } + public sealed class NoOpHandler : IEventHandler { public ValueTask HandleAsync(TestEvent message, CancellationToken cancellationToken) => ValueTask.CompletedTask; } + + public sealed class BlockingHandler(HandlerGate gate) : IEventHandler + { + public async ValueTask HandleAsync(TestEvent message, CancellationToken cancellationToken) + { + gate.Started.TrySetResult(); + try + { + await Task.Delay(Timeout.InfiniteTimeSpan, cancellationToken); + } + catch (OperationCanceledException) + { + gate.Cancelled = true; + if (gate.FailAfterCancellation) + { + throw new InvalidOperationException("Handler failed after shutdown cancellation."); + } + + throw; + } + } + } + + public sealed class CompletingHandler(HandlerGate gate) : IEventHandler + { + public async ValueTask HandleAsync(TestEvent message, CancellationToken cancellationToken) + { + // completes successfully once the endpoint starts stopping + var stopping = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + await using var registration = cancellationToken.Register(() => stopping.TrySetResult()); + gate.Started.TrySetResult(); + await stopping.Task; + } + } } diff --git a/src/Mocha/test/Mocha.Transport.RabbitMQ.Tests/Behaviors/ShutdownTests.cs b/src/Mocha/test/Mocha.Transport.RabbitMQ.Tests/Behaviors/ShutdownTests.cs new file mode 100644 index 00000000000..7117164d9a1 --- /dev/null +++ b/src/Mocha/test/Mocha.Transport.RabbitMQ.Tests/Behaviors/ShutdownTests.cs @@ -0,0 +1,65 @@ +using Microsoft.Extensions.DependencyInjection; +using Mocha.Transport.RabbitMQ.Tests.Helpers; + +namespace Mocha.Transport.RabbitMQ.Tests.Behaviors; + +[Collection("RabbitMQ")] +public sealed class ShutdownTests(RabbitMQFixture fixture) +{ + [Fact] + public async Task StopAsync_Should_CancelHandlerAndRequeueMessage_When_EndpointStops() + { + // arrange + await using var vhost = await fixture.CreateVhostAsync(); + var gate = new HandlerGate(); + await using var bus = await new ServiceCollection() + .AddSingleton(vhost.ConnectionFactory) + .AddSingleton(gate) + .AddMessageBus() + .AddEventHandler() + .AddRabbitMQ(t => t.Endpoint("shutdown").Handler()) + .BuildTestBusAsync(); + var runtime = (MessagingRuntime)bus.Provider.GetRequiredService(); + using var scope = bus.Provider.CreateScope(); + await scope.ServiceProvider.GetRequiredService() + .PublishAsync(new TestEvent(), TestContext.Current.CancellationToken); + await gate.Started.Task.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken); + + // act + await runtime.StopAsync(CancellationToken.None).AsTask() + .WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken); + + // assert + await using var connection = await vhost.ConnectionFactory.CreateConnectionAsync(TestContext.Current.CancellationToken); + await using var channel = await connection.CreateChannelAsync(cancellationToken: TestContext.Current.CancellationToken); + var message = await channel.BasicGetAsync("shutdown", autoAck: true, TestContext.Current.CancellationToken); + Assert.True(gate.Cancelled); + Assert.True(message?.Redelivered); + Assert.False(runtime.IsStarted); + } + + public sealed class TestEvent; + + public sealed class HandlerGate + { + public TaskCompletionSource Started { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously); + public bool Cancelled { get; set; } + } + + public sealed class BlockingHandler(HandlerGate gate) : IEventHandler + { + public async ValueTask HandleAsync(TestEvent message, CancellationToken cancellationToken) + { + gate.Started.TrySetResult(); + try + { + await Task.Delay(Timeout.InfiniteTimeSpan, cancellationToken); + } + catch (OperationCanceledException) + { + gate.Cancelled = true; + throw; + } + } + } +}