vfxreview/RenderWorker/VFXReviewWorker/ApiClient.cs

using System.Net;
using System.Net.Http.Headers;
using System.Net.Http.Json;
using System.Text.Json;
using Microsoft.Extensions.Logging;

namespace VFXReviewWorker;

/// <summary>
/// HTTP client for the VFXReview /api/ext pipeline endpoints. Workers talk
/// only HTTP — never a database (§2.1 principle 4).
/// </summary>
public sealed class ApiClient
{
    private readonly HttpClient _http;
    private readonly WorkerOptions _options;
    private readonly ILogger<ApiClient> _log;

    // Presigned uploads: a separate client so the URL is not sent with our
    // Authorization header, and a generous per-request timeout for media.
    private static readonly HttpClient _upload = new() { Timeout = TimeSpan.FromMinutes(30) };

    public ApiClient(WorkerOptions options, ILogger<ApiClient> log)
    {
        _options = options;
        _log = log;
        _http = new HttpClient
        {
            BaseAddress = new Uri(options.ServerUrl.TrimEnd('/') + "/"),
            Timeout = TimeSpan.FromSeconds(30),
        };
        _http.DefaultRequestHeaders.Authorization = new AuthenticationHeaderValue("Bearer", options.ApiKey);
    }

    private static readonly JsonSerializerOptions Json = WorkerOptions.JsonOpts;

    // E6
    public async Task<RegisterResponse> RegisterAsync(WorkerOptions o, CancellationToken ct)
    {
        var res = await _http.PostAsJsonAsync("api/ext/workers/register", new
        {
            name = o.MachineName,
            hostname = Dns.GetHostName(),
            workerVersion = o.WorkerVersion,
            aeVersion = o.AeVersion,
            // The server hands out BLENDER_RENDER work on the strength of the
            // blender entry, so it reports the resolved path rather than a bare
            // flag: a machine claiming Blender jobs it cannot run is worse than
            // one that never claims them.
            capabilities = new
            {
                maxConcurrentJobs = 1,
                tools = new
                {
                    ffmpeg = o.FfmpegPath != null,
                    oiiotool = o.OiiotoolPath != null,
                    blender = !string.IsNullOrWhiteSpace(o.BlenderPath) && File.Exists(o.BlenderPath) ? o.BlenderPath : null,
                },
            },
        }, Json, ct);
        res.EnsureSuccessStatusCode();
        return (await res.Content.ReadFromJsonAsync<RegisterResponse>(Json, ct))!;
    }

    // E7
    public async Task<HeartbeatResponse> HeartbeatAsync(string machineId, string? currentJobId, CancellationToken ct)
    {
        var res = await _http.PostAsJsonAsync($"api/ext/workers/{machineId}/heartbeat", new
        {
            cpuPercent = (double?)null,
            memPercent = (double?)null,
            diskFreeGb = GetDiskFreeGb(),
            currentJobId,
        }, Json, ct);
        res.EnsureSuccessStatusCode();
        return (await res.Content.ReadFromJsonAsync<HeartbeatResponse>(Json, ct))!;
    }

    // E8 — null when the queue is empty / machine outside its render window (204)
    public async Task<ClaimedJob?> ClaimAsync(string machineId, string[] types, CancellationToken ct)
    {
        var res = await _http.PostAsJsonAsync("api/ext/render/jobs/claim", new
        {
            machineId,
            types,
        }, Json, ct);
        if (res.StatusCode == HttpStatusCode.NoContent) return null;
        res.EnsureSuccessStatusCode();
        var body = await res.Content.ReadFromJsonAsync<ClaimResponse>(Json, ct);
        return body?.Job;
    }

    // E9 — best-effort by contract: a progress report is a nicety, and losing
    // one must never abort work in flight. Transport failures return null here
    // rather than throwing at the call site; the lease and the reaper are what
    // actually protect the job.
    public async Task<ProgressResponse?> ProgressAsync(string jobId, string machineId, double progress,
        int? currentFrame, int? totalFrames, int? etaSeconds, string? logTail, CancellationToken ct)
    {
        try
        {
            using var req = new HttpRequestMessage(HttpMethod.Patch, $"api/ext/render/jobs/{jobId}/progress")
            {
                Content = JsonContent.Create(new { machineId, progress, currentFrame, totalFrames, etaSeconds, logTail }, options: Json),
            };
            var res = await _http.SendAsync(req, ct);
            if (!res.IsSuccessStatusCode) return null;
            return await res.Content.ReadFromJsonAsync<ProgressResponse>(Json, ct);
        }
        catch (OperationCanceledException) when (ct.IsCancellationRequested)
        {
            throw; // a real shutdown still stops us
        }
        catch (Exception ex)
        {
            _log.LogDebug(ex, "Progress report failed for job {JobId} (non-fatal)", jobId);
            return null;
        }
    }

    // E10 / E11 — lifecycle reports, sent via the spooler for durability.
    // Returns true when the server accepted (2xx) OR permanently rejected (409
    // contradictory / 404 gone) — both mean "stop retrying".
    public async Task<bool> SendLifecycleAsync(string kind, string jobId, Dictionary<string, object?> payload, CancellationToken ct)
    {
        var path = kind switch
        {
            "fail" => $"api/ext/render/jobs/{jobId}/fail",
            "complete-render" => $"api/ext/render/jobs/{jobId}/complete-render",
            "validation" => $"api/ext/render/jobs/{jobId}/validation",
            "delivery" => $"api/ext/render/jobs/{jobId}/delivery",
            "finalize" => $"api/ext/render/jobs/{jobId}/finalize",
            "attach-highres" => $"api/ext/render/jobs/{jobId}/attach-highres",
            _ => throw new ArgumentOutOfRangeException(nameof(kind), kind, "unknown lifecycle report kind"),
        };
        var res = await _http.PostAsJsonAsync(path, payload, Json, ct);
        if (res.IsSuccessStatusCode) return true;
        if (res.StatusCode is HttpStatusCode.Conflict or HttpStatusCode.NotFound or HttpStatusCode.UnprocessableEntity)
        {
            // Conflict/NotFound: contradictory or vanished report — nothing to retry.
            // UnprocessableEntity: the server rejected the payload shape itself, so
            // resending the same bytes can only ever produce the same 422 — retrying
            // just spams the log forever instead of surfacing the real problem.
            _log.LogWarning("Server permanently rejected {Kind} for job {JobId}: {Status} {Body}",
                kind, jobId, (int)res.StatusCode, await res.Content.ReadAsStringAsync(ct));
            return true;
        }
        _log.LogWarning("Lifecycle report {Kind} for {JobId} failed with {Status}; will retry", kind, jobId, (int)res.StatusCode);
        return false;
    }

    // E20 + PUT upload. Returns the storage key, or null when the upload failed.
    // onProgress receives the bytes sent so far, for status lines in the app.
    public async Task<string?> UploadArtifactAsync(
        string jobId, string machineId, string kind, string filePath, string contentType, CancellationToken ct,
        Action<long>? onProgress = null)
    {
        try
        {
            if (!File.Exists(filePath))
            {
                _log.LogWarning("Artifact {Kind} not found at {Path}", kind, filePath);
                return null;
            }
            var size = new FileInfo(filePath).Length;
            var sw = System.Diagnostics.Stopwatch.StartNew();
            // Uploads can run for many minutes on a slow uplink; without these two
            // lines the log goes silent and a busy worker looks hung.
            if (kind != "log")
                _log.LogInformation("Uploading {Kind} for job {JobId}: {File} ({MB:F1} MB)",
                    kind, jobId, Path.GetFileName(filePath), size / 1048576.0);

            var presignRes = await _http.PostAsJsonAsync($"api/ext/render/jobs/{jobId}/artifact-presign", new
            {
                machineId,
                kind,
                fileName = Path.GetFileName(filePath),
                contentType,
                multipart = size >= MultipartThreshold,
                size,
            }, Json, ct);
            if (!presignRes.IsSuccessStatusCode)
            {
                _log.LogWarning("Presign for {Kind} failed: {Status}", kind, (int)presignRes.StatusCode);
                return null;
            }
            var presign = await presignRes.Content.ReadFromJsonAsync<PresignResponse>(Json, ct);
            if (presign is null) return null;

            if (presign.UploadId is not null && presign.PartUrls is { Count: > 0 } && presign.PartSize > 0)
            {
                if (!await UploadMultipartAsync(jobId, machineId, kind, filePath, size, presign, onProgress, ct))
                    return null;
            }
            else
            {
                // Single PUT: small files, or a server that predates multipart.
                await using var stream = File.OpenRead(filePath);
                var content = new ProgressStreamContent(stream, onProgress);
                content.Headers.ContentType = new MediaTypeHeaderValue(contentType);
                var putRes = await _upload.PutAsync(presign.PresignedUrl, content, ct);
                if (!putRes.IsSuccessStatusCode)
                {
                    _log.LogWarning("Upload of {Kind} failed after {Sec:F0}s: {Status}", kind, sw.Elapsed.TotalSeconds, (int)putRes.StatusCode);
                    return null;
                }
            }
            if (kind != "log")
                _log.LogInformation("Uploaded {Kind} for job {JobId} in {Sec:F0}s ({Speed:F1} MB/s)",
                    kind, jobId, sw.Elapsed.TotalSeconds, size / 1048576.0 / Math.Max(0.001, sw.Elapsed.TotalSeconds));
            return presign.Key;
        }
        catch (Exception ex)
        {
            _log.LogWarning(ex, "Artifact upload ({Kind}) failed for job {JobId}", kind, jobId);
            return null;
        }
    }

    // Files at least this big go up as parallel multipart parts (E20b).
    private const long MultipartThreshold = 16L * 1024 * 1024;
    private const int PartAttempts = 4;

    /// <summary>
    /// Uploads the parts concurrently (UploadConcurrency connections), each
    /// retried on its own so a dropped connection costs one part, not the
    /// file; then asks the server to stitch them. Aborts on failure so the
    /// stored parts don't linger in the bucket.
    /// </summary>
    private async Task<bool> UploadMultipartAsync(string jobId, string machineId, string kind, string filePath,
        long size, PresignResponse presign, Action<long>? onProgress, CancellationToken ct)
    {
        var partUrls = presign.PartUrls!;
        var partSize = presign.PartSize;
        var etags = new string?[partUrls.Count];
        var sentPerPart = new long[partUrls.Count];
        long sentTotal = 0;
        var next = -1;

        using var handle = File.OpenHandle(filePath, FileMode.Open, FileAccess.Read, FileShare.Read, FileOptions.Asynchronous);
        using var failed = CancellationTokenSource.CreateLinkedTokenSource(ct);

        async Task Worker()
        {
            int i;
            while ((i = Interlocked.Increment(ref next)) < partUrls.Count && !failed.IsCancellationRequested)
            {
                var offset = i * partSize;
                var len = (int)Math.Min(partSize, size - offset);
                var buffer = new byte[len];
                var read = 0;
                while (read < len) read += await RandomAccess.ReadAsync(handle, buffer.AsMemory(read), offset + read, failed.Token);

                for (var attempt = 1; ; attempt++)
                {
                    var part = i;
                    // Progress restarts from zero for a retried part.
                    Interlocked.Add(ref sentTotal, -Interlocked.Exchange(ref sentPerPart[part], 0));
                    var content = new ProgressStreamContent(new MemoryStream(buffer, writable: false), sent =>
                    {
                        var delta = sent - Interlocked.Exchange(ref sentPerPart[part], sent);
                        onProgress?.Invoke(Interlocked.Add(ref sentTotal, delta));
                    });
                    try
                    {
                        var res = await _upload.PutAsync(partUrls[i], content, failed.Token);
                        var etag = res.Headers.ETag?.Tag
                                   ?? (res.Headers.TryGetValues("ETag", out var v) ? v.FirstOrDefault() : null);
                        if (res.IsSuccessStatusCode && etag is not null)
                        {
                            etags[i] = etag;
                            break;
                        }
                        _log.LogWarning("Part {Part}/{Count} of {Kind} for job {JobId} failed: {Status} (attempt {N}/{Max})",
                            i + 1, partUrls.Count, kind, jobId, (int)res.StatusCode, attempt, PartAttempts);
                    }
                    catch (Exception ex) when (!failed.IsCancellationRequested)
                    {
                        _log.LogWarning("Part {Part}/{Count} of {Kind} for job {JobId} failed: {Error} (attempt {N}/{Max})",
                            i + 1, partUrls.Count, kind, jobId, ex.Message, attempt, PartAttempts);
                    }
                    if (attempt >= PartAttempts)
                    {
                        failed.Cancel(); // no point finishing the other parts
                        return;
                    }
                    await Task.Delay(TimeSpan.FromSeconds(5 * attempt), failed.Token);
                }
            }
        }

        var concurrency = Math.Clamp(_options.UploadConcurrency, 1, 64);
        try
        {
            await Task.WhenAll(Enumerable.Range(0, Math.Min(concurrency, partUrls.Count)).Select(_ => Task.Run(Worker)));
        }
        catch (OperationCanceledException)
        {
            // a part gave up, or shutdown — either way the upload is aborted below
        }

        var complete = !failed.IsCancellationRequested && etags.All(e => e is not null);
        if (complete)
        {
            var res = await _http.PostAsJsonAsync($"api/ext/render/jobs/{jobId}/artifact-complete", new
            {
                machineId,
                key = presign.Key,
                uploadId = presign.UploadId,
                parts = etags.Select((etag, i) => new { partNumber = i + 1, etag }),
            }, Json, ct);
            if (res.IsSuccessStatusCode) return true;
            _log.LogWarning("Completing multipart {Kind} for job {JobId} failed: {Status} {Body}",
                kind, jobId, (int)res.StatusCode, await res.Content.ReadAsStringAsync(ct));
        }

        try
        {
            await _http.PostAsJsonAsync($"api/ext/render/jobs/{jobId}/artifact-complete", new
            {
                machineId,
                key = presign.Key,
                uploadId = presign.UploadId,
                abort = true,
            }, Json, CancellationToken.None);
        }
        catch { /* parts expire on their own */ }
        return false;
    }

    public Task<string?> UploadLogAsync(string jobId, string machineId, string logFilePath, CancellationToken ct) =>
        UploadArtifactAsync(jobId, machineId, "log", logFilePath, "text/plain", ct);

    /// <summary>
    /// Reports the deployed worker version so the web UI can show what is
    /// running where. Fire-and-forget — never worth failing a startup over.
    /// </summary>
    public async Task ReportVersionAsync(WorkerOptions o, CancellationToken ct)
    {
        try
        {
            await _http.PostAsJsonAsync("api/ext/tools/ping", new
            {
                tool = "RENDER_WORKER",
                machineName = o.MachineName,
                version = o.WorkerVersion,
                hostUser = Environment.UserName,
                hostApp = $"{Environment.OSVersion.VersionString} ({Dns.GetHostName()})",
                details = new
                {
                    aeVersion = o.AeVersion,
                    aerenderPath = o.AerenderPath,
                    ffmpeg = o.FfmpegPath is not null,
                    plateSearchRoots = o.PlateSearchRoots,
                },
            }, Json, ct);
        }
        catch (Exception ex)
        {
            _log.LogDebug(ex, "Version ping failed (non-fatal)");
        }
    }

    // E3 — used by crash-recovery reconcile (§7.5)
    public async Task<ExportDetail?> GetExportAsync(string exportId, CancellationToken ct)
    {
        var res = await _http.GetAsync($"api/ext/exports/{exportId}", ct);
        if (!res.IsSuccessStatusCode) return null;
        var body = await res.Content.ReadFromJsonAsync<ExportDetailResponse>(Json, ct);
        return body?.Export;
    }

    private static double? GetDiskFreeGb()
    {
        try
        {
            var root = Path.GetPathRoot(Environment.SystemDirectory);
            if (root is null) return null;
            return Math.Round(new DriveInfo(root).AvailableFreeSpace / 1_073_741_824.0, 1);
        }
        catch
        {
            return null;
        }
    }
}