MCVIngenieros.Healthchecks/__/HealthCheckContainer.cs

using MCVIngenieros.Healthchecks.Abstracts;
using MCVIngenieros.Healthchecks.Options;
using Microsoft.Extensions.Options;
using System.Collections.Concurrent;
using System.Collections.Immutable;
using System.ComponentModel;
using System.Reflection;

namespace MCVIngenieros.Healthchecks.__;

// Health check container implementation
[EditorBrowsable(EditorBrowsableState.Never)]
public sealed class HealthCheckContainer : IHealthCheckContainer, IDisposable
{
    private readonly IServiceProvider provider;
    private readonly ILogger<HealthCheckContainer> _logger;
    private readonly TimeSpan _cacheDuration;
    private readonly TimeSpan? _timeout;
    private readonly static ConcurrentDictionary<string, Type> _checkTypes = new();
    private readonly ConcurrentDictionary<string, (HealthCheckResult Result, DateTime LastUpdate)> _cache = new();
    private readonly CancellationTokenSource _cts = new();
    private readonly Task _backgroundTask;
    private readonly Assembly[]? _assembliesToScan;

    public HealthCheckContainer(IServiceProvider provider, ILogger<HealthCheckContainer> logger, IOptions<HealthCheckContainerOptions> options)
    {
        this.provider = provider;
        _logger = logger;
        _cacheDuration = options.Value.CacheDuration;
        _timeout = options.Value.Timeout;

        _assembliesToScan = options.Value.AssembliesToScan?.ToArray() ?? null;
        _assembliesToScan = _assembliesToScan?.Distinct().ToArray();

        LoadHealthChecks();

        // Start background refresh
        _backgroundTask = Task.Run(() => BackgroundRefreshAsync(_cts.Token));
    }

    public void LoadHealthChecks()
    {
        var checks = provider.GetServices<IHealthCheck>();
        foreach (var check in checks)
        {
            var type = check.GetType();
            if (_checkTypes.TryAdd(type.Name, type))
            {
                _logger.LogInformation($"Loaded health check: {type.Name}");
            }
            else
            {
                _logger.LogInformation($"Health check {type.Name} is already registered. Skipping duplicate.");
            }
        }
    }

    // Discover all IHealthCheck implementations in loaded assemblies
    public void DiscoverHealthChecks()
    {
        if (_assembliesToScan == null)
            return;

        foreach (var assembly in _assembliesToScan)
        {
            var checkTypes = assembly.GetTypes()
            .Where(t => typeof(IHealthCheck).IsAssignableFrom(t) && !t.IsInterface && !t.IsAbstract);

            assembly.GetReferencedAssemblies()
                .Select(Assembly.Load)
                .ToList()
                .ForEach(a =>
                {
                    var types = a.GetTypes()
                        .Where(t => typeof(IHealthCheck).IsAssignableFrom(t) && !t.IsInterface && !t.IsAbstract);
                    checkTypes = checkTypes.Concat(types);
                });

            checkTypes = checkTypes.Distinct();

            foreach (var type in checkTypes)
            {
                if (_checkTypes.TryAdd(type.Name, type))
                {
                    _logger.LogInformation($"Discovered health check: {type.Name}");
                }
                else
                {
                    _logger.LogInformation($"Health check {type.Name} is already registered. Skipping duplicate.");
                }
            }
        }
    }

    public IReadOnlyList<string> GetAllCheckNames() => [.. _checkTypes.Keys];

    public async Task<List<HealthCheckResult>> RunAllChecks(CancellationToken cancellationToken = default)
    {
        _logger.LogInformation("Running all health checks...");
        var tasks = _checkTypes.Keys.Select(name => RunCheck(name, cancellationToken));
        var results = await Task.WhenAll(tasks);
        _logger.LogInformation("Completed running all health checks.");
        return results.Where(r => r != null).ToList()!;
    }

    public Task<List<HealthCheckResult>> GetAllChecksResults(CancellationToken cancellationToken = default)
    {
        var results = new List<HealthCheckResult>();
        foreach (var item in _cache)
        {
            results.Add(item.Value.Result);
        }
        return Task.FromResult(results);
    }

    public async Task<HealthCheckResult?> RunCheck(string name, CancellationToken cancellationToken = default)
    {
        var now = DateTime.UtcNow;
        _logger.LogInformation($"Running health check: {name}");
        if (!_checkTypes.TryGetValue(name, out var type))
        {
            _logger.LogWarning($"Health check {name} not found.");
            return null;
        }
        if (_cache.TryGetValue(name, out var cached) && now - cached.LastUpdate < _cacheDuration)
        {
            _logger.LogInformation($"Returning cached result for {name}.");
            return cached.Result;
        }
        var result = await RunCheckInstanceAsync(type, cancellationToken);
        _cache[name] = (result, DateTime.UtcNow);
        _logger.LogInformation($"Health check {name} completed. Result: {(result.IsHealthy ? "Healthy" : "Unhealthy")}");
        return result;
    }

    public async Task<List<HealthCheckResult>> RunChecks(IEnumerable<string> names, CancellationToken cancellationToken = default)
    {
        _logger.LogInformation($"Running health checks: {string.Join(", ", names)}");
        var tasks = names.Select(name => RunCheck(name, cancellationToken));
        var results = await Task.WhenAll(tasks);
        _logger.LogInformation($"Completed running specified health checks. Runned {results.Count(r => r != null)} checks.");
        return results.Where(r => r != null).ToList()!;
    }

    // Background refresh of all checks
    private async Task BackgroundRefreshAsync(CancellationToken cancellationToken)
    {
        while (!cancellationToken.IsCancellationRequested)
        {
            try
            {
                var tasks = _checkTypes.Keys.Select(name => RunCheck(name, cancellationToken));
                await Task.WhenAll(tasks);
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "Error during background health check refresh");
            }
            await Task.Delay(_cacheDuration, cancellationToken);
        }
    }

    // Run a single check instance with timeout, retry, circuit breaker, fallback
    private async Task<HealthCheckResult> RunCheckInstanceAsync(Type type, CancellationToken cancellationToken)
    {
        _logger.LogInformation($"Starting health check: {type.Name}");
        var name = type.Name;
        var logger = provider.GetService(typeof(ILogger<>).MakeGenericType(type)) as ILogger;
        try
        {
            IHealthCheck? check = ActivatorUtilities.CreateInstance(provider, type) as IHealthCheck
                    ?? throw new InvalidOperationException($"Could not create health check: {name}");
            var retry = check.RetryAttempts ?? HealthChecksConfigs.Default.RetryAttempts;
            TimeSpan timeout = (TimeSpan)((check.Timeout ?? _timeout) ?? HealthChecksConfigs.Default.Timeout!);
            TimeSpan delay = (TimeSpan)(check.RetryDelay ?? HealthChecksConfigs.Default.RetryDelay!);

            for (int attempt = 0; attempt <= retry; attempt++)
            {
                try
                {
                    using var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
                    cts.CancelAfter(timeout);
                    var result = await check.CheckAsync(cts.Token);
                    _logger.LogInformation($"Health check {name} completed. Result: {(result.IsHealthy ? "Healthy" : "Unhealthy")}");
                    result.Name = name;
                    return result;
                }
                catch (Exception ex)
                {
                    logger?.LogError(ex, $"Health check {name} failed on attempt {attempt + 1} -- Exception was not controlled.");
                    if (attempt == retry)
                    {
                        return new HealthCheckResult(false)
                        {
                            Name = name,
                            Details = $"Exception: {ex.GetBaseException().Message} -- StackTrace: {ex.GetBaseException().StackTrace}",
                            Severity = HealthCheckSeverity.Critical
                        };
                    }
                    await Task.Delay(delay, cancellationToken);
                }
            }
        }
        catch (Exception ex)
        {
            _logger.LogError(ex.GetBaseException(), $"Error creating instance of health check {name}");
            return new HealthCheckResult(false) { Name = name, Details = $"Error creating instance: {ex.GetBaseException().Message}", Severity = HealthCheckSeverity.Critical };
        }
        return new HealthCheckResult(false) { Name = name, Details = "Unknown error", Severity = HealthCheckSeverity.Critical };
    }

    public void Dispose()
    {
        _cts.Cancel();
        _backgroundTask?.Dispose();
        _cts.Dispose();
    }

    public IReadOnlyDictionary<string, IHealthCheck> GetAllChecks()
    {
        var dict = new Dictionary<string, IHealthCheck>();
        foreach (var item in _checkTypes.ToImmutableDictionary())
        {
            IHealthCheck? check = (IHealthCheck?)ActivatorUtilities.CreateInstance(provider, item.Value);
            if (check == null) continue;
            dict[item.Key] = check;
        }
        return dict.ToImmutableDictionary();
    }
}