clusterflux-public/scripts/cancellation-smoke.js

#!/usr/bin/env node

const assert = require("assert");
const cp = require("child_process");
const net = require("net");
const path = require("path");

const repo = path.resolve(__dirname, "..");

function waitForJsonLine(child) {
  return new Promise((resolve, reject) => {
    let buffer = "";
    child.stdout.on("data", (chunk) => {
      buffer += chunk.toString();
      const newline = buffer.indexOf("\n");
      if (newline < 0) return;
      try {
        resolve(JSON.parse(buffer.slice(0, newline).trim()));
      } catch (error) {
        reject(error);
      }
    });
    child.once("exit", (code) => {
      reject(new Error(`process exited before JSON line with code ${code}`));
    });
  });
}

function send(addr, message) {
  return new Promise((resolve, reject) => {
    const socket = net.connect(addr.port, addr.host, () => {
      socket.write(`${JSON.stringify(message)}\n`);
    });
    let buffer = "";
    socket.on("data", (chunk) => {
      buffer += chunk.toString();
      const newline = buffer.indexOf("\n");
      if (newline < 0) return;
      socket.end();
      try {
        resolve(JSON.parse(buffer.slice(0, newline)));
      } catch (error) {
        reject(error);
      }
    });
    socket.on("error", reject);
  });
}

function spawnNode(addr) {
  const child = cp.spawn(
    "cargo",
    [
      "run",
      "-q",
      "-p",
      "disasmer-node",
      "--bin",
      "disasmer-node",
      "--",
      "--coordinator",
      `${addr.host}:${addr.port}`,
      "--tenant",
      "tenant",
      "--project-id",
      "project",
      "--node",
      "node-cancel",
      "--process",
      "vp-cancel",
      "--task",
      "compile-linux",
      "--artifact",
      "/vfs/artifacts/cancelled-output.txt",
      "--emit-ready",
      "--control-poll-ms",
      "5000",
    ],
    { cwd: repo }
  );

  let buffer = "";
  let rawStdout = "";
  let stderr = "";
  let exitCode = null;
  let exitSignal = null;
  const messages = [];
  const waiters = [];

  child.stdout.on("data", (chunk) => {
    rawStdout += chunk.toString();
    buffer += chunk.toString();
    while (true) {
      const newline = buffer.indexOf("\n");
      if (newline < 0) return;
      const line = buffer.slice(0, newline).trim();
      buffer = buffer.slice(newline + 1);
      if (!line) continue;
      let message;
      try {
        message = JSON.parse(line);
      } catch (error) {
        rejectWaiters(new Error(`node emitted non-JSON line: ${line}\n${stderr}`));
        return;
      }
      messages.push(message);
      flush();
    }
  });
  child.stderr.on("data", (chunk) => {
    stderr += chunk.toString();
  });
  child.on("exit", (code, signal) => {
    exitCode = code;
    exitSignal = signal;
    flush();
    rejectWaiters(
      new Error(
        `node process exited before expected message with code ${code} signal ${signal}\nstdout:\n${rawStdout}\nstderr:\n${stderr}`
      )
    );
  });

  function flush() {
    for (const waiter of [...waiters]) {
      const message = messages.find(waiter.predicate);
      if (!message) continue;
      clearTimeout(waiter.timer);
      waiters.splice(waiters.indexOf(waiter), 1);
      waiter.resolve(message);
    }
  }

  function waitForMessage(predicate, timeoutMs = 120000) {
    const existing = messages.find(predicate);
    if (existing) return Promise.resolve(existing);
    if (exitCode !== null) {
      return Promise.reject(
        new Error(
          `node process already exited with code ${exitCode} signal ${exitSignal}\nstdout:\n${rawStdout}\nstderr:\n${stderr}`
        )
      );
    }
    return new Promise((resolve, reject) => {
      const timer = setTimeout(() => {
        child.kill("SIGKILL");
        reject(new Error(`timed out waiting for node message\n${stderr}`));
      }, timeoutMs);
      waiters.push({ predicate, resolve, reject, timer });
    });
  }

  function rejectWaiters(error) {
    for (const waiter of [...waiters]) {
      clearTimeout(waiter.timer);
      waiters.splice(waiters.indexOf(waiter), 1);
      waiter.reject(error);
    }
  }

  function waitForExit() {
    if (child.exitCode !== null) {
      if (child.exitCode === 0) return Promise.resolve();
      return Promise.reject(new Error(`node process failed with code ${child.exitCode}\n${stderr}`));
    }
    return new Promise((resolve, reject) => {
      child.on("exit", (code) => {
        if (code === 0) {
          resolve();
        } else {
          reject(new Error(`node process failed with code ${code}\n${stderr}`));
        }
      });
    });
  }

  return { child, waitForMessage, waitForExit };
}

(async () => {
  const coordinator = cp.spawn(
    "cargo",
    [
      "run",
      "-q",
      "-p",
      "disasmer-coordinator",
      "--bin",
      "disasmer-coordinator",
      "--",
      "--listen",
      "127.0.0.1:0",
    ],
    { cwd: repo }
  );

  let node;
  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");

    node = spawnNode(addr);
    const nodeReady = await node.waitForMessage((message) => message.node_status === "ready");
    assert.strictEqual(nodeReady.process, "vp-cancel");
    assert.strictEqual(nodeReady.task, "compile-linux");

    const cancel = await send(addr, {
      type: "cancel_task",
      tenant: "tenant",
      project: "project",
      process: "vp-cancel",
      node: "node-cancel",
      task: "compile-linux",
    });
    assert.strictEqual(cancel.type, "task_cancellation_requested");
    assert.strictEqual(cancel.process, "vp-cancel");
    assert.strictEqual(cancel.task, "compile-linux");
    assert.strictEqual(cancel.node, "node-cancel");

    const report = await node.waitForMessage((message) => message.node_status === "cancelled");
    assert.strictEqual(report.terminal_state, "cancelled");
    assert.strictEqual(report.status_code, null);
    assert.strictEqual(report.large_bytes_uploaded, false);
    assert.strictEqual(report.coordinator_response.type, "task_recorded");
    await node.waitForExit();

    const events = await send(addr, {
      type: "list_task_events",
      tenant: "tenant",
      project: "project",
      actor_user: "user",
      process: "vp-cancel"
    });
    assert.strictEqual(events.type, "task_events");
    assert.strictEqual(events.events.length, 1);
    assert.strictEqual(events.events[0].node, "node-cancel");
    assert.strictEqual(events.events[0].process, "vp-cancel");
    assert.strictEqual(events.events[0].task, "compile-linux");
    assert.strictEqual(events.events[0].terminal_state, "cancelled");
    assert.strictEqual(events.events[0].status_code, null);

    const control = await send(addr, {
      type: "poll_task_control",
      tenant: "tenant",
      project: "project",
      process: "vp-cancel",
      node: "node-cancel",
      task: "compile-linux",
    });
    assert.strictEqual(control.type, "task_control");
    assert.strictEqual(control.cancel_requested, false);
  } finally {
    if (node && node.child.exitCode === null) node.child.kill("SIGKILL");
    coordinator.kill("SIGTERM");
  }

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