vfxreview/RenderWorker/VFXReviewWorker/JobExecutor.cs

using System.Text.Json;
using Microsoft.Extensions.Logging;

namespace VFXReviewWorker;

/// <summary>
/// Executes one claimed job end-to-end: writes crash-recovery state, runs
/// aerender, uploads the full log, and enqueues the durable complete/fail
/// report. One job at a time per machine (§7.2).
/// </summary>
public sealed class JobExecutor
{
    private readonly WorkerOptions _options;
    private readonly ApiClient _api;
    private readonly AerenderRunner _runner;
    private readonly PreviewStage _preview;
    private readonly ReportSpooler _spooler;
    private readonly ILogger<JobExecutor> _log;

    public JobExecutor(WorkerOptions options, ApiClient api, AerenderRunner runner, PreviewStage preview, ReportSpooler spooler, ILogger<JobExecutor> log)
    {
        _options = options;
        _api = api;
        _runner = runner;
        _preview = preview;
        _spooler = spooler;
        _log = log;
    }

    public async Task RunAsync(ClaimedJob job, string machineId, ServerConfig config, CancellationToken cancelJob, CancellationToken shutdown)
    {
        if (job.Type == "PREVIEW_ONLY")
        {
            await RunPreviewAsync(job, machineId, config, cancelJob, shutdown);
            return;
        }

        var manifest = job.AsRenderManifest();
        var mapper = new PathMapper(_options.PathMappings);
        var outputDirLocal = mapper.Map(manifest.OutputDir);

        WriteState(new CurrentJobState
        {
            JobId = job.Id,
            ExportId = job.ExportId,
            MachineId = machineId,
            OutputDirLocal = outputDirLocal,
            ClaimedAt = DateTimeOffset.UtcNow,
        });

        try
        {
            var result = await _runner.RunAsync(job, manifest, machineId, outputDirLocal, config, cancelJob, shutdown);
            var logKey = await _api.UploadLogAsync(job.Id, machineId, result.FullLogPath, shutdown);

            if (result.Success)
            {
                long totalBytes = 0;
                int fileCount = 0;
                var prefix = AerenderRunner.OutputPatternPrefix(manifest.OutputPattern);
                try
                {
                    foreach (var f in Directory.EnumerateFiles(outputDirLocal))
                    {
                        if (!Path.GetFileName(f).StartsWith(prefix, StringComparison.OrdinalIgnoreCase)) continue;
                        fileCount++;
                        totalBytes += new FileInfo(f).Length;
                    }
                }
                catch { /* stats are best-effort in Phase 2; validation owns this in Phase 3 */ }

                if (fileCount == 0)
                {
                    // aerender can exit 0 having rendered nothing (bad frame
                    // range, template mismatch, silently swallowed error).
                    // Until Phase 3 validation exists, an empty output dir must
                    // never become READY_FOR_QC.
                    _log.LogError("Job {JobId}: aerender exited 0 but no \"{Prefix}*\" files exist in {Dir} — reporting failure",
                        job.Id, prefix, outputDirLocal);
                    _spooler.Enqueue("fail", job.Id, new Dictionary<string, object?>
                    {
                        ["machineId"] = machineId,
                        ["stage"] = "RENDER",
                        ["exitCode"] = 0,
                        ["errorMessage"] = "aerender exited 0 but produced no output files in " + outputDirLocal,
                        ["logTail"] = result.LogTail,
                        ["logFileKey"] = logKey,
                        ["retryable"] = false,
                    });
                    return;
                }

                _log.LogInformation("Job {JobId} complete: {Frames} frames in {Sec:F0}s ({Files} files, {Bytes} bytes)",
                    job.Id, manifest.TotalFrames, result.RenderSeconds, fileCount, totalBytes);
                _spooler.Enqueue("complete-render", job.Id, new Dictionary<string, object?>
                {
                    ["machineId"] = machineId,
                    ["renderSeconds"] = (int)result.RenderSeconds,
                    ["logFileKey"] = logKey,
                    ["logTail"] = result.LogTail,
                    ["exrFileCount"] = fileCount,
                    ["exrTotalBytes"] = totalBytes,
                });
            }
            else if (result.ErrorMessage == "Cancelled")
            {
                // Ack the cancel so the server clears the CANCEL_JOB command
                _log.LogInformation("Job {JobId} cancelled — acknowledging", job.Id);
                _spooler.Enqueue("fail", job.Id, new Dictionary<string, object?>
                {
                    ["machineId"] = machineId,
                    ["stage"] = "RENDER",
                    ["errorMessage"] = "Cancelled by server",
                    ["logTail"] = result.LogTail,
                    ["logFileKey"] = logKey,
                    ["retryable"] = false,
                });
            }
            else
            {
                _log.LogError("Job {JobId} failed (exit {Exit}, retryable {Retryable}): {Error}",
                    job.Id, result.ExitCode, result.Retryable, result.ErrorMessage);
                _spooler.Enqueue("fail", job.Id, new Dictionary<string, object?>
                {
                    ["machineId"] = machineId,
                    ["stage"] = "RENDER",
                    ["exitCode"] = result.ExitCode,
                    ["errorMessage"] = result.ErrorMessage,
                    ["logTail"] = result.LogTail,
                    ["logFileKey"] = logKey,
                    ["retryable"] = result.Retryable,
                });
            }
        }
        catch (Exception ex)
        {
            _log.LogError(ex, "Unexpected executor error for job {JobId}", job.Id);
            _spooler.Enqueue("fail", job.Id, new Dictionary<string, object?>
            {
                ["machineId"] = machineId,
                ["stage"] = "RENDER",
                ["errorMessage"] = $"Worker exception: {ex.Message}",
                ["retryable"] = true,
            });
        }
        finally
        {
            ClearState();
        }
    }

    /// <summary>
    /// PREVIEW_ONLY job: rebuild the shot around the rendered EXRs with slate
    /// and burn-ins, render the delivery MOV + review MP4, upload the review
    /// media, and finalize (which registers the internal-only Version).
    /// </summary>
    private async Task RunPreviewAsync(ClaimedJob job, string machineId, ServerConfig config, CancellationToken cancelJob, CancellationToken shutdown)
    {
        var manifest = job.AsPreviewManifest();

        WriteState(new CurrentJobState
        {
            JobId = job.Id,
            ExportId = job.ExportId,
            MachineId = machineId,
            OutputDirLocal = new PathMapper(_options.PathMappings).Map(manifest.OutputDir),
            ClaimedAt = DateTimeOffset.UtcNow,
        });

        try
        {
            var result = await _preview.RunAsync(job, manifest, machineId, config, cancelJob, shutdown);

            if (!result.Success)
            {
                _log.LogError("Preview job {JobId} failed (retryable {Retryable}): {Error}",
                    job.Id, result.Retryable, result.ErrorMessage);
                _spooler.Enqueue("fail", job.Id, new Dictionary<string, object?>
                {
                    ["machineId"] = machineId,
                    ["stage"] = "PREVIEW",
                    ["errorMessage"] = result.ErrorMessage,
                    ["retryable"] = result.Retryable,
                });
                return; // temp AEP is deliberately kept for debugging
            }

            var previewKey = await _api.UploadArtifactAsync(
                job.Id, machineId, "preview", result.Mp4Path!, "video/mp4", shutdown);
            if (previewKey is null)
            {
                // The renders are on the SAN and fine — only the upload failed,
                // so this is worth retrying without rebuilding anything.
                _spooler.Enqueue("fail", job.Id, new Dictionary<string, object?>
                {
                    ["machineId"] = machineId,
                    ["stage"] = "PREVIEW",
                    ["errorMessage"] = "Preview MP4 upload failed",
                    ["retryable"] = true,
                });
                return;
            }

            string? thumbKey = null;
            if (result.ThumbnailPath is not null)
            {
                thumbKey = await _api.UploadArtifactAsync(
                    job.Id, machineId, "thumbnail", result.ThumbnailPath, "image/jpeg", shutdown);
            }

            _spooler.Enqueue("finalize", job.Id, new Dictionary<string, object?>
            {
                ["machineId"] = machineId,
                ["artifacts"] = new Dictionary<string, object?>
                {
                    ["deliveryMovPath"] = manifest.MovOutput,
                    ["previewMovKey"] = previewKey,
                    ["thumbnailKey"] = thumbKey,
                },
                ["renderStats"] = new Dictionary<string, object?>
                {
                    ["previewSeconds"] = (int)result.Seconds,
                },
                ["media"] = new Dictionary<string, object?>
                {
                    ["width"] = result.Build?.Width,
                    ["height"] = result.Build?.Height,
                    ["frameCount"] = result.Build?.FrameCount,
                    ["fps"] = manifest.Fps,
                    ["fileName"] = Path.GetFileName(result.Mp4Path!),
                },
            });

            // Only clean up once the outputs are safely reported.
            CleanScratch(job.Id, result.TempAep);
        }
        catch (Exception ex)
        {
            _log.LogError(ex, "Unexpected preview error for job {JobId}", job.Id);
            _spooler.Enqueue("fail", job.Id, new Dictionary<string, object?>
            {
                ["machineId"] = machineId,
                ["stage"] = "PREVIEW",
                ["errorMessage"] = $"Worker exception: {ex.Message}",
                ["retryable"] = true,
            });
        }
        finally
        {
            ClearState();
        }
    }

    private static void CleanScratch(string jobId, string? tempAep)
    {
        PreviewStage.TryDelete(Path.Combine(WorkerOptions.ScratchDir, $"job_{jobId}_context.json"));
        PreviewStage.TryDelete(Path.Combine(WorkerOptions.ScratchDir, $"job_{jobId}_result.json"));
        PreviewStage.TryDelete(Path.Combine(WorkerOptions.ScratchDir, $"job_{jobId}_run.jsx"));
        PreviewStage.TryDelete(Path.Combine(WorkerOptions.ScratchDir, $"job_{jobId}_thumb.jpg"));
        if (tempAep is not null) PreviewStage.TryDelete(tempAep);
    }

    /// <summary>
    /// Crash recovery (§7.5): on service start, if current-job.json names a
    /// job, ask the server what became of it. Still ours (CLAIMED/RUNNING) →
    /// we crashed mid-render; report a retryable fail so the server requeues
    /// (never resume a partial aerender). Reassigned/finished → discard state;
    /// the next attempt's PrepareOutputDir clears any orphaned frames.
    /// </summary>
    public async Task ReconcileAsync(string machineId, CancellationToken ct)
    {
        var state = ReadState();
        if (state is null) return;
        _log.LogWarning("Found crash-recovery state for job {JobId} — reconciling", state.JobId);
        try
        {
            var stillOurs = false;
            if (!string.IsNullOrEmpty(state.ExportId))
            {
                var export = await _api.GetExportAsync(state.ExportId!, ct);
                var job = export?.RenderJobs.FirstOrDefault(j => j.Id == state.JobId);
                stillOurs = job?.Status is "CLAIMED" or "RUNNING";
            }
            if (stillOurs)
            {
                _log.LogWarning("Server still shows job {JobId} as ours — reporting crash for requeue", state.JobId);
                _spooler.Enqueue("fail", state.JobId, new Dictionary<string, object?>
                {
                    ["machineId"] = machineId,
                    ["stage"] = "RENDER",
                    ["errorMessage"] = "Worker restarted mid-render (crash/power loss); never resuming a partial render",
                    ["retryable"] = true,
                });
            }
        }
        catch (Exception ex)
        {
            _log.LogWarning(ex, "Reconcile failed for job {JobId}; discarding local state (reaper covers the lease)", state.JobId);
        }
        ClearState();
    }

    private static void WriteState(CurrentJobState state)
    {
        Directory.CreateDirectory(WorkerOptions.DataDir);
        File.WriteAllText(WorkerOptions.StateFilePath, JsonSerializer.Serialize(state, WorkerOptions.JsonOpts));
    }

    private static CurrentJobState? ReadState()
    {
        try
        {
            if (!File.Exists(WorkerOptions.StateFilePath)) return null;
            return JsonSerializer.Deserialize<CurrentJobState>(File.ReadAllText(WorkerOptions.StateFilePath), WorkerOptions.JsonOpts);
        }
        catch
        {
            return null;
        }
    }

    private static void ClearState()
    {
        try { File.Delete(WorkerOptions.StateFilePath); } catch { }
    }
}