scadaproj/ZB.MOM.WW.Health/src/ZB.MOM.WW.Health.EntityFrameworkCore/DatabaseHealthCheck.cs

using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Diagnostics.HealthChecks;

namespace ZB.MOM.WW.Health.EntityFrameworkCore;

/// <summary>
/// Health check that verifies database reachability through an EF Core <typeparamref name="TContext"/>.
/// </summary>
/// <remarks>
/// <para>
/// The default probe calls
/// <see cref="Microsoft.EntityFrameworkCore.Infrastructure.DatabaseFacade.CanConnectAsync(CancellationToken)"/>
/// (the ScadaBridge pattern): <see cref="HealthStatus.Healthy"/> when it returns <c>true</c>,
/// <see cref="HealthStatus.Unhealthy"/> when it returns <c>false</c> or throws. Supplying
/// <see cref="DatabaseHealthCheckOptions{TContext}.ProbeQuery"/> swaps in a stricter query-based probe
/// (the OtOpcUa "query <c>Deployments</c>" pattern): the result is <see cref="HealthStatus.Healthy"/>
/// unless the delegate throws, in which case it is <see cref="HealthStatus.Unhealthy"/>. No exception
/// escapes <see cref="CheckHealthAsync"/>.
/// </para>
/// <para>
/// The context is resolved from the application <see cref="IServiceProvider"/>: an
/// <see cref="IDbContextFactory{TContext}"/> is used when one is registered (each probe gets a fresh,
/// disposed context); otherwise a scoped <typeparamref name="TContext"/> is resolved from a new DI
/// scope. Recommended registration tag: <c>ZbHealthTags.Ready</c> (applied by the registrant).
/// </para>
/// <para>
/// The scoped-resolution path is safe for <c>AddDbContextPool</c>: disposing the
/// <see cref="IServiceScope"/> returns the pooled context to the pool rather than destroying it,
/// so no pooled instance is prematurely discarded.
/// </para>
/// </remarks>
/// <typeparam name="TContext">The EF Core <see cref="DbContext"/> to probe.</typeparam>
public sealed class DatabaseHealthCheck<TContext> : IHealthCheck
    where TContext : DbContext
{
    private readonly IServiceProvider _serviceProvider;
    private readonly DatabaseHealthCheckOptions<TContext> _options;

    /// <summary>Initializes a new <see cref="DatabaseHealthCheck{TContext}"/>.</summary>
    /// <param name="serviceProvider">
    /// Application service provider used to resolve <typeparamref name="TContext"/> — preferring a
    /// registered <see cref="IDbContextFactory{TContext}"/>, otherwise a scoped instance.
    /// </param>
    /// <param name="options">
    /// Probe override and timeout. When <c>null</c>, the default <c>CanConnectAsync</c> probe with a
    /// 10 s timeout is used.
    /// </param>
    public DatabaseHealthCheck(
        IServiceProvider serviceProvider,
        DatabaseHealthCheckOptions<TContext>? options = null)
    {
        _serviceProvider = serviceProvider ?? throw new ArgumentNullException(nameof(serviceProvider));
        _options = options ?? new DatabaseHealthCheckOptions<TContext>();
    }

    /// <inheritdoc />
    public async Task<HealthCheckResult> CheckHealthAsync(
        HealthCheckContext context,
        CancellationToken cancellationToken = default)
    {
        using var timeoutCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
        timeoutCts.CancelAfter(_options.Timeout);

        try
        {
            var result = await ProbeAsync(timeoutCts.Token).ConfigureAwait(false);
            // Eagerly release the pending timer on the happy path so the OS timer
            // resource is not held for the full timeout duration.
            timeoutCts.CancelAfter(Timeout.InfiniteTimeSpan);
            return result;
        }
        catch (OperationCanceledException ex)
            when (timeoutCts.IsCancellationRequested && !cancellationToken.IsCancellationRequested)
        {
            return HealthCheckResult.Unhealthy($"Database probe timed out after {_options.Timeout}.", ex);
        }
        catch (Exception ex)
        {
            return HealthCheckResult.Unhealthy("Database connection failed.", ex);
        }
    }

    private async Task<HealthCheckResult> ProbeAsync(CancellationToken cancellationToken)
    {
        var factory = _serviceProvider.GetService<IDbContextFactory<TContext>>();
        if (factory is not null)
        {
            await using var db = await factory.CreateDbContextAsync(cancellationToken).ConfigureAwait(false);
            return await RunProbeAsync(db, cancellationToken).ConfigureAwait(false);
        }

        await using var scope = _serviceProvider.CreateAsyncScope();
        var scoped = scope.ServiceProvider.GetRequiredService<TContext>();
        return await RunProbeAsync(scoped, cancellationToken).ConfigureAwait(false);
    }

    private async Task<HealthCheckResult> RunProbeAsync(TContext db, CancellationToken cancellationToken)
    {
        if (_options.ProbeQuery is { } probeQuery)
        {
            await probeQuery(db, cancellationToken).ConfigureAwait(false);
            return HealthCheckResult.Healthy("Database query probe succeeded.");
        }

        var canConnect = await db.Database.CanConnectAsync(cancellationToken).ConfigureAwait(false);
        return canConnect
            ? HealthCheckResult.Healthy("Database connection is available.")
            : HealthCheckResult.Unhealthy("Database connection failed.");
    }
}