101 lines
4.7 KiB
C#
101 lines
4.7 KiB
C#
using System.Net.Http.Headers;
|
|
using System.Security.Cryptography;
|
|
using System.Text;
|
|
using System.Text.Json;
|
|
|
|
namespace Zcbot.WindowsNode;
|
|
|
|
internal sealed class JobOutputUploader(NodeConfig config)
|
|
{
|
|
private static readonly JsonSerializerOptions JsonOptions = new() { WriteIndented = true };
|
|
|
|
internal async Task UploadAsync(RecoverableJob job)
|
|
{
|
|
var jobDirectory = Path.Combine(
|
|
NodePaths.ForCurrentMachine().JobsDirectory, job.JobId.ToString("D"));
|
|
var completionPath = Path.Combine(jobDirectory, "upload-complete.json");
|
|
if (File.Exists(completionPath)) return;
|
|
var terminalPath = Path.Combine(jobDirectory, "terminal.json");
|
|
using var terminal = JsonDocument.Parse(await File.ReadAllBytesAsync(terminalPath));
|
|
if (terminal.RootElement.GetProperty("status").GetString() != "succeeded") return;
|
|
var manifest = terminal.RootElement.GetProperty("artifact_manifest").Clone();
|
|
|
|
using var client = new HttpClient { BaseAddress = config.ServerUrl };
|
|
client.DefaultRequestHeaders.Authorization =
|
|
new AuthenticationHeaderValue("Bearer", config.NodeToken);
|
|
client.DefaultRequestHeaders.Add("X-Node-Id", config.NodeId.ToString());
|
|
client.DefaultRequestHeaders.Add("X-Lease-Id", job.LeaseId.ToString());
|
|
client.DefaultRequestHeaders.Add("X-Request-Digest", job.RequestDigest);
|
|
foreach (var artifact in manifest.EnumerateArray())
|
|
{
|
|
var localId = artifact.GetProperty("artifact_id").GetString()!;
|
|
var filename = artifact.GetProperty("filename").GetString()!;
|
|
var expectedSize = artifact.GetProperty("size_bytes").GetInt64();
|
|
var expectedDigest = artifact.GetProperty("sha256").GetString()!;
|
|
var path = Path.Combine(jobDirectory, "output", filename);
|
|
var info = new FileInfo(path);
|
|
if (!info.Exists || info.Length != expectedSize)
|
|
{
|
|
throw new InvalidDataException($"Output artifact is missing or changed: {localId}.");
|
|
}
|
|
await using (var verify = new FileStream(
|
|
path, FileMode.Open, FileAccess.Read, FileShare.Read, 64 * 1024,
|
|
FileOptions.Asynchronous | FileOptions.SequentialScan))
|
|
{
|
|
var digest = Convert.ToHexString(
|
|
await SHA256.HashDataAsync(verify)).ToLowerInvariant();
|
|
if (digest != expectedDigest)
|
|
{
|
|
throw new InvalidDataException($"Output artifact digest changed: {localId}.");
|
|
}
|
|
}
|
|
await using var stream = new FileStream(
|
|
path, FileMode.Open, FileAccess.Read, FileShare.Read, 64 * 1024,
|
|
FileOptions.Asynchronous | FileOptions.SequentialScan);
|
|
using var content = new StreamContent(stream);
|
|
content.Headers.ContentType = new MediaTypeHeaderValue("application/octet-stream");
|
|
content.Headers.ContentLength = expectedSize;
|
|
content.Headers.Add("X-Content-SHA256", expectedDigest);
|
|
content.Headers.Add("X-Content-Length", expectedSize.ToString());
|
|
using var response = await client.PutAsync(
|
|
$"/v1/compute/jobs/{job.JobId:D}/outputs/{Uri.EscapeDataString(localId)}",
|
|
content);
|
|
response.EnsureSuccessStatusCode();
|
|
}
|
|
|
|
using var completeContent = new StringContent(
|
|
JsonSerializer.Serialize(new { artifact_manifest = manifest }),
|
|
Encoding.UTF8,
|
|
"application/json");
|
|
using var completeResponse = await client.PostAsync(
|
|
$"/v1/compute/jobs/{job.JobId:D}/outputs/complete", completeContent);
|
|
completeResponse.EnsureSuccessStatusCode();
|
|
var responseBody = await completeResponse.Content.ReadAsByteArrayAsync();
|
|
AtomicWrite(completionPath, responseBody);
|
|
}
|
|
|
|
private static void AtomicWrite(string path, byte[] responseBody)
|
|
{
|
|
using var response = JsonDocument.Parse(responseBody);
|
|
var content = JsonSerializer.SerializeToUtf8Bytes(new
|
|
{
|
|
completed_at = DateTimeOffset.UtcNow,
|
|
response = response.RootElement,
|
|
}, JsonOptions);
|
|
var temporary = path + ".tmp-" + Guid.NewGuid().ToString("N");
|
|
try
|
|
{
|
|
using var stream = new FileStream(
|
|
temporary, FileMode.CreateNew, FileAccess.Write, FileShare.None,
|
|
4096, FileOptions.WriteThrough);
|
|
stream.Write(content);
|
|
stream.Flush(flushToDisk: true);
|
|
File.Move(temporary, path, overwrite: false);
|
|
}
|
|
finally
|
|
{
|
|
if (File.Exists(temporary)) File.Delete(temporary);
|
|
}
|
|
}
|
|
}
|