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

Information
Class: AsyncResponse.Transports.SQS.QueuedSqsMessageDispatcher
Assembly: AsyncResponse.Transports.SQS
File(s): /_/src/Transports/AsyncResponse.Transports.SQS/SqsMessageDispatcher.cs
Line coverage
99%
Covered lines: 120
Uncovered lines: 1
Coverable lines: 121
Total lines: 463
Line coverage: 99.1%
Branch coverage
92%
Covered branches: 13
Total branches: 14
Branch coverage: 92.8%
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%
get_CanAcceptMore()100%11100%
get_FreeCapacity()100%11100%
WaitForCapacityAsync()75%4475%
HandleAsync()100%44100%
DisposeAsync()100%22100%
RunWorkerAsync()100%44100%

File(s)

/_/src/Transports/AsyncResponse.Transports.SQS/SqsMessageDispatcher.cs

#LineLine coverage
 1using Microsoft.Extensions.Logging;
 2using System.Diagnostics;
 3using System.Threading.Channels;
 4
 5namespace AsyncResponse.Transports.SQS;
 6
 7internal enum SqsSubscriberRole
 8{
 9    Worker,
 10    ResponseIngress
 11}
 12
 13internal abstract class SqsMessageDispatcher : IAsyncDisposable
 14{
 15    private readonly Func<SqsTransportDelivery, CancellationToken, Task> _handler;
 16    private readonly SqsAsyncResponseOptions _transportOptions;
 17    private readonly SqsSubscriberOptions _subscriberOptions;
 18    private readonly string _queue;
 19    private readonly SqsSubscriberRole _role;
 20
 21    protected SqsMessageDispatcher(
 22        Func<SqsTransportDelivery, CancellationToken, Task> handler,
 23        SqsAsyncResponseOptions transportOptions,
 24        SqsSubscriberOptions subscriberOptions,
 25        ILogger logger,
 26        string queue,
 27        SqsSubscriberRole role)
 28    {
 29        _handler = handler;
 30        _transportOptions = transportOptions;
 31        _subscriberOptions = subscriberOptions;
 32        Logger = logger;
 33        _queue = queue;
 34        _role = role;
 35    }
 36
 37    protected ILogger Logger { get; }
 38    protected TimeSpan? RedeliveryDelay => _subscriberOptions.RedeliveryDelay;
 39
 40    /// <summary>Creates the dispatcher configured by the subscriber options.</summary>
 41    public static SqsMessageDispatcher Create(
 42        Func<SqsTransportDelivery, CancellationToken, Task> handler,
 43        SqsAsyncResponseOptions transportOptions,
 44        SqsSubscriberOptions subscriberOptions,
 45        ILogger logger,
 46        string queue,
 47        SqsSubscriberRole role)
 48    {
 49        SqsOptionsValidator.ValidateSubscriber(transportOptions, subscriberOptions, role);
 50
 51        return subscriberOptions.AckMode == SqsAckMode.AckAfterHandlerCompletes
 52            ? new AwaitingSqsMessageDispatcher(
 53                handler,
 54                transportOptions,
 55                subscriberOptions,
 56                logger,
 57                queue,
 58                role)
 59            : new QueuedSqsMessageDispatcher(
 60                handler,
 61                transportOptions,
 62                subscriberOptions,
 63                logger,
 64                queue,
 65                role);
 66    }
 67
 68    /// <summary>Validates the supplied subscriber options.</summary>
 69    public static void ValidateOptions(
 70        SqsAsyncResponseOptions transportOptions,
 71        SqsSubscriberOptions subscriberOptions,
 72        SqsSubscriberRole role)
 73        => SqsOptionsValidator.ValidateSubscriber(transportOptions, subscriberOptions, role);
 74
 75    /// <summary>Handles the delivered message.</summary>
 76    public abstract Task HandleAsync(
 77        SqsTransportDelivery delivery,
 78        CancellationToken subscriberCancellationToken);
 79
 80    /// <summary>
 81    /// Whether the dispatcher can accept more deliveries right now. Awaiting dispatchers always can
 82    /// (handlers run inline); the queued dispatcher returns <c>false</c> while its bounded queue is
 83    /// saturated so the receive loop stops pulling messages instead of receiving and releasing them —
 84    /// SQS counts every receive toward the queue's redrive policy.
 85    /// </summary>
 86    public virtual bool CanAcceptMore => true;
 87
 88    /// <summary>
 89    /// Number of deliveries the dispatcher can accept right now. The receive loop requests at most
 90    /// this many messages per receive in early-ACK mode so a burst never overflows the background queue.
 91    /// </summary>
 92    public virtual int FreeCapacity => int.MaxValue;
 93
 94    /// <summary>
 95    /// Waits until the dispatcher can accept at least one more delivery. Completes immediately for
 96    /// awaiting dispatchers; the queued dispatcher waits for a background worker to free a slot.
 97    /// </summary>
 98    public virtual ValueTask WaitForCapacityAsync(CancellationToken cancellationToken) => ValueTask.CompletedTask;
 99
 100    /// <summary>Releases resources held by this instance.</summary>
 101    public virtual ValueTask DisposeAsync() => ValueTask.CompletedTask;
 102
 103    protected async Task ExecuteHandlerAsync(
 104        SqsTransportDelivery delivery,
 105        CancellationToken cancellationToken,
 106        bool logFailures = true)
 107    {
 108        using var activity = AsyncResponseDiagnostics.StartActivity(
 109            "asyncresponse.sqs.receive",
 110            ActivityKind.Consumer);
 111        activity?.SetTag("asyncresponse.transport", "aws_sqs");
 112        activity?.SetTag("asyncresponse.sqs.role", _role.ToString());
 113        activity?.SetTag("asyncresponse.sqs.ack_mode", _subscriberOptions.AckMode.ToString());
 114        activity?.SetTag("messaging.system", "aws_sqs");
 115        activity?.SetTag("messaging.destination.name", _queue);
 116        activity?.SetTag("messaging.message.id", delivery.MessageId);
 117        activity?.SetTag("messaging.aws_sqs.receive_count", delivery.ReceiveCount);
 118
 119        if (TryReadCorrelationId(delivery) is { } correlationId)
 120            AsyncResponseDiagnostics.SetCorrelationId(activity, correlationId);
 121
 122        try
 123        {
 124            await _handler(delivery, cancellationToken).ConfigureAwait(false);
 125        }
 126        catch (Exception ex)
 127        {
 128            if (logFailures)
 129                Logger.LogError(ex, "SQS message handling failed for message {MessageId}.", delivery.MessageId);
 130            AsyncResponseDiagnostics.SetError(activity, ex);
 131            throw;
 132        }
 133    }
 134
 135    /// <summary>
 136    /// Best-effort <c>ChangeMessageVisibility</c>: the receipt handle may already be expired or the
 137    /// message deleted by a competing consumer, and either way SQS redelivery still owns the retry.
 138    /// </summary>
 139    protected async ValueTask TryChangeVisibilityAsync(SqsTransportDelivery delivery, TimeSpan delay)
 140    {
 141        try
 142        {
 143            // Settlement deliberately ignores cancellation (as every sibling transport does).
 144            await delivery.ChangeVisibilityAsync(delay, CancellationToken.None).ConfigureAwait(false);
 145        }
 146        catch (Exception ex)
 147        {
 148            Logger.LogWarning(
 149                ex,
 150                "Failed to change visibility of SQS message {MessageId} on {Queue}; it stays invisible until the visibil
 151                delivery.MessageId,
 152                _queue);
 153        }
 154    }
 155
 156    protected async ValueTask NotifyBackgroundFailureAsync(
 157        SqsTransportDelivery delivery,
 158        Exception exception,
 159        string queue,
 160        SqsSubscriberRole role)
 161    {
 162        var callback = _subscriberOptions.OnBackgroundFailure;
 163        if (callback is null)
 164            return;
 165
 166        try
 167        {
 168            var context = new SqsBackgroundFailureContext(
 169                queue,
 170                role.ToString(),
 171                delivery.MessageId,
 172                delivery.ReceiveCount,
 173                TryReadCorrelationId(delivery),
 174                exception);
 175            await callback(context).ConfigureAwait(false);
 176        }
 177        catch (Exception callbackException)
 178        {
 179            Logger.LogError(
 180                callbackException,
 181                "SQS background failure callback failed for already-deleted message {MessageId} on {Queue}.",
 182                delivery.MessageId,
 183                queue);
 184        }
 185    }
 186
 187    private string? TryReadCorrelationId(SqsTransportDelivery delivery)
 188        => !string.IsNullOrWhiteSpace(_transportOptions.CorrelationIdAttribute)
 189            && delivery.MessageAttributes.TryGetValue(_transportOptions.CorrelationIdAttribute, out var value)
 190            && !string.IsNullOrWhiteSpace(value)
 191                ? value
 192                : null;
 193}
 194
 195internal sealed class AwaitingSqsMessageDispatcher(
 196    Func<SqsTransportDelivery, CancellationToken, Task> handler,
 197    SqsAsyncResponseOptions transportOptions,
 198    SqsSubscriberOptions subscriberOptions,
 199    ILogger logger,
 200    string queue,
 201    SqsSubscriberRole role)
 202    : SqsMessageDispatcher(handler, transportOptions, subscriberOptions, logger, queue, role)
 203{
 204    /// <summary>Handles the delivered message.</summary>
 205    public override async Task HandleAsync(
 206        SqsTransportDelivery delivery,
 207        CancellationToken subscriberCancellationToken)
 208    {
 209        try
 210        {
 211            await ExecuteHandlerAsync(delivery, subscriberCancellationToken).ConfigureAwait(false);
 212        }
 213        catch (OperationCanceledException) when (subscriberCancellationToken.IsCancellationRequested)
 214        {
 215            // Host shutdown, not a handler failure: shortening visibility would hasten a
 216            // redelivery of work that never ran as if it had failed. Leave the message
 217            // untouched — its visibility timeout lapses on its own and at-least-once
 218            // redelivery applies after restart (parity with the RabbitMQ/Redis/Kafka/DB
 219            // dispatchers).
 220            throw;
 221        }
 222        catch (Exception)
 223        {
 224            // SQS has no explicit NACK or dead-letter call: leaving the message undeleted lets it
 225            // reappear when its visibility timeout expires, ApproximateReceiveCount increments, and
 226            // the queue's redrive policy dead-letters it after maxReceiveCount receives.
 227            if (RedeliveryDelay is { } redeliveryDelay)
 228                await TryChangeVisibilityAsync(delivery, redeliveryDelay).ConfigureAwait(false);
 229            return;
 230        }
 231
 232        // The delete sits outside the handler's try/catch: a transient DeleteMessage failure after
 233        // a successful handler must not be misread as a handler failure — shortening visibility
 234        // here would hasten a duplicate of work whose side effects already completed and burn
 235        // receives toward the redrive policy. Swallow and log instead; the message reappears when
 236        // its visibility timeout expires and at-least-once redelivery applies.
 237        try
 238        {
 239            await delivery.DeleteAsync().ConfigureAwait(false);
 240        }
 241        catch (Exception ex)
 242        {
 243            Logger.LogWarning(
 244                ex,
 245                "Failed to delete SQS message {MessageId} after a successful handler; it may be redelivered after its vi
 246                delivery.MessageId);
 247        }
 248    }
 249}
 250
 251internal sealed class QueuedSqsMessageDispatcher : SqsMessageDispatcher
 252{
 253    private readonly Channel<SqsTransportDelivery> _queue;
 254    private readonly Task[] _workers;
 24255    private readonly CancellationTokenSource _drainCancellation = new();
 256    private readonly TimeSpan _drainTimeout;
 257    private readonly int _capacity;
 258    private readonly string _queueName;
 259    private readonly SqsSubscriberRole _role;
 260    private int _pendingCount;
 261    private int _runningCount;
 262    private int _disposeStarted;
 263
 264    /// <summary>Creates an ACK-after-enqueue dispatcher with a bounded background queue.</summary>
 265    public QueuedSqsMessageDispatcher(
 266        Func<SqsTransportDelivery, CancellationToken, Task> handler,
 267        SqsAsyncResponseOptions transportOptions,
 268        SqsSubscriberOptions subscriberOptions,
 269        ILogger logger,
 270        string queue,
 271        SqsSubscriberRole role)
 24272        : base(handler, transportOptions, subscriberOptions, logger, queue, role)
 273    {
 24274        _drainTimeout = subscriberOptions.BackgroundDrainTimeout;
 24275        _capacity = subscriberOptions.BackgroundQueueCapacity;
 24276        _queueName = queue;
 24277        _role = role;
 24278        _queue = Channel.CreateBounded<SqsTransportDelivery>(new BoundedChannelOptions(subscriberOptions.BackgroundQueue
 24279        {
 24280            AllowSynchronousContinuations = false,
 24281            FullMode = BoundedChannelFullMode.Wait,
 24282            SingleReader = subscriberOptions.BackgroundWorkerCount == 1,
 24283            SingleWriter = false
 24284        });
 285
 24286        _workers = Enumerable.Range(0, subscriberOptions.BackgroundWorkerCount)
 60287            .Select(workerIndex => Task.Run(() => RunWorkerAsync(workerIndex)))
 24288            .ToArray();
 289
 24290        Logger.LogInformation(
 24291            "Created SQS ACK-after-enqueue dispatcher for {Queue} with {WorkerCount} worker(s), queue capacity {QueueCap
 24292            _queueName,
 24293            subscriberOptions.BackgroundWorkerCount,
 24294            subscriberOptions.BackgroundQueueCapacity,
 24295            _drainTimeout);
 24296    }
 297
 70298    internal int PendingCount => Volatile.Read(ref _pendingCount);
 68299    internal int RunningCount => Volatile.Read(ref _runningCount);
 300
 50301    public override bool CanAcceptMore => Volatile.Read(ref _pendingCount) < _capacity;
 302
 15303    public override int FreeCapacity => Math.Max(0, _capacity - Volatile.Read(ref _pendingCount));
 304
 305    public override async ValueTask WaitForCapacityAsync(CancellationToken cancellationToken)
 306    {
 307        // WaitToWriteAsync completes when the bounded channel has room (or the channel is completed
 308        // during dispose, in which case there is nothing left to gate).
 44309        while (!CanAcceptMore)
 310        {
 31311            if (!await _queue.Writer.WaitToWriteAsync(cancellationToken).ConfigureAwait(false))
 0312                return;
 313        }
 13314    }
 315
 316    /// <summary>Handles the delivered message.</summary>
 317    public override async Task HandleAsync(
 318        SqsTransportDelivery delivery,
 319        CancellationToken subscriberCancellationToken)
 320    {
 42321        Interlocked.Increment(ref _pendingCount);
 42322        if (!_queue.Writer.TryWrite(delivery))
 323        {
 4324            Interlocked.Decrement(ref _pendingCount);
 4325            Logger.LogWarning(
 4326                "SQS background queue rejected message {MessageId} for {Queue}; leaving it to redeliver via its visibili
 4327                delivery.MessageId,
 4328                _queueName,
 4329                PendingCount,
 4330                RunningCount);
 331            // Do not release visibility to zero here: SQS counts every receive toward the queue's
 332            // redrive policy, so an instantly re-receivable message that keeps hitting a full queue
 333            // would cross maxReceiveCount and dead-letter without ever being processed. Let the
 334            // visibility timeout lapse naturally (or shorten it via RedeliveryDelay when configured)
 335            // so redelivery lands after capacity has had time to free.
 4336            if (RedeliveryDelay is { } redeliveryDelay)
 2337                await TryChangeVisibilityAsync(delivery, redeliveryDelay).ConfigureAwait(false);
 4338            return;
 339        }
 340
 341        // The delivery now belongs to a background worker, which decrements _pendingCount when it
 342        // dequeues. Do not touch the counter or release visibility here, even if the delete below
 343        // fails — the message is already executing in-process and releasing it would trigger a
 344        // duplicate execution via redelivery.
 345        try
 346        {
 38347            await delivery.DeleteAsync().ConfigureAwait(false);
 36348        }
 2349        catch (Exception ex)
 350        {
 2351            Logger.LogError(
 2352                ex,
 2353                "Failed to delete SQS message {MessageId} for {Queue} after enqueue; it is being processed but SQS will 
 2354                delivery.MessageId,
 2355                _queueName);
 2356        }
 42357    }
 358
 359    /// <summary>Releases resources held by this instance.</summary>
 360    public override async ValueTask DisposeAsync()
 361    {
 26362        if (Interlocked.Exchange(ref _disposeStarted, 1) != 0)
 2363            return;
 364
 24365        Logger.LogInformation(
 24366            "Draining SQS ACK-after-enqueue dispatcher for {Queue}. Pending={PendingCount}, Running={RunningCount}.",
 24367            _queueName,
 24368            PendingCount,
 24369            RunningCount);
 24370        _queue.Writer.TryComplete();
 371
 372        try
 373        {
 24374            await Task.WhenAll(_workers).WaitAsync(_drainTimeout).ConfigureAwait(false);
 18375            _drainCancellation.Dispose();
 18376        }
 4377        catch (TimeoutException ex)
 378        {
 4379            _drainCancellation.Cancel();
 4380            Logger.LogWarning(
 4381                ex,
 4382                "Timed out while draining SQS ACK-after-enqueue dispatcher for {Queue}. Pending={PendingCount}, Running=
 4383                _queueName,
 4384                PendingCount,
 4385                RunningCount);
 386
 4387            _ = Task.WhenAll(_workers).ContinueWith(
 4388                _ => _drainCancellation.Dispose(),
 4389                CancellationToken.None,
 4390                TaskContinuationOptions.ExecuteSynchronously,
 4391                TaskScheduler.Default);
 4392        }
 2393        catch (Exception ex)
 394        {
 395            // A worker faulted outside its own handler guard (DB/NATS dispatcher parity). WhenAll
 396            // only completes once every worker has finished, so the source is safe to dispose here
 397            // — and the fault must not escape DisposeAsync and mask the real shutdown path.
 2398            Logger.LogDebug(ex, "SQS ACK-after-enqueue dispatcher drain for {Queue} ended with an error.", _queueName);
 2399            _drainCancellation.Dispose();
 2400        }
 26401    }
 402
 403    private async Task RunWorkerAsync(int workerIndex)
 404    {
 134405        await foreach (var delivery in _queue.Reader.ReadAllAsync().ConfigureAwait(false))
 406        {
 38407            Interlocked.Decrement(ref _pendingCount);
 408
 409            // Once the drain budget has lapsed, STOP executing (DB/Redis/Pub-Sub parity). The
 410            // token below cannot stop the real handler — it is
 411            // `_ingress.HandleWorkerMessageAsync(payload)`, whose target takes no
 412            // CancellationToken — so past the budget the loop kept starting fresh work beyond the
 413            // host's shutdown budget, and every entry still queued at process exit vanished with
 414            // no record (deleted at enqueue, so the redrive policy never sees it again). No DLQ
 415            // write is possible for a deleted message; OnBackgroundFailure is the record.
 38416            if (_drainCancellation.IsCancellationRequested)
 417            {
 2418                var lapsed = new OperationCanceledException(
 2419                    "The ACK-after-enqueue drain budget lapsed before this already-deleted message was handled.");
 2420                Logger.LogWarning(
 2421                    "SQS background handler for already-deleted message {MessageId} on {Queue} was not started: the drai
 2422                    delivery.MessageId,
 2423                    _queueName);
 2424                await NotifyBackgroundFailureAsync(delivery, lapsed, _queueName, _role).ConfigureAwait(false);
 2425                continue;
 426            }
 427
 36428            Interlocked.Increment(ref _runningCount);
 429
 430            try
 431            {
 36432                Logger.LogDebug(
 36433                    "SQS background worker {WorkerIndex} handling message {MessageId} for {Queue}. Pending={PendingCount
 36434                    workerIndex,
 36435                    delivery.MessageId,
 36436                    _queueName,
 36437                    PendingCount,
 36438                    RunningCount);
 36439                await ExecuteHandlerAsync(
 36440                    delivery,
 36441                    _drainCancellation.Token,
 36442                    logFailures: false).ConfigureAwait(false);
 30443            }
 6444            catch (Exception ex)
 445            {
 6446                Logger.LogError(
 6447                    ex,
 6448                    "SQS background handler failed for already-deleted message {MessageId} on {Queue}.",
 6449                    delivery.MessageId,
 6450                    _queueName);
 4451                await NotifyBackgroundFailureAsync(
 4452                    delivery,
 4453                    ex,
 4454                    _queueName,
 4455                    _role).ConfigureAwait(false);
 456            }
 457            finally
 458            {
 36459                Interlocked.Decrement(ref _runningCount);
 460            }
 34461        }
 28462    }
 463}