Vectrum
Sign In
555 lines · 21.6 KB · TypeScript
import { createHash, randomBytes } from "node:crypto";
import { count, db, filter, findById, findFirst, newId, orNull, withLock, type RecordModel } from "./db";
import type { RunnerIdentity } from "./runners";
import type { JobSpec, StepSpec } from "./workflows";

export type RunStatus = "queued" | "in_progress" | "success" | "failure" | "cancelled";
export type JobStatus = "waiting" | "queued" | "in_progress" | "success" | "failure" | "cancelled" | "skipped";
export type StepStatus = "pending" | "in_progress" | "success" | "failure" | "cancelled" | "skipped";
export type RunEvent = "push" | "workflow_dispatch";

export const JOB_TOKEN_PREFIX = "vjt_";

/** A running job whose runner hasn't been heard from for this long is given up on. */
const STALE_AFTER_MS = 90_000;
/** Output kept per step; the rest is dropped with a note. */
const MAX_STEP_LOG = 4 * 1024 * 1024;

const FINISHED = new Set<string>(["success", "failure", "cancelled", "skipped"]);
export const isFinished = (status: string) => FINISHED.has(status);

const sha256 = (s: string) => createHash("sha256").update(s).digest("hex");

export type Run = {
  id: string;
  repoId: string;
  number: number;
  workflowPath: string;
  workflowName: string;
  event: RunEvent;
  ref: string;
  sha: string;
  title: string;
  actor: string | null;
  status: RunStatus;
  error: string | null;
  createdAt: number;
  startedAt: number | null;
  finishedAt: number | null;
};

function toRun(r: RecordModel): Run {
  return {
    id: r.id,
    repoId: r.repo,
    number: r.number,
    workflowPath: r.workflow_path,
    workflowName: r.workflow_name,
    event: r.event,
    ref: r.ref,
    sha: r.sha,
    title: r.title,
    actor: r.expand?.actor?.username ?? null,
    status: r.status,
    error: orNull(r.error),
    createdAt: r.created_at,
    startedAt: orNull(r.started_at),
    finishedAt: orNull(r.finished_at),
  };
}

export type Step = { index: number; name: string; status: StepStatus; exitCode: number | null; startedAt: number | null; finishedAt: number | null };

export type Job = {
  id: string;
  runId: string;
  key: string;
  name: string;
  runsOn: string[];
  needs: string[];
  status: JobStatus;
  runnerName: string | null;
  cancelRequested: boolean;
  createdAt: number;
  startedAt: number | null;
  finishedAt: number | null;
};

function toJob(r: RecordModel): Job {
  return {
    id: r.id,
    runId: r.run,
    key: r.key,
    name: r.name,
    runsOn: r.runs_on,
    needs: r.needs,
    status: r.status,
    runnerName: orNull(r.runner_name),
    cancelRequested: r.cancel_requested,
    createdAt: r.created_at,
    startedAt: orNull(r.started_at),
    finishedAt: orNull(r.finished_at),
  };
}

/** What a job stores of its workflow definition, to hand to the runner that picks it up. */
type StoredSpec = { env: Record<string, string>; timeoutMinutes: number; steps: StepSpec[] };

export const CHECKOUT_STEP = "Check out repository";

const collection = async (name: string) => (await db()).collection(name);

/**
 * Runs, jobs and steps change in read-modify-write sequences: from the web UI, from runners
 * reporting in and from sweeps for lost runners. They all take this lock, so they don't interleave.
 * The helpers under "Progress" expect it to be held already (taking it again would deadlock).
 */
const withRunsLock = <T>(fn: () => Promise<T>) => withLock("runs", fn);

// ---------------------------------------------------------------------------
// Creating runs

export async function createRun(input: {
  repoId: string;
  workflowPath: string;
  workflowName: string;
  event: RunEvent;
  ref: string;
  sha: string;
  title: string;
  actorId: string | null;
  /** The parsed jobs, or why the workflow file can't run (the run is then created as failed). */
  jobs: JobSpec[] | { error: string };
}) {
  return withRunsLock(async () => {
    const now = Date.now();
    const last = await findFirst("workflow_runs", filter("repo = {:repo}", { repo: input.repoId }), { sort: "-number", fields: "number" });
    const number = (last?.number ?? 0) + 1;
    const error = "error" in input.jobs ? input.jobs.error : null;
    // The run, its jobs and their steps are created together, or not at all.
    const batch = (await db()).createBatch();
    const runId = newId();
    batch.collection("workflow_runs").create({
      id: runId,
      repo: input.repoId,
      number,
      workflow_path: input.workflowPath,
      workflow_name: input.workflowName,
      event: input.event,
      ref: input.ref,
      sha: input.sha,
      title: input.title.slice(0, 300),
      actor: input.actorId ?? "",
      status: error ? "failure" : "queued",
      error: error ?? "",
      created_at: now,
      finished_at: error ? now : 0,
    });
    if (Array.isArray(input.jobs)) {
      input.jobs.forEach((job, position) => {
        const jobId = newId();
        const spec: StoredSpec = { env: job.env, timeoutMinutes: job.timeoutMinutes, steps: job.steps };
        batch.collection("workflow_jobs").create({
          id: jobId,
          run: runId,
          position,
          key: job.key,
          name: job.name,
          runs_on: job.runsOn,
          needs: job.needs,
          spec,
          status: job.needs.length ? "waiting" : "queued",
          created_at: now,
        });
        [CHECKOUT_STEP, ...job.steps.map((s) => s.name)].forEach((name, idx) => {
          batch.collection("job_steps").create({ job: jobId, idx, name, status: "pending", exit_code: null });
        });
      });
    }
    await batch.send();
    return number;
  });
}

// ---------------------------------------------------------------------------
// Reading runs

const pathFilter = (repoId: string, workflowPath?: string | null) =>
  workflowPath ? filter("repo = {:repoId} && workflow_path = {:workflowPath}", { repoId, workflowPath }) : filter("repo = {:repoId}", { repoId });

/** `offset` is a multiple of `limit` (pages of runs). */
export async function listRuns(repoId: string, opts: { workflowPath?: string | null; limit: number; offset: number }) {
  const page = Math.floor(opts.offset / opts.limit) + 1;
  const { items } = await (await collection("workflow_runs")).getList(page, opts.limit, {
    filter: pathFilter(repoId, opts.workflowPath),
    sort: "-number",
    expand: "actor",
    skipTotal: true,
  });
  return items.map(toRun);
}

export async function countRuns(repoId: string, workflowPath?: string | null) {
  return count("workflow_runs", pathFilter(repoId, workflowPath));
}

/** Workflows that have run, with the name they had last, so ones deleted since still show up in the filter. */
export async function listRunWorkflows(repoId: string) {
  const rows = await (await collection("workflow_runs")).getFullList({
    filter: filter("repo = {:repoId}", { repoId }),
    sort: "-number",
    fields: "workflow_path,workflow_name",
  });
  const latest = new Map<string, string>();
  for (const r of rows) if (!latest.has(r.workflow_path)) latest.set(r.workflow_path, r.workflow_name);
  return [...latest]
    .map(([path, name]) => ({ path, name }))
    .sort((a, b) => a.name.toLowerCase().localeCompare(b.name.toLowerCase()));
}

export async function getRun(repoId: string, number: number) {
  if (!Number.isInteger(number)) return null;
  const row = await findFirst("workflow_runs", filter("repo = {:repoId} && number = {:number}", { repoId, number }), { expand: "actor" });
  return row ? toRun(row) : null;
}

export async function listJobs(runId: string) {
  const rows = await (await collection("workflow_jobs")).getFullList({ filter: filter("run = {:runId}", { runId }), sort: "position,@rowid" });
  return rows.map(toJob);
}

export async function listSteps(jobId: string): Promise<Step[]> {
  const rows = await (await collection("job_steps")).getFullList({ filter: filter("job = {:jobId}", { jobId }), sort: "idx" });
  return rows.map((r) => ({
    index: r.idx,
    name: r.name,
    status: r.status,
    exitCode: typeof r.exit_code === "number" ? r.exit_code : null,
    startedAt: orNull(r.started_at),
    finishedAt: orNull(r.finished_at),
  }));
}

/** Each step's output, joined from its chunks. */
export async function readJobLogs(jobId: string) {
  const rows = await (await collection("job_logs")).getFullList({
    filter: filter("job = {:jobId}", { jobId }),
    sort: "step,@rowid",
    fields: "step,text",
  });
  const logs = new Map<number, string>();
  for (const { step, text } of rows) logs.set(step, (logs.get(step) ?? "") + text);
  return logs;
}

export type CommitStatus = "pending" | "success" | "failure";

/**
 * The combined state of each commit's latest runs (one per workflow), for status marks in commit
 * lists, with the newest of those runs to link to.
 */
export async function commitStatuses(repoId: string, shas: string[]) {
  const out = new Map<string, { status: CommitStatus; runNumber: number }>();
  const latest = new Map<string, { sha: string; status: RunStatus; number: number }>();
  // A few dozen commits at a time keep the filter short.
  for (let i = 0; i < shas.length; i += 40) {
    const chunk = shas.slice(i, i + 40);
    const rows = await (await collection("workflow_runs")).getFullList({
      filter: `${filter("repo = {:repoId}", { repoId })} && (${chunk.map((sha, j) => filter(`sha = {:s${j}}`, { [`s${j}`]: sha })).join(" || ")})`,
      sort: "-number",
      fields: "sha,workflow_path,status,number",
    });
    for (const r of rows) {
      const key = `${r.sha}\0${r.workflow_path}`;
      if (!latest.has(key)) latest.set(key, { sha: r.sha, status: r.status, number: r.number });
    }
  }
  const rank: Record<CommitStatus, number> = { success: 0, failure: 1, pending: 2 };
  for (const { sha, status, number } of latest.values()) {
    const prev = out.get(sha);
    const now: CommitStatus = status === "success" ? "success" : status === "failure" || status === "cancelled" ? "failure" : "pending";
    // Pending wins (something still runs), then failure, then success.
    out.set(sha, {
      status: prev && rank[prev.status] > rank[now] ? prev.status : now,
      runNumber: Math.max(prev?.runNumber ?? 0, number),
    });
  }
  return out;
}

// ---------------------------------------------------------------------------
// Progress

const statusIn = (statuses: string[]) => `(${statuses.map((s) => filter("status = {:s}", { s })).join(" || ")})`;

/** Moves a job's steps in one of the `from` states to `to`, closing out the ones that started. */
async function setSteps(jobId: string, from: StepStatus[], to: StepStatus, now = Date.now()) {
  const steps = await collection("job_steps");
  const rows = await steps.getFullList({ filter: `${filter("job = {:jobId}", { jobId })} && ${statusIn(from)}` });
  for (const step of rows) {
    await steps.update(step.id, { status: to, finished_at: step.finished_at || (step.started_at ? now : 0) });
  }
}

/**
 * Moves a run forward after one of its jobs changed: jobs whose needs all succeeded are queued, jobs
 * whose needs failed are skipped, and the run's own status follows its jobs.
 */
async function advanceRun(runId: string) {
  const now = Date.now();
  const jobsCollection = await collection("workflow_jobs");
  const jobs = await jobsCollection.getFullList({ filter: filter("run = {:runId}", { runId }), fields: "id,key,needs,status" });
  const statusOf = new Map(jobs.map((j) => [j.key as string, j.status as JobStatus]));
  for (let changed = true; changed; ) {
    changed = false;
    for (const job of jobs.filter((j) => statusOf.get(j.key) === "waiting")) {
      const needs = (job.needs as string[]).map((k) => statusOf.get(k));
      if (needs.every((s) => s === "success")) {
        await jobsCollection.update(job.id, { status: "queued" });
        statusOf.set(job.key, "queued");
        changed = true;
      } else if (needs.some((s) => s && s !== "success" && isFinished(s))) {
        await jobsCollection.update(job.id, { status: "skipped", finished_at: now });
        await setSteps(job.id, ["pending"], "skipped");
        statusOf.set(job.key, "skipped");
        changed = true;
      }
    }
  }

  const statuses = [...statusOf.values()];
  let status: RunStatus;
  if (statuses.every(isFinished)) {
    status = statuses.includes("failure") ? "failure" : statuses.includes("cancelled") ? "cancelled" : "success";
  } else {
    status = statuses.some((s) => s !== "queued" && s !== "waiting") ? "in_progress" : "queued";
  }
  const runs = await collection("workflow_runs");
  const run = await runs.getOne(runId, { fields: "started_at,finished_at" });
  await runs.update(runId, {
    status,
    started_at: status === "queued" ? run.started_at : run.started_at || now,
    finished_at: isFinished(status) ? run.finished_at || now : 0,
  });
}

/** A job's steps by index, to update as output and step states come in. */
type StepRecords = Map<number, RecordModel>;

async function stepRecords(jobId: string): Promise<StepRecords> {
  const rows = await (await collection("job_steps")).getFullList({ filter: filter("job = {:jobId}", { jobId }) });
  return new Map(rows.map((r) => [r.idx as number, r]));
}

async function appendLog(jobId: string, steps: StepRecords, step: number, text: string) {
  const record = steps.get(step);
  if (!text || !record) return;
  const size: number = record.log_size;
  if (size >= MAX_STEP_LOG) return;
  const room = MAX_STEP_LOG - size;
  const chunk = text.length > room ? `${text.slice(0, room)}\n[Output truncated: steps keep at most 4 MB of logs.]\n` : text;
  await (await collection("job_logs")).create({ job: jobId, step, text: chunk });
  record.log_size = size + chunk.length;
  await (await collection("job_steps")).update(record.id, { log_size: record.log_size });
}

/** Ends a job: unfinished steps are closed out, its token stops working and the run moves on. */
async function finishJob(jobId: string, runId: string, status: "success" | "failure" | "cancelled") {
  const now = Date.now();
  await (await collection("workflow_jobs")).update(jobId, { status, finished_at: now, token_hash: "" });
  await setSteps(jobId, ["in_progress"], status, now);
  await setSteps(jobId, ["pending"], "skipped", now);
  await advanceRun(runId);
}

/**
 * A job its runner lost (it crashed, lost its network or was deleted). If the runner never started
 * it, it goes back in the queue for another runner; otherwise it fails with `why` in its log.
 */
async function abandonJob(job: RecordModel, why: string) {
  const started = await findFirst("job_steps", filter("job = {:jobId} && started_at > 0", { jobId: job.id }), { sort: "-idx", fields: "idx" });
  if (!started && !job.cancel_requested) {
    await (await collection("workflow_jobs")).update(job.id, {
      status: "queued",
      runner: "",
      runner_name: "",
      token_hash: "",
      started_at: 0,
      heartbeat_at: 0,
    });
    await advanceRun(job.run);
    return;
  }
  await appendLog(job.id, await stepRecords(job.id), started?.idx ?? 0, `\n${why}\n`);
  await finishJob(job.id, job.run, job.cancel_requested ? "cancelled" : "failure");
}

async function sweep() {
  const stale = await (await collection("workflow_jobs")).getFullList({
    filter: filter("status = 'in_progress' && heartbeat_at < {:before}", { before: Date.now() - STALE_AFTER_MS }),
    fields: "id,run,cancel_requested",
  });
  for (const job of stale) await abandonJob(job, "The runner stopped responding, so the job was stopped.");
}

/** Gives up on running jobs whose runner went quiet. */
export async function sweepStaleJobs() {
  await withRunsLock(sweep);
}

export async function cancelRun(runId: string) {
  await withRunsLock(async () => {
    const now = Date.now();
    const jobs = await collection("workflow_jobs");
    const pending = await jobs.getFullList({ filter: filter("run = {:runId} && (status = 'waiting' || status = 'queued')", { runId }), fields: "id" });
    for (const { id } of pending) {
      await jobs.update(id, { status: "cancelled", finished_at: now });
      await setSteps(id, ["pending"], "skipped", now);
    }
    // Running jobs are stopped by their runner, which learns of it on its next update.
    const running = await jobs.getFullList({ filter: filter("run = {:runId} && status = 'in_progress'", { runId }), fields: "id" });
    for (const { id } of running) await jobs.update(id, { cancel_requested: true });
    await advanceRun(runId);
  });
}

// ---------------------------------------------------------------------------
// The runner side

export type ClaimedJob = Awaited<ReturnType<typeof claimJob>>;

/**
 * Hands the oldest queued job this runner can run to it: in a repository it serves, needing only
 * labels it has. Returns null when there's nothing to do.
 */
export async function claimJob(runner: RunnerIdentity, origin: string) {
  const labels = new Set(runner.labels.map((l) => l.toLowerCase()));
  return withRunsLock(async () => {
    await sweep();
    const jobs = await collection("workflow_jobs");

    // One job at a time: a runner that asks again while it still has one has lost it.
    const lost = await jobs.getFullList({
      filter: filter("runner = {:runner} && status = 'in_progress'", { runner: runner.id }),
      fields: "id,run,cancel_requested",
    });
    for (const job of lost) await abandonJob(job, "The runner restarted while running this job.");

    // Its own repository, or (user runners) any the user owns.
    const scope = runner.repoId
      ? filter("run.repo = {:repo}", { repo: runner.repoId })
      : runner.userId
        ? filter("run.repo.owner = {:user}", { user: runner.userId })
        : null;
    if (!scope) return null;
    const { items: candidates } = await jobs.getList(1, 200, {
      filter: `status = 'queued' && ${scope}`,
      sort: "created_at,position,@rowid",
      fields: "id,runs_on",
      skipTotal: true,
    });
    const match = candidates.find((c) => (c.runs_on as string[]).every((l) => labels.has(l.toLowerCase())));
    if (!match) return null;

    const token = `${JOB_TOKEN_PREFIX}${randomBytes(30).toString("base64url")}`;
    const now = Date.now();
    const job = await jobs.update(match.id, {
      status: "in_progress",
      runner: runner.id,
      runner_name: runner.name,
      token_hash: sha256(token),
      started_at: now,
      heartbeat_at: now,
    });
    await advanceRun(job.run);

    const run = await (await collection("workflow_runs")).getOne(job.run, { expand: "actor,repo.owner" });
    const repo = run.expand?.repo as RecordModel;
    const owner: string = repo.expand?.owner?.username;
    const spec = job.spec as StoredSpec;
    const repository = `${owner}/${repo.name}`;
    const ref: string = run.ref;
    const refName = ref.replace(/^refs\/(heads|tags)\//, "");

    return {
      id: job.id as string,
      token,
      repository,
      cloneUrl: `${origin}/${repository}.git`,
      sha: run.sha as string,
      ref,
      timeoutMinutes: spec.timeoutMinutes,
      // Variables every step sees, under the workflow's and job's own env.
      env: {
        CI: "true",
        VECTRUM_ACTIONS: "true",
        VECTRUM_SERVER_URL: origin,
        VECTRUM_REPOSITORY: repository,
        VECTRUM_REPOSITORY_OWNER: owner,
        VECTRUM_SHA: run.sha,
        VECTRUM_REF: ref,
        VECTRUM_REF_NAME: refName,
        VECTRUM_REF_TYPE: ref.startsWith("refs/tags/") ? "tag" : "branch",
        VECTRUM_EVENT_NAME: run.event,
        VECTRUM_WORKFLOW: run.workflow_name,
        VECTRUM_RUN_ID: run.id,
        VECTRUM_RUN_NUMBER: String(run.number),
        VECTRUM_JOB: job.key,
        VECTRUM_ACTOR: run.expand?.actor?.username ?? "",
        ...spec.env,
      } as Record<string, string>,
      steps: spec.steps.map((s, i) => ({ index: i + 1, ...s })),
    };
  });
}

type StepUpdate = { index: number; status: StepStatus; exitCode?: number | null };

/**
 * A runner reporting on its job: step states, new output, and at the end the job's result.
 * Returns whether the job should be cancelled, or null if this runner doesn't hold the job.
 */
export async function updateJob(
  runner: RunnerIdentity,
  jobId: string,
  update: { steps?: StepUpdate[]; logs?: { step: number; text: string }[]; status?: "success" | "failure" | "cancelled" },
) {
  return withRunsLock(async () => {
    const job = await findById("workflow_jobs", jobId);
    if (!job || job.runner !== runner.id || job.status !== "in_progress") return null;
    const now = Date.now();
    await (await collection("workflow_jobs")).update(jobId, { heartbeat_at: now });

    const steps = await stepRecords(jobId);
    for (const log of update.logs ?? []) await appendLog(jobId, steps, log.step, log.text);
    for (const s of update.steps ?? []) {
      const step = steps.get(s.index);
      if (!step) continue;
      const unfinished = s.status === "pending" || s.status === "in_progress" || (s.status === "skipped" && !step.started_at);
      await (await collection("job_steps")).update(step.id, {
        status: s.status,
        exit_code: s.exitCode ?? step.exit_code,
        started_at: s.status === "skipped" ? step.started_at : step.started_at || now,
        finished_at: unfinished ? 0 : now,
      });
    }
    if (update.status) await finishJob(jobId, job.run, update.status);
    return { cancelled: job.cancel_requested as boolean };
  });
}

/** The repository a job's checkout token may read, while the job is running. */
export async function authenticateJobToken(secret: string) {
  if (!secret.startsWith(JOB_TOKEN_PREFIX)) return null;
  const job = await findFirst("workflow_jobs", filter("token_hash = {:hash} && status = 'in_progress'", { hash: sha256(secret) }), {
    expand: "run",
    fields: "expand.run.repo",
  });
  const repoId = job?.expand?.run?.repo as string | undefined;
  return repoId ? { repoId } : null;
}

/** Jobs in a repository waiting for a runner to pick them up. */
export async function queuedJobCount(repoId: string) {
  return count("workflow_jobs", filter("run.repo = {:repoId} && status = 'queued'", { repoId }));
}