< Summary - AsyncResponse (Release / net8.0+net10.0 / unit+integration)

Information
Class: AsyncResponse.Transports.RabbitMQ.QueuedRabbitMqMessageDispatcher
Assembly: AsyncResponse.Transports.RabbitMQ
File(s): /home/runner/work/AsyncResponse/AsyncResponse/src/Transports/AsyncResponse.Transports.RabbitMQ/RabbitMqMessageDispatcher.cs
Line coverage
100%
Covered lines: 94
Uncovered lines: 0
Coverable lines: 94
Total lines: 449
Line coverage: 100%
Branch coverage
100%
Covered branches: 4
Total branches: 4
Branch coverage: 100%
Method coverage

Feature is only available for sponsors

Upgrade to PRO version

Metrics

MethodBranch coverage Crap Score Cyclomatic complexity Line coverage
.ctor(...)100%11100%
get_PendingCount()100%11100%
get_RunningCount()100%11100%
HandleAsync()100%22100%
DisposeAsync()100%22100%
RunWorkerAsync()100%11100%

File(s)

/home/runner/work/AsyncResponse/AsyncResponse/src/Transports/AsyncResponse.Transports.RabbitMQ/RabbitMqMessageDispatcher.cs

#LineLine coverage
 1using Microsoft.Extensions.Logging;
 2using RabbitMQ.Client;
 3using System.Collections;
 4using System.Diagnostics;
 5using System.Threading.Channels;
 6
 7namespace AsyncResponse.Transports.RabbitMQ;
 8
 9internal enum RabbitMqSubscriberRole
 10{
 11    Worker,
 12    ResponseIngress
 13}
 14
 15internal abstract class RabbitMqMessageDispatcher : IAsyncDisposable
 16{
 17    private readonly Func<RabbitMqDelivery, CancellationToken, Task> _handler;
 18    private readonly RabbitMqSubscriberOptions _subscriberOptions;
 19    private readonly string _queue;
 20    private readonly RabbitMqSubscriberRole _role;
 21
 22    /// <summary>Runs the RabbitMqMessageDispatcher operation.</summary>
 23    protected RabbitMqMessageDispatcher(
 24        Func<RabbitMqDelivery, CancellationToken, Task> handler,
 25        RabbitMqAsyncResponseOptions transportOptions,
 26        RabbitMqSubscriberOptions subscriberOptions,
 27        ILogger logger,
 28        string queue,
 29        RabbitMqSubscriberRole role)
 30    {
 31        _handler = handler;
 32        TransportOptions = transportOptions;
 33        _subscriberOptions = subscriberOptions;
 34        Logger = logger;
 35        _queue = queue;
 36        _role = role;
 37    }
 38
 39    protected RabbitMqAsyncResponseOptions TransportOptions { get; }
 40    protected ILogger Logger { get; }
 41
 42    /// <summary>
 43    /// Maximum delivery attempts before a failing <see cref="RabbitMqAckMode.AckAfterHandlerCompletes"/> handler
 44    /// rejects without requeue. <c>0</c> means unlimited (requeue forever).
 45    /// </summary>
 46    protected int MaxDeliveryAttempts => _subscriberOptions.MaxDeliveryAttempts;
 47
 48    /// <summary>
 49    /// Resolves the 1-based delivery attempt for a message from the broker's <c>x-death</c> count and the
 50    /// <c>redelivered</c> flag. A message seen for the first time is attempt 1.
 51    /// </summary>
 52    internal static int ResolveDeliveryAttempt(RabbitMqDelivery delivery)
 53    {
 54        var priorAttempts = Math.Max(ReadDeathCount(delivery.BasicProperties), delivery.Redelivered ? 1L : 0L);
 55        var attempt = priorAttempts + 1;
 56        return attempt > int.MaxValue ? int.MaxValue : (int)attempt;
 57    }
 58
 59    private static long ReadDeathCount(IReadOnlyBasicProperties properties)
 60    {
 61        if (properties.Headers is null
 62            || !properties.Headers.TryGetValue("x-death", out var raw)
 63            || raw is not IEnumerable entries)
 64        {
 65            return 0;
 66        }
 67
 68        long max = 0;
 69        foreach (var entry in entries)
 70        {
 71            if (entry is not IDictionary fields || !fields.Contains("count") || fields["count"] is not { } countValue)
 72                continue;
 73
 74            try
 75            {
 76                max = Math.Max(max, Convert.ToInt64(countValue));
 77            }
 78            catch (Exception ex) when (ex is FormatException or InvalidCastException or OverflowException)
 79            {
 80                // Ignore malformed x-death entries; fall back to the redelivered flag.
 81            }
 82        }
 83
 84        return max;
 85    }
 86
 87    /// <summary>Creates the configured dispatcher.</summary>
 88    public static RabbitMqMessageDispatcher Create(
 89        Func<RabbitMqDelivery, CancellationToken, Task> handler,
 90        RabbitMqAsyncResponseOptions transportOptions,
 91        RabbitMqSubscriberOptions subscriberOptions,
 92        ILogger logger,
 93        string queue,
 94        RabbitMqSubscriberRole role)
 95    {
 96        ValidateOptions(transportOptions, subscriberOptions, role);
 97
 98        return subscriberOptions.AckMode == RabbitMqAckMode.AckAfterHandlerCompletes
 99            ? new AwaitingRabbitMqMessageDispatcher(
 100                handler,
 101                transportOptions,
 102                subscriberOptions,
 103                logger,
 104                queue,
 105                role)
 106            : new QueuedRabbitMqMessageDispatcher(
 107                handler,
 108                transportOptions,
 109                subscriberOptions,
 110                logger,
 111                queue,
 112                role);
 113    }
 114
 115    /// <summary>Validates the supplied options.</summary>
 116    public static void ValidateOptions(
 117        RabbitMqAsyncResponseOptions transportOptions,
 118        RabbitMqSubscriberOptions subscriberOptions,
 119        RabbitMqSubscriberRole role)
 120    {
 121        var optionPath = role is RabbitMqSubscriberRole.Worker
 122            ? $"{nameof(RabbitMqAsyncResponseOptions)}.{nameof(RabbitMqAsyncResponseOptions.WorkerSubscriber)}"
 123            : $"{nameof(RabbitMqAsyncResponseOptions)}.{nameof(RabbitMqAsyncResponseOptions.ResponseSubscriber)}";
 124
 125        if (StringComparer.Ordinal.Equals(transportOptions.WorkerQueue, transportOptions.ResponseQueue))
 126        {
 127            throw new InvalidOperationException(
 128                $"{nameof(RabbitMqAsyncResponseOptions)}.{nameof(RabbitMqAsyncResponseOptions.WorkerQueue)} and " +
 129                $"{nameof(RabbitMqAsyncResponseOptions.ResponseQueue)} must be distinct so worker and response subscribe
 130        }
 131
 132        if (subscriberOptions.PrefetchCount == 0)
 133            throw new InvalidOperationException($"{optionPath}.{nameof(RabbitMqSubscriberOptions.PrefetchCount)} must be
 134
 135        switch (subscriberOptions.AckMode)
 136        {
 137            case RabbitMqAckMode.AckAfterHandlerCompletes:
 138                return;
 139
 140            case RabbitMqAckMode.AckAfterEnqueue:
 141                if (subscriberOptions.BackgroundWorkerCount <= 0)
 142                {
 143                    throw new InvalidOperationException(
 144                        $"{optionPath}.{nameof(RabbitMqSubscriberOptions.BackgroundWorkerCount)} must be explicitly conf
 145                        $"when {nameof(RabbitMqSubscriberOptions.AckMode)} is {nameof(RabbitMqAckMode.AckAfterEnqueue)}.
 146                }
 147
 148                if (subscriberOptions.BackgroundQueueCapacity <= 0)
 149                {
 150                    throw new InvalidOperationException(
 151                        $"{optionPath}.{nameof(RabbitMqSubscriberOptions.BackgroundQueueCapacity)} must be explicitly co
 152                        $"when {nameof(RabbitMqSubscriberOptions.AckMode)} is {nameof(RabbitMqAckMode.AckAfterEnqueue)}.
 153                }
 154
 155                if (subscriberOptions.BackgroundDrainTimeout <= TimeSpan.Zero)
 156                    throw new InvalidOperationException($"{optionPath}.{nameof(RabbitMqSubscriberOptions.BackgroundDrain
 157
 158                if (transportOptions.ShutdownTimeout <= TimeSpan.Zero)
 159                    throw new InvalidOperationException($"{nameof(RabbitMqAsyncResponseOptions)}.{nameof(RabbitMqAsyncRe
 160
 161                // RabbitMQ spends the background drain plus the bounded connection close
 162                // (ShutdownTimeout) at shutdown; both must fit inside the host budget.
 163                ShutdownBudgetValidator.Validate(
 164                    "RabbitMQ",
 165                    $"{nameof(RabbitMqAsyncResponseOptions)}.{nameof(RabbitMqAsyncResponseOptions.HostShutdownTimeout)}"
 166                    transportOptions.HostShutdownTimeout,
 167                    ($"{optionPath}.{nameof(RabbitMqSubscriberOptions.BackgroundDrainTimeout)}", subscriberOptions.Backg
 168                    ($"{nameof(RabbitMqAsyncResponseOptions)}.{nameof(RabbitMqAsyncResponseOptions.ShutdownTimeout)}", t
 169
 170                return;
 171
 172            default:
 173                throw new InvalidOperationException(
 174                    $"{optionPath}.{nameof(RabbitMqSubscriberOptions.AckMode)} has unsupported value '{subscriberOptions
 175        }
 176    }
 177
 178    /// <summary>Handles the delivered message.</summary>
 179    public abstract Task HandleAsync(
 180        RabbitMqDelivery delivery,
 181        IRabbitMqChannel channel,
 182        CancellationToken subscriberCancellationToken);
 183
 184    /// <summary>Releases resources held by this instance.</summary>
 185    public virtual ValueTask DisposeAsync() => ValueTask.CompletedTask;
 186
 187    /// <summary>Runs the ExecuteHandlerAsync operation.</summary>
 188    protected async Task ExecuteHandlerAsync(
 189        RabbitMqDelivery delivery,
 190        CancellationToken cancellationToken,
 191        bool logFailures = true)
 192    {
 193        using var activity = AsyncResponseDiagnostics.StartActivity(
 194            "asyncresponse.rabbitmq.receive",
 195            ActivityKind.Consumer);
 196        activity?.SetTag("asyncresponse.transport", "rabbitmq");
 197        activity?.SetTag("asyncresponse.rabbitmq.role", _role.ToString());
 198        activity?.SetTag("asyncresponse.rabbitmq.ack_mode", _subscriberOptions.AckMode.ToString());
 199        activity?.SetTag("messaging.system", "rabbitmq");
 200        activity?.SetTag("messaging.destination.name", _queue);
 201        activity?.SetTag("messaging.rabbitmq.exchange", delivery.Exchange);
 202        activity?.SetTag("messaging.rabbitmq.routing_key", delivery.RoutingKey);
 203        activity?.SetTag("messaging.rabbitmq.delivery_tag", delivery.DeliveryTag);
 204        activity?.SetTag("messaging.message.id", delivery.BasicProperties.MessageId);
 205
 206        if (!string.IsNullOrWhiteSpace(delivery.BasicProperties.CorrelationId))
 207            AsyncResponseDiagnostics.SetCorrelationId(activity, delivery.BasicProperties.CorrelationId);
 208
 209        try
 210        {
 211            await _handler(delivery, cancellationToken).ConfigureAwait(false);
 212        }
 213        catch (Exception ex)
 214        {
 215            if (logFailures)
 216                Logger.LogError(ex, "RabbitMQ message handling failed for delivery {DeliveryTag}.", delivery.DeliveryTag
 217            AsyncResponseDiagnostics.SetError(activity, ex);
 218            throw;
 219        }
 220    }
 221
 222    /// <summary>Runs the NotifyBackgroundFailureAsync operation.</summary>
 223    protected async ValueTask NotifyBackgroundFailureAsync(
 224        RabbitMqDelivery delivery,
 225        Exception exception,
 226        string queue,
 227        RabbitMqSubscriberRole role)
 228    {
 229        var callback = _subscriberOptions.OnBackgroundFailure;
 230        if (callback is null)
 231            return;
 232
 233        try
 234        {
 235            await callback(new RabbitMqBackgroundFailureContext(
 236                queue,
 237                role.ToString(),
 238                delivery.Exchange,
 239                delivery.RoutingKey,
 240                delivery.DeliveryTag,
 241                exception)).ConfigureAwait(false);
 242        }
 243        catch (Exception callbackException)
 244        {
 245            Logger.LogError(
 246                callbackException,
 247                "RabbitMQ background failure callback failed for already-ACKed delivery {DeliveryTag} on {Queue}.",
 248                delivery.DeliveryTag,
 249                queue);
 250        }
 251    }
 252}
 253
 254internal sealed class AwaitingRabbitMqMessageDispatcher(
 255    Func<RabbitMqDelivery, CancellationToken, Task> handler,
 256    RabbitMqAsyncResponseOptions transportOptions,
 257    RabbitMqSubscriberOptions subscriberOptions,
 258    ILogger logger,
 259    string queue,
 260    RabbitMqSubscriberRole role)
 261    : RabbitMqMessageDispatcher(handler, transportOptions, subscriberOptions, logger, queue, role)
 262{
 263    /// <summary>Handles the delivered message.</summary>
 264    public override async Task HandleAsync(
 265        RabbitMqDelivery delivery,
 266        IRabbitMqChannel channel,
 267        CancellationToken subscriberCancellationToken)
 268    {
 269        try
 270        {
 271            await ExecuteHandlerAsync(delivery, subscriberCancellationToken).ConfigureAwait(false);
 272            await channel.BasicAckAsync(delivery.DeliveryTag, subscriberCancellationToken).ConfigureAwait(false);
 273        }
 274        catch
 275        {
 276            // Requeue for redelivery, unless a delivery cap is configured and this delivery has reached it —
 277            // then reject without requeue so the broker dead-letters (or drops) it instead of hot-looping.
 278            var requeue = MaxDeliveryAttempts <= 0
 279                || ResolveDeliveryAttempt(delivery) < MaxDeliveryAttempts;
 280            await channel.BasicNackAsync(delivery.DeliveryTag, requeue, CancellationToken.None).ConfigureAwait(false);
 281        }
 282    }
 283}
 284
 285internal sealed class QueuedRabbitMqMessageDispatcher : RabbitMqMessageDispatcher
 286{
 287    private readonly Channel<RabbitMqDelivery> _queue;
 288    private readonly Task[] _workers;
 3289    private readonly CancellationTokenSource _drainCancellation = new();
 290    private readonly TimeSpan _drainTimeout;
 291    private readonly string _queueName;
 292    private readonly RabbitMqSubscriberRole _role;
 293    private int _pendingCount;
 294    private int _runningCount;
 295    private int _disposeStarted;
 296
 297    /// <summary>Runs the QueuedRabbitMqMessageDispatcher operation.</summary>
 298    public QueuedRabbitMqMessageDispatcher(
 299        Func<RabbitMqDelivery, CancellationToken, Task> handler,
 300        RabbitMqAsyncResponseOptions transportOptions,
 301        RabbitMqSubscriberOptions subscriberOptions,
 302        ILogger logger,
 303        string queue,
 304        RabbitMqSubscriberRole role)
 3305        : base(handler, transportOptions, subscriberOptions, logger, queue, role)
 306    {
 3307        _drainTimeout = subscriberOptions.BackgroundDrainTimeout;
 3308        _queueName = queue;
 3309        _role = role;
 3310        _queue = Channel.CreateBounded<RabbitMqDelivery>(new BoundedChannelOptions(subscriberOptions.BackgroundQueueCapa
 3311        {
 3312            AllowSynchronousContinuations = false,
 3313            FullMode = BoundedChannelFullMode.Wait,
 3314            SingleReader = subscriberOptions.BackgroundWorkerCount == 1,
 3315            SingleWriter = false
 3316        });
 317
 3318        _workers = Enumerable.Range(0, subscriberOptions.BackgroundWorkerCount)
 3319            .Select(workerIndex => Task.Run(() => RunWorkerAsync(workerIndex)))
 3320            .ToArray();
 321
 3322        Logger.LogInformation(
 3323            "Created RabbitMQ ACK-after-enqueue dispatcher for {Queue} with {WorkerCount} worker(s), queue capacity {Que
 3324            _queueName,
 3325            subscriberOptions.BackgroundWorkerCount,
 3326            subscriberOptions.BackgroundQueueCapacity,
 3327            _drainTimeout);
 3328    }
 329
 3330    internal int PendingCount => Volatile.Read(ref _pendingCount);
 3331    internal int RunningCount => Volatile.Read(ref _runningCount);
 332
 333    /// <summary>Handles the delivered message.</summary>
 334    public override async Task HandleAsync(
 335        RabbitMqDelivery delivery,
 336        IRabbitMqChannel channel,
 337        CancellationToken subscriberCancellationToken)
 338    {
 339        // The client owns the delivery body's memory only until the consumer callback returns
 340        // ("Accessing the body at a later point is unsafe as its memory can be already
 341        // released" — RabbitMQ.Client v7). This dispatcher hands the delivery to background
 342        // workers that read the body after the callback, so materialize a private copy now.
 343        // The awaiting dispatcher consumes the body inside the callback and stays zero-copy.
 3344        delivery = delivery with { Body = delivery.Body.ToArray() };
 345
 3346        Interlocked.Increment(ref _pendingCount);
 3347        if (!_queue.Writer.TryWrite(delivery))
 348        {
 2349            Interlocked.Decrement(ref _pendingCount);
 2350            Logger.LogWarning(
 2351                "RabbitMQ background queue rejected delivery {DeliveryTag} for {Queue}; returning NACK. Pending={Pending
 2352                delivery.DeliveryTag,
 2353                _queueName,
 2354                PendingCount,
 2355                RunningCount);
 2356            await channel.BasicNackAsync(delivery.DeliveryTag, requeue: true, subscriberCancellationToken).ConfigureAwai
 3357            return;
 358        }
 359
 360        // The delivery now belongs to a background worker, which decrements _pendingCount when it dequeues.
 361        // Do not touch the counter or NACK here, even if the ACK below fails — the message is already
 362        // executing in-process and a NACK would trigger a duplicate execution via requeue.
 363        try
 364        {
 3365            await channel.BasicAckAsync(delivery.DeliveryTag, subscriberCancellationToken).ConfigureAwait(false);
 3366        }
 2367        catch (Exception ex)
 368        {
 2369            Logger.LogError(
 2370                ex,
 2371                "Failed to ACK RabbitMQ delivery {DeliveryTag} for {Queue} after enqueue; it is being processed but the 
 2372                delivery.DeliveryTag,
 2373                _queueName);
 3374        }
 3375    }
 376
 377    /// <summary>Releases resources held by this instance.</summary>
 378    public override async ValueTask DisposeAsync()
 379    {
 3380        if (Interlocked.Exchange(ref _disposeStarted, 1) != 0)
 3381            return;
 382
 3383        Logger.LogInformation(
 3384            "Draining RabbitMQ ACK-after-enqueue dispatcher for {Queue}. Pending={PendingCount}, Running={RunningCount}.
 3385            _queueName,
 3386            PendingCount,
 3387            RunningCount);
 3388        _queue.Writer.TryComplete();
 389
 390        try
 391        {
 3392            await Task.WhenAll(_workers).WaitAsync(_drainTimeout).ConfigureAwait(false);
 3393            _drainCancellation.Dispose();
 3394        }
 2395        catch (TimeoutException ex)
 396        {
 2397            _drainCancellation.Cancel();
 2398            Logger.LogWarning(
 2399                ex,
 2400                "Timed out while draining RabbitMQ ACK-after-enqueue dispatcher for {Queue}. Pending={PendingCount}, Run
 2401                _queueName,
 2402                PendingCount,
 2403                RunningCount);
 404
 405            // The workers are still running and read _drainCancellation.Token each loop, so disposing it now
 406            // would throw ObjectDisposedException inside them. Dispose once they actually finish, off the
 407            // shutdown path, so the source is not leaked either.
 2408            _ = Task.WhenAll(_workers).ContinueWith(
 3409                _ => _drainCancellation.Dispose(),
 2410                CancellationToken.None,
 2411                TaskContinuationOptions.ExecuteSynchronously,
 2412                TaskScheduler.Default);
 3413        }
 3414    }
 415
 416    private async Task RunWorkerAsync(int workerIndex)
 417    {
 3418        await foreach (var delivery in _queue.Reader.ReadAllAsync().ConfigureAwait(false))
 419        {
 3420            Interlocked.Decrement(ref _pendingCount);
 3421            Interlocked.Increment(ref _runningCount);
 422
 423            try
 424            {
 3425                await ExecuteHandlerAsync(
 3426                    delivery,
 3427                    _drainCancellation.Token,
 3428                    logFailures: false).ConfigureAwait(false);
 3429            }
 3430            catch (Exception ex)
 431            {
 2432                Logger.LogError(
 2433                    ex,
 2434                    "RabbitMQ background handler failed for already-ACKed delivery {DeliveryTag} on {Queue}.",
 2435                    delivery.DeliveryTag,
 2436                    _queueName);
 2437                await NotifyBackgroundFailureAsync(
 2438                    delivery,
 2439                    ex,
 2440                    _queueName,
 2441                    _role).ConfigureAwait(false);
 442            }
 443            finally
 444            {
 3445                Interlocked.Decrement(ref _runningCount);
 446            }
 3447        }
 3448    }
 449}