clusterflux-public/scripts/artifact-export-smoke.js

#!/usr/bin/env node

const assert = require("assert");
const cp = require("child_process");
const fs = require("fs");
const os = require("os");
const path = require("path");
const { nodeIdentity, signedNodeRequest } = require("./node-signing");
const {
  ensureRootlessPodman,
  flagshipNodeCapabilities,
  launchFlagship,
  repo,
  runFlagshipWorker,
  send,
  waitForJsonLine,
  waitForNodeStatus,
} = require("./real-flagship-harness");

const sourceNode = "node-export-source";
const sourceIdentity = nodeIdentity("artifact-export-smoke", sourceNode);

function nodeCapabilities() {
  return flagshipNodeCapabilities();
}

function runJson(command, args, options = {}) {
  return new Promise((resolve, reject) => {
    const child = cp.spawn(command, args, { cwd: repo, ...options });
    let stdout = "";
    let stderr = "";
    child.stdout.on("data", (chunk) => {
      stdout += chunk.toString();
    });
    child.stderr.on("data", (chunk) => {
      stderr += chunk.toString();
    });
    child.once("error", reject);
    child.once("exit", (code) => {
      if (code !== 0) {
        reject(
          new Error(
            `${command} ${args.join(" ")} failed with code ${code}\n${stderr}\n${stdout}`
          )
        );
        return;
      }
      try {
        resolve(JSON.parse(stdout));
      } catch (error) {
        reject(new Error(`${command} did not return JSON\n${stdout}\n${error.message}`));
      }
    });
  });
}

async function reportNode(
  addr,
  node,
  identity,
  { directConnectivity = true, online = true, artifacts = [] } = {}
) {
  const response = await send(addr, signedNodeRequest(node, identity, "report_node_capabilities", {
    type: "report_node_capabilities",
    tenant: "tenant",
    project: "project",
    node,
    capabilities: nodeCapabilities(),
    cached_environment_digests: [],
    dependency_cache_digests: [],
    source_snapshots: [],
    artifact_locations: artifacts,
    direct_connectivity: directConnectivity,
    online,
  }));
  assert.strictEqual(response.type, "node_capabilities_recorded");
  assert.strictEqual(response.node, node);
}

(async () => {
  ensureRootlessPodman();
  const coordinator = cp.spawn(
    "cargo",
    [
      "run",
      "-q",
      "-p",
      "clusterflux-coordinator",
      "--bin",
      "clusterflux-coordinator",
      "--",
      "--listen",
      "127.0.0.1:0",
      "--allow-local-trusted-loopback",
    ],
    {
      cwd: repo,
      env: { ...process.env, CLUSTERFLUX_NODE_STALE_AFTER_SECONDS: "1" },
    }
  );

  let worker;
  try {
    const ready = await waitForJsonLine(coordinator);
    const [host, portText] = ready.listen.split(":");
    const addr = { host, port: Number(portText) };
    assert.strictEqual((await send(addr, { type: "ping" })).type, "pong");

    worker = await runFlagshipWorker(addr, sourceNode, sourceIdentity);
    const workerReady = await worker.ready;
    assert.strictEqual(workerReady.node_status, "ready");
    assert.strictEqual(workerReady.mode, "worker");
    const firstNodeTaskCompletion = waitForNodeStatus(worker.child, "completed");
    const { compileEvent, process: virtualProcess } = await launchFlagship(addr);
    const workerCompletion = await firstNodeTaskCompletion;
    assert.strictEqual(workerCompletion.node_status, "completed");
    assert.strictEqual(
      workerCompletion.task_assignment_response.task_spec.task_definition,
      "prepare_source"
    );
    assert.strictEqual(
      workerCompletion.virtual_thread,
      workerCompletion.task_assignment_response.task_spec.task_instance
    );
    assert.strictEqual(compileEvent.status_code, 0);
    assert.match(
      compileEvent.artifact_path,
      /^\/vfs\/artifacts\/hello-clusterflux-[0-9a-f]{64}$/
    );
    const artifact = compileEvent.artifact_path.slice("/vfs/artifacts/".length);

    await reportNode(addr, sourceNode, sourceIdentity, {
      artifacts: [artifact],
    });

    const receiverIdentity = nodeIdentity("artifact-export-smoke", "node-export-receiver");
    const attachedReceiver = await send(addr, {
      type: "attach_node",
      tenant: "tenant",
      project: "project",
      node: "node-export-receiver",
      public_key: receiverIdentity.publicKey,
    });
    assert.strictEqual(attachedReceiver.type, "node_attached");
    await reportNode(addr, "node-export-receiver", receiverIdentity);

    const exportPlan = await send(addr, {
      type: "export_artifact_to_node",
      tenant: "tenant",
      project: "project",
      actor_user: "user",
      artifact,
      receiver_node: "node-export-receiver",
      direct_connectivity: true,
      failure_reason: "",
    });
    assert.strictEqual(exportPlan.type, "artifact_export_plan");
    assert.strictEqual(exportPlan.source_node, "node-export-source");
    assert.strictEqual(exportPlan.receiver_node, "node-export-receiver");
    assert.strictEqual(exportPlan.plan.transport, "NativeQuic");
    assert.strictEqual(exportPlan.plan.scope.tenant, "tenant");
    assert.strictEqual(exportPlan.plan.scope.project, "project");
    assert.strictEqual(exportPlan.plan.scope.process, virtualProcess);
    assert.deepStrictEqual(exportPlan.plan.scope.object, { Artifact: artifact });
    assert.strictEqual(exportPlan.plan.source.node, "node-export-source");
    assert.strictEqual(exportPlan.plan.destination.node, "node-export-receiver");
    assert.strictEqual(exportPlan.plan.coordinator_assisted_rendezvous, true);
    assert.strictEqual(exportPlan.plan.coordinator_bulk_relay_allowed, false);
    assert.match(exportPlan.plan.authorization_digest, /^sha256:[0-9a-f]{64}$/);
    assert.strictEqual(exportPlan.artifact_size_bytes, compileEvent.artifact_size_bytes);

    const temp = fs.mkdtempSync(path.join(os.tmpdir(), "clusterflux-artifact-export-"));
    const exportPath = path.join(temp, "hello-clusterflux");
    const cliExport = await runJson("cargo", [
      "run",
      "-q",
      "-p",
      "clusterflux-cli",
      "--bin",
      "clusterflux",
      "--",
      "artifact",
      "export",
      "--coordinator",
      `${addr.host}:${addr.port}`,
      "--tenant",
      "tenant",
      "--project-id",
      "project",
      "--user",
      "user",
      "--json",
      artifact,
      "--receiver-node",
      "node-export-receiver",
      "--to",
      exportPath,
    ]);
    assert.strictEqual(cliExport.command, "artifact export");
    assert.strictEqual(cliExport.export_plan.local_bytes_written_by_cli, true);
    assert.strictEqual(cliExport.export_plan.local_export_status, "local_bytes_written");
    assert.strictEqual(
      cliExport.export_plan.bytes_written,
      compileEvent.artifact_size_bytes
    );
    assert.strictEqual(cliExport.local_export.stream.content_material_returned_in_report, false);
    assert.strictEqual(
      cliExport.local_export.verified_digest,
      compileEvent.artifact_digest
    );
    fs.chmodSync(exportPath, 0o755);
    assert.strictEqual(
      cp.execFileSync(exportPath, { encoding: "utf8" }),
      "hello from a real Clusterflux build\n"
    );

    const crossTenant = await send(addr, {
      type: "export_artifact_to_node",
      tenant: "other",
      project: "project",
      actor_user: "user",
      artifact,
      receiver_node: "node-export-receiver",
      direct_connectivity: true,
      failure_reason: "",
    });
    assert.strictEqual(crossTenant.type, "error");
    assert.match(crossTenant.message, /tenant mismatch/);

    await reportNode(addr, sourceNode, sourceIdentity, { artifacts: [artifact] });
    const failedDirect = await send(addr, {
      type: "export_artifact_to_node",
      tenant: "tenant",
      project: "project",
      actor_user: "user",
      artifact,
      receiver_node: "node-export-receiver",
      direct_connectivity: false,
      failure_reason: "nat traversal failed",
    });
    assert.strictEqual(failedDirect.type, "error");
    assert.match(failedDirect.message, /nat traversal failed/);
    assert.match(failedDirect.message, /coordinator bulk relay is disabled/);

    await reportNode(addr, "node-export-receiver", receiverIdentity, { online: false });
    await new Promise((resolve) => setTimeout(resolve, 2100));
    await reportNode(addr, sourceNode, sourceIdentity, { artifacts: [artifact] });
    const offlineReceiver = await send(addr, {
      type: "export_artifact_to_node",
      tenant: "tenant",
      project: "project",
      actor_user: "user",
      artifact,
      receiver_node: "node-export-receiver",
      direct_connectivity: true,
      failure_reason: "",
    });
    assert.strictEqual(offlineReceiver.type, "error");
    assert.match(offlineReceiver.message, /offline/);
  } finally {
    worker?.child.kill("SIGTERM");
    coordinator.kill("SIGTERM");
  }

  console.log("Artifact export smoke passed");
})().catch((error) => {
  console.error(error.stack || error.message);
  process.exit(1);
});