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

Information
Class: AsyncResponse.Transports.SQS.QueuedSqsMessageDispatcher
Assembly: AsyncResponse.Transports.SQS
File(s): /home/runner/work/AsyncResponse/AsyncResponse/src/Transports/AsyncResponse.Transports.SQS/SqsMessageDispatcher.cs
Line coverage
99%
Covered lines: 107
Uncovered lines: 1
Coverable lines: 108
Total lines: 408
Line coverage: 99%
Branch coverage
100%
Covered branches: 8
Total branches: 8
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%
get_CanAcceptMore()100%11100%
get_FreeCapacity()100%11100%
WaitForCapacityAsync()100%2275%
HandleAsync()100%44100%
DisposeAsync()100%22100%
RunWorkerAsync()100%11100%

File(s)

/home/runner/work/AsyncResponse/AsyncResponse/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            await delivery.ChangeVisibilityAsync(delay).ConfigureAwait(false);
 144        }
 145        catch (Exception ex)
 146        {
 147            Logger.LogWarning(
 148                ex,
 149                "Failed to change visibility of SQS message {MessageId} on {Queue}; it stays invisible until the visibil
 150                delivery.MessageId,
 151                _queue);
 152        }
 153    }
 154
 155    protected async ValueTask NotifyBackgroundFailureAsync(
 156        SqsTransportDelivery delivery,
 157        Exception exception,
 158        string queue,
 159        SqsSubscriberRole role)
 160    {
 161        var callback = _subscriberOptions.OnBackgroundFailure;
 162        if (callback is null)
 163            return;
 164
 165        try
 166        {
 167            var context = new SqsBackgroundFailureContext(
 168                queue,
 169                role.ToString(),
 170                delivery.MessageId,
 171                delivery.ReceiveCount,
 172                TryReadCorrelationId(delivery),
 173                exception);
 174            await callback(context).ConfigureAwait(false);
 175        }
 176        catch (Exception callbackException)
 177        {
 178            Logger.LogError(
 179                callbackException,
 180                "SQS background failure callback failed for already-deleted message {MessageId} on {Queue}.",
 181                delivery.MessageId,
 182                queue);
 183        }
 184    }
 185
 186    private string? TryReadCorrelationId(SqsTransportDelivery delivery)
 187        => !string.IsNullOrWhiteSpace(_transportOptions.CorrelationIdAttribute)
 188            && delivery.MessageAttributes.TryGetValue(_transportOptions.CorrelationIdAttribute, out var value)
 189            && !string.IsNullOrWhiteSpace(value)
 190                ? value
 191                : null;
 192}
 193
 194internal sealed class AwaitingSqsMessageDispatcher(
 195    Func<SqsTransportDelivery, CancellationToken, Task> handler,
 196    SqsAsyncResponseOptions transportOptions,
 197    SqsSubscriberOptions subscriberOptions,
 198    ILogger logger,
 199    string queue,
 200    SqsSubscriberRole role)
 201    : SqsMessageDispatcher(handler, transportOptions, subscriberOptions, logger, queue, role)
 202{
 203    /// <summary>Handles the delivered message.</summary>
 204    public override async Task HandleAsync(
 205        SqsTransportDelivery delivery,
 206        CancellationToken subscriberCancellationToken)
 207    {
 208        try
 209        {
 210            await ExecuteHandlerAsync(delivery, subscriberCancellationToken).ConfigureAwait(false);
 211            await delivery.DeleteAsync().ConfigureAwait(false);
 212        }
 213        catch (Exception)
 214        {
 215            // SQS has no explicit NACK or dead-letter call: leaving the message undeleted lets it
 216            // reappear when its visibility timeout expires, ApproximateReceiveCount increments, and
 217            // the queue's redrive policy dead-letters it after maxReceiveCount receives.
 218            if (RedeliveryDelay is { } redeliveryDelay)
 219                await TryChangeVisibilityAsync(delivery, redeliveryDelay).ConfigureAwait(false);
 220        }
 221    }
 222}
 223
 224internal sealed class QueuedSqsMessageDispatcher : SqsMessageDispatcher
 225{
 226    private readonly Channel<SqsTransportDelivery> _queue;
 227    private readonly Task[] _workers;
 3228    private readonly CancellationTokenSource _drainCancellation = new();
 229    private readonly TimeSpan _drainTimeout;
 230    private readonly int _capacity;
 231    private readonly string _queueName;
 232    private readonly SqsSubscriberRole _role;
 233    private int _pendingCount;
 234    private int _runningCount;
 235    private int _disposeStarted;
 236
 237    /// <summary>Creates an ACK-after-enqueue dispatcher with a bounded background queue.</summary>
 238    public QueuedSqsMessageDispatcher(
 239        Func<SqsTransportDelivery, CancellationToken, Task> handler,
 240        SqsAsyncResponseOptions transportOptions,
 241        SqsSubscriberOptions subscriberOptions,
 242        ILogger logger,
 243        string queue,
 244        SqsSubscriberRole role)
 3245        : base(handler, transportOptions, subscriberOptions, logger, queue, role)
 246    {
 3247        _drainTimeout = subscriberOptions.BackgroundDrainTimeout;
 3248        _capacity = subscriberOptions.BackgroundQueueCapacity;
 3249        _queueName = queue;
 3250        _role = role;
 3251        _queue = Channel.CreateBounded<SqsTransportDelivery>(new BoundedChannelOptions(subscriberOptions.BackgroundQueue
 3252        {
 3253            AllowSynchronousContinuations = false,
 3254            FullMode = BoundedChannelFullMode.Wait,
 3255            SingleReader = subscriberOptions.BackgroundWorkerCount == 1,
 3256            SingleWriter = false
 3257        });
 258
 3259        _workers = Enumerable.Range(0, subscriberOptions.BackgroundWorkerCount)
 3260            .Select(workerIndex => Task.Run(() => RunWorkerAsync(workerIndex)))
 3261            .ToArray();
 262
 3263        Logger.LogInformation(
 3264            "Created SQS ACK-after-enqueue dispatcher for {Queue} with {WorkerCount} worker(s), queue capacity {QueueCap
 3265            _queueName,
 3266            subscriberOptions.BackgroundWorkerCount,
 3267            subscriberOptions.BackgroundQueueCapacity,
 3268            _drainTimeout);
 3269    }
 270
 3271    internal int PendingCount => Volatile.Read(ref _pendingCount);
 3272    internal int RunningCount => Volatile.Read(ref _runningCount);
 273
 3274    public override bool CanAcceptMore => Volatile.Read(ref _pendingCount) < _capacity;
 275
 3276    public override int FreeCapacity => Math.Max(0, _capacity - Volatile.Read(ref _pendingCount));
 277
 278    public override async ValueTask WaitForCapacityAsync(CancellationToken cancellationToken)
 279    {
 280        // WaitToWriteAsync completes when the bounded channel has room (or the channel is completed
 281        // during dispose, in which case there is nothing left to gate).
 3282        while (!CanAcceptMore)
 283        {
 2284            if (!await _queue.Writer.WaitToWriteAsync(cancellationToken).ConfigureAwait(false))
 0285                return;
 286        }
 3287    }
 288
 289    /// <summary>Handles the delivered message.</summary>
 290    public override async Task HandleAsync(
 291        SqsTransportDelivery delivery,
 292        CancellationToken subscriberCancellationToken)
 293    {
 3294        Interlocked.Increment(ref _pendingCount);
 3295        if (!_queue.Writer.TryWrite(delivery))
 296        {
 2297            Interlocked.Decrement(ref _pendingCount);
 2298            Logger.LogWarning(
 2299                "SQS background queue rejected message {MessageId} for {Queue}; leaving it to redeliver via its visibili
 2300                delivery.MessageId,
 2301                _queueName,
 2302                PendingCount,
 2303                RunningCount);
 304            // Do not release visibility to zero here: SQS counts every receive toward the queue's
 305            // redrive policy, so an instantly re-receivable message that keeps hitting a full queue
 306            // would cross maxReceiveCount and dead-letter without ever being processed. Let the
 307            // visibility timeout lapse naturally (or shorten it via RedeliveryDelay when configured)
 308            // so redelivery lands after capacity has had time to free.
 2309            if (RedeliveryDelay is { } redeliveryDelay)
 2310                await TryChangeVisibilityAsync(delivery, redeliveryDelay).ConfigureAwait(false);
 2311            return;
 312        }
 313
 314        // The delivery now belongs to a background worker, which decrements _pendingCount when it
 315        // dequeues. Do not touch the counter or release visibility here, even if the delete below
 316        // fails — the message is already executing in-process and releasing it would trigger a
 317        // duplicate execution via redelivery.
 318        try
 319        {
 3320            await delivery.DeleteAsync().ConfigureAwait(false);
 3321        }
 2322        catch (Exception ex)
 323        {
 2324            Logger.LogError(
 2325                ex,
 2326                "Failed to delete SQS message {MessageId} for {Queue} after enqueue; it is being processed but SQS will 
 2327                delivery.MessageId,
 2328                _queueName);
 2329        }
 3330    }
 331
 332    /// <summary>Releases resources held by this instance.</summary>
 333    public override async ValueTask DisposeAsync()
 334    {
 3335        if (Interlocked.Exchange(ref _disposeStarted, 1) != 0)
 2336            return;
 337
 3338        Logger.LogInformation(
 3339            "Draining SQS ACK-after-enqueue dispatcher for {Queue}. Pending={PendingCount}, Running={RunningCount}.",
 3340            _queueName,
 3341            PendingCount,
 3342            RunningCount);
 3343        _queue.Writer.TryComplete();
 344
 345        try
 346        {
 3347            await Task.WhenAll(_workers).WaitAsync(_drainTimeout).ConfigureAwait(false);
 3348            _drainCancellation.Dispose();
 3349        }
 2350        catch (TimeoutException ex)
 351        {
 2352            _drainCancellation.Cancel();
 2353            Logger.LogWarning(
 2354                ex,
 2355                "Timed out while draining SQS ACK-after-enqueue dispatcher for {Queue}. Pending={PendingCount}, Running=
 2356                _queueName,
 2357                PendingCount,
 2358                RunningCount);
 359
 2360            _ = Task.WhenAll(_workers).ContinueWith(
 2361                _ => _drainCancellation.Dispose(),
 2362                CancellationToken.None,
 2363                TaskContinuationOptions.ExecuteSynchronously,
 2364                TaskScheduler.Default);
 2365        }
 3366    }
 367
 368    private async Task RunWorkerAsync(int workerIndex)
 369    {
 3370        await foreach (var delivery in _queue.Reader.ReadAllAsync().ConfigureAwait(false))
 371        {
 3372            Interlocked.Decrement(ref _pendingCount);
 3373            Interlocked.Increment(ref _runningCount);
 374
 375            try
 376            {
 3377                Logger.LogDebug(
 3378                    "SQS background worker {WorkerIndex} handling message {MessageId} for {Queue}. Pending={PendingCount
 3379                    workerIndex,
 3380                    delivery.MessageId,
 3381                    _queueName,
 3382                    PendingCount,
 3383                    RunningCount);
 3384                await ExecuteHandlerAsync(
 3385                    delivery,
 3386                    _drainCancellation.Token,
 3387                    logFailures: false).ConfigureAwait(false);
 3388            }
 2389            catch (Exception ex)
 390            {
 2391                Logger.LogError(
 2392                    ex,
 2393                    "SQS background handler failed for already-deleted message {MessageId} on {Queue}.",
 2394                    delivery.MessageId,
 2395                    _queueName);
 2396                await NotifyBackgroundFailureAsync(
 2397                    delivery,
 2398                    ex,
 2399                    _queueName,
 2400                    _role).ConfigureAwait(false);
 401            }
 402            finally
 403            {
 3404                Interlocked.Decrement(ref _runningCount);
 405            }
 3406        }
 3407    }
 408}