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

Information
Class: AsyncResponse.Transports.NATS.NatsMessageDispatcher
Assembly: AsyncResponse.Transports.NATS
File(s): /_/src/Transports/AsyncResponse.Transports.NATS/NatsMessageDispatcher.cs
Line coverage
96%
Covered lines: 230
Uncovered lines: 8
Coverable lines: 238
Total lines: 474
Line coverage: 96.6%
Branch coverage
97%
Covered branches: 43
Total branches: 44
Branch coverage: 97.7%
Method coverage

Feature is only available for sponsors

Upgrade to PRO version

Metrics

MethodBranch coverage Crap Score Cyclomatic complexity Line coverage
.ctor(...)100%44100%
HandleAsync()87.5%8882.22%
ExecuteHandlerAsync()100%1414100%
HandleEarlyAckAsync()100%22100%
BackgroundWorkerLoopAsync()100%44100%
HandleFailureAsync()100%66100%
NakQuietlyAsync()100%11100%
DeadLetterAsync()100%22100%
InvokeBackgroundFailureAsync()100%22100%
SanitizeHeaderValue(...)100%11100%
DisposeAsync()100%22100%

File(s)

/_/src/Transports/AsyncResponse.Transports.NATS/NatsMessageDispatcher.cs

#LineLine coverage
 1using Microsoft.Extensions.Logging;
 2using System.Diagnostics;
 3using System.Threading.Channels;
 4
 5namespace AsyncResponse.Transports.NATS;
 6
 7internal enum NatsSubscriberRole
 8{
 9    Worker,
 10    ResponseIngress
 11}
 12
 13/// <summary>
 14/// Applies the acknowledgement, redelivery, and dead-letter policy to JetStream deliveries.
 15/// <list type="bullet">
 16/// <item><description><see cref="NatsAckMode.AckAfterHandlerCompletes"/>: run the handler, then ACK;
 17/// on failure NAK for redelivery until <see cref="NatsSubscriberOptions.MaxDeliveryAttempts"/>, then
 18/// dead-letter and terminate.</description></item>
 19/// <item><description><see cref="NatsAckMode.AckAfterEnqueue"/>: enqueue to a bounded background
 20/// queue and ACK immediately; background handler failures are dead-lettered and reported.</description></item>
 21/// </list>
 22/// </summary>
 23internal sealed class NatsMessageDispatcher : IAsyncDisposable
 24{
 25    private readonly Func<NatsJobDelivery, CancellationToken, Task> _handler;
 26    private readonly INatsJetStreamTransport _jetStream;
 27    private readonly NatsAsyncResponseTransportOptions _options;
 28    private readonly NatsSubscriberOptions _subscriberOptions;
 29    private readonly NatsTransportSubjectSchema _schema;
 30    private readonly ILogger _logger;
 31    private readonly NatsSubscriberRole _role;
 32    private readonly string _consumer;
 33
 34    private readonly Channel<NatsJobDelivery>? _backgroundQueue;
 35    private readonly Task[]? _backgroundWorkers;
 36    private readonly CancellationTokenSource? _backgroundCts;
 37
 38    /// <summary>Runs the NatsMessageDispatcher operation.</summary>
 45839    public NatsMessageDispatcher(
 45840        Func<NatsJobDelivery, CancellationToken, Task> handler,
 45841        INatsJetStreamTransport jetStream,
 45842        NatsAsyncResponseTransportOptions options,
 45843        NatsSubscriberOptions subscriberOptions,
 45844        NatsTransportSubjectSchema schema,
 45845        ILogger logger,
 45846        NatsSubscriberRole role,
 45847        string consumer)
 48    {
 45849        NatsTransportOptionsValidator.ValidateSubscriber(options, subscriberOptions, role.ToString());
 50
 45851        _handler = handler;
 45852        _jetStream = jetStream;
 45853        _options = options;
 45854        _subscriberOptions = subscriberOptions;
 45855        _schema = schema;
 45856        _logger = logger;
 45857        _role = role;
 45858        _consumer = consumer;
 59
 45860        if (subscriberOptions.AckMode is NatsAckMode.AckAfterEnqueue)
 61        {
 2862            _backgroundQueue = Channel.CreateBounded<NatsJobDelivery>(new BoundedChannelOptions(subscriberOptions.Backgr
 2863            {
 2864                SingleReader = false,
 2865                SingleWriter = true,
 2866                FullMode = BoundedChannelFullMode.Wait
 2867            });
 2868            _backgroundCts = new CancellationTokenSource();
 2869            _backgroundWorkers = new Task[subscriberOptions.BackgroundWorkerCount];
 12470            for (var i = 0; i < _backgroundWorkers.Length; i++)
 6871                _backgroundWorkers[i] = Task.Run(() => BackgroundWorkerLoopAsync(_backgroundCts.Token));
 72        }
 45873    }
 74
 75    /// <summary>Handles the delivered message.</summary>
 76    public async Task HandleAsync(NatsJobDelivery delivery, CancellationToken cancellationToken)
 77    {
 78        // Pre-execution cap, BEFORE either ack mode (DB/Redis dispatcher parity).
 79        // HandleFailureAsync below is the only other place the cap is consulted, and it runs only
 80        // when the handler THREW — so a delivery that ends any other way (the process dies
 81        // mid-handler, the host is killed, a NAK fails) never reaches it. The consumer is created
 82        // with MaxDeliver = -1 on the premise that THIS dispatcher bounds attempts, so without
 83        // this check such a message redelivered forever after each AckWait, killing each replica
 84        // in turn, and was never dead-lettered. Settlement uses CancellationToken.None for the
 85        // usual reason: burying a poison message must not be abandoned half-done by a shutdown.
 50986        var cap = _subscriberOptions.MaxDeliveryAttempts;
 50987        if (cap > 0 && delivery.NumDelivered > cap)
 88        {
 289            _logger.LogError(
 290                "Message on subject {Subject} ({Role}) arrived on delivery {NumDelivered} with a cap of {MaxDeliveryAtte
 291                delivery.Subject,
 292                _role,
 293                delivery.NumDelivered,
 294                cap);
 95
 296            var shouldTerminate = await DeadLetterAsync(
 297                delivery,
 298                new InvalidOperationException(
 299                    $"Message exceeded {cap} delivery attempts without settling (delivery {delivery.NumDelivered})."),
 2100                CancellationToken.None).ConfigureAwait(false);
 2101            if (shouldTerminate)
 102            {
 103                // Guarded like the failure path's Term: a thrown settlement would unwind the
 104                // consume loop while the un-termed message redelivers and is dead-lettered again.
 105                try
 106                {
 2107                    await delivery.TermAsync().ConfigureAwait(false);
 2108                }
 0109                catch (Exception ex)
 110                {
 0111                    _logger.LogWarning(
 0112                        ex,
 0113                        "Failed to TERM NATS message on subject {Subject} ({Role}) after dead-lettering; it may redelive
 0114                        delivery.Subject,
 0115                        _role);
 0116                }
 117            }
 118            else
 119            {
 0120                await NakQuietlyAsync(delivery).ConfigureAwait(false);
 121            }
 122
 2123            return;
 124        }
 125
 507126        if (_subscriberOptions.AckMode is NatsAckMode.AckAfterEnqueue)
 127        {
 46128            await HandleEarlyAckAsync(delivery, cancellationToken).ConfigureAwait(false);
 46129            return;
 130        }
 131
 132        try
 133        {
 461134            await ExecuteHandlerAsync(delivery, cancellationToken).ConfigureAwait(false);
 437135        }
 2136        catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
 137        {
 138            // Host shutdown, not a handler failure: NAK would burn a delivery attempt on work
 139            // that never ran, and at the attempt cap the failure path would dead-letter — or,
 140            // with dead-lettering disabled, TERMINATE — healthy work. Leave the delivery
 141            // unsettled; AckWait lapses on its own and at-least-once redelivery applies after
 142            // restart (parity with the RabbitMQ/Redis/Kafka/DB dispatchers).
 2143            throw;
 144        }
 22145        catch (Exception ex)
 146        {
 22147            await HandleFailureAsync(delivery, ex, cancellationToken).ConfigureAwait(false);
 22148            return;
 149        }
 150
 151        // The ACK sits outside the handler's try/catch: a transient ack failure after a successful
 152        // handler must not be misread as a handler failure — NAK/dead-letter here would redeliver
 153        // (or bury) work whose side effects already completed. Swallow and log instead; the ack
 154        // window lapses on its own and at-least-once redelivery applies.
 155        try
 156        {
 437157            await delivery.AckAsync().ConfigureAwait(false);
 435158        }
 2159        catch (Exception ex)
 160        {
 2161            _logger.LogWarning(
 2162                ex,
 2163                "Failed to ACK NATS message on subject {Subject} ({Role}) after a successful handler; it may be redelive
 2164                delivery.Subject,
 2165                _role);
 2166        }
 507167    }
 168
 169    // Single choke point for handler execution so both ACK modes emit the consumer receive span.
 170    private async Task ExecuteHandlerAsync(NatsJobDelivery delivery, CancellationToken cancellationToken)
 171    {
 501172        using var activity = AsyncResponseDiagnostics.StartActivity(
 501173            "asyncresponse.nats.receive",
 501174            ActivityKind.Consumer);
 501175        activity?.SetTag("asyncresponse.transport", "nats");
 501176        activity?.SetTag("asyncresponse.nats.role", _role.ToString());
 501177        activity?.SetTag("asyncresponse.nats.ack_mode", _subscriberOptions.AckMode.ToString());
 501178        activity?.SetTag("messaging.system", "nats");
 501179        activity?.SetTag("messaging.destination.name", delivery.Subject);
 501180        activity?.SetTag("messaging.nats.num_delivered", delivery.NumDelivered);
 181
 501182        if (delivery.Headers.TryGetValue(_options.CorrelationIdHeader, out var correlationId))
 105183            AsyncResponseDiagnostics.SetCorrelationId(activity, correlationId);
 184
 185        try
 186        {
 501187            await _handler(delivery, cancellationToken).ConfigureAwait(false);
 469188        }
 32189        catch (Exception ex)
 190        {
 32191            AsyncResponseDiagnostics.SetError(activity, ex);
 32192            throw;
 193        }
 469194    }
 195
 196    private async Task HandleEarlyAckAsync(NatsJobDelivery delivery, CancellationToken cancellationToken)
 197    {
 198        // Accept into the background queue and ACK. If the queue is saturated, wait for a worker to
 199        // free a slot instead of NAKing: the wait blocks the consume loop, so the subscriber stops
 200        // pulling new messages until capacity frees rather than churning NAK/redeliver cycles.
 46201        if (!_backgroundQueue!.Writer.TryWrite(delivery))
 202        {
 203            try
 204            {
 8205                _logger.LogDebug("Background queue full for {Role}; pausing the consume loop until capacity frees.", _ro
 8206                await _backgroundQueue.Writer.WriteAsync(delivery, cancellationToken).ConfigureAwait(false);
 4207            }
 4208            catch (Exception ex) when (ex is OperationCanceledException or ChannelClosedException)
 209            {
 210                // Subscriber stopping or dispatcher disposing while parked: the delivery was never
 211                // enqueued, so NAK so JetStream redelivers elsewhere; if the NAK itself fails the
 212                // AckWait lapses to the same effect.
 4213                _logger.LogDebug("Background queue unavailable for {Role} during shutdown; NAKing message for redelivery
 214                try
 215                {
 4216                    await delivery.NakAsync(_subscriberOptions.RedeliveryDelay).ConfigureAwait(false);
 2217                }
 2218                catch (Exception nakException)
 219                {
 2220                    _logger.LogWarning(
 2221                        nakException,
 2222                        "Failed to NAK NATS message on subject {Subject} ({Role}) while stopping; AckWait will lapse and
 2223                        delivery.Subject,
 2224                        _role);
 2225                }
 226
 4227                return;
 228            }
 229        }
 230
 231        // The ACK sits outside the enqueue try/catch, and never NAKs or escapes: the delivery is
 232        // already owned by a background worker, so a NAK would redeliver a job that is being
 233        // executed, and a thrown ack failure would unwind the consume loop and rebuild the whole
 234        // subscriber — draining the workers mid-handler while the un-ACKed message redelivers
 235        // after AckWait and runs again. Swallow and log; at-least-once redelivery applies.
 236        try
 237        {
 42238            await delivery.AckAsync().ConfigureAwait(false);
 38239        }
 4240        catch (Exception ex)
 241        {
 4242            _logger.LogWarning(
 4243                ex,
 4244                "Failed to ACK NATS message on subject {Subject} ({Role}) after enqueueing it for background execution; 
 4245                delivery.Subject,
 4246                _role);
 4247        }
 46248    }
 249
 250    private async Task BackgroundWorkerLoopAsync(CancellationToken cancellationToken)
 251    {
 252        // Token-less ReadAllAsync: on shutdown the queue is completed and fully drained, so every
 253        // already-ACKed delivery is either attempted or explicitly dead-lettered below — never
 254        // silently dropped.
 152255        await foreach (var delivery in _backgroundQueue!.Reader.ReadAllAsync().ConfigureAwait(false))
 256        {
 257            // Once the drain budget has lapsed, STOP executing (DB/Redis/Pub-Sub parity). The
 258            // token below cannot stop the real handler — it is `_ingress.HandleWorkerMessageAsync
 259            // (payload)`, whose target takes no CancellationToken — so past the budget the loop
 260            // kept starting fresh work beyond the host's shutdown budget, and every entry still
 261            // queued at process exit vanished with no record (ACKed at enqueue, so JetStream never
 262            // redelivers it). Route the rest through the dead-letter/OnBackgroundFailure path.
 42263            if (_backgroundCts!.IsCancellationRequested)
 264            {
 2265                var lapsed = new OperationCanceledException(
 2266                    "The ACK-after-enqueue drain budget lapsed before this already-ACKed message was handled.");
 2267                _logger.LogWarning(
 2268                    "NATS background handler for already-ACKed message on subject {Subject} ({Role}) was not started: th
 2269                    delivery.Subject,
 2270                    _role);
 2271                await DeadLetterAsync(delivery, lapsed, CancellationToken.None).ConfigureAwait(false);
 2272                await InvokeBackgroundFailureAsync(delivery, lapsed).ConfigureAwait(false);
 2273                continue;
 274            }
 275
 276            try
 277            {
 40278                await ExecuteHandlerAsync(delivery, cancellationToken).ConfigureAwait(false);
 32279            }
 4280            catch (OperationCanceledException ex) when (_backgroundCts!.IsCancellationRequested)
 281            {
 282                // The drain budget lapsed with this already-ACKed message still unprocessed:
 283                // JetStream will not redeliver it, so surface the drop through OnBackgroundFailure
 284                // instead of dead-lettering a never-run job as a handler failure (Kafka/Redis
 285                // dispatcher parity).
 4286                _logger.LogWarning(
 4287                    "NATS background handler for already-ACKed message on subject {Subject} ({Role}) was canceled during
 4288                    delivery.Subject,
 4289                    _role);
 4290                await InvokeBackgroundFailureAsync(delivery, ex).ConfigureAwait(false);
 4291            }
 4292            catch (Exception ex)
 293            {
 4294                _logger.LogError(ex, "Background handler failed for {Role} on subject {Subject} after early ACK.", _role
 4295                await DeadLetterAsync(delivery, ex, CancellationToken.None).ConfigureAwait(false);
 4296                await InvokeBackgroundFailureAsync(delivery, ex).ConfigureAwait(false);
 4297            }
 40298        }
 34299    }
 300
 301    private async Task HandleFailureAsync(NatsJobDelivery delivery, Exception exception, CancellationToken cancellationT
 302    {
 22303        var maxAttempts = _subscriberOptions.MaxDeliveryAttempts;
 22304        if (maxAttempts > 0 && delivery.NumDelivered >= maxAttempts)
 305        {
 13306            _logger.LogError(
 13307                exception,
 13308                "Message on subject {Subject} ({Role}) failed after {Attempts} attempts; dead-lettering.",
 13309                delivery.Subject,
 13310                _role,
 13311                delivery.NumDelivered);
 312
 313            // CancellationToken.None like every other settlement in this package (the
 314            // pre-execution cap and the early-ACK failure path already pin it): burying a poison
 315            // message must not be abandoned by a shutdown — with the stopping token, a handler
 316            // failing on its LAST attempt during a stop had the DLQ publish throw on the cancelled
 317            // token, and the message was NAKed back instead of buried.
 13318            var shouldTerminate = await DeadLetterAsync(delivery, exception, CancellationToken.None).ConfigureAwait(fals
 13319            if (shouldTerminate)
 320            {
 321                // Guarded like both ack sites: TermAsync is the same JetStream request/reply as
 322                // Ack/Nak and can throw, and a thrown settlement would unwind the consume loop and
 323                // rebuild the whole subscriber — while the un-termed message redelivers after
 324                // AckWait and dead-letters AGAIN, forever. Swallow and log; the duplicate
 325                // dead-letter on redelivery is the bounded at-least-once outcome.
 326                try
 327                {
 11328                    await delivery.TermAsync().ConfigureAwait(false);
 9329                }
 2330                catch (Exception ex)
 331                {
 2332                    _logger.LogWarning(
 2333                        ex,
 2334                        "Failed to TERM NATS message on subject {Subject} ({Role}) after dead-lettering; it may redelive
 2335                        delivery.Subject,
 2336                        _role);
 2337                }
 338            }
 339            else
 340            {
 2341                _logger.LogWarning(
 2342                    exception,
 2343                    "Dead-letter publish failed for subject {Subject} ({Role}); NAKing so the message can be retried.",
 2344                    delivery.Subject,
 2345                    _role);
 2346                await NakQuietlyAsync(delivery).ConfigureAwait(false);
 347            }
 348        }
 349        else
 350        {
 9351            _logger.LogWarning(
 9352                exception,
 9353                "Message on subject {Subject} ({Role}) failed on attempt {Attempt}; NAKing for redelivery.",
 9354                delivery.Subject,
 9355                _role,
 9356                delivery.NumDelivered);
 9357            await NakQuietlyAsync(delivery).ConfigureAwait(false);
 358        }
 22359    }
 360
 361    /// <summary>
 362    /// NAKs with the configured redelivery delay, swallowing settlement failures like the ack
 363    /// sites: a thrown NAK would unwind the consume loop and rebuild the subscriber, and the only
 364    /// consequence of a lost NAK is that redelivery waits for AckWait instead of the delay.
 365    /// </summary>
 366    private async Task NakQuietlyAsync(NatsJobDelivery delivery)
 367    {
 368        try
 369        {
 11370            await delivery.NakAsync(_subscriberOptions.RedeliveryDelay).ConfigureAwait(false);
 9371        }
 2372        catch (Exception ex)
 373        {
 2374            _logger.LogWarning(
 2375                ex,
 2376                "Failed to NAK NATS message on subject {Subject} ({Role}); redelivery falls back to AckWait.",
 2377                delivery.Subject,
 2378                _role);
 2379        }
 11380    }
 381
 382    private async Task<bool> DeadLetterAsync(NatsJobDelivery delivery, Exception exception, CancellationToken cancellati
 383    {
 21384        if (!_options.DeadLetterEnabled)
 385        {
 2386            _logger.LogError(
 2387                exception,
 2388                "Message on subject {Subject} ({Role}) is unprocessable and dead-lettering is disabled; it will be dropp
 2389                delivery.Subject,
 2390                _role);
 2391            return true;
 392        }
 393
 19394        var headers = new Dictionary<string, string>(delivery.Headers, StringComparer.OrdinalIgnoreCase)
 19395        {
 19396            ["AR-DeadLetter-Reason"] = SanitizeHeaderValue(exception.Message),
 19397            ["AR-DeadLetter-Source-Subject"] = delivery.Subject,
 19398            ["AR-DeadLetter-Role"] = _role.ToString()
 19399        };
 400
 401        // The inbound Nats-Msg-Id belongs to the LIVE publish, not to this one. Carrying it over
 402        // makes a second dead-letter of the same message inside the DLQ stream's duplicate window
 403        // — reachable whenever the Term below fails and the message redelivers after AckWait — a
 404        // deduplicated publish, which the caller reads as a DLQ failure and answers with a NAK,
 405        // looping until the window passes.
 19406        headers.Remove("Nats-Msg-Id");
 407
 408        try
 409        {
 19410            await _jetStream.PublishAsync(_schema.DeadLetterSubject, delivery.Payload, headers, cancellationToken).Confi
 17411            _logger.LogInformation("Dead-lettered message from subject {Subject} ({Role}) to {DeadLetterSubject}.", deli
 17412            return true;
 413        }
 2414        catch (Exception ex)
 415        {
 2416            _logger.LogError(ex, "Failed to dead-letter message from subject {Subject} ({Role}).", delivery.Subject, _ro
 2417            return false;
 418        }
 21419    }
 420
 421    private async Task InvokeBackgroundFailureAsync(NatsJobDelivery delivery, Exception exception)
 422    {
 12423        if (_subscriberOptions.OnBackgroundFailure is null)
 2424            return;
 425
 426        try
 427        {
 10428            delivery.Headers.TryGetValue(_options.CorrelationIdHeader, out var correlationId);
 10429            var context = new NatsBackgroundFailureContext(delivery.Subject, _consumer, _role.ToString(), delivery.NumDe
 10430            await _subscriberOptions.OnBackgroundFailure(context).ConfigureAwait(false);
 8431        }
 2432        catch (Exception ex)
 433        {
 2434            _logger.LogError(ex, "OnBackgroundFailure callback threw for {Role}.", _role);
 2435        }
 12436    }
 437
 438    private static string SanitizeHeaderValue(string value)
 19439        => value.Replace('\r', ' ').Replace('\n', ' ');
 440
 441    /// <summary>Releases resources held by this instance.</summary>
 442    public async ValueTask DisposeAsync()
 443    {
 456444        if (_backgroundQueue is null)
 428445            return;
 446
 28447        _backgroundQueue.Writer.TryComplete();
 448        try
 449        {
 28450            await Task.WhenAll(_backgroundWorkers!).WaitAsync(_subscriberOptions.BackgroundDrainTimeout).ConfigureAwait(
 20451            _backgroundCts!.Dispose();
 20452        }
 453        catch (TimeoutException)
 454        {
 6455            _logger.LogWarning("Background handlers for {Role} did not drain within {Timeout}.", _role, _subscriberOptio
 6456            await _backgroundCts!.CancelAsync().ConfigureAwait(false);
 457
 458            // The workers are still running and observe _backgroundCts.Token inside ReadAllAsync, so disposing
 459            // it now would throw ObjectDisposedException inside them. Dispose once they actually finish, off
 460            // the shutdown path, so the source is not leaked either.
 6461            _ = Task.WhenAll(_backgroundWorkers!).ContinueWith(
 6462                _ => _backgroundCts.Dispose(),
 6463                CancellationToken.None,
 6464                TaskContinuationOptions.ExecuteSynchronously,
 6465                TaskScheduler.Default);
 466        }
 2467        catch (Exception ex)
 468        {
 469            // WhenAll only completes once every worker has finished, so the source is safe to dispose here.
 2470            _logger.LogDebug(ex, "Background worker drain for {Role} ended with an error.", _role);
 2471            _backgroundCts!.Dispose();
 2472        }
 456473    }
 474}