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

Information
Class: AsyncResponse.Transports.PostgreSQL.PostgreSqlTransportStore
Assembly: AsyncResponse.Transports.PostgreSQL
File(s): /home/runner/work/AsyncResponse/AsyncResponse/src/Transports/AsyncResponse.Transports.PostgreSQL/PostgreSqlTransportStore.cs
Line coverage
94%
Covered lines: 235
Uncovered lines: 14
Coverable lines: 249
Total lines: 439
Line coverage: 94.3%
Branch coverage
97%
Covered branches: 41
Total branches: 42
Branch coverage: 97.6%
Method coverage

Feature is only available for sponsors

Upgrade to PRO version

Metrics

MethodBranch coverage Crap Score Cyclomatic complexity Line coverage
.ctor(...)100%11100%
EnsureCreatedAsync()100%66100%
PublishAsync()100%11100%
TryClaimAsync()100%11100%
ClaimBatchAsync()100%66100%
InsertAsync()100%88100%
AckAsync()100%11100%
RenewLeaseAsync()100%210%
NakAsync()100%11100%
DeadLetterAsync()75%8881.4%
ExecuteListenAsync()100%22100%
PruneDeadLettersIfDueAsync()100%44100%
ShouldPruneDeadLetters()100%22100%
DeserializeHeaders(...)100%22100%
Sanitize(...)100%11100%
SchemaAdvisoryLockKey(...)100%22100%
Quote(...)100%11100%
IndexName(...)100%22100%
.cctor()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>
 23internal sealed record PostgreSqlTransportDelivery(
 24    Guid Id,
 25    string Queue,
 26    string Payload,
 27    IReadOnlyDictionary<string, string> Headers,
 28    int Attempt,
 29    Func<ValueTask> AckAsync,
 30    Func<TimeSpan, ValueTask> NakAsync,
 31    Func<Exception, bool, CancellationToken, ValueTask<bool>> DeadLetterAsync,
 32    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;
 340    private readonly SemaphoreSlim _ensureGate = new(1, 1);
 41    private bool _created;
 42    private readonly long _schemaLockKey;
 43    private long _lastDeadLetterPruneTicks;
 44
 345    public PostgreSqlTransportStore(
 346        NpgsqlDataSource dataSource,
 347        IOptions<PostgreSqlAsyncResponseTransportOptions> options,
 348        ILogger<PostgreSqlTransportStore>? logger = null)
 49    {
 350        _dataSource = dataSource;
 351        _options = options.Value;
 352        _logger = logger;
 353        PostgreSqlTransportOptionsValidator.ValidateCommon(_options);
 354        Schema = Quote(_options.SchemaName);
 355        MessageTable = $"{Schema}.{Quote(_options.MessageTable)}";
 356        _schemaLockKey = SchemaAdvisoryLockKey(_options.SchemaName);
 357    }
 58
 59    public string Schema { get; }
 60    public string MessageTable { get; }
 61
 62    public async Task EnsureCreatedAsync(CancellationToken cancellationToken = default)
 63    {
 364        if (_created || !_options.AutoCreateSchema)
 365            return;
 66
 367        await _ensureGate.WaitAsync(cancellationToken).ConfigureAwait(false);
 68        try
 69        {
 370            if (_created)
 171                return;
 72
 373            await using var connection = await _dataSource.OpenConnectionAsync(cancellationToken).ConfigureAwait(false);
 174            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.
 181            await using (var lockCommand = connection.CreateCommand())
 82            {
 183                lockCommand.Transaction = transaction;
 184                lockCommand.CommandText = "SELECT pg_advisory_xact_lock(@lock_key);";
 185                lockCommand.Parameters.AddWithValue("lock_key", _schemaLockKey);
 186                await lockCommand.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
 87            }
 88
 189            await using var command = connection.CreateCommand();
 190            command.Transaction = transaction;
 191            command.CommandText =
 192                $"""
 193                CREATE SCHEMA IF NOT EXISTS {Schema};
 194
 195                CREATE TABLE IF NOT EXISTS {MessageTable} (
 196                    id uuid PRIMARY KEY,
 197                    queue text NOT NULL,
 198                    payload_json jsonb NOT NULL,
 199                    headers_json jsonb NOT NULL DEFAULT jsonb_build_object(),
 1100                    created_at timestamptz NOT NULL DEFAULT now(),
 1101                    available_at timestamptz NOT NULL DEFAULT now(),
 1102                    locked_until timestamptz NULL,
 1103                    lock_id uuid NULL,
 1104                    attempts integer NOT NULL DEFAULT 0,
 1105                    dead_letter_reason text NULL
 1106                );
 1107                CREATE INDEX IF NOT EXISTS {Quote(IndexName(_options.MessageTable, "claim"))}
 1108                    ON {MessageTable} (queue, available_at, locked_until, created_at);
 1109                CREATE INDEX IF NOT EXISTS {Quote(IndexName(_options.MessageTable, "created"))}
 1110                    ON {MessageTable} (created_at);
 1111                """;
 1112            await command.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
 1113            await transaction.CommitAsync(cancellationToken).ConfigureAwait(false);
 1114            _created = true;
 1115        }
 116        finally
 117        {
 3118            _ensureGate.Release();
 119        }
 3120    }
 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    {
 3133        await InsertAsync(id, queue, payload, headers, deadLetterReason: null, notify: true, cancellationToken).Configur
 1134        await PruneDeadLettersIfDueAsync(cancellationToken).ConfigureAwait(false);
 1135    }
 136
 137    public async Task<PostgreSqlTransportDelivery?> TryClaimAsync(string queue, TimeSpan lockTimeout, CancellationToken 
 138    {
 1139        await EnsureCreatedAsync(cancellationToken).ConfigureAwait(false);
 1140        var lockId = Guid.NewGuid();
 141
 1142        await using var connection = await _dataSource.OpenConnectionAsync(cancellationToken).ConfigureAwait(false);
 1143        await using var command = connection.CreateCommand();
 1144        command.CommandText =
 1145            $"""
 1146            WITH next AS (
 1147                SELECT id
 1148                FROM {MessageTable}
 1149                WHERE queue = @queue
 1150                  AND available_at <= now()
 1151                  AND (locked_until IS NULL OR locked_until <= now())
 1152                ORDER BY created_at
 1153                FOR UPDATE SKIP LOCKED
 1154                LIMIT 1
 1155            )
 1156            UPDATE {MessageTable} AS message
 1157            SET attempts = message.attempts + 1,
 1158                locked_until = now() + @lock_timeout,
 1159                lock_id = @lock_id
 1160            FROM next
 1161            WHERE message.id = next.id
 1162            RETURNING message.id, message.queue, message.payload_json::text, message.headers_json::text, message.attempt
 1163            """;
 1164        command.Parameters.AddWithValue("queue", queue);
 1165        command.Parameters.AddWithValue("lock_timeout", lockTimeout);
 1166        command.Parameters.AddWithValue("lock_id", lockId);
 167
 1168        await using var reader = await command.ExecuteReaderAsync(cancellationToken).ConfigureAwait(false);
 1169        if (!await reader.ReadAsync(cancellationToken).ConfigureAwait(false))
 1170            return null;
 171
 1172        var id = reader.GetGuid(0);
 1173        var payload = reader.GetString(2);
 1174        var headerJson = reader.GetString(3);
 1175        var attempt = reader.GetInt32(4);
 1176        var headers = DeserializeHeaders(headerJson);
 177
 1178        return new PostgreSqlTransportDelivery(
 1179            id,
 1180            reader.GetString(1),
 1181            payload,
 1182            headers,
 1183            attempt,
 1184            () => AckAsync(id, lockId),
 1185            delay => NakAsync(id, lockId, delay),
 1186            (exception, deleteOriginal, token) => DeadLetterAsync(id, lockId, queue, payload, headers, exception, delete
 1187            () => RenewLeaseAsync(id, lockId, lockTimeout));
 1188    }
 189
 190    public async IAsyncEnumerable<PostgreSqlTransportDelivery> ClaimBatchAsync(
 191        string queue,
 192        int batchSize,
 193        TimeSpan lockTimeout,
 194        [EnumeratorCancellation] CancellationToken cancellationToken)
 195    {
 1196        for (var i = 0; i < batchSize; i++)
 197        {
 1198            var delivery = await TryClaimAsync(queue, lockTimeout, cancellationToken).ConfigureAwait(false);
 1199            if (delivery is null)
 1200                yield break;
 1201            yield return delivery;
 202        }
 1203    }
 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    {
 3214        await EnsureCreatedAsync(cancellationToken).ConfigureAwait(false);
 3215        await using var connection = await _dataSource.OpenConnectionAsync(cancellationToken).ConfigureAwait(false);
 1216        await using var command = connection.CreateCommand();
 1217        command.CommandText =
 1218            $"""
 1219            INSERT INTO {MessageTable} (id, queue, payload_json, headers_json, dead_letter_reason)
 1220            VALUES (@id, @queue, @payload_json, @headers_json, @dead_letter_reason)
 1221            ON CONFLICT (id) DO NOTHING;
 1222            """;
 1223        if (notify)
 1224            command.CommandText += "SELECT pg_notify(@channel, @payload);";
 225
 1226        command.Parameters.AddWithValue("id", id);
 1227        command.Parameters.AddWithValue("queue", queue);
 1228        command.Parameters.Add("payload_json", NpgsqlDbType.Jsonb).Value = payload;
 1229        command.Parameters.Add("headers_json", NpgsqlDbType.Jsonb).Value = AsyncResponseJson.Serialize(headers ?? EmptyH
 1230        command.Parameters.AddWithValue("dead_letter_reason", (object?)deadLetterReason ?? DBNull.Value);
 1231        if (notify)
 232        {
 1233            command.Parameters.AddWithValue("channel", _options.NotificationChannel);
 1234            command.Parameters.AddWithValue("payload", queue);
 235        }
 236
 1237        await command.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
 1238    }
 239
 240    private async ValueTask AckAsync(Guid id, Guid lockId)
 241    {
 1242        await using var connection = await _dataSource.OpenConnectionAsync(CancellationToken.None).ConfigureAwait(false)
 1243        await using var command = connection.CreateCommand();
 1244        command.CommandText = $"DELETE FROM {MessageTable} WHERE id = @id AND lock_id = @lock_id;";
 1245        command.Parameters.AddWithValue("id", id);
 1246        command.Parameters.AddWithValue("lock_id", lockId);
 1247        await command.ExecuteNonQueryAsync(CancellationToken.None).ConfigureAwait(false);
 1248    }
 249
 250    private async ValueTask<bool> RenewLeaseAsync(Guid id, Guid lockId, TimeSpan lockTimeout)
 251    {
 0252        await using var connection = await _dataSource.OpenConnectionAsync(CancellationToken.None).ConfigureAwait(false)
 0253        await using var command = connection.CreateCommand();
 0254        command.CommandText =
 0255            $"""
 0256            UPDATE {MessageTable}
 0257            SET locked_until = now() + @lock_timeout
 0258            WHERE id = @id AND lock_id = @lock_id;
 0259            """;
 0260        command.Parameters.AddWithValue("id", id);
 0261        command.Parameters.AddWithValue("lock_id", lockId);
 0262        command.Parameters.AddWithValue("lock_timeout", lockTimeout);
 0263        return await command.ExecuteNonQueryAsync(CancellationToken.None).ConfigureAwait(false) > 0;
 0264    }
 265
 266    private async ValueTask NakAsync(Guid id, Guid lockId, TimeSpan delay)
 267    {
 1268        await using var connection = await _dataSource.OpenConnectionAsync(CancellationToken.None).ConfigureAwait(false)
 1269        await using var command = connection.CreateCommand();
 1270        command.CommandText =
 1271            $"""
 1272            UPDATE {MessageTable}
 1273            SET available_at = now() + @delay,
 1274                locked_until = NULL,
 1275                lock_id = NULL
 1276            WHERE id = @id AND lock_id = @lock_id;
 1277            SELECT pg_notify(@channel, @payload);
 1278            """;
 1279        command.Parameters.AddWithValue("id", id);
 1280        command.Parameters.AddWithValue("lock_id", lockId);
 1281        command.Parameters.AddWithValue("delay", delay);
 1282        command.Parameters.AddWithValue("channel", _options.NotificationChannel);
 1283        command.Parameters.AddWithValue("payload", "retry");
 1284        await command.ExecuteNonQueryAsync(CancellationToken.None).ConfigureAwait(false);
 1285    }
 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    {
 3297        if (!_options.DeadLetterEnabled)
 298        {
 1299            if (deleteOriginal)
 1300                await AckAsync(id, lockId).ConfigureAwait(false);
 1301            return true;
 302        }
 303
 3304        var deadHeaders = new Dictionary<string, string>(headers, StringComparer.OrdinalIgnoreCase)
 3305        {
 3306            ["AR-DeadLetter-Reason"] = Sanitize(exception.Message),
 3307            ["AR-DeadLetter-Source-Queue"] = sourceQueue
 3308        };
 309
 310        try
 311        {
 3312            if (!deleteOriginal)
 313            {
 3314                await InsertAsync(Guid.NewGuid(), _options.DeadLetterQueue, payload, deadHeaders, exception.Message, not
 0315                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.
 3321            await EnsureCreatedAsync(cancellationToken).ConfigureAwait(false);
 3322            await using var connection = await _dataSource.OpenConnectionAsync(cancellationToken).ConfigureAwait(false);
 1323            await using var transaction = await connection.BeginTransactionAsync(cancellationToken).ConfigureAwait(false
 1324            await using var command = connection.CreateCommand();
 1325            command.Transaction = transaction;
 1326            command.CommandText =
 1327                $"""
 1328                INSERT INTO {MessageTable} (id, queue, payload_json, headers_json, dead_letter_reason)
 1329                VALUES (@id, @queue, @payload_json, @headers_json, @dead_letter_reason)
 1330                ON CONFLICT (id) DO NOTHING;
 1331                DELETE FROM {MessageTable} WHERE id = @source_id AND lock_id = @lock_id;
 1332                """;
 1333            command.Parameters.AddWithValue("id", Guid.NewGuid());
 1334            command.Parameters.AddWithValue("queue", _options.DeadLetterQueue);
 1335            command.Parameters.Add("payload_json", NpgsqlDbType.Jsonb).Value = payload;
 1336            command.Parameters.Add("headers_json", NpgsqlDbType.Jsonb).Value = AsyncResponseJson.Serialize(deadHeaders);
 1337            command.Parameters.AddWithValue("dead_letter_reason", exception.Message);
 1338            command.Parameters.AddWithValue("source_id", id);
 1339            command.Parameters.AddWithValue("lock_id", lockId);
 1340            await command.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
 1341            await transaction.CommitAsync(cancellationToken).ConfigureAwait(false);
 1342            return true;
 1343        }
 2344        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.
 2348            _logger?.LogError(
 2349                ex,
 2350                "Failed to write PostgreSQL dead-letter row for message {MessageId} from queue {SourceQueue}.",
 2351                id,
 2352                sourceQueue);
 2353            return false;
 354        }
 3355    }
 356
 357    public async Task ExecuteListenAsync(Func<Task> onNotification, CancellationToken cancellationToken)
 358    {
 1359        await EnsureCreatedAsync(cancellationToken).ConfigureAwait(false);
 1360        await using var connection = await _dataSource.OpenConnectionAsync(cancellationToken).ConfigureAwait(false);
 1361        connection.Notification += (_, _) => _ = onNotification();
 1362        await using (var command = connection.CreateCommand())
 363        {
 1364            command.CommandText = $"LISTEN {Quote(_options.NotificationChannel)};";
 1365            await command.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
 366        }
 367
 1368        while (!cancellationToken.IsCancellationRequested)
 1369            await connection.WaitAsync(cancellationToken).ConfigureAwait(false);
 1370    }
 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    {
 1379        if (_options.DeadLetterRetention is not { } retention || !ShouldPruneDeadLetters())
 1380            return;
 381
 1382        await using var connection = await _dataSource.OpenConnectionAsync(cancellationToken).ConfigureAwait(false);
 1383        await using var command = connection.CreateCommand();
 1384        command.CommandText = $"DELETE FROM {MessageTable} WHERE queue = @queue AND created_at < now() - @retention;";
 1385        command.Parameters.AddWithValue("queue", _options.DeadLetterQueue);
 1386        command.Parameters.AddWithValue("retention", retention);
 1387        await command.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
 1388    }
 389
 390    private bool ShouldPruneDeadLetters()
 391    {
 1392        var now = DateTime.UtcNow.Ticks;
 1393        var last = Interlocked.Read(ref _lastDeadLetterPruneTicks);
 1394        return now - last >= DeadLetterPruneThrottle.Ticks
 1395            && Interlocked.CompareExchange(ref _lastDeadLetterPruneTicks, now, last) == last;
 396    }
 397
 398    private static IReadOnlyDictionary<string, string> DeserializeHeaders(string json)
 399    {
 1400        var parsed = AsyncResponseJson.Deserialize<Dictionary<string, string>>(json);
 1401        return parsed is null
 1402            ? EmptyHeaders
 1403            : new Dictionary<string, string>(parsed, StringComparer.OrdinalIgnoreCase);
 404    }
 405
 3406    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;
 3417        var hash = offset;
 3418        foreach (var b in Encoding.UTF8.GetBytes($"asyncresponse:ddl:{schemaName}"))
 419        {
 3420            hash ^= b;
 3421            hash *= prime;
 422        }
 423
 3424        return unchecked((long)hash);
 425    }
 426
 3427    private static string Quote(string identifier) => "\"" + identifier + "\"";
 428
 429    private static string IndexName(string table, string suffix)
 430    {
 1431        var name = $"{table}_{suffix}_idx";
 1432        return name.Length <= 63 ? name : name[..63];
 433    }
 434
 1435    private static readonly TimeSpan DeadLetterPruneThrottle = TimeSpan.FromMinutes(1);
 436
 1437    private static readonly IReadOnlyDictionary<string, string> EmptyHeaders =
 1438        new Dictionary<string, string>(0, StringComparer.OrdinalIgnoreCase);
 439}