Repository Public
1/** Leases: claim (SKIP LOCKED queue), fenced heartbeat, release. */2import type { RunStatus } from '@aos/shared';3import { getKernel, query, withTx, type Kernel } from './db.js';4import { announceRunStatus, RUN_COLS, type RunRow } from './runs.js';56export const DEFAULT_LEASE_MS = 90_000;7export const DEFAULT_HEARTBEAT_MS = 30_000;89export interface ClaimedRun {10 run: RunRow;11 /** Fencing token of this claim; every later state write is guarded by it. */12 token: string;13 prevStatus: RunStatus;14 missionBudgetUsd: number;15}1617/**18 * Claims the oldest claimable run: queued (and due), or running with an expired lease.19 * One statement locks the row with FOR UPDATE SKIP LOCKED, so two workers can never claim the same20 * run; the fencing token is incremented, which invalidates every write of a previous holder.21 */22export async function claimRun(workerId: string, opts: { leaseMs?: number; kernel?: Kernel } = {}): Promise<ClaimedRun | null> {23 const k = opts.kernel ?? getKernel();24 const leaseMs = opts.leaseMs ?? DEFAULT_LEASE_MS;25 return withTx(k, async (tx) => {26 const cols = RUN_COLS.split(', ').map((c) => `r.${c}`).join(', ');27 const row = await tx.one<RunRow & { prev_status: RunStatus; mission_budget_usd: unknown }>(28 `WITH c AS (29 SELECT id, status AS prev_status FROM runs30 WHERE (status = 'queued' AND (wake_at IS NULL OR wake_at <= now()))31 OR (status = 'running' AND (lease_expires_at IS NULL OR lease_expires_at < now()))32 ORDER BY created_at, id33 FOR UPDATE SKIP LOCKED LIMIT 1)34 UPDATE runs r SET status = 'running', lease_owner = $1, lease_expires_at = now() + $2 * interval '1 millisecond',35 fencing_token = r.fencing_token + 1, started_at = COALESCE(r.started_at, now()), wake_at = NULL36 FROM c WHERE r.id = c.id37 RETURNING ${cols}, c.prev_status, (SELECT budget_usd FROM missions m WHERE m.id = r.mission_id) AS mission_budget_usd`,38 [workerId, leaseMs],39 );40 if (!row) return null;41 if (row.prev_status !== 'running') await announceRunStatus(tx, row, 'running', { worker: workerId });42 return { run: row, token: String(row.fencing_token), prevStatus: row.prev_status, missionBudgetUsd: Number(row.mission_budget_usd) };43 });44}4546export type HeartbeatResult = { ok: true; cancelRequested: boolean } | { ok: false };4748/** Extends the lease if we still hold it. `ok:false` ⇒ the lease is lost for good. Throws on storage errors. */49export async function heartbeat(k: Kernel, runId: string, token: string, leaseMs: number): Promise<HeartbeatResult> {50 const rows = await query<{ cancel_requested: boolean }>(51 k,52 `UPDATE runs SET lease_expires_at = now() + $3 * interval '1 millisecond'53 WHERE id = $1 AND fencing_token = $2::bigint AND status = 'running' RETURNING cancel_requested`,54 [runId, token, leaseMs],55 );56 const r = rows[0];57 return r ? { ok: true, cancelRequested: r.cancel_requested } : { ok: false };58}5960/** Graceful hand-over: the run stays 'running' but its lease expires now, so any worker may resume it. */61export async function releaseLease(k: Kernel, runId: string, token: string): Promise<void> {62 await query(63 k,64 `UPDATE runs SET lease_expires_at = now(), lease_owner = NULL WHERE id = $1 AND fencing_token = $2::bigint AND status = 'running'`,65 [runId, token],66 );67 await query(k, `SELECT pg_notify('aos_wake', 'q')`);68}6970/**71 * Per-run heartbeat. A refused heartbeat (fencing) reports `lost` immediately; storage errors are72 * retried at a quarter of the interval until the lease would have expired by our own (monotonic) clock,73 * then the lease is reported lost too — the worker must stop touching the run before another claims it.74 */75export class Heartbeat {76 private timer: NodeJS.Timeout | null = null;77 private stopped = false;78 private busy = false;79 private lastOk = performance.now();8081 constructor(82 private readonly k: Kernel,83 private readonly runId: string,84 private readonly token: string,85 private readonly leaseMs: number,86 private readonly intervalMs: number,87 private readonly on: { lost: (why: string) => void; cancel: () => void },88 ) {}8990 start(): void {91 this.lastOk = performance.now();92 this.schedule(this.intervalMs);93 }94 stop(): void {95 this.stopped = true;96 if (this.timer) clearTimeout(this.timer);97 this.timer = null;98 }99 /** Beat now (used when a cancel doorbell rings). */100 poke(): void {101 if (!this.stopped) void this.beat();102 }103 private schedule(ms: number): void {104 if (this.stopped) return;105 if (this.timer) clearTimeout(this.timer);106 this.timer = setTimeout(() => void this.beat(), ms);107 this.timer.unref();108 }109 private async beat(): Promise<void> {110 if (this.stopped || this.busy) return;111 this.busy = true;112 let next = this.intervalMs;113 try {114 const sent = performance.now();115 const r = await heartbeat(this.k, this.runId, this.token, this.leaseMs);116 if (this.stopped) return;117 if (!r.ok) {118 this.stop();119 this.on.lost('heartbeat refused (fencing token changed or run no longer running)');120 return;121 }122 this.lastOk = sent;123 if (r.cancelRequested) this.on.cancel();124 } catch (err) {125 if (this.stopped) return;126 // Safety margin: give up one heartbeat interval's quarter before the lease can expire.127 if (performance.now() - this.lastOk >= this.leaseMs - Math.min(this.intervalMs, this.leaseMs / 4)) {128 this.stop();129 this.on.lost(`heartbeat failed until lease expiry: ${(err as Error).message}`);130 return;131 }132 next = Math.max(50, Math.floor(this.intervalMs / 4));133 } finally {134 this.busy = false;135 }136 this.schedule(next);137 }138}139