w4c-workflows-api/Services/Execution/WorkerHostService.cs

using System.Diagnostics;
using Microsoft.EntityFrameworkCore;
using w4c_workflows.Data;
using w4c_workflows.Models;
using w4c_workflows.Services.Messaging;
using w4c_workflows.Services.Nodes;

namespace w4c_workflows.Services.Execution;

/// <summary>
/// The worker loop: a long-lived consumer that pulls <c>task.run</c> jobs off
/// every tenant's jobs stream, executes them with the matching
/// <see cref="IScriptExecutor"/>, publishes a <c>task.result</c> event back, and
/// acks the job. Redelivery of crashed jobs is handled by claiming pending
/// messages past an idle threshold (at-least-once; idempotency is the control
/// plane's job via correlation_id + attempt).
///
/// The worker is not pinned to one static tenant: it enumerates the tenants that
/// have <em>workflows</em> (and therefore may have queued jobs) on each tick. This
/// lets a single control-plane process also act as the worker for its own tenants
/// (the self-hosted / single-node case), while still being usable as a dedicated
/// <c>--worker</c> deployment that fans out across many tenants.
/// </summary>
public class WorkerHostService : BackgroundService
{
    private readonly IJobQueue _jobs;
    private readonly IEventBus _events;
    private readonly RuntimeRegistry _runtimes;
    private readonly RemoteServerExecutor _remote;
    private readonly GraphJobExecutor _graphJobs;
    private readonly IServiceScopeFactory _scopeFactory;
    private readonly ILogger<WorkerHostService> _logger;

    private readonly string _consumer;
    private readonly string _resultStream;
    private readonly int _batchSize;
    private readonly TimeSpan _pollDelay;
    private readonly TimeSpan _claimIdle;

    public WorkerHostService(
        IJobQueue jobs,
        IEventBus events,
        RuntimeRegistry runtimes,
        RemoteServerExecutor remote,
        GraphJobExecutor graphJobs,
        IServiceScopeFactory scopeFactory,
        IConfiguration config,
        ILogger<WorkerHostService> logger)
    {
        _jobs = jobs;
        _events = events;
        _runtimes = runtimes;
        _remote = remote;
        _graphJobs = graphJobs;
        _scopeFactory = scopeFactory;
        _logger = logger;

        _consumer = $"{Environment.MachineName}-{Guid.NewGuid():N}"[..28];
        _resultStream = config["Workflows:ResultStream"] ?? Streams.Results;
        _batchSize = ParseInt(config["Workflows:WorkerBatchSize"], 10);
        _pollDelay = TimeSpan.FromMilliseconds(ParseInt(config["Workflows:WorkerPollDelayMs"], 500));
        _claimIdle = TimeSpan.FromSeconds(ParseInt(config["Workflows:WorkerClaimIdleSeconds"], 30));
    }

    protected override async Task ExecuteAsync(CancellationToken stoppingToken)
    {
        _logger.LogInformation("Worker starting as consumer {Consumer}", _consumer);

        while (!stoppingToken.IsCancellationRequested)
        {
            try
            {
                var processed = await ProcessBatchAsync(stoppingToken);
                if (processed == 0)
                    await Task.Delay(_pollDelay, stoppingToken);
            }
            catch (OperationCanceledException) when (stoppingToken.IsCancellationRequested)
            {
                break;
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "Worker loop error");
                try { await Task.Delay(_pollDelay, stoppingToken); } catch (OperationCanceledException) { break; }
            }
        }

        _logger.LogInformation("Worker stopped");
    }

    /// <summary>
    /// Scans the tenants that have compiled workflows (the only tenants that can
    /// have queued job messages) and processes their jobs streams. Ensures each
    /// stream's consumer group exists before reading so a fresh worker picks up
    /// only new messages (backlog replay is intentionally not attempted — a
    /// stalled run is recovered by the control-plane timeout sweep instead).
    /// </summary>
    private async Task<int> ProcessBatchAsync(CancellationToken ct)
    {
        var tenants = await ListTenantIdsAsync(ct);
        var processed = 0;

        foreach (var tenant in tenants)
        {
            if (ct.IsCancellationRequested)
                break;

            if (processed >= _batchSize)
                break;

            await _jobs.EnsureGroupAsync(tenant, ct);

            var messages = await _jobs.ReadGroupAsync(tenant, _consumer, _batchSize, ct);
            if (messages.Count == 0)
                messages = await _jobs.ClaimPendingAsync(tenant, _consumer, _claimIdle, _batchSize, ct);

            foreach (var message in messages)
            {
                ct.ThrowIfCancellationRequested();
                await ExecuteJobAsync(tenant, message, ct);
                processed++;
                if (processed >= _batchSize)
                    break;
            }
        }

        return processed;
    }

    private async Task<IReadOnlyList<string>> ListTenantIdsAsync(CancellationToken ct)
    {
        using var scope = _scopeFactory.CreateScope();
        var db = scope.ServiceProvider.GetRequiredService<WorkflowsDbContext>();
        return await db.Workflows
            .Where(w => w.Status == WorkflowStatus.Compiled)
            .Select(w => w.TenantId)
            .Distinct()
            .OrderBy(t => t)
            .ToListAsync(ct);
    }

    private async Task ExecuteJobAsync(string tenant, StreamMessage message, CancellationToken ct)
    {
        // Node-mode runs arrive as one `graph.run` job covering the whole graph;
        // script tasks arrive as `task.run`. Dispatch on the wire type before
        // parsing the script contract, whose required fields a graph job lacks.
        if (message.Fields.TryGetValue("type", out var kind)
            && string.Equals(kind, GraphRunInvocation.TypeValue, StringComparison.Ordinal))
        {
            await _graphJobs.ExecuteAsync(tenant, message, ct);
            return;
        }

        TaskInvocation invocation;
        try
        {
            invocation = TaskInvocation.FromFields(message.Fields);
        }
        catch (Exception ex)
        {
            _logger.LogWarning("Dead-lettering malformed job {MessageId}: {Error}", message.Id, ex.Message);
            await _jobs.DeadLetterAsync(tenant, message.Id, message.Fields, "malformed", ct);
            return;
        }

        if (invocation.Type != TaskInvocation.TypeValue)
        {
            // Foreign message on our stream — ack and move on.
            await _jobs.AckAsync(tenant, message.Id, ct);
            return;
        }

        // A job's tenant must match the stream it was read from; a mismatch means
        // the control plane enqueued under a different tenant key and must be
        // dead-lettered rather than silently dropped.
        if (!string.Equals(invocation.TenantId, tenant, StringComparison.Ordinal))
        {
            _logger.LogWarning("Job {MessageId} tenant {JobTenant} != stream tenant {StreamTenant}; dead-lettering",
                message.Id, invocation.TenantId, tenant);
            await _jobs.DeadLetterAsync(tenant, message.Id, message.Fields, "tenant_mismatch", ct);
            return;
        }

        var sw = Stopwatch.StartNew();
        ExecutionResult result;

        // S8: a task with a `server` reference runs remotely over SSH (via the
        // w4c-webapi server-console exec endpoint). The local runtime availability
        // probe is irrelevant — the target server owns its own runtime. Without a
        // `server` field the local subprocess path is unchanged (the default).
        if (!string.IsNullOrWhiteSpace(invocation.Server))
        {
            result = await _remote.ExecuteAsync(invocation, ct);
        }
        else
        {
            var executor = _runtimes.Resolve(invocation.Language);
            if (executor == null || !executor.IsAvailable())
            {
                var reason = executor == null
                    ? $"unknown language '{invocation.Language}'"
                    : $"runtime for language '{invocation.Language}' is unavailable on this worker";
                await PublishResultAsync(invocation,
                    new ExecutionResult(false, -1, string.Empty, string.Empty, null, reason, TimeSpan.Zero), ct);
                await _jobs.AckAsync(tenant, message.Id, ct);
                return;
            }

            try
            {
                result = await executor.ExecuteAsync(invocation, ct);
            }
            catch (Exception ex) when (ex is not OperationCanceledException)
            {
                result = new ExecutionResult(false, -1, string.Empty, string.Empty, null, ex.Message, sw.Elapsed);
            }
        }
        sw.Stop();

        await PublishResultAsync(invocation, result, ct);
        await _jobs.AckAsync(tenant, message.Id, ct);

        _logger.LogDebug(
            "Executed task {TaskKey} (run {RunId}, attempt {Attempt}, tenant {Tenant}) → {Status} in {Ms}ms",
            invocation.TaskKey, invocation.RunId, invocation.Attempt, tenant,
            result.Success ? "succeeded" : "failed", sw.ElapsedMilliseconds);
    }

    private Task PublishResultAsync(TaskInvocation invocation, ExecutionResult result, CancellationToken ct)
        => _events.PublishAsync(invocation.TenantId, _resultStream, TaskResultMessage.ToFields(invocation, result), ct);

    private static int ParseInt(string? text, int fallback)
        => int.TryParse(text, out var value) && value > 0 ? value : fallback;
}