AiDOPWarehouse/server/Plugins/Admin.NET.Plugin.AiDOP/DataPlatform/MdpRebuild/ModuleRebuildWorker.cs

using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;

namespace Admin.NET.Plugin.AiDOP.DataPlatform.MdpRebuild;

public sealed class ModuleRebuildWorker : BackgroundService
{
    private static readonly TimeSpan IdleDelay = TimeSpan.FromSeconds(5);

    private readonly IServiceScopeFactory _scopeFactory;
    private readonly ModuleRebuildQueue _queue;
    private readonly ILogger _logger;

    public ModuleRebuildWorker(
        IServiceScopeFactory scopeFactory,
        ModuleRebuildQueue queue,
        ILoggerFactory loggerFactory)
    {
        _scopeFactory = scopeFactory;
        _queue = queue;
        _logger = loggerFactory.CreateLogger(nameof(ModuleRebuildWorker));
    }

    protected override async Task ExecuteAsync(CancellationToken stoppingToken)
    {
        try { await Task.Delay(TimeSpan.FromSeconds(8), stoppingToken); }
        catch (OperationCanceledException) { return; }

        await FailStaleAsync(stoppingToken);

        var running = new List<Task>();
        while (!stoppingToken.IsCancellationRequested)
        {
            try
            {
                await FillSlotsAsync(running, stoppingToken);
            }
            catch (OperationCanceledException) when (stoppingToken.IsCancellationRequested)
            {
                break;
            }
            catch (Exception ex)
            {
                _logger.LogWarning(ex, "[ModuleRebuildWorker] run failed");
            }

            running.RemoveAll(t => t.IsCompleted);
            if (running.Count == 0)
            {
                try
                {
                    using var linked = CancellationTokenSource.CreateLinkedTokenSource(stoppingToken);
                    linked.CancelAfter(IdleDelay);
                    try { await _queue.Reader.ReadAsync(linked.Token); }
                    catch (OperationCanceledException) when (!stoppingToken.IsCancellationRequested) { }
                }
                catch (OperationCanceledException) { break; }
            }
            else
            {
                var pulse = WaitPulseAsync(stoppingToken);
                var completed = await Task.WhenAny(running.Append(pulse));
                if (completed != pulse)
                    running.Remove(completed);
            }
        }

        try { await Task.WhenAll(running); }
        catch (Exception ex) { _logger.LogWarning(ex, "[ModuleRebuildWorker] drain failed"); }
    }

    private async Task FillSlotsAsync(List<Task> running, CancellationToken stoppingToken)
    {
        int max = 2;
        try
        {
            using var probe = _scopeFactory.CreateScope();
            max = probe.ServiceProvider.GetRequiredService<IModuleRebuildCapability>().MaxParallelScopes;
        }
        catch
        {
            max = 2;
        }

        while (running.Count < max && !stoppingToken.IsCancellationRequested)
        {
            var started = await TryStartOneAsync(stoppingToken);
            if (started == null)
                break;
            running.Add(started);
        }
    }

    private async Task<Task> TryStartOneAsync(CancellationToken stoppingToken)
    {
        using var probe = _scopeFactory.CreateScope();
        var capability = probe.ServiceProvider.GetRequiredService<IModuleRebuildCapability>();
        var enabled = MdpRebuildScope.RebuildModules.Where(capability.IsEnabled).ToArray();
        if (enabled.Length == 0)
            return null;
        var job = await probe.ServiceProvider.GetRequiredService<IModuleRebuildJobStore>()
            .ClaimNextQueuedAsync(enabled, stoppingToken);
        if (job == null)
            return null;
        return ExecuteClaimedAsync(job, stoppingToken);
    }

    private async Task ExecuteClaimedAsync(Admin.NET.Plugin.AiDOP.Entity.SmartOps.AdoModuleDashboardRebuildJob job, CancellationToken stoppingToken)
    {
        using var scope = _scopeFactory.CreateScope();
        var svc = scope.ServiceProvider.GetRequiredService<ModuleRebuildService>();
        var runLock = scope.ServiceProvider.GetRequiredService<IModuleRebuildLock>();
        var handler = scope.ServiceProvider.GetServices<IModuleRebuildHandler>()
            .FirstOrDefault(x => string.Equals(x.ModuleCode, job.ModuleCode, StringComparison.OrdinalIgnoreCase));
        if (handler == null)
        {
            job.Status = ModuleRebuildStatus.Failed;
            job.CurrentStage = ModuleRebuildStages.Failed;
            job.ErrorMessage = $"未注册 {job.ModuleCode} 重算处理器";
            job.FinishedAt = DateTime.Now;
            job.UpdateTime = DateTime.Now;
            await scope.ServiceProvider.GetRequiredService<IModuleRebuildJobStore>().UpdateAsync(job, CancellationToken.None);
            return;
        }

        // 进度写入已带 heartbeat；后台另开连接会与单例 ISqlSugarClient 抢同一条 MySQL 连接。
        await svc.RunClaimedAsync(job, handler, runLock, stoppingToken);
    }

    private async Task WaitPulseAsync(CancellationToken stoppingToken)
    {
        using var linked = CancellationTokenSource.CreateLinkedTokenSource(stoppingToken);
        linked.CancelAfter(IdleDelay);
        try { await _queue.Reader.ReadAsync(linked.Token); }
        catch (OperationCanceledException) { }
    }

    private async Task FailStaleAsync(CancellationToken stoppingToken)
    {
        try
        {
            using var scope = _scopeFactory.CreateScope();
            await scope.ServiceProvider.GetRequiredService<ModuleRebuildService>().FailStaleAsync(stoppingToken);
        }
        catch (Exception ex)
        {
            _logger.LogWarning(ex, "[ModuleRebuildWorker] fail-stale failed");
        }
    }
}