zcbot/windows-node/Zcbot.WindowsNode/WorkspaceStore.cs

using System.Text.Json;

namespace Zcbot.WindowsNode;

internal sealed class WorkspaceStore(NodePaths paths)
{
    private static readonly JsonSerializerOptions JsonOptions = new() { WriteIndented = true };

    internal string CurrentDirectory(RecoverableJob job) =>
        Path.Combine(paths.WorkspacesDirectory, Binding(job).WorkspaceId.ToString("D"), "current");

    internal void Prepare(RecoverableJob job, string capability)
    {
        var root = WorkspaceRoot(job);
        var current = Path.Combine(root, "current");
        var rollback = Path.Combine(root, "rollback");
        Directory.CreateDirectory(root);
        var binding = Binding(job);
        if (binding.Mode == "continue")
        {
            var metadata = ReadMetadata(root);
            if (metadata.Capability != capability
                || metadata.HeadJobId != binding.SourceJobId)
            {
                throw new InvalidDataException("Workspace head does not match the requested source job.");
            }
            if (!Directory.Exists(current))
            {
                if (Directory.Exists(rollback)) Directory.Move(rollback, current);
                else throw new InvalidDataException("Workspace current state is missing.");
            }
            if (Directory.Exists(rollback)) Directory.Delete(rollback, recursive: true);
            Directory.Move(current, rollback);
        }
        else if (File.Exists(Path.Combine(root, "workspace.json")) || Directory.Exists(current))
        {
            throw new InvalidDataException("A new workspace already contains state.");
        }

        var jobDirectory = Path.Combine(paths.JobsDirectory, job.JobId.ToString("D"));
        AtomicJson(Path.Combine(jobDirectory, "workspace.json"), new
        {
            workspace_id = binding.WorkspaceId,
            source_job_id = binding.SourceJobId,
            mode = binding.Mode,
            local_base_path = binding.Mode == "continue" ? rollback : null,
        });
    }

    internal long Promote(RecoverableJob job, string capability, string stateFilename)
    {
        var root = WorkspaceRoot(job);
        var current = Path.Combine(root, "current");
        var output = Path.Combine(paths.JobsDirectory, job.JobId.ToString("D"), "output");
        var statePath = Path.Combine(output, stateFilename);
        if (!File.Exists(statePath))
        {
            throw new InvalidDataException("Adapter did not create the workspace state output.");
        }
        if (Directory.Exists(current))
        {
            throw new InvalidDataException("Workspace current destination already exists.");
        }
        Directory.Move(output, current);
        return CommitCurrent(job, capability);
    }

    internal long EnsureSucceeded(
        RecoverableJob job, string capability, string stateFilename)
    {
        var current = CurrentDirectory(job);
        if (File.Exists(Path.Combine(current, stateFilename)))
        {
            return CommitCurrent(job, capability);
        }
        var output = Path.Combine(paths.JobsDirectory, job.JobId.ToString("D"), "output");
        if (File.Exists(Path.Combine(output, stateFilename)))
        {
            return Promote(job, capability, stateFilename);
        }
        Restore(job);
        throw new InvalidDataException("Successful workspace output cannot be recovered.");
    }

    private long CommitCurrent(RecoverableJob job, string capability)
    {
        var root = WorkspaceRoot(job);
        var sizeBytes = Directory.EnumerateFiles(root, "*", SearchOption.AllDirectories)
            .Where(item => !Path.GetFileName(item).Equals(
                "workspace.json", StringComparison.OrdinalIgnoreCase))
            .Sum(item => new FileInfo(item).Length);
        AtomicJson(Path.Combine(root, "workspace.json"), new
        {
            workspace_id = Binding(job).WorkspaceId,
            capability,
            head_job_id = job.JobId,
            size_bytes = sizeBytes,
            updated_at = DateTimeOffset.UtcNow,
        });
        AtomicJson(
            Path.Combine(paths.JobsDirectory, job.JobId.ToString("D"), "workspace-result.json"),
            new { size_bytes = sizeBytes });
        return sizeBytes;
    }

    internal void Restore(RecoverableJob job)
    {
        if (Binding(job).Mode != "continue") return;
        var root = WorkspaceRoot(job);
        var current = Path.Combine(root, "current");
        var rollback = Path.Combine(root, "rollback");
        if (!Directory.Exists(current) && Directory.Exists(rollback))
        {
            Directory.Move(rollback, current);
        }
    }

    private string WorkspaceRoot(RecoverableJob job) =>
        Path.Combine(paths.WorkspacesDirectory, Binding(job).WorkspaceId.ToString("D"));

    private static WorkspaceBinding Binding(RecoverableJob job) =>
        job.Workspace
        ?? throw new InvalidDataException("Job has no persistent workspace binding.");

    private static WorkspaceMetadata ReadMetadata(string root)
    {
        var path = Path.Combine(root, "workspace.json");
        using var document = JsonDocument.Parse(File.ReadAllBytes(path));
        var value = document.RootElement;
        if (!value.TryGetProperty("capability", out var capability)
            || !value.TryGetProperty("head_job_id", out var head)
            || !Guid.TryParse(head.GetString(), out var headJobId))
        {
            throw new InvalidDataException("Workspace metadata is invalid.");
        }
        return new WorkspaceMetadata(capability.GetString() ?? "", headJobId);
    }

    private static void AtomicJson(string path, object value)
    {
        Directory.CreateDirectory(Path.GetDirectoryName(path)!);
        var temporary = path + ".tmp-" + Guid.NewGuid().ToString("N");
        try
        {
            using (var stream = new FileStream(
                temporary,
                FileMode.CreateNew,
                FileAccess.Write,
                FileShare.None,
                4096,
                FileOptions.WriteThrough))
            {
                var content = JsonSerializer.SerializeToUtf8Bytes(value, JsonOptions);
                stream.Write(content);
                stream.Flush(flushToDisk: true);
            }
            File.Move(temporary, path, overwrite: true);
        }
        finally
        {
            if (File.Exists(temporary)) File.Delete(temporary);
        }
    }

    private sealed record WorkspaceMetadata(string Capability, Guid HeadJobId);
}