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

Information
Class: AsyncResponse.Transports.PostgreSQL.PostgreSqlTransportDelivery
Assembly: AsyncResponse.Transports.PostgreSQL
File(s): /home/runner/work/AsyncResponse/AsyncResponse/src/Transports/AsyncResponse.Transports.PostgreSQL/PostgreSqlTransportStore.cs
Line coverage
100%
Covered lines: 10
Uncovered lines: 0
Coverable lines: 10
Total lines: 439
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.PostgreSQL/PostgreSqlTransportStore.cs

#LineLine coverage
 1using Microsoft.Extensions.Logging;
 2using Microsoft.Extensions.Options;
 3using Npgsql;
 4using NpgsqlTypes;
 5using System.Runtime.CompilerServices;
 6using System.Text;
 7using System.Text.Json;
 8
 9namespace AsyncResponse.Transports.PostgreSQL;
 10
 11internal enum PostgreSqlSubscriberRole
 12{
 13    Worker,
 14    ResponseIngress
 15}
 16
 17/// <summary>A claimed PostgreSQL transport row, decoupled from Npgsql types for dispatch tests.</summary>
 18/// <remarks>
 19/// <c>RenewAsync</c> extends the claim's lease (<c>locked_until</c>) by the original lock timeout,
 20/// fenced on the claim's <c>lock_id</c>; it returns <c>false</c> when the fence no longer matches
 21/// (the lease lapsed and another subscriber re-claimed the row).
 22/// </remarks>
 323internal sealed record PostgreSqlTransportDelivery(
 324    Guid Id,
 325    string Queue,
 326    string Payload,
 327    IReadOnlyDictionary<string, string> Headers,
 328    int Attempt,
 329    Func<ValueTask> AckAsync,
 330    Func<TimeSpan, ValueTask> NakAsync,
 331    Func<Exception, bool, CancellationToken, ValueTask<bool>> DeadLetterAsync,
 332    Func<ValueTask<bool>> RenewAsync);
 33
 34/// <summary>Small SQL adapter for the PostgreSQL transport queue table.</summary>
 35internal sealed class PostgreSqlTransportStore
 36{
 37    private readonly NpgsqlDataSource _dataSource;
 38    private readonly PostgreSqlAsyncResponseTransportOptions _options;
 39    private readonly ILogger<PostgreSqlTransportStore>? _logger;
 40    private readonly SemaphoreSlim _ensureGate = new(1, 1);
 41    private bool _created;
 42    private readonly long _schemaLockKey;
 43    private long _lastDeadLetterPruneTicks;
 44
 45    public PostgreSqlTransportStore(
 46        NpgsqlDataSource dataSource,
 47        IOptions<PostgreSqlAsyncResponseTransportOptions> options,
 48        ILogger<PostgreSqlTransportStore>? logger = null)
 49    {
 50        _dataSource = dataSource;
 51        _options = options.Value;
 52        _logger = logger;
 53        PostgreSqlTransportOptionsValidator.ValidateCommon(_options);
 54        Schema = Quote(_options.SchemaName);
 55        MessageTable = $"{Schema}.{Quote(_options.MessageTable)}";
 56        _schemaLockKey = SchemaAdvisoryLockKey(_options.SchemaName);
 57    }
 58
 59    public string Schema { get; }
 60    public string MessageTable { get; }
 61
 62    public async Task EnsureCreatedAsync(CancellationToken cancellationToken = default)
 63    {
 64        if (_created || !_options.AutoCreateSchema)
 65            return;
 66
 67        await _ensureGate.WaitAsync(cancellationToken).ConfigureAwait(false);
 68        try
 69        {
 70            if (_created)
 71                return;
 72
 73            await using var connection = await _dataSource.OpenConnectionAsync(cancellationToken).ConfigureAwait(false);
 74            await using var transaction = await connection.BeginTransactionAsync(cancellationToken).ConfigureAwait(false
 75
 76            // Serialize schema creation across processes. CREATE ... IF NOT EXISTS is not atomic against a
 77            // concurrent create of the same object: two instances starting together both pass the existence
 78            // check and collide on the system catalog ("duplicate key ... pg_type_typname_nsp_index"). A
 79            // transaction-scoped advisory lock (keyed by schema, shared with the channel store) lets one
 80            // instance build the schema while the rest wait and then find it already present.
 81            await using (var lockCommand = connection.CreateCommand())
 82            {
 83                lockCommand.Transaction = transaction;
 84                lockCommand.CommandText = "SELECT pg_advisory_xact_lock(@lock_key);";
 85                lockCommand.Parameters.AddWithValue("lock_key", _schemaLockKey);
 86                await lockCommand.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
 87            }
 88
 89            await using var command = connection.CreateCommand();
 90            command.Transaction = transaction;
 91            command.CommandText =
 92                $"""
 93                CREATE SCHEMA IF NOT EXISTS {Schema};
 94
 95                CREATE TABLE IF NOT EXISTS {MessageTable} (
 96                    id uuid PRIMARY KEY,
 97                    queue text NOT NULL,
 98                    payload_json jsonb NOT NULL,
 99                    headers_json jsonb NOT NULL DEFAULT jsonb_build_object(),
 100                    created_at timestamptz NOT NULL DEFAULT now(),
 101                    available_at timestamptz NOT NULL DEFAULT now(),
 102                    locked_until timestamptz NULL,
 103                    lock_id uuid NULL,
 104                    attempts integer NOT NULL DEFAULT 0,
 105                    dead_letter_reason text NULL
 106                );
 107                CREATE INDEX IF NOT EXISTS {Quote(IndexName(_options.MessageTable, "claim"))}
 108                    ON {MessageTable} (queue, available_at, locked_until, created_at);
 109                CREATE INDEX IF NOT EXISTS {Quote(IndexName(_options.MessageTable, "created"))}
 110                    ON {MessageTable} (created_at);
 111                """;
 112            await command.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
 113            await transaction.CommitAsync(cancellationToken).ConfigureAwait(false);
 114            _created = true;
 115        }
 116        finally
 117        {
 118            _ensureGate.Release();
 119        }
 120    }
 121
 122    /// <summary>
 123    /// Publishes a queue row. The caller supplies the id so a retried publish is idempotent
 124    /// (<c>ON CONFLICT DO NOTHING</c>) rather than inserting a duplicate job.
 125    /// </summary>
 126    public async Task PublishAsync(
 127        Guid id,
 128        string queue,
 129        string payload,
 130        IReadOnlyDictionary<string, string>? headers,
 131        CancellationToken cancellationToken)
 132    {
 133        await InsertAsync(id, queue, payload, headers, deadLetterReason: null, notify: true, cancellationToken).Configur
 134        await PruneDeadLettersIfDueAsync(cancellationToken).ConfigureAwait(false);
 135    }
 136
 137    public async Task<PostgreSqlTransportDelivery?> TryClaimAsync(string queue, TimeSpan lockTimeout, CancellationToken 
 138    {
 139        await EnsureCreatedAsync(cancellationToken).ConfigureAwait(false);
 140        var lockId = Guid.NewGuid();
 141
 142        await using var connection = await _dataSource.OpenConnectionAsync(cancellationToken).ConfigureAwait(false);
 143        await using var command = connection.CreateCommand();
 144        command.CommandText =
 145            $"""
 146            WITH next AS (
 147                SELECT id
 148                FROM {MessageTable}
 149                WHERE queue = @queue
 150                  AND available_at <= now()
 151                  AND (locked_until IS NULL OR locked_until <= now())
 152                ORDER BY created_at
 153                FOR UPDATE SKIP LOCKED
 154                LIMIT 1
 155            )
 156            UPDATE {MessageTable} AS message
 157            SET attempts = message.attempts + 1,
 158                locked_until = now() + @lock_timeout,
 159                lock_id = @lock_id
 160            FROM next
 161            WHERE message.id = next.id
 162            RETURNING message.id, message.queue, message.payload_json::text, message.headers_json::text, message.attempt
 163            """;
 164        command.Parameters.AddWithValue("queue", queue);
 165        command.Parameters.AddWithValue("lock_timeout", lockTimeout);
 166        command.Parameters.AddWithValue("lock_id", lockId);
 167
 168        await using var reader = await command.ExecuteReaderAsync(cancellationToken).ConfigureAwait(false);
 169        if (!await reader.ReadAsync(cancellationToken).ConfigureAwait(false))
 170            return null;
 171
 172        var id = reader.GetGuid(0);
 173        var payload = reader.GetString(2);
 174        var headerJson = reader.GetString(3);
 175        var attempt = reader.GetInt32(4);
 176        var headers = DeserializeHeaders(headerJson);
 177
 178        return new PostgreSqlTransportDelivery(
 179            id,
 180            reader.GetString(1),
 181            payload,
 182            headers,
 183            attempt,
 184            () => AckAsync(id, lockId),
 185            delay => NakAsync(id, lockId, delay),
 186            (exception, deleteOriginal, token) => DeadLetterAsync(id, lockId, queue, payload, headers, exception, delete
 187            () => RenewLeaseAsync(id, lockId, lockTimeout));
 188    }
 189
 190    public async IAsyncEnumerable<PostgreSqlTransportDelivery> ClaimBatchAsync(
 191        string queue,
 192        int batchSize,
 193        TimeSpan lockTimeout,
 194        [EnumeratorCancellation] CancellationToken cancellationToken)
 195    {
 196        for (var i = 0; i < batchSize; i++)
 197        {
 198            var delivery = await TryClaimAsync(queue, lockTimeout, cancellationToken).ConfigureAwait(false);
 199            if (delivery is null)
 200                yield break;
 201            yield return delivery;
 202        }
 203    }
 204
 205    private async Task InsertAsync(
 206        Guid id,
 207        string queue,
 208        string payload,
 209        IReadOnlyDictionary<string, string>? headers,
 210        string? deadLetterReason,
 211        bool notify,
 212        CancellationToken cancellationToken)
 213    {
 214        await EnsureCreatedAsync(cancellationToken).ConfigureAwait(false);
 215        await using var connection = await _dataSource.OpenConnectionAsync(cancellationToken).ConfigureAwait(false);
 216        await using var command = connection.CreateCommand();
 217        command.CommandText =
 218            $"""
 219            INSERT INTO {MessageTable} (id, queue, payload_json, headers_json, dead_letter_reason)
 220            VALUES (@id, @queue, @payload_json, @headers_json, @dead_letter_reason)
 221            ON CONFLICT (id) DO NOTHING;
 222            """;
 223        if (notify)
 224            command.CommandText += "SELECT pg_notify(@channel, @payload);";
 225
 226        command.Parameters.AddWithValue("id", id);
 227        command.Parameters.AddWithValue("queue", queue);
 228        command.Parameters.Add("payload_json", NpgsqlDbType.Jsonb).Value = payload;
 229        command.Parameters.Add("headers_json", NpgsqlDbType.Jsonb).Value = AsyncResponseJson.Serialize(headers ?? EmptyH
 230        command.Parameters.AddWithValue("dead_letter_reason", (object?)deadLetterReason ?? DBNull.Value);
 231        if (notify)
 232        {
 233            command.Parameters.AddWithValue("channel", _options.NotificationChannel);
 234            command.Parameters.AddWithValue("payload", queue);
 235        }
 236
 237        await command.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
 238    }
 239
 240    private async ValueTask AckAsync(Guid id, Guid lockId)
 241    {
 242        await using var connection = await _dataSource.OpenConnectionAsync(CancellationToken.None).ConfigureAwait(false)
 243        await using var command = connection.CreateCommand();
 244        command.CommandText = $"DELETE FROM {MessageTable} WHERE id = @id AND lock_id = @lock_id;";
 245        command.Parameters.AddWithValue("id", id);
 246        command.Parameters.AddWithValue("lock_id", lockId);
 247        await command.ExecuteNonQueryAsync(CancellationToken.None).ConfigureAwait(false);
 248    }
 249
 250    private async ValueTask<bool> RenewLeaseAsync(Guid id, Guid lockId, TimeSpan lockTimeout)
 251    {
 252        await using var connection = await _dataSource.OpenConnectionAsync(CancellationToken.None).ConfigureAwait(false)
 253        await using var command = connection.CreateCommand();
 254        command.CommandText =
 255            $"""
 256            UPDATE {MessageTable}
 257            SET locked_until = now() + @lock_timeout
 258            WHERE id = @id AND lock_id = @lock_id;
 259            """;
 260        command.Parameters.AddWithValue("id", id);
 261        command.Parameters.AddWithValue("lock_id", lockId);
 262        command.Parameters.AddWithValue("lock_timeout", lockTimeout);
 263        return await command.ExecuteNonQueryAsync(CancellationToken.None).ConfigureAwait(false) > 0;
 264    }
 265
 266    private async ValueTask NakAsync(Guid id, Guid lockId, TimeSpan delay)
 267    {
 268        await using var connection = await _dataSource.OpenConnectionAsync(CancellationToken.None).ConfigureAwait(false)
 269        await using var command = connection.CreateCommand();
 270        command.CommandText =
 271            $"""
 272            UPDATE {MessageTable}
 273            SET available_at = now() + @delay,
 274                locked_until = NULL,
 275                lock_id = NULL
 276            WHERE id = @id AND lock_id = @lock_id;
 277            SELECT pg_notify(@channel, @payload);
 278            """;
 279        command.Parameters.AddWithValue("id", id);
 280        command.Parameters.AddWithValue("lock_id", lockId);
 281        command.Parameters.AddWithValue("delay", delay);
 282        command.Parameters.AddWithValue("channel", _options.NotificationChannel);
 283        command.Parameters.AddWithValue("payload", "retry");
 284        await command.ExecuteNonQueryAsync(CancellationToken.None).ConfigureAwait(false);
 285    }
 286
 287    private async ValueTask<bool> DeadLetterAsync(
 288        Guid id,
 289        Guid lockId,
 290        string sourceQueue,
 291        string payload,
 292        IReadOnlyDictionary<string, string> headers,
 293        Exception exception,
 294        bool deleteOriginal,
 295        CancellationToken cancellationToken)
 296    {
 297        if (!_options.DeadLetterEnabled)
 298        {
 299            if (deleteOriginal)
 300                await AckAsync(id, lockId).ConfigureAwait(false);
 301            return true;
 302        }
 303
 304        var deadHeaders = new Dictionary<string, string>(headers, StringComparer.OrdinalIgnoreCase)
 305        {
 306            ["AR-DeadLetter-Reason"] = Sanitize(exception.Message),
 307            ["AR-DeadLetter-Source-Queue"] = sourceQueue
 308        };
 309
 310        try
 311        {
 312            if (!deleteOriginal)
 313            {
 314                await InsertAsync(Guid.NewGuid(), _options.DeadLetterQueue, payload, deadHeaders, exception.Message, not
 315                return true;
 316            }
 317
 318            // The DLQ insert and the original-row delete must commit atomically: split across two
 319            // connections, a crash between them leaves the original row to be redelivered and
 320            // dead-lettered again, duplicating the DLQ entry.
 321            await EnsureCreatedAsync(cancellationToken).ConfigureAwait(false);
 322            await using var connection = await _dataSource.OpenConnectionAsync(cancellationToken).ConfigureAwait(false);
 323            await using var transaction = await connection.BeginTransactionAsync(cancellationToken).ConfigureAwait(false
 324            await using var command = connection.CreateCommand();
 325            command.Transaction = transaction;
 326            command.CommandText =
 327                $"""
 328                INSERT INTO {MessageTable} (id, queue, payload_json, headers_json, dead_letter_reason)
 329                VALUES (@id, @queue, @payload_json, @headers_json, @dead_letter_reason)
 330                ON CONFLICT (id) DO NOTHING;
 331                DELETE FROM {MessageTable} WHERE id = @source_id AND lock_id = @lock_id;
 332                """;
 333            command.Parameters.AddWithValue("id", Guid.NewGuid());
 334            command.Parameters.AddWithValue("queue", _options.DeadLetterQueue);
 335            command.Parameters.Add("payload_json", NpgsqlDbType.Jsonb).Value = payload;
 336            command.Parameters.Add("headers_json", NpgsqlDbType.Jsonb).Value = AsyncResponseJson.Serialize(deadHeaders);
 337            command.Parameters.AddWithValue("dead_letter_reason", exception.Message);
 338            command.Parameters.AddWithValue("source_id", id);
 339            command.Parameters.AddWithValue("lock_id", lockId);
 340            await command.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
 341            await transaction.CommitAsync(cancellationToken).ConfigureAwait(false);
 342            return true;
 343        }
 344        catch (Exception ex)
 345        {
 346            // Callers decide the redelivery consequence from the false return; log the cause here so
 347            // a failing dead-letter write is never silent.
 348            _logger?.LogError(
 349                ex,
 350                "Failed to write PostgreSQL dead-letter row for message {MessageId} from queue {SourceQueue}.",
 351                id,
 352                sourceQueue);
 353            return false;
 354        }
 355    }
 356
 357    public async Task ExecuteListenAsync(Func<Task> onNotification, CancellationToken cancellationToken)
 358    {
 359        await EnsureCreatedAsync(cancellationToken).ConfigureAwait(false);
 360        await using var connection = await _dataSource.OpenConnectionAsync(cancellationToken).ConfigureAwait(false);
 361        connection.Notification += (_, _) => _ = onNotification();
 362        await using (var command = connection.CreateCommand())
 363        {
 364            command.CommandText = $"LISTEN {Quote(_options.NotificationChannel)};";
 365            await command.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
 366        }
 367
 368        while (!cancellationToken.IsCancellationRequested)
 369            await connection.WaitAsync(cancellationToken).ConfigureAwait(false);
 370    }
 371
 372    /// <summary>
 373    /// Opportunistically deletes dead-letter rows older than the configured retention. No-op unless
 374    /// <see cref="PostgreSqlAsyncResponseTransportOptions.DeadLetterRetention"/> is set, and throttled
 375    /// so the DELETE runs at most once per minute regardless of publish rate.
 376    /// </summary>
 377    private async Task PruneDeadLettersIfDueAsync(CancellationToken cancellationToken)
 378    {
 379        if (_options.DeadLetterRetention is not { } retention || !ShouldPruneDeadLetters())
 380            return;
 381
 382        await using var connection = await _dataSource.OpenConnectionAsync(cancellationToken).ConfigureAwait(false);
 383        await using var command = connection.CreateCommand();
 384        command.CommandText = $"DELETE FROM {MessageTable} WHERE queue = @queue AND created_at < now() - @retention;";
 385        command.Parameters.AddWithValue("queue", _options.DeadLetterQueue);
 386        command.Parameters.AddWithValue("retention", retention);
 387        await command.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
 388    }
 389
 390    private bool ShouldPruneDeadLetters()
 391    {
 392        var now = DateTime.UtcNow.Ticks;
 393        var last = Interlocked.Read(ref _lastDeadLetterPruneTicks);
 394        return now - last >= DeadLetterPruneThrottle.Ticks
 395            && Interlocked.CompareExchange(ref _lastDeadLetterPruneTicks, now, last) == last;
 396    }
 397
 398    private static IReadOnlyDictionary<string, string> DeserializeHeaders(string json)
 399    {
 400        var parsed = AsyncResponseJson.Deserialize<Dictionary<string, string>>(json);
 401        return parsed is null
 402            ? EmptyHeaders
 403            : new Dictionary<string, string>(parsed, StringComparer.OrdinalIgnoreCase);
 404    }
 405
 406    private static string Sanitize(string value) => value.Replace('\r', ' ').Replace('\n', ' ');
 407
 408    /// <summary>
 409    /// Stable 64-bit advisory-lock key for serializing schema creation. Must be byte-for-byte identical
 410    /// to the channel store's algorithm/discriminator so that, for a shared schema, the channel and
 411    /// transport take the same lock and never race each other on CREATE SCHEMA.
 412    /// </summary>
 413    internal static long SchemaAdvisoryLockKey(string schemaName)
 414    {
 415        const ulong offset = 14695981039346656037UL;
 416        const ulong prime = 1099511628211UL;
 417        var hash = offset;
 418        foreach (var b in Encoding.UTF8.GetBytes($"asyncresponse:ddl:{schemaName}"))
 419        {
 420            hash ^= b;
 421            hash *= prime;
 422        }
 423
 424        return unchecked((long)hash);
 425    }
 426
 427    private static string Quote(string identifier) => "\"" + identifier + "\"";
 428
 429    private static string IndexName(string table, string suffix)
 430    {
 431        var name = $"{table}_{suffix}_idx";
 432        return name.Length <= 63 ? name : name[..63];
 433    }
 434
 435    private static readonly TimeSpan DeadLetterPruneThrottle = TimeSpan.FromMinutes(1);
 436
 437    private static readonly IReadOnlyDictionary<string, string> EmptyHeaders =
 438        new Dictionary<string, string>(0, StringComparer.OrdinalIgnoreCase);
 439}