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

Information
Class: AsyncResponse.Transports.SQS.AwaitingSqsMessageDispatcher
Assembly: AsyncResponse.Transports.SQS
File(s): /_/src/Transports/AsyncResponse.Transports.SQS/SqsMessageDispatcher.cs
Line coverage
100%
Covered lines: 17
Uncovered lines: 0
Coverable lines: 17
Total lines: 463
Line coverage: 100%
Branch coverage
100%
Covered branches: 2
Total branches: 2
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%
HandleAsync()100%22100%

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)
 418202    : 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        {
 451211            await ExecuteHandlerAsync(delivery, subscriberCancellationToken).ConfigureAwait(false);
 435212        }
 2213        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).
 2220            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.
 14227            if (RedeliveryDelay is { } redeliveryDelay)
 8228                await TryChangeVisibilityAsync(delivery, redeliveryDelay).ConfigureAwait(false);
 14229            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        {
 435239            await delivery.DeleteAsync().ConfigureAwait(false);
 433240        }
 2241        catch (Exception ex)
 242        {
 2243            Logger.LogWarning(
 2244                ex,
 2245                "Failed to delete SQS message {MessageId} after a successful handler; it may be redelivered after its vi
 2246                delivery.MessageId);
 2247        }
 449248    }
 249}
 250
 251internal sealed class QueuedSqsMessageDispatcher : SqsMessageDispatcher
 252{
 253    private readonly Channel<SqsTransportDelivery> _queue;
 254    private readonly Task[] _workers;
 255    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)
 272        : base(handler, transportOptions, subscriberOptions, logger, queue, role)
 273    {
 274        _drainTimeout = subscriberOptions.BackgroundDrainTimeout;
 275        _capacity = subscriberOptions.BackgroundQueueCapacity;
 276        _queueName = queue;
 277        _role = role;
 278        _queue = Channel.CreateBounded<SqsTransportDelivery>(new BoundedChannelOptions(subscriberOptions.BackgroundQueue
 279        {
 280            AllowSynchronousContinuations = false,
 281            FullMode = BoundedChannelFullMode.Wait,
 282            SingleReader = subscriberOptions.BackgroundWorkerCount == 1,
 283            SingleWriter = false
 284        });
 285
 286        _workers = Enumerable.Range(0, subscriberOptions.BackgroundWorkerCount)
 287            .Select(workerIndex => Task.Run(() => RunWorkerAsync(workerIndex)))
 288            .ToArray();
 289
 290        Logger.LogInformation(
 291            "Created SQS ACK-after-enqueue dispatcher for {Queue} with {WorkerCount} worker(s), queue capacity {QueueCap
 292            _queueName,
 293            subscriberOptions.BackgroundWorkerCount,
 294            subscriberOptions.BackgroundQueueCapacity,
 295            _drainTimeout);
 296    }
 297
 298    internal int PendingCount => Volatile.Read(ref _pendingCount);
 299    internal int RunningCount => Volatile.Read(ref _runningCount);
 300
 301    public override bool CanAcceptMore => Volatile.Read(ref _pendingCount) < _capacity;
 302
 303    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).
 309        while (!CanAcceptMore)
 310        {
 311            if (!await _queue.Writer.WaitToWriteAsync(cancellationToken).ConfigureAwait(false))
 312                return;
 313        }
 314    }
 315
 316    /// <summary>Handles the delivered message.</summary>
 317    public override async Task HandleAsync(
 318        SqsTransportDelivery delivery,
 319        CancellationToken subscriberCancellationToken)
 320    {
 321        Interlocked.Increment(ref _pendingCount);
 322        if (!_queue.Writer.TryWrite(delivery))
 323        {
 324            Interlocked.Decrement(ref _pendingCount);
 325            Logger.LogWarning(
 326                "SQS background queue rejected message {MessageId} for {Queue}; leaving it to redeliver via its visibili
 327                delivery.MessageId,
 328                _queueName,
 329                PendingCount,
 330                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.
 336            if (RedeliveryDelay is { } redeliveryDelay)
 337                await TryChangeVisibilityAsync(delivery, redeliveryDelay).ConfigureAwait(false);
 338            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        {
 347            await delivery.DeleteAsync().ConfigureAwait(false);
 348        }
 349        catch (Exception ex)
 350        {
 351            Logger.LogError(
 352                ex,
 353                "Failed to delete SQS message {MessageId} for {Queue} after enqueue; it is being processed but SQS will 
 354                delivery.MessageId,
 355                _queueName);
 356        }
 357    }
 358
 359    /// <summary>Releases resources held by this instance.</summary>
 360    public override async ValueTask DisposeAsync()
 361    {
 362        if (Interlocked.Exchange(ref _disposeStarted, 1) != 0)
 363            return;
 364
 365        Logger.LogInformation(
 366            "Draining SQS ACK-after-enqueue dispatcher for {Queue}. Pending={PendingCount}, Running={RunningCount}.",
 367            _queueName,
 368            PendingCount,
 369            RunningCount);
 370        _queue.Writer.TryComplete();
 371
 372        try
 373        {
 374            await Task.WhenAll(_workers).WaitAsync(_drainTimeout).ConfigureAwait(false);
 375            _drainCancellation.Dispose();
 376        }
 377        catch (TimeoutException ex)
 378        {
 379            _drainCancellation.Cancel();
 380            Logger.LogWarning(
 381                ex,
 382                "Timed out while draining SQS ACK-after-enqueue dispatcher for {Queue}. Pending={PendingCount}, Running=
 383                _queueName,
 384                PendingCount,
 385                RunningCount);
 386
 387            _ = Task.WhenAll(_workers).ContinueWith(
 388                _ => _drainCancellation.Dispose(),
 389                CancellationToken.None,
 390                TaskContinuationOptions.ExecuteSynchronously,
 391                TaskScheduler.Default);
 392        }
 393        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.
 398            Logger.LogDebug(ex, "SQS ACK-after-enqueue dispatcher drain for {Queue} ended with an error.", _queueName);
 399            _drainCancellation.Dispose();
 400        }
 401    }
 402
 403    private async Task RunWorkerAsync(int workerIndex)
 404    {
 405        await foreach (var delivery in _queue.Reader.ReadAllAsync().ConfigureAwait(false))
 406        {
 407            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.
 416            if (_drainCancellation.IsCancellationRequested)
 417            {
 418                var lapsed = new OperationCanceledException(
 419                    "The ACK-after-enqueue drain budget lapsed before this already-deleted message was handled.");
 420                Logger.LogWarning(
 421                    "SQS background handler for already-deleted message {MessageId} on {Queue} was not started: the drai
 422                    delivery.MessageId,
 423                    _queueName);
 424                await NotifyBackgroundFailureAsync(delivery, lapsed, _queueName, _role).ConfigureAwait(false);
 425                continue;
 426            }
 427
 428            Interlocked.Increment(ref _runningCount);
 429
 430            try
 431            {
 432                Logger.LogDebug(
 433                    "SQS background worker {WorkerIndex} handling message {MessageId} for {Queue}. Pending={PendingCount
 434                    workerIndex,
 435                    delivery.MessageId,
 436                    _queueName,
 437                    PendingCount,
 438                    RunningCount);
 439                await ExecuteHandlerAsync(
 440                    delivery,
 441                    _drainCancellation.Token,
 442                    logFailures: false).ConfigureAwait(false);
 443            }
 444            catch (Exception ex)
 445            {
 446                Logger.LogError(
 447                    ex,
 448                    "SQS background handler failed for already-deleted message {MessageId} on {Queue}.",
 449                    delivery.MessageId,
 450                    _queueName);
 451                await NotifyBackgroundFailureAsync(
 452                    delivery,
 453                    ex,
 454                    _queueName,
 455                    _role).ConfigureAwait(false);
 456            }
 457            finally
 458            {
 459                Interlocked.Decrement(ref _runningCount);
 460            }
 461        }
 462    }
 463}