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

Information
Class: AsyncResponse.Transports.RabbitMQ.RabbitMqMessageDispatcher
Assembly: AsyncResponse.Transports.RabbitMQ
File(s): /home/runner/work/AsyncResponse/AsyncResponse/src/Transports/AsyncResponse.Transports.RabbitMQ/RabbitMqMessageDispatcher.cs
Line coverage
100%
Covered lines: 120
Uncovered lines: 0
Coverable lines: 120
Total lines: 449
Line coverage: 100%
Branch coverage
83%
Covered branches: 50
Total branches: 60
Branch coverage: 83.3%
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_MaxDeliveryAttempts()100%11100%
ResolveDeliveryAttempt(...)50%22100%
ReadDeathCount(...)100%1414100%
Create(...)100%22100%
ValidateOptions(...)100%1818100%
DisposeAsync()100%11100%
ExecuteHandlerAsync()59.09%2222100%
NotifyBackgroundFailureAsync()100%22100%

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>
 323    protected RabbitMqMessageDispatcher(
 324        Func<RabbitMqDelivery, CancellationToken, Task> handler,
 325        RabbitMqAsyncResponseOptions transportOptions,
 326        RabbitMqSubscriberOptions subscriberOptions,
 327        ILogger logger,
 328        string queue,
 329        RabbitMqSubscriberRole role)
 30    {
 331        _handler = handler;
 332        TransportOptions = transportOptions;
 333        _subscriberOptions = subscriberOptions;
 334        Logger = logger;
 335        _queue = queue;
 336        _role = role;
 337    }
 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>
 246    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    {
 254        var priorAttempts = Math.Max(ReadDeathCount(delivery.BasicProperties), delivery.Redelivered ? 1L : 0L);
 255        var attempt = priorAttempts + 1;
 256        return attempt > int.MaxValue ? int.MaxValue : (int)attempt;
 57    }
 58
 59    private static long ReadDeathCount(IReadOnlyBasicProperties properties)
 60    {
 261        if (properties.Headers is null
 262            || !properties.Headers.TryGetValue("x-death", out var raw)
 263            || raw is not IEnumerable entries)
 64        {
 265            return 0;
 66        }
 67
 268        long max = 0;
 269        foreach (var entry in entries)
 70        {
 271            if (entry is not IDictionary fields || !fields.Contains("count") || fields["count"] is not { } countValue)
 72                continue;
 73
 74            try
 75            {
 276                max = Math.Max(max, Convert.ToInt64(countValue));
 277            }
 278            catch (Exception ex) when (ex is FormatException or InvalidCastException or OverflowException)
 79            {
 80                // Ignore malformed x-death entries; fall back to the redelivered flag.
 281            }
 82        }
 83
 384        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    {
 396        ValidateOptions(transportOptions, subscriberOptions, role);
 97
 398        return subscriberOptions.AckMode == RabbitMqAckMode.AckAfterHandlerCompletes
 399            ? new AwaitingRabbitMqMessageDispatcher(
 3100                handler,
 3101                transportOptions,
 3102                subscriberOptions,
 3103                logger,
 3104                queue,
 3105                role)
 3106            : new QueuedRabbitMqMessageDispatcher(
 3107                handler,
 3108                transportOptions,
 3109                subscriberOptions,
 3110                logger,
 3111                queue,
 3112                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    {
 3121        var optionPath = role is RabbitMqSubscriberRole.Worker
 3122            ? $"{nameof(RabbitMqAsyncResponseOptions)}.{nameof(RabbitMqAsyncResponseOptions.WorkerSubscriber)}"
 3123            : $"{nameof(RabbitMqAsyncResponseOptions)}.{nameof(RabbitMqAsyncResponseOptions.ResponseSubscriber)}";
 124
 3125        if (StringComparer.Ordinal.Equals(transportOptions.WorkerQueue, transportOptions.ResponseQueue))
 126        {
 3127            throw new InvalidOperationException(
 3128                $"{nameof(RabbitMqAsyncResponseOptions)}.{nameof(RabbitMqAsyncResponseOptions.WorkerQueue)} and " +
 3129                $"{nameof(RabbitMqAsyncResponseOptions.ResponseQueue)} must be distinct so worker and response subscribe
 130        }
 131
 3132        if (subscriberOptions.PrefetchCount == 0)
 3133            throw new InvalidOperationException($"{optionPath}.{nameof(RabbitMqSubscriberOptions.PrefetchCount)} must be
 134
 3135        switch (subscriberOptions.AckMode)
 136        {
 137            case RabbitMqAckMode.AckAfterHandlerCompletes:
 3138                return;
 139
 140            case RabbitMqAckMode.AckAfterEnqueue:
 3141                if (subscriberOptions.BackgroundWorkerCount <= 0)
 142                {
 3143                    throw new InvalidOperationException(
 3144                        $"{optionPath}.{nameof(RabbitMqSubscriberOptions.BackgroundWorkerCount)} must be explicitly conf
 3145                        $"when {nameof(RabbitMqSubscriberOptions.AckMode)} is {nameof(RabbitMqAckMode.AckAfterEnqueue)}.
 146                }
 147
 3148                if (subscriberOptions.BackgroundQueueCapacity <= 0)
 149                {
 3150                    throw new InvalidOperationException(
 3151                        $"{optionPath}.{nameof(RabbitMqSubscriberOptions.BackgroundQueueCapacity)} must be explicitly co
 3152                        $"when {nameof(RabbitMqSubscriberOptions.AckMode)} is {nameof(RabbitMqAckMode.AckAfterEnqueue)}.
 153                }
 154
 3155                if (subscriberOptions.BackgroundDrainTimeout <= TimeSpan.Zero)
 3156                    throw new InvalidOperationException($"{optionPath}.{nameof(RabbitMqSubscriberOptions.BackgroundDrain
 157
 3158                if (transportOptions.ShutdownTimeout <= TimeSpan.Zero)
 3159                    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.
 3163                ShutdownBudgetValidator.Validate(
 3164                    "RabbitMQ",
 3165                    $"{nameof(RabbitMqAsyncResponseOptions)}.{nameof(RabbitMqAsyncResponseOptions.HostShutdownTimeout)}"
 3166                    transportOptions.HostShutdownTimeout,
 3167                    ($"{optionPath}.{nameof(RabbitMqSubscriberOptions.BackgroundDrainTimeout)}", subscriberOptions.Backg
 3168                    ($"{nameof(RabbitMqAsyncResponseOptions)}.{nameof(RabbitMqAsyncResponseOptions.ShutdownTimeout)}", t
 169
 3170                return;
 171
 172            default:
 3173                throw new InvalidOperationException(
 3174                    $"{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>
 3185    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    {
 3193        using var activity = AsyncResponseDiagnostics.StartActivity(
 3194            "asyncresponse.rabbitmq.receive",
 3195            ActivityKind.Consumer);
 3196        activity?.SetTag("asyncresponse.transport", "rabbitmq");
 3197        activity?.SetTag("asyncresponse.rabbitmq.role", _role.ToString());
 3198        activity?.SetTag("asyncresponse.rabbitmq.ack_mode", _subscriberOptions.AckMode.ToString());
 3199        activity?.SetTag("messaging.system", "rabbitmq");
 3200        activity?.SetTag("messaging.destination.name", _queue);
 3201        activity?.SetTag("messaging.rabbitmq.exchange", delivery.Exchange);
 3202        activity?.SetTag("messaging.rabbitmq.routing_key", delivery.RoutingKey);
 3203        activity?.SetTag("messaging.rabbitmq.delivery_tag", delivery.DeliveryTag);
 3204        activity?.SetTag("messaging.message.id", delivery.BasicProperties.MessageId);
 205
 3206        if (!string.IsNullOrWhiteSpace(delivery.BasicProperties.CorrelationId))
 3207            AsyncResponseDiagnostics.SetCorrelationId(activity, delivery.BasicProperties.CorrelationId);
 208
 209        try
 210        {
 3211            await _handler(delivery, cancellationToken).ConfigureAwait(false);
 3212        }
 2213        catch (Exception ex)
 214        {
 2215            if (logFailures)
 2216                Logger.LogError(ex, "RabbitMQ message handling failed for delivery {DeliveryTag}.", delivery.DeliveryTag
 2217            AsyncResponseDiagnostics.SetError(activity, ex);
 3218            throw;
 219        }
 3220    }
 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    {
 2229        var callback = _subscriberOptions.OnBackgroundFailure;
 2230        if (callback is null)
 2231            return;
 232
 233        try
 234        {
 2235            await callback(new RabbitMqBackgroundFailureContext(
 2236                queue,
 2237                role.ToString(),
 2238                delivery.Exchange,
 2239                delivery.RoutingKey,
 2240                delivery.DeliveryTag,
 2241                exception)).ConfigureAwait(false);
 2242        }
 2243        catch (Exception callbackException)
 244        {
 2245            Logger.LogError(
 2246                callbackException,
 2247                "RabbitMQ background failure callback failed for already-ACKed delivery {DeliveryTag} on {Queue}.",
 2248                delivery.DeliveryTag,
 2249                queue);
 2250        }
 2251    }
 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;
 289    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)
 305        : base(handler, transportOptions, subscriberOptions, logger, queue, role)
 306    {
 307        _drainTimeout = subscriberOptions.BackgroundDrainTimeout;
 308        _queueName = queue;
 309        _role = role;
 310        _queue = Channel.CreateBounded<RabbitMqDelivery>(new BoundedChannelOptions(subscriberOptions.BackgroundQueueCapa
 311        {
 312            AllowSynchronousContinuations = false,
 313            FullMode = BoundedChannelFullMode.Wait,
 314            SingleReader = subscriberOptions.BackgroundWorkerCount == 1,
 315            SingleWriter = false
 316        });
 317
 318        _workers = Enumerable.Range(0, subscriberOptions.BackgroundWorkerCount)
 319            .Select(workerIndex => Task.Run(() => RunWorkerAsync(workerIndex)))
 320            .ToArray();
 321
 322        Logger.LogInformation(
 323            "Created RabbitMQ ACK-after-enqueue dispatcher for {Queue} with {WorkerCount} worker(s), queue capacity {Que
 324            _queueName,
 325            subscriberOptions.BackgroundWorkerCount,
 326            subscriberOptions.BackgroundQueueCapacity,
 327            _drainTimeout);
 328    }
 329
 330    internal int PendingCount => Volatile.Read(ref _pendingCount);
 331    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.
 344        delivery = delivery with { Body = delivery.Body.ToArray() };
 345
 346        Interlocked.Increment(ref _pendingCount);
 347        if (!_queue.Writer.TryWrite(delivery))
 348        {
 349            Interlocked.Decrement(ref _pendingCount);
 350            Logger.LogWarning(
 351                "RabbitMQ background queue rejected delivery {DeliveryTag} for {Queue}; returning NACK. Pending={Pending
 352                delivery.DeliveryTag,
 353                _queueName,
 354                PendingCount,
 355                RunningCount);
 356            await channel.BasicNackAsync(delivery.DeliveryTag, requeue: true, subscriberCancellationToken).ConfigureAwai
 357            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        {
 365            await channel.BasicAckAsync(delivery.DeliveryTag, subscriberCancellationToken).ConfigureAwait(false);
 366        }
 367        catch (Exception ex)
 368        {
 369            Logger.LogError(
 370                ex,
 371                "Failed to ACK RabbitMQ delivery {DeliveryTag} for {Queue} after enqueue; it is being processed but the 
 372                delivery.DeliveryTag,
 373                _queueName);
 374        }
 375    }
 376
 377    /// <summary>Releases resources held by this instance.</summary>
 378    public override async ValueTask DisposeAsync()
 379    {
 380        if (Interlocked.Exchange(ref _disposeStarted, 1) != 0)
 381            return;
 382
 383        Logger.LogInformation(
 384            "Draining RabbitMQ ACK-after-enqueue dispatcher for {Queue}. Pending={PendingCount}, Running={RunningCount}.
 385            _queueName,
 386            PendingCount,
 387            RunningCount);
 388        _queue.Writer.TryComplete();
 389
 390        try
 391        {
 392            await Task.WhenAll(_workers).WaitAsync(_drainTimeout).ConfigureAwait(false);
 393            _drainCancellation.Dispose();
 394        }
 395        catch (TimeoutException ex)
 396        {
 397            _drainCancellation.Cancel();
 398            Logger.LogWarning(
 399                ex,
 400                "Timed out while draining RabbitMQ ACK-after-enqueue dispatcher for {Queue}. Pending={PendingCount}, Run
 401                _queueName,
 402                PendingCount,
 403                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.
 408            _ = Task.WhenAll(_workers).ContinueWith(
 409                _ => _drainCancellation.Dispose(),
 410                CancellationToken.None,
 411                TaskContinuationOptions.ExecuteSynchronously,
 412                TaskScheduler.Default);
 413        }
 414    }
 415
 416    private async Task RunWorkerAsync(int workerIndex)
 417    {
 418        await foreach (var delivery in _queue.Reader.ReadAllAsync().ConfigureAwait(false))
 419        {
 420            Interlocked.Decrement(ref _pendingCount);
 421            Interlocked.Increment(ref _runningCount);
 422
 423            try
 424            {
 425                await ExecuteHandlerAsync(
 426                    delivery,
 427                    _drainCancellation.Token,
 428                    logFailures: false).ConfigureAwait(false);
 429            }
 430            catch (Exception ex)
 431            {
 432                Logger.LogError(
 433                    ex,
 434                    "RabbitMQ background handler failed for already-ACKed delivery {DeliveryTag} on {Queue}.",
 435                    delivery.DeliveryTag,
 436                    _queueName);
 437                await NotifyBackgroundFailureAsync(
 438                    delivery,
 439                    ex,
 440                    _queueName,
 441                    _role).ConfigureAwait(false);
 442            }
 443            finally
 444            {
 445                Interlocked.Decrement(ref _runningCount);
 446            }
 447        }
 448    }
 449}