DodoSSH/src/DodoSSH.Api/Setup/Persistence.cs

using DodoSSH.Infrastructure;
using Microsoft.EntityFrameworkCore;
using Microsoft.EntityFrameworkCore.Diagnostics;
using Microsoft.Extensions.Diagnostics.HealthChecks;
using Microsoft.Extensions.Options;
using Npgsql;

namespace DodoSSH.Api.Setup;

/// <summary>Database wiring.</summary>
internal static class Persistence
{
    private const string ConnectionStringName = "Postgres";

    /// <summary>
    /// Advisory lock every instance takes before migrating.
    /// </summary>
    /// <remarks>
    /// The value is arbitrary and only has to be the same in every replica and different from whatever
    /// else might lock in this database. Spelling "MIGR" in the low bytes makes the <c>pg_locks</c> row
    /// recognisable to whoever is looking at a stuck deployment at the time.
    /// </remarks>
    private const long MigrationLockKey = 0x0D0D_0555_4D49_4752;

    internal static IServiceCollection AddDodoPersistence(
        this IServiceCollection services,
        IConfiguration configuration)
    {
        var connectionString = configuration.GetConnectionString(ConnectionStringName)
            ?? throw new InvalidOperationException(
                $"ConnectionStrings:{ConnectionStringName} is not configured.");

        services.AddDbContext<DodoDbContext>((provider, options) => options
            .UseNpgsql(connectionString, npgsql => npgsql
                .MigrationsHistoryTable("__EFMigrationsHistory", DodoDbContext.SchemaName)
                // Retry on transient faults only. A write transaction replayed by the strategy runs on
                // the same context, which still holds what the failed attempt tracked — every such
                // transaction goes through FreshAttempts.ExecuteFreshAsync, which starts a retry from
                // the database instead.
                .EnableRetryOnFailure(3))
            .UseSnakeCaseNamingConvention()
            // None in production. The seam a test uses to make a commit fail the way a dropped
            // connection would, which is the only way to watch a retry happen.
            .AddInterceptors(provider.GetServices<IInterceptor>()));

        services.AddHealthChecks()
            .AddCheck<DatabaseHealthCheck>(
                "postgres",
                tags: [HealthChecks.ReadyTag, HealthChecks.StartupTag]);

        return services;
    }

    /// <summary>
    /// Brings the schema up to date, before the first request rather than during it.
    /// </summary>
    /// <remarks>
    /// <para>
    /// Ordered ahead of <c>RunAsync</c> deliberately: a migration that ran in the background would let
    /// the first requests through against a half-applied schema, and those are precisely the requests
    /// that provision accounts. Failing to migrate therefore fails to start, which is the loudest signal
    /// available and the one an orchestrator already knows how to act on.
    /// </para>
    /// <para>
    /// See <see cref="DatabaseOptions.AutoMigrate"/> for the deployments that should switch this off.
    /// </para>
    /// </remarks>
    internal static async Task<WebApplication> MigrateDodoDatabaseAsync(this WebApplication app)
    {
        ArgumentNullException.ThrowIfNull(app);

        var logger = app.Services.GetRequiredService<ILoggerFactory>().CreateLogger("DodoSSH.Startup");

        if (!app.Services.GetRequiredService<IOptions<DatabaseOptions>>().Value.AutoMigrate)
        {
            StartupLog.AutoMigrateDisabled(logger);

            return app;
        }

        var cancellationToken = app.Lifetime.ApplicationStopping;

        var scope = app.Services.CreateAsyncScope();
        await using var scopeHandle = scope.ConfigureAwait(false);

        var database = scope.ServiceProvider.GetRequiredService<DodoDbContext>().Database;

        // A connection of its own, held open across the whole migration. A session-level advisory lock
        // belongs to the connection that took it and EF opens and closes one per command, so borrowing
        // its connection would drop the lock somewhere in the middle.
        var guard = await MigrationGuard
            .AcquireAsync(database.GetConnectionString(), logger, cancellationToken)
            .ConfigureAwait(false);

        await using var guardHandle = guard.ConfigureAwait(false);

        var pending = await database.GetPendingMigrationsAsync(cancellationToken).ConfigureAwait(false);
        var pendingList = pending.ToList();

        // Reported even when there is nothing to do. Two replicas starting together both say something,
        // and "up to date" from the second is what shows the lock worked rather than that it was skipped.
        if (pendingList.Count == 0)
        {
            StartupLog.SchemaUpToDate(logger);

            return app;
        }

        // Named rather than joined inside the call, matching the health check below: the analyzer cannot
        // see through a static wrapper to a guard, and a join over a handful of names once per process
        // start is not worth the shape needed to convince it otherwise.
        var pendingDescription = string.Join(", ", pendingList);

        StartupLog.ApplyingMigrations(logger, pendingList.Count, pendingDescription);

        await database.MigrateAsync(cancellationToken).ConfigureAwait(false);

        StartupLog.MigrationsApplied(logger, pendingList.Count);

        return app;
    }

    /// <summary>
    /// Holds the advisory lock that makes concurrent starts safe.
    /// </summary>
    /// <remarks>
    /// Without it, two replicas rolled out together run <c>MigrateAsync</c> at the same moment: each
    /// reads the same empty history table, both apply the same migration, and the second fails on an
    /// object that already exists — a crash loop on the day of a schema change, which is the worst day
    /// for one. PostgreSQL grants the lock to one of them and blocks the other until it is done, at
    /// which point the second finds nothing pending and starts serving.
    /// </remarks>
    private sealed class MigrationGuard(NpgsqlConnection? connection) : IAsyncDisposable
    {
        internal static async Task<MigrationGuard> AcquireAsync(
            string? connectionString,
            ILogger logger,
            CancellationToken cancellationToken)
        {
            var connection = new NpgsqlConnection(connectionString);

            try
            {
                await connection.OpenAsync(cancellationToken).ConfigureAwait(false);
            }
            catch (PostgresException exception)
                when (string.Equals(
                    exception.SqlState,
                    PostgresErrorCodes.InvalidCatalogName,
                    StringComparison.Ordinal))
            {
                // Nothing to lock against yet: the database itself is missing, and EF's own Migrate
                // creates it. Racing here would have one replica win the CREATE DATABASE and the other
                // restart into the ordinary locked path, which is a worse first boot than it looks and
                // still converges — whereas refusing to start would leave a fresh deployment stuck on a
                // step this setting exists to remove.
                await connection.DisposeAsync().ConfigureAwait(false);

                StartupLog.MigratingUnlockedDatabase(logger);

                return new MigrationGuard(null);
            }
            catch
            {
                await connection.DisposeAsync().ConfigureAwait(false);
                throw;
            }

            var command = new NpgsqlCommand("SELECT pg_advisory_lock($1)", connection);
            await using var commandHandle = command.ConfigureAwait(false);

            command.Parameters.Add(new NpgsqlParameter { Value = MigrationLockKey });

            await command.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);

            return new MigrationGuard(connection);
        }

        /// <inheritdoc />
        public async ValueTask DisposeAsync()
        {
            if (connection is null)
            {
                return;
            }

            // Released explicitly rather than left to the connection closing. Both free the lock, but
            // only this one does it at a point in the code where a failure is visible — and an
            // unreleased migration lock is a deployment that hangs on its next start with no clue why.
            try
            {
                var command = new NpgsqlCommand("SELECT pg_advisory_unlock($1)", connection);
                await using var commandHandle = command.ConfigureAwait(false);

                command.Parameters.Add(new NpgsqlParameter { Value = MigrationLockKey });

                await command.ExecuteNonQueryAsync().ConfigureAwait(false);
            }
            finally
            {
                await connection.DisposeAsync().ConfigureAwait(false);
            }
        }
    }
}

/// <summary>
/// Reports the database reachable and the schema current.
/// </summary>
/// <remarks>
/// Pending migrations make the instance unready rather than crashing it. Startup normally applies them
/// — see <see cref="Persistence.MigrateDodoDatabaseAsync"/> — so reaching this check with migrations
/// still pending means either that <see cref="DatabaseOptions.AutoMigrate"/> is off and whatever owns
/// the schema has not run, or that the schema moved under a live instance. The right answer to both is
/// to stop taking traffic and say so loudly rather than to attempt a repair mid-flight.
/// </remarks>
internal sealed class DatabaseHealthCheck(DodoDbContext context, ILogger<DatabaseHealthCheck> logger)
    : IHealthCheck
{
    /// <inheritdoc />
    public async Task<HealthCheckResult> CheckHealthAsync(
        HealthCheckContext healthCheckContext,
        CancellationToken cancellationToken = default)
    {
        try
        {
            if (!await context.Database.CanConnectAsync(cancellationToken).ConfigureAwait(false))
            {
                return HealthCheckResult.Unhealthy("Cannot connect to PostgreSQL.");
            }

            var pending = await context.Database
                .GetPendingMigrationsAsync(cancellationToken)
                .ConfigureAwait(false);

            var pendingList = pending.ToList();
            if (pendingList.Count > 0)
            {
                // Computed once and used for both the log and the health result, so this is not
                // work done only for logging.
                var pendingDescription = string.Join(", ", pendingList);

                StartupLog.PendingMigrations(logger, pendingList.Count, pendingDescription);

                return HealthCheckResult.Unhealthy(
                    $"{pendingList.Count} migration(s) pending: {pendingDescription}");
            }

            return HealthCheckResult.Healthy();
        }
        catch (Exception exception) when (exception is not OperationCanceledException)
        {
            return HealthCheckResult.Unhealthy("PostgreSQL health check failed.", exception);
        }
    }
}