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

Information
Class: AsyncResponse.Transports.MongoDB.MongoDbTransportDelivery
Assembly: AsyncResponse.Transports.MongoDB
File(s): /home/runner/work/AsyncResponse/AsyncResponse/src/Transports/AsyncResponse.Transports.MongoDB/MongoDbTransportStore.cs
Line coverage
100%
Covered lines: 10
Uncovered lines: 0
Coverable lines: 10
Total lines: 454
Line coverage: 100%
Branch coverage
N/A
Covered branches: 0
Total branches: 0
Branch coverage: N/A
Method coverage

Feature is only available for sponsors

Upgrade to PRO version

Metrics

MethodBranch coverage Crap Score Cyclomatic complexity Line coverage
.ctor(...)100%11100%

File(s)

/home/runner/work/AsyncResponse/AsyncResponse/src/Transports/AsyncResponse.Transports.MongoDB/MongoDbTransportStore.cs

#LineLine coverage
 1using Microsoft.Extensions.Logging;
 2using Microsoft.Extensions.Options;
 3using MongoDB.Bson;
 4using MongoDB.Bson.Serialization.Attributes;
 5using MongoDB.Bson.Serialization.Options;
 6using MongoDB.Driver;
 7using System.Runtime.CompilerServices;
 8using System.Security.Cryptography;
 9using System.Text;
 10
 11namespace AsyncResponse.Transports.MongoDB;
 12
 13internal enum MongoDbSubscriberRole
 14{
 15    Worker,
 16    ResponseIngress
 17}
 18
 19/// <summary>A claimed MongoDB transport document, decoupled from driver types for dispatch tests.</summary>
 20/// <remarks>
 21/// <c>RenewAsync</c> extends the claim's lease (<c>locked_until</c>) by the original lock timeout,
 22/// fenced on the claim's <c>lock_id</c>; it returns <c>false</c> when the fence no longer matches
 23/// (the lease lapsed and another subscriber re-claimed the document).
 24/// </remarks>
 325internal sealed record MongoDbTransportDelivery(
 326    Guid Id,
 327    string Queue,
 328    string Payload,
 329    IReadOnlyDictionary<string, string> Headers,
 330    int Attempt,
 331    Func<ValueTask> AckAsync,
 332    Func<TimeSpan, ValueTask> NakAsync,
 333    Func<Exception, bool, CancellationToken, ValueTask<bool>> DeadLetterAsync,
 334    Func<ValueTask<bool>> RenewAsync);
 35
 36/// <summary>Small document adapter for the MongoDB transport queue collection.</summary>
 37internal sealed class MongoDbTransportStore : IDisposable
 38{
 39    private readonly IMongoCollection<MongoTransportMessageDocument> _messages;
 40    private readonly MongoDbAsyncResponseTransportOptions _options;
 41    private readonly ILogger<MongoDbTransportStore>? _logger;
 42    private readonly SemaphoreSlim _ensureGate = new(1, 1);
 43    private readonly IMongoClient? _ownedClient;
 44    private bool _created;
 45    private long _lastDeadLetterPruneTicks;
 46
 47    public MongoDbTransportStore(
 48        IMongoDatabase database,
 49        IOptions<MongoDbAsyncResponseTransportOptions> options,
 50        IMongoClient? ownedClient = null,
 51        ILogger<MongoDbTransportStore>? logger = null)
 52    {
 53        _options = options.Value;
 54        _logger = logger;
 55        MongoDbTransportOptionsValidator.ValidateCommon(_options);
 56        _messages = database.GetCollection<MongoTransportMessageDocument>(_options.MessageCollection);
 57        _ownedClient = ownedClient;
 58    }
 59
 60    public async Task EnsureCreatedAsync(CancellationToken cancellationToken = default)
 61    {
 62        if (_created || !_options.AutoCreateIndexes)
 63            return;
 64
 65        await _ensureGate.WaitAsync(cancellationToken).ConfigureAwait(false);
 66        try
 67        {
 68            if (_created)
 69                return;
 70
 71            await _messages.Indexes.CreateOneAsync(
 72                new CreateIndexModel<MongoTransportMessageDocument>(
 73                    Builders<MongoTransportMessageDocument>.IndexKeys
 74                        .Ascending(item => item.Queue)
 75                        .Ascending(item => item.AvailableAtUtc)
 76                        .Ascending(item => item.CreatedAtUtc),
 77                    new CreateIndexOptions { Name = $"{_options.MessageCollection}_claim_idx" }),
 78                cancellationToken: cancellationToken).ConfigureAwait(false);
 79            await _messages.Indexes.CreateOneAsync(
 80                new CreateIndexModel<MongoTransportMessageDocument>(
 81                    Builders<MongoTransportMessageDocument>.IndexKeys.Ascending(item => item.CreatedAtUtc),
 82                    new CreateIndexOptions { Name = $"{_options.MessageCollection}_created_idx" }),
 83                cancellationToken: cancellationToken).ConfigureAwait(false);
 84            _created = true;
 85        }
 86        finally
 87        {
 88            _ensureGate.Release();
 89        }
 90    }
 91
 92    /// <summary>
 93    /// Publishes a queue document. The caller supplies the id so a retried publish is idempotent —
 94    /// a duplicate-key insert is treated as success rather than enqueuing the same job twice.
 95    /// </summary>
 96    public async Task PublishAsync(
 97        Guid id,
 98        string queue,
 99        string payload,
 100        IReadOnlyDictionary<string, string>? headers,
 101        CancellationToken cancellationToken)
 102    {
 103        await InsertAsync(id, queue, payload, headers, deadLetterReason: null, cancellationToken).ConfigureAwait(false);
 104        await PruneDeadLettersIfDueAsync(cancellationToken).ConfigureAwait(false);
 105    }
 106
 107    public async Task<MongoDbTransportDelivery?> TryClaimAsync(string queue, TimeSpan lockTimeout, CancellationToken can
 108    {
 109        await EnsureCreatedAsync(cancellationToken).ConfigureAwait(false);
 110        var lockId = Guid.NewGuid();
 111
 112        // findOneAndUpdate is atomic per document: of all competing consumers, exactly one observes
 113        // the document unlocked and stamps its lock_id/locked_until in the same server-side step.
 114        var claimed = await _messages.FindOneAndUpdateAsync(
 115            BuildClaimFilter(queue),
 116            BuildClaimUpdate(lockId, lockTimeout),
 117            new FindOneAndUpdateOptions<MongoTransportMessageDocument>
 118            {
 119                Sort = Builders<MongoTransportMessageDocument>.Sort.Ascending(item => item.CreatedAtUtc),
 120                ReturnDocument = ReturnDocument.After
 121            },
 122            cancellationToken).ConfigureAwait(false);
 123        if (claimed is null)
 124            return null;
 125
 126        var headers = claimed.Headers is null
 127            ? EmptyHeaders
 128            : new Dictionary<string, string>(claimed.Headers, StringComparer.OrdinalIgnoreCase);
 129
 130        return new MongoDbTransportDelivery(
 131            claimed.Id,
 132            claimed.Queue,
 133            claimed.Payload,
 134            headers,
 135            claimed.Attempts,
 136            () => AckAsync(claimed.Id, lockId),
 137            delay => NakAsync(claimed.Id, lockId, delay),
 138            (exception, deleteOriginal, token) => DeadLetterAsync(claimed.Id, lockId, queue, claimed.Payload, headers, e
 139            () => RenewLeaseAsync(claimed.Id, lockId, lockTimeout));
 140    }
 141
 142    /// <summary>
 143    /// Claim filter: available and not (still) locked, evaluated against the server clock
 144    /// (<c>$$NOW</c>) so publisher/consumer clock skew never fences messages in or out.
 145    /// A missing <c>locked_until</c> compares as null and therefore as expired.
 146    /// </summary>
 147    internal static FilterDefinition<MongoTransportMessageDocument> BuildClaimFilter(string queue)
 148        => Builders<MongoTransportMessageDocument>.Filter.Eq(item => item.Queue, queue)
 149           & new BsonDocumentFilterDefinition<MongoTransportMessageDocument>(new BsonDocument(
 150               "$expr",
 151               new BsonDocument("$and", new BsonArray
 152               {
 153                   new BsonDocument("$lte", new BsonArray { "$available_at", "$$NOW" }),
 154                   new BsonDocument("$or", new BsonArray
 155                   {
 156                       new BsonDocument("$eq", new BsonArray { "$locked_until", BsonNull.Value }),
 157                       new BsonDocument("$lte", new BsonArray { "$locked_until", "$$NOW" })
 158                   })
 159               })));
 160
 161    internal static UpdateDefinition<MongoTransportMessageDocument> BuildClaimUpdate(Guid lockId, TimeSpan lockTimeout)
 162        => Builders<MongoTransportMessageDocument>.Update.Pipeline(new[]
 163        {
 164            new BsonDocument("$set", new BsonDocument
 165            {
 166                ["attempts"] = new BsonDocument("$add", new BsonArray
 167                {
 168                    new BsonDocument("$ifNull", new BsonArray { "$attempts", 0 }),
 169                    1
 170                }),
 171                ["locked_until"] = new BsonDocument("$add", new BsonArray { "$$NOW", lockTimeout.TotalMilliseconds }),
 172                ["lock_id"] = new BsonBinaryData(lockId, GuidRepresentation.Standard)
 173            })
 174        });
 175
 176    public async IAsyncEnumerable<MongoDbTransportDelivery> ClaimBatchAsync(
 177        string queue,
 178        int batchSize,
 179        TimeSpan lockTimeout,
 180        [EnumeratorCancellation] CancellationToken cancellationToken)
 181    {
 182        for (var i = 0; i < batchSize; i++)
 183        {
 184            var delivery = await TryClaimAsync(queue, lockTimeout, cancellationToken).ConfigureAwait(false);
 185            if (delivery is null)
 186                yield break;
 187            yield return delivery;
 188        }
 189    }
 190
 191    private async Task InsertAsync(
 192        Guid id,
 193        string queue,
 194        string payload,
 195        IReadOnlyDictionary<string, string>? headers,
 196        string? deadLetterReason,
 197        CancellationToken cancellationToken)
 198    {
 199        await EnsureCreatedAsync(cancellationToken).ConfigureAwait(false);
 200        var now = DateTime.UtcNow;
 201        var document = new MongoTransportMessageDocument
 202        {
 203            Id = id,
 204            Queue = queue,
 205            Payload = payload,
 206            Headers = headers is null
 207                ? new Dictionary<string, string>(0, StringComparer.OrdinalIgnoreCase)
 208                : new Dictionary<string, string>(headers, StringComparer.OrdinalIgnoreCase),
 209            CreatedAtUtc = now,
 210            // "Available immediately on arrival": InsertOne cannot evaluate $$NOW, and stamping the
 211            // client clock here would let client-ahead-of-server skew hide a fresh message from the
 212            // server-clock claim filter until the skew elapsed. Epoch expresses what the SQL stores'
 213            // "available_at DEFAULT now()" expresses; a NAK re-stamps a real server-relative time.
 214            AvailableAtUtc = DateTime.UnixEpoch,
 215            Attempts = 0,
 216            DeadLetterReason = deadLetterReason
 217        };
 218        try
 219        {
 220            await _messages.InsertOneAsync(document, cancellationToken: cancellationToken).ConfigureAwait(false);
 221        }
 222        catch (MongoWriteException ex) when (ex.WriteError?.Category == ServerErrorCategory.DuplicateKey)
 223        {
 224            // A retried publish found the document already inserted: idempotent success.
 225        }
 226    }
 227
 228    private async ValueTask AckAsync(Guid id, Guid lockId)
 229    {
 230        await _messages.DeleteOneAsync(
 231            Builders<MongoTransportMessageDocument>.Filter.Eq(item => item.Id, id)
 232            & Builders<MongoTransportMessageDocument>.Filter.Eq(item => item.LockId, lockId),
 233            CancellationToken.None).ConfigureAwait(false);
 234    }
 235
 236    private async ValueTask NakAsync(Guid id, Guid lockId, TimeSpan delay)
 237    {
 238        await _messages.UpdateOneAsync(
 239            Builders<MongoTransportMessageDocument>.Filter.Eq(item => item.Id, id)
 240            & Builders<MongoTransportMessageDocument>.Filter.Eq(item => item.LockId, lockId),
 241            BuildNakUpdate(delay),
 242            options: null,
 243            CancellationToken.None).ConfigureAwait(false);
 244    }
 245
 246    private async ValueTask<bool> RenewLeaseAsync(Guid id, Guid lockId, TimeSpan lockTimeout)
 247    {
 248        var result = await _messages.UpdateOneAsync(
 249            Builders<MongoTransportMessageDocument>.Filter.Eq(item => item.Id, id)
 250            & Builders<MongoTransportMessageDocument>.Filter.Eq(item => item.LockId, lockId),
 251            BuildRenewUpdate(lockTimeout),
 252            options: null,
 253            CancellationToken.None).ConfigureAwait(false);
 254        return result.MatchedCount > 0;
 255    }
 256
 257    /// <summary>
 258    /// Fenced lease renewal: extends <c>locked_until</c> from the server clock (<c>$$NOW</c>) only
 259    /// while the claim's <c>lock_id</c> fence still matches, mirroring the claim update.
 260    /// </summary>
 261    internal static UpdateDefinition<MongoTransportMessageDocument> BuildRenewUpdate(TimeSpan lockTimeout)
 262        => Builders<MongoTransportMessageDocument>.Update.Pipeline(new[]
 263        {
 264            new BsonDocument("$set", new BsonDocument
 265            {
 266                ["locked_until"] = new BsonDocument("$add", new BsonArray { "$$NOW", lockTimeout.TotalMilliseconds })
 267            })
 268        });
 269
 270    internal static UpdateDefinition<MongoTransportMessageDocument> BuildNakUpdate(TimeSpan delay)
 271        => Builders<MongoTransportMessageDocument>.Update.Pipeline(new[]
 272        {
 273            new BsonDocument("$set", new BsonDocument
 274            {
 275                ["available_at"] = new BsonDocument("$add", new BsonArray { "$$NOW", delay.TotalMilliseconds }),
 276                ["locked_until"] = BsonNull.Value,
 277                ["lock_id"] = BsonNull.Value
 278            })
 279        });
 280
 281    private async ValueTask<bool> DeadLetterAsync(
 282        Guid id,
 283        Guid lockId,
 284        string sourceQueue,
 285        string payload,
 286        IReadOnlyDictionary<string, string> headers,
 287        Exception exception,
 288        bool deleteOriginal,
 289        CancellationToken cancellationToken)
 290    {
 291        if (!_options.DeadLetterEnabled)
 292        {
 293            if (deleteOriginal)
 294                await AckAsync(id, lockId).ConfigureAwait(false);
 295            return true;
 296        }
 297
 298        var deadHeaders = new Dictionary<string, string>(headers, StringComparer.OrdinalIgnoreCase)
 299        {
 300            ["AR-DeadLetter-Reason"] = Sanitize(exception.Message),
 301            ["AR-DeadLetter-Source-Queue"] = sourceQueue
 302        };
 303
 304        try
 305        {
 306            if (!deleteOriginal)
 307            {
 308                await InsertAsync(Guid.NewGuid(), _options.DeadLetterQueue, payload, deadHeaders, exception.Message, can
 309                return true;
 310            }
 311
 312            // MongoDB has no cross-document transaction we can rely on here (the transport must also
 313            // work on standalone servers), so the DLQ insert uses an id derived deterministically
 314            // from the source document: if a crash lands between the insert and the delete, the
 315            // redelivered message dead-letters onto the same id and the duplicate-key insert is
 316            // swallowed — the DLQ never accumulates copies of one poison message.
 317            await InsertAsync(DeadLetterId(id), _options.DeadLetterQueue, payload, deadHeaders, exception.Message, cance
 318            await AckAsync(id, lockId).ConfigureAwait(false);
 319            return true;
 320        }
 321        catch (Exception ex)
 322        {
 323            // Callers decide the redelivery consequence from the false return; log the cause here so
 324            // a failing dead-letter write is never silent.
 325            _logger?.LogError(
 326                ex,
 327                "Failed to write MongoDB dead-letter document for message {MessageId} from queue {SourceQueue}.",
 328                id,
 329                sourceQueue);
 330            return false;
 331        }
 332    }
 333
 334    /// <summary>
 335    /// Watches the queue collection with a change stream and invokes
 336    /// <paramref name="onNotification"/> whenever a document is inserted into
 337    /// <paramref name="queue"/>. Runs until cancellation or a stream error; callers treat the wake
 338    /// as an optimization over <see cref="MongoDbSubscriberOptions.EmptyPollDelay"/> polling.
 339    /// </summary>
 340    public async Task WatchQueueAsync(string queue, Func<Task> onNotification, CancellationToken cancellationToken)
 341    {
 342        await EnsureCreatedAsync(cancellationToken).ConfigureAwait(false);
 343        using var cursor = await _messages.WatchAsync(
 344            BuildQueueWatchPipeline(queue),
 345            cancellationToken: cancellationToken).ConfigureAwait(false);
 346        while (await cursor.MoveNextAsync(cancellationToken).ConfigureAwait(false))
 347        {
 348            foreach (var _ in cursor.Current)
 349                await onNotification().ConfigureAwait(false);
 350        }
 351    }
 352
 353    /// <summary>Change-stream pipeline for queue wakes: a <c>$match</c> on inserts into one logical queue.</summary>
 354    internal static PipelineDefinition<ChangeStreamDocument<MongoTransportMessageDocument>, ChangeStreamDocument<MongoTr
 355        => new EmptyPipelineDefinition<ChangeStreamDocument<MongoTransportMessageDocument>>()
 356            .Match(new BsonDocument("$and", new BsonArray
 357            {
 358                new BsonDocument("operationType", "insert"),
 359                new BsonDocument("fullDocument.queue", queue)
 360            }));
 361
 362    /// <summary>Returns <c>true</c> when the server rejected the change stream itself (not a transient cursor error).</
 363    internal static bool IsChangeStreamUnsupported(Exception exception)
 364        => exception is MongoCommandException commandException
 365           && (commandException.Code == 40573
 366               || commandException.Message.Contains("only supported on replica sets", StringComparison.OrdinalIgnoreCase
 367
 368    /// <summary>
 369    /// Opportunistically deletes dead-letter documents older than the configured retention. No-op
 370    /// unless <see cref="MongoDbAsyncResponseTransportOptions.DeadLetterRetention"/> is set, and
 371    /// throttled so the delete runs at most once per minute regardless of publish rate.
 372    /// </summary>
 373    private async Task PruneDeadLettersIfDueAsync(CancellationToken cancellationToken)
 374    {
 375        if (_options.DeadLetterRetention is not { } retention || !ShouldPruneDeadLetters())
 376            return;
 377
 378        await _messages.DeleteManyAsync(
 379            Builders<MongoTransportMessageDocument>.Filter.Eq(item => item.Queue, _options.DeadLetterQueue)
 380            & Builders<MongoTransportMessageDocument>.Filter.Lt(item => item.CreatedAtUtc, DateTime.UtcNow.Subtract(rete
 381            cancellationToken).ConfigureAwait(false);
 382    }
 383
 384    private bool ShouldPruneDeadLetters()
 385    {
 386        var now = DateTime.UtcNow.Ticks;
 387        var last = Interlocked.Read(ref _lastDeadLetterPruneTicks);
 388        return now - last >= DeadLetterPruneThrottle.Ticks
 389            && Interlocked.CompareExchange(ref _lastDeadLetterPruneTicks, now, last) == last;
 390    }
 391
 392    /// <summary>
 393    /// Deterministic dead-letter document id for a source message: the same poison message always
 394    /// maps to the same DLQ id, making the insert-then-delete pair idempotent under crash-redelivery.
 395    /// </summary>
 396    internal static Guid DeadLetterId(Guid sourceId)
 397    {
 398        var hash = SHA256.HashData(Encoding.UTF8.GetBytes($"asyncresponse:deadletter:{sourceId:N}"));
 399        return new Guid(hash.AsSpan(0, 16));
 400    }
 401
 402    private static string Sanitize(string value) => value.Replace('\r', ' ').Replace('\n', ' ');
 403
 404    private static readonly TimeSpan DeadLetterPruneThrottle = TimeSpan.FromMinutes(1);
 405
 406    private static readonly IReadOnlyDictionary<string, string> EmptyHeaders =
 407        new Dictionary<string, string>(0, StringComparer.OrdinalIgnoreCase);
 408
 409    /// <summary>Disposes the Mongo client when the store created (and therefore owns) it.</summary>
 410    public void Dispose()
 411    {
 412        _ensureGate.Dispose();
 413        (_ownedClient as IDisposable)?.Dispose();
 414    }
 415}
 416
 417internal sealed class MongoTransportMessageDocument
 418{
 419    [BsonId]
 420    [BsonElement("_id")]
 421    [BsonGuidRepresentation(GuidRepresentation.Standard)]
 422    public Guid Id { get; set; }
 423
 424    [BsonElement("queue")]
 425    public string Queue { get; set; } = "";
 426
 427    [BsonElement("payload")]
 428    public string Payload { get; set; } = "";
 429
 430    // Array-of-documents representation keeps arbitrary header names (dots, dollars) legal as values
 431    // rather than as BSON field names.
 432    [BsonElement("headers")]
 433    [BsonDictionaryOptions(DictionaryRepresentation.ArrayOfDocuments)]
 434    public Dictionary<string, string>? Headers { get; set; }
 435
 436    [BsonElement("created_at")]
 437    public DateTime CreatedAtUtc { get; set; }
 438
 439    [BsonElement("available_at")]
 440    public DateTime AvailableAtUtc { get; set; }
 441
 442    [BsonElement("locked_until")]
 443    public DateTime? LockedUntilUtc { get; set; }
 444
 445    [BsonElement("lock_id")]
 446    [BsonGuidRepresentation(GuidRepresentation.Standard)]
 447    public Guid? LockId { get; set; }
 448
 449    [BsonElement("attempts")]
 450    public int Attempts { get; set; }
 451
 452    [BsonElement("dead_letter_reason")]
 453    public string? DeadLetterReason { get; set; }
 454}