Vectrum
Sign In
179 lines · 7.2 KB · TypeScript
import { createHash, randomBytes } from "node:crypto";
import { ClientResponseError } from "pocketbase";
import { db, filter, findFirst, orNull, type RecordModel } from "./db";

export const RUNNER_TOKEN_PREFIX = "vrn_";

/** A runner counts as online if it polled for work this recently. */
const ONLINE_WINDOW_MS = 60_000;

/** Where a runner takes jobs from: one repository, or every repository a user owns. */
export type RunnerScope = { repoId: string } | { userId: string };

export type Runner = {
  id: string;
  repoId: string | null;
  userId: string | null;
  name: string;
  /** The labels set in the UI. */
  labels: string[];
  /** Labels every runner has: self-hosted and, once it has connected, its OS and architecture. */
  systemLabels: string[];
  hint: string;
  hostname: string | null;
  os: string | null;
  arch: string | null;
  version: string | null;
  createdAt: number;
  lastSeenAt: number | null;
  expiresAt: number | null;
  expired: boolean;
  online: boolean;
  /** The job it is running, if any. */
  busy: { runNumber: number; jobName: string; repo: string } | null;
};

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

export class RunnerError extends Error {}

export function systemLabels(row: { os: string | null; arch: string | null }) {
  return ["self-hosted", ...(row.os ? [row.os] : []), ...(row.arch ? [row.arch] : [])];
}

const isOnline = (lastSeenAt: number) => !!lastSeenAt && Date.now() - lastSeenAt < ONLINE_WINDOW_MS;
const isExpired = (expiresAt: number) => !!expiresAt && expiresAt < Date.now();

/** The job each of these runners is running, if any. */
async function busyJobs(runnerIds: string[]) {
  const busy = new Map<string, NonNullable<Runner["busy"]>>();
  if (runnerIds.length === 0) return busy;
  const jobs = await (await db()).collection("workflow_jobs").getFullList({
    filter: `status = 'in_progress' && (${runnerIds.map((id, i) => filter(`runner = {:r${i}}`, { [`r${i}`]: id })).join(" || ")})`,
    expand: "run.repo.owner",
  });
  for (const job of jobs) {
    const run = job.expand?.run as RecordModel;
    const repo = run.expand?.repo as RecordModel;
    busy.set(job.runner, { runNumber: run.number, jobName: job.name, repo: `${repo.expand?.owner?.username}/${repo.name}` });
  }
  return busy;
}

function toRunner(row: RecordModel, busy: Map<string, NonNullable<Runner["busy"]>>): Runner {
  const os = orNull<string>(row.os);
  const arch = orNull<string>(row.arch);
  return {
    id: row.id,
    repoId: orNull(row.repo),
    userId: orNull(row.user),
    name: row.name,
    labels: row.labels ?? [],
    systemLabels: systemLabels({ os, arch }),
    hint: row.hint,
    hostname: orNull(row.hostname),
    os,
    arch,
    version: orNull(row.version),
    createdAt: row.created_at,
    lastSeenAt: orNull(row.last_seen_at),
    expiresAt: orNull(row.expires_at),
    expired: isExpired(row.expires_at),
    online: isOnline(row.last_seen_at),
    busy: busy.get(row.id) ?? null,
  };
}

const scopeFilter = (scope: RunnerScope) =>
  "repoId" in scope ? filter("repo = {:repoId}", { repoId: scope.repoId }) : filter("user = {:userId}", { userId: scope.userId });

/** Splits "gpu, docker linux" into labels; they're matched case-insensitively, so stored lowercase. */
export function parseLabels(input: string) {
  const labels = [...new Set(input.split(/[\s,]+/).filter(Boolean).map((l) => l.toLowerCase()))];
  const bad = labels.find((l) => !/^[a-z0-9][a-z0-9._-]{0,49}$/.test(l));
  if (bad) throw new RunnerError(`“${bad}” isn't a valid label. Use letters, digits, '.', '-' and '_'.`);
  if (labels.length > 20) throw new RunnerError("A runner can have at most 20 labels.");
  return labels;
}

/** Registers a runner and returns its secret token; only its hash is stored. */
export async function createRunner(
  scope: RunnerScope,
  input: { name: string; labels: string; expiresInDays: number | null },
) {
  const name = input.name.trim().slice(0, 100);
  if (!name) throw new RunnerError("Give the runner a name, e.g. the machine it runs on.");
  const labels = parseLabels(input.labels);
  const days = input.expiresInDays;
  if (days !== null && !(Number.isInteger(days) && days > 0 && days <= 3650)) {
    throw new RunnerError("Choose when the runner’s token expires.");
  }
  const secret = `${RUNNER_TOKEN_PREFIX}${randomBytes(30).toString("base64url")}`;
  const now = Date.now();
  await (await db()).collection("runners").create({
    repo: "repoId" in scope ? scope.repoId : "",
    user: "userId" in scope ? scope.userId : "",
    name,
    labels,
    token_hash: sha256(secret),
    hint: secret.slice(-4),
    created_at: now,
    expires_at: days ? now + days * 86400e3 : 0,
  });
  return secret;
}

export async function listRunners(scope: RunnerScope) {
  const rows = await (await db()).collection("runners").getFullList({ filter: scopeFilter(scope) });
  const busy = await busyJobs(rows.map((r) => r.id));
  return rows.map((r) => toRunner(r, busy)).sort((a, b) => a.name.toLowerCase().localeCompare(b.name.toLowerCase()));
}

export async function deleteRunner(scope: RunnerScope, runnerId: string) {
  const row = await findFirst("runners", `${filter("id = {:runnerId}", { runnerId })} && ${scopeFilter(scope)}`, { fields: "id" });
  if (row) await (await db()).collection("runners").delete(row.id);
}

export type RunnerIdentity = { id: string; name: string; repoId: string | null; userId: string | null; labels: string[] };

/** Resolves a runner token from the runner API, or null if unknown or expired. */
export async function authenticateRunner(secret: string): Promise<RunnerIdentity | null> {
  if (!secret.startsWith(RUNNER_TOKEN_PREFIX)) return null;
  const row = await findFirst("runners", filter("token_hash = {:hash}", { hash: sha256(secret) }));
  if (!row || isExpired(row.expires_at)) return null;
  return {
    id: row.id,
    name: row.name,
    repoId: orNull(row.repo),
    userId: orNull(row.user),
    labels: [...(row.labels ?? []), ...systemLabels({ os: orNull(row.os), arch: orNull(row.arch) })],
  };
}

const clean = (v: unknown, max = 100) => (typeof v === "string" && v.trim() ? v.trim().slice(0, max) : null);

/** Records that a runner is alive, with what it reports about its machine. */
export async function touchRunner(runnerId: string, info: Record<string, unknown>) {
  const reported = {
    hostname: clean(info.hostname),
    os: clean(info.os, 20)?.toLowerCase(),
    arch: clean(info.arch, 20)?.toLowerCase(),
    version: clean(info.version, 20),
  };
  // Only what the runner reported; the rest stays as it was.
  const known = Object.fromEntries(Object.entries(reported).filter(([, v]) => v));
  try {
    await (await db()).collection("runners").update(runnerId, { last_seen_at: Date.now(), ...known });
  } catch (e) {
    // Deleted while it was polling.
    if (!(e instanceof ClientResponseError && e.status === 404)) throw e;
  }
}

/** How many runners could pick up jobs in a repository, and how many of those are online. */
export async function runnerSummary(repoId: string, ownerId: string) {
  const rows = await (await db())
    .collection("runners")
    .getFullList({ filter: filter("repo = {:repoId} || user = {:ownerId}", { repoId, ownerId }), fields: "last_seen_at" });
  return { total: rows.length, online: rows.filter((r) => isOnline(r.last_seen_at)).length };
}