chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,34 @@
|
||||
import { Counter, Gauge, type Registry } from "prom-client";
|
||||
|
||||
/** Prometheus metrics for dequeue backpressure. */
|
||||
export class BackpressureMetrics {
|
||||
/** 1 while backpressure is engaged (computed signal, set even in dry-run). */
|
||||
readonly engaged: Gauge<string>;
|
||||
/** 1 when running in dry-run (gates inert). */
|
||||
readonly dryRun: Gauge<string>;
|
||||
/** Dequeue attempts the gate skipped - or would have, in dry-run (labelled). */
|
||||
readonly skipsTotal: Counter<string>;
|
||||
|
||||
constructor(opts: { register: Registry; prefix?: string }) {
|
||||
const prefix = opts.prefix ?? "supervisor_backpressure";
|
||||
|
||||
this.engaged = new Gauge({
|
||||
name: `${prefix}_engaged`,
|
||||
help: "1 while dequeue backpressure is engaged (computed signal, regardless of dry-run)",
|
||||
registers: [opts.register],
|
||||
});
|
||||
|
||||
this.dryRun = new Gauge({
|
||||
name: `${prefix}_dry_run`,
|
||||
help: "1 when dequeue backpressure is in dry-run mode (gates inert)",
|
||||
registers: [opts.register],
|
||||
});
|
||||
|
||||
this.skipsTotal = new Counter({
|
||||
name: `${prefix}_skipped_dequeues_total`,
|
||||
help: "Dequeue attempts skipped by backpressure (or would be, in dry-run)",
|
||||
labelNames: ["dry_run"],
|
||||
registers: [opts.register],
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,353 @@
|
||||
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
|
||||
import { Registry } from "prom-client";
|
||||
import { BackpressureMonitor, type BackpressureSignalSource } from "./backpressureMonitor.js";
|
||||
import { BackpressureMetrics } from "./backpressureMetrics.js";
|
||||
|
||||
function countingSource(verdict: { engaged: boolean } | null): {
|
||||
source: BackpressureSignalSource;
|
||||
reads: () => number;
|
||||
} {
|
||||
let reads = 0;
|
||||
return {
|
||||
source: {
|
||||
read: async () => {
|
||||
reads++;
|
||||
return verdict;
|
||||
},
|
||||
},
|
||||
reads: () => reads,
|
||||
};
|
||||
}
|
||||
|
||||
describe("BackpressureMonitor", () => {
|
||||
beforeEach(() => {
|
||||
vi.useFakeTimers();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.useRealTimers();
|
||||
});
|
||||
|
||||
it("when disabled, never skips dequeue and never reads the signal source", () => {
|
||||
// Even though the source would report "engaged", a disabled monitor must be
|
||||
// a complete no-op: this is the backwards-compatibility guarantee.
|
||||
const { source, reads } = countingSource({ engaged: true });
|
||||
const monitor = new BackpressureMonitor({ enabled: false, source });
|
||||
|
||||
monitor.start();
|
||||
|
||||
expect(monitor.shouldSkipDequeue()).toBe(false);
|
||||
expect(reads()).toBe(0);
|
||||
|
||||
monitor.stop();
|
||||
});
|
||||
|
||||
it("when enabled and the source reports engaged, skips dequeue after a refresh", async () => {
|
||||
const { source } = countingSource({ engaged: true });
|
||||
const monitor = new BackpressureMonitor({ enabled: true, source, refreshIntervalMs: 1000 });
|
||||
|
||||
monitor.start();
|
||||
await vi.advanceTimersByTimeAsync(0); // flush the initial async read
|
||||
|
||||
expect(monitor.shouldSkipDequeue()).toBe(true);
|
||||
|
||||
monitor.stop();
|
||||
});
|
||||
|
||||
it("when enabled and the source reports clear, does not skip dequeue", async () => {
|
||||
const { source } = countingSource({ engaged: false });
|
||||
const monitor = new BackpressureMonitor({ enabled: true, source, refreshIntervalMs: 1000 });
|
||||
|
||||
monitor.start();
|
||||
await vi.advanceTimersByTimeAsync(0);
|
||||
|
||||
expect(monitor.shouldSkipDequeue()).toBe(false);
|
||||
|
||||
monitor.stop();
|
||||
});
|
||||
|
||||
it("fails open (stops skipping) when the source throws", async () => {
|
||||
let call = 0;
|
||||
const source: BackpressureSignalSource = {
|
||||
read: async () => {
|
||||
call++;
|
||||
if (call === 1) {
|
||||
return { engaged: true };
|
||||
}
|
||||
throw new Error("signal source unreachable");
|
||||
},
|
||||
};
|
||||
const monitor = new BackpressureMonitor({ enabled: true, source, refreshIntervalMs: 1000 });
|
||||
|
||||
monitor.start();
|
||||
await vi.advanceTimersByTimeAsync(0);
|
||||
expect(monitor.shouldSkipDequeue()).toBe(true); // engaged from the first read
|
||||
|
||||
await vi.advanceTimersByTimeAsync(1000); // next refresh throws
|
||||
expect(monitor.shouldSkipDequeue()).toBe(false); // fail-open: a dead source must not pin the brake
|
||||
|
||||
monitor.stop();
|
||||
});
|
||||
|
||||
it("fails open when the source reports unknown (null)", async () => {
|
||||
const { source } = countingSource(null);
|
||||
const monitor = new BackpressureMonitor({ enabled: true, source, refreshIntervalMs: 1000 });
|
||||
|
||||
monitor.start();
|
||||
await vi.advanceTimersByTimeAsync(0);
|
||||
|
||||
expect(monitor.shouldSkipDequeue()).toBe(false);
|
||||
|
||||
monitor.stop();
|
||||
});
|
||||
|
||||
it("fails open when the cached verdict goes stale (older than max age)", async () => {
|
||||
// Source stops updating (e.g. hangs) after the first read; the verdict ages out.
|
||||
const source: BackpressureSignalSource = {
|
||||
read: async () => ({ engaged: true, ts: Date.now() }),
|
||||
};
|
||||
const monitor = new BackpressureMonitor({
|
||||
enabled: true,
|
||||
source,
|
||||
refreshIntervalMs: 1_000_000, // effectively only the initial read fires
|
||||
maxVerdictAgeMs: 15_000,
|
||||
});
|
||||
|
||||
monitor.start();
|
||||
await vi.advanceTimersByTimeAsync(0);
|
||||
expect(monitor.shouldSkipDequeue()).toBe(true);
|
||||
|
||||
await vi.advanceTimersByTimeAsync(15_001); // verdict now older than max age
|
||||
expect(monitor.shouldSkipDequeue()).toBe(false);
|
||||
|
||||
monitor.stop();
|
||||
});
|
||||
|
||||
it("does not read the source on the hot path (reads are driven by the refresh tick)", async () => {
|
||||
const { source, reads } = countingSource({ engaged: true });
|
||||
const monitor = new BackpressureMonitor({ enabled: true, source, refreshIntervalMs: 1000 });
|
||||
|
||||
monitor.start();
|
||||
await vi.advanceTimersByTimeAsync(0);
|
||||
expect(reads()).toBe(1); // just the initial refresh
|
||||
|
||||
for (let i = 0; i < 1000; i++) {
|
||||
monitor.shouldSkipDequeue();
|
||||
}
|
||||
|
||||
expect(reads()).toBe(1); // hot-path calls performed zero I/O
|
||||
|
||||
monitor.stop();
|
||||
});
|
||||
|
||||
it("does not start an overlapping refresh while one is in flight", async () => {
|
||||
let reads = 0;
|
||||
const source: BackpressureSignalSource = {
|
||||
// Never resolves - simulates a hung read.
|
||||
read: () => {
|
||||
reads++;
|
||||
return new Promise<{ engaged: boolean } | null>(() => {});
|
||||
},
|
||||
};
|
||||
const monitor = new BackpressureMonitor({ enabled: true, source, refreshIntervalMs: 1000 });
|
||||
|
||||
monitor.start();
|
||||
await vi.advanceTimersByTimeAsync(3000); // several intervals while the first read hangs
|
||||
|
||||
expect(reads).toBe(1); // in-flight guard prevents stacking
|
||||
|
||||
monitor.stop();
|
||||
});
|
||||
|
||||
it("stops refreshing after stop()", async () => {
|
||||
const { source, reads } = countingSource({ engaged: true });
|
||||
const monitor = new BackpressureMonitor({ enabled: true, source, refreshIntervalMs: 1000 });
|
||||
|
||||
monitor.start();
|
||||
await vi.advanceTimersByTimeAsync(0);
|
||||
const readsAtStop = reads();
|
||||
|
||||
monitor.stop();
|
||||
await vi.advanceTimersByTimeAsync(5000);
|
||||
|
||||
expect(reads()).toBe(readsAtStop);
|
||||
});
|
||||
|
||||
it("isEngaged reflects the hard engaged state (the signal for freezing scale-up)", async () => {
|
||||
const { source } = countingSource({ engaged: true });
|
||||
const monitor = new BackpressureMonitor({ enabled: true, source, refreshIntervalMs: 1000 });
|
||||
|
||||
monitor.start();
|
||||
await vi.advanceTimersByTimeAsync(0);
|
||||
|
||||
expect(monitor.isEngaged()).toBe(true);
|
||||
|
||||
monitor.stop();
|
||||
});
|
||||
|
||||
it("isEngaged is false when clear and when stale", async () => {
|
||||
const source: BackpressureSignalSource = {
|
||||
read: async () => ({ engaged: true, ts: Date.now() }),
|
||||
};
|
||||
const monitor = new BackpressureMonitor({
|
||||
enabled: true,
|
||||
source,
|
||||
refreshIntervalMs: 1_000_000,
|
||||
maxVerdictAgeMs: 15_000,
|
||||
});
|
||||
|
||||
monitor.start();
|
||||
await vi.advanceTimersByTimeAsync(0);
|
||||
expect(monitor.isEngaged()).toBe(true);
|
||||
|
||||
await vi.advanceTimersByTimeAsync(15_001); // stale → fail-open
|
||||
expect(monitor.isEngaged()).toBe(false);
|
||||
|
||||
monitor.stop();
|
||||
});
|
||||
|
||||
it("ramps the dequeue gate after release instead of resuming instantly", async () => {
|
||||
let engaged = true;
|
||||
let rnd = 0.5;
|
||||
const source: BackpressureSignalSource = { read: async () => ({ engaged }) };
|
||||
const monitor = new BackpressureMonitor({
|
||||
enabled: true,
|
||||
source,
|
||||
refreshIntervalMs: 1000,
|
||||
rampMs: 10_000,
|
||||
random: () => rnd,
|
||||
});
|
||||
|
||||
monitor.start();
|
||||
await vi.advanceTimersByTimeAsync(0);
|
||||
expect(monitor.shouldSkipDequeue()).toBe(true); // hard engaged
|
||||
|
||||
// Release: the next refresh observes the clear verdict and starts the ramp.
|
||||
engaged = false;
|
||||
await vi.advanceTimersByTimeAsync(1000);
|
||||
expect(monitor.isEngaged()).toBe(false);
|
||||
|
||||
// Just after release (progress ~0): skip probability ~1, so skip regardless.
|
||||
rnd = 0.99;
|
||||
expect(monitor.shouldSkipDequeue()).toBe(true);
|
||||
|
||||
// Halfway through the ramp (progress 0.5): skip probability 0.5.
|
||||
await vi.advanceTimersByTimeAsync(5000);
|
||||
rnd = 0.4;
|
||||
expect(monitor.shouldSkipDequeue()).toBe(true); // 0.4 < 0.5 → skip
|
||||
rnd = 0.6;
|
||||
expect(monitor.shouldSkipDequeue()).toBe(false); // 0.6 ≥ 0.5 → allow
|
||||
|
||||
// Past the ramp window: never skip.
|
||||
await vi.advanceTimersByTimeAsync(5000);
|
||||
rnd = 0.0;
|
||||
expect(monitor.shouldSkipDequeue()).toBe(false);
|
||||
|
||||
monitor.stop();
|
||||
});
|
||||
|
||||
it("fails open on an engaged verdict with no timestamp when staleness is enforced", async () => {
|
||||
// A verdict claiming engaged but carrying no ts can't be checked for freshness;
|
||||
// when maxVerdictAgeMs is set we must not trust it (else a dead producer could
|
||||
// pin the brake forever).
|
||||
const { source } = countingSource({ engaged: true }); // no ts
|
||||
const monitor = new BackpressureMonitor({
|
||||
enabled: true,
|
||||
source,
|
||||
refreshIntervalMs: 1000,
|
||||
maxVerdictAgeMs: 15_000,
|
||||
});
|
||||
|
||||
monitor.start();
|
||||
await vi.advanceTimersByTimeAsync(0);
|
||||
|
||||
expect(monitor.computeEngaged()).toBe(false);
|
||||
expect(monitor.shouldSkipDequeue()).toBe(false);
|
||||
|
||||
monitor.stop();
|
||||
});
|
||||
|
||||
it("in dry-run, the gates are inert but computeEngaged still reflects the real signal", async () => {
|
||||
const { source } = countingSource({ engaged: true });
|
||||
const monitor = new BackpressureMonitor({
|
||||
enabled: true,
|
||||
source,
|
||||
refreshIntervalMs: 1000,
|
||||
dryRun: true,
|
||||
});
|
||||
|
||||
monitor.start();
|
||||
await vi.advanceTimersByTimeAsync(0);
|
||||
|
||||
expect(monitor.computeEngaged()).toBe(true); // real signal, for observability/metrics
|
||||
expect(monitor.isEngaged()).toBe(false); // inert: no scale-up freeze
|
||||
expect(monitor.shouldSkipDequeue()).toBe(false); // inert: no dequeue skip
|
||||
|
||||
monitor.stop();
|
||||
});
|
||||
|
||||
it("logs on verdict transitions", async () => {
|
||||
let engaged = true;
|
||||
const source: BackpressureSignalSource = { read: async () => ({ engaged }) };
|
||||
const logs: Array<{ message: string; meta?: Record<string, unknown> }> = [];
|
||||
const logger = {
|
||||
info: (message: string, meta?: Record<string, unknown>) => logs.push({ message, meta }),
|
||||
};
|
||||
const monitor = new BackpressureMonitor({
|
||||
enabled: true,
|
||||
source,
|
||||
refreshIntervalMs: 1000,
|
||||
logger,
|
||||
});
|
||||
|
||||
monitor.start();
|
||||
await vi.advanceTimersByTimeAsync(0);
|
||||
expect(logs.some((l) => l.meta?.engaged === true)).toBe(true);
|
||||
|
||||
engaged = false;
|
||||
await vi.advanceTimersByTimeAsync(1000);
|
||||
expect(logs.some((l) => l.meta?.engaged === false)).toBe(true);
|
||||
|
||||
monitor.stop();
|
||||
});
|
||||
|
||||
it("records prometheus metrics", async () => {
|
||||
const { source } = countingSource({ engaged: true });
|
||||
const register = new Registry();
|
||||
const metrics = new BackpressureMetrics({ register });
|
||||
const monitor = new BackpressureMonitor({
|
||||
enabled: true,
|
||||
source,
|
||||
refreshIntervalMs: 1000,
|
||||
metrics,
|
||||
});
|
||||
|
||||
monitor.start();
|
||||
await vi.advanceTimersByTimeAsync(0);
|
||||
|
||||
expect(await register.metrics()).toContain("supervisor_backpressure_engaged 1");
|
||||
|
||||
monitor.shouldSkipDequeue();
|
||||
expect(await register.metrics()).toMatch(
|
||||
/supervisor_backpressure_skipped_dequeues_total\{dry_run="false"\} [1-9]/
|
||||
);
|
||||
|
||||
monitor.stop();
|
||||
});
|
||||
|
||||
it("resumes instantly when no ramp is configured", async () => {
|
||||
let engaged = true;
|
||||
const source: BackpressureSignalSource = { read: async () => ({ engaged }) };
|
||||
const monitor = new BackpressureMonitor({ enabled: true, source, refreshIntervalMs: 1000 });
|
||||
|
||||
monitor.start();
|
||||
await vi.advanceTimersByTimeAsync(0);
|
||||
expect(monitor.shouldSkipDequeue()).toBe(true);
|
||||
|
||||
engaged = false;
|
||||
await vi.advanceTimersByTimeAsync(1000);
|
||||
expect(monitor.shouldSkipDequeue()).toBe(false); // no ramp → instant resume
|
||||
|
||||
monitor.stop();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,179 @@
|
||||
import type { BackpressureMetrics } from "./backpressureMetrics.js";
|
||||
|
||||
export interface BackpressureLogger {
|
||||
info(message: string, meta?: Record<string, unknown>): void;
|
||||
}
|
||||
|
||||
export type BackpressureVerdict = {
|
||||
engaged: boolean;
|
||||
/** Epoch ms the verdict was produced. Used for consumer-side staleness fail-open. */
|
||||
ts?: number;
|
||||
};
|
||||
|
||||
/**
|
||||
* Source of the current backpressure verdict. `read()` returns `null` when the
|
||||
* verdict is unknown (missing/unreadable) - the monitor treats unknown as
|
||||
* "not engaged" (fail-open).
|
||||
*/
|
||||
export interface BackpressureSignalSource {
|
||||
read(): Promise<BackpressureVerdict | null>;
|
||||
}
|
||||
|
||||
export type BackpressureMonitorOptions = {
|
||||
enabled: boolean;
|
||||
source: BackpressureSignalSource;
|
||||
refreshIntervalMs?: number;
|
||||
/**
|
||||
* If set, a cached verdict older than this is treated as unknown (fail-open).
|
||||
* Guards against the source silently going stale (e.g. hanging reads).
|
||||
*/
|
||||
maxVerdictAgeMs?: number;
|
||||
/**
|
||||
* If set, after backpressure releases the dequeue gate stays partially engaged
|
||||
* for this long, skipping a linearly-decaying fraction of attempts so the
|
||||
* aggregate dequeue rate ramps from ~0 to full instead of snapping to full and
|
||||
* re-flooding a freshly-recovered cluster. 0/unset = instant resume.
|
||||
*/
|
||||
rampMs?: number;
|
||||
/** Injectable RNG for the resume ramp; defaults to Math.random. */
|
||||
random?: () => number;
|
||||
/**
|
||||
* When true, the gates are inert (never skip dequeues, never freeze scale-up).
|
||||
* computeEngaged() still reflects the real signal so it can be observed.
|
||||
*/
|
||||
dryRun?: boolean;
|
||||
logger?: BackpressureLogger;
|
||||
metrics?: BackpressureMetrics;
|
||||
};
|
||||
|
||||
const DEFAULT_REFRESH_INTERVAL_MS = 1000;
|
||||
|
||||
export class BackpressureMonitor {
|
||||
private verdict: BackpressureVerdict | null = null;
|
||||
private timer?: ReturnType<typeof setInterval>;
|
||||
private refreshInFlight = false;
|
||||
private wasEngaged = false;
|
||||
private releasedAt?: number;
|
||||
|
||||
constructor(private readonly opts: BackpressureMonitorOptions) {
|
||||
this.opts.metrics?.dryRun.set(this.opts.dryRun ? 1 : 0);
|
||||
}
|
||||
|
||||
start(): void {
|
||||
if (!this.opts.enabled) {
|
||||
return;
|
||||
}
|
||||
|
||||
void this.refreshTick();
|
||||
this.timer = setInterval(
|
||||
() => void this.refreshTick(),
|
||||
this.opts.refreshIntervalMs ?? DEFAULT_REFRESH_INTERVAL_MS
|
||||
);
|
||||
}
|
||||
|
||||
/** Skip a tick if the previous refresh is still in flight, so slow/hung reads can't stack. */
|
||||
private async refreshTick(): Promise<void> {
|
||||
if (this.refreshInFlight) {
|
||||
return;
|
||||
}
|
||||
this.refreshInFlight = true;
|
||||
try {
|
||||
await this.refresh();
|
||||
} finally {
|
||||
this.refreshInFlight = false;
|
||||
}
|
||||
}
|
||||
|
||||
stop(): void {
|
||||
if (this.timer) {
|
||||
clearInterval(this.timer);
|
||||
this.timer = undefined;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Raw hard backpressure state: true while the (fresh) verdict says engaged,
|
||||
* ignoring dry-run. Used for observability/metrics so the real signal is
|
||||
* visible even when the gates are inert.
|
||||
*/
|
||||
computeEngaged(): boolean {
|
||||
const verdict = this.verdict;
|
||||
if (verdict?.engaged !== true) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// When staleness enforcement is on, an engaged verdict must carry a fresh
|
||||
// timestamp. A missing or stale ts can't be trusted (a dead producer could
|
||||
// otherwise pin the brake forever), so fail open.
|
||||
const maxAge = this.opts.maxVerdictAgeMs;
|
||||
if (maxAge !== undefined) {
|
||||
if (verdict.ts === undefined || Date.now() - verdict.ts > maxAge) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
/**
|
||||
* Effective hard state: the signal for freezing consumer-pool scale-up. Inert
|
||||
* (false) in dry-run. Hot-path read, no I/O.
|
||||
*/
|
||||
isEngaged(): boolean {
|
||||
return this.opts.dryRun ? false : this.computeEngaged();
|
||||
}
|
||||
|
||||
/** Hot-path read: synchronous, never performs I/O. Inert (false) in dry-run. */
|
||||
shouldSkipDequeue(): boolean {
|
||||
const wouldSkip = this.computeShouldSkip();
|
||||
if (wouldSkip) {
|
||||
this.opts.metrics?.skipsTotal.inc({ dry_run: this.opts.dryRun ? "true" : "false" });
|
||||
}
|
||||
return this.opts.dryRun ? false : wouldSkip;
|
||||
}
|
||||
|
||||
private computeShouldSkip(): boolean {
|
||||
if (this.computeEngaged()) {
|
||||
return true;
|
||||
}
|
||||
|
||||
// Post-release ramp: skip a linearly-decaying fraction of attempts so the
|
||||
// aggregate dequeue rate climbs back to full over rampMs rather than snapping.
|
||||
const rampMs = this.opts.rampMs;
|
||||
if (rampMs && this.releasedAt !== undefined) {
|
||||
const elapsed = Date.now() - this.releasedAt;
|
||||
if (elapsed < rampMs) {
|
||||
const skipProbability = 1 - elapsed / rampMs;
|
||||
return (this.opts.random ?? Math.random)() < skipProbability;
|
||||
}
|
||||
}
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
private async refresh(): Promise<void> {
|
||||
try {
|
||||
this.verdict = await this.opts.source.read();
|
||||
} catch {
|
||||
// Fail-open: a dead/unreachable source must never pin the brake. Treat as
|
||||
// unknown (no verdict) so dequeue resumes as if backpressure were off.
|
||||
this.verdict = null;
|
||||
}
|
||||
|
||||
// Track the engaged→released transition to anchor the resume ramp. Use the
|
||||
// staleness-aware state so a stale verdict doesn't pin wasEngaged / the gauge.
|
||||
const nowEngaged = this.computeEngaged();
|
||||
this.opts.metrics?.engaged.set(nowEngaged ? 1 : 0);
|
||||
|
||||
if (nowEngaged !== this.wasEngaged) {
|
||||
this.opts.logger?.info("backpressure verdict changed", {
|
||||
engaged: nowEngaged,
|
||||
dryRun: !!this.opts.dryRun,
|
||||
});
|
||||
}
|
||||
if (this.wasEngaged && !nowEngaged) {
|
||||
this.releasedAt = Date.now();
|
||||
}
|
||||
this.wasEngaged = nowEngaged;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,95 @@
|
||||
import { describe, it, expect } from "vitest";
|
||||
import { parsePodCount, K8sPodCountSignalSource } from "./k8sPodCountSignalSource.js";
|
||||
|
||||
describe("parsePodCount", () => {
|
||||
it("reads the pods object count", () => {
|
||||
const text = [
|
||||
"# HELP apiserver_storage_objects Number of stored objects",
|
||||
"# TYPE apiserver_storage_objects gauge",
|
||||
'apiserver_storage_objects{resource="pods"} 8421',
|
||||
'apiserver_storage_objects{resource="configmaps"} 17',
|
||||
].join("\n");
|
||||
expect(parsePodCount(text)).toBe(8421);
|
||||
});
|
||||
|
||||
it("is tolerant of extra labels in any order", () => {
|
||||
const text = 'apiserver_storage_objects{group="",resource="pods",extra="x"} 12';
|
||||
expect(parsePodCount(text)).toBe(12);
|
||||
});
|
||||
|
||||
it("parses scientific notation", () => {
|
||||
const text = 'apiserver_storage_objects{resource="pods"} 1.2e+04';
|
||||
expect(parsePodCount(text)).toBe(12000);
|
||||
});
|
||||
|
||||
it("throws when the pods metric is absent", () => {
|
||||
const text = 'apiserver_storage_objects{resource="configmaps"} 17';
|
||||
expect(() => parsePodCount(text)).toThrow(/not found/);
|
||||
});
|
||||
|
||||
it("throws on a non-finite value (e.g. 1e999)", () => {
|
||||
const text = 'apiserver_storage_objects{resource="pods"} 1e999';
|
||||
expect(() => parsePodCount(text)).toThrow();
|
||||
});
|
||||
|
||||
it("throws on a negative value", () => {
|
||||
const text = 'apiserver_storage_objects{resource="pods"} -5';
|
||||
expect(() => parsePodCount(text)).toThrow();
|
||||
});
|
||||
});
|
||||
|
||||
function metrics(count: number): string {
|
||||
return `apiserver_storage_objects{resource="pods"} ${count}`;
|
||||
}
|
||||
|
||||
describe("K8sPodCountSignalSource", () => {
|
||||
it("engages at the engage threshold and reports the count", async () => {
|
||||
const counts: number[] = [];
|
||||
const source = new K8sPodCountSignalSource({
|
||||
fetchMetrics: async () => metrics(10000),
|
||||
engageThreshold: 10000,
|
||||
releaseThreshold: 5000,
|
||||
reportPodCount: (c) => counts.push(c),
|
||||
});
|
||||
const verdict = await source.read();
|
||||
expect(verdict.engaged).toBe(true);
|
||||
expect(typeof verdict.ts).toBe("number");
|
||||
expect(counts).toEqual([10000]);
|
||||
});
|
||||
|
||||
it("does not engage below the engage threshold", async () => {
|
||||
const source = new K8sPodCountSignalSource({
|
||||
fetchMetrics: async () => metrics(9999),
|
||||
engageThreshold: 10000,
|
||||
releaseThreshold: 5000,
|
||||
});
|
||||
expect((await source.read()).engaged).toBe(false);
|
||||
});
|
||||
|
||||
it("stays engaged in the hysteresis band, releases only below release threshold", async () => {
|
||||
let count = 10000;
|
||||
const source = new K8sPodCountSignalSource({
|
||||
fetchMetrics: async () => metrics(count),
|
||||
engageThreshold: 10000,
|
||||
releaseThreshold: 5000,
|
||||
});
|
||||
expect((await source.read()).engaged).toBe(true); // engage
|
||||
count = 7000;
|
||||
expect((await source.read()).engaged).toBe(true); // band -> still engaged
|
||||
count = 4999;
|
||||
expect((await source.read()).engaged).toBe(false); // below release -> off
|
||||
count = 7000;
|
||||
expect((await source.read()).engaged).toBe(false); // band again -> stays off
|
||||
});
|
||||
|
||||
it("propagates scrape failures (monitor fails open on throw)", async () => {
|
||||
const source = new K8sPodCountSignalSource({
|
||||
fetchMetrics: async () => {
|
||||
throw new Error("connection refused");
|
||||
},
|
||||
engageThreshold: 10000,
|
||||
releaseThreshold: 5000,
|
||||
});
|
||||
await expect(source.read()).rejects.toThrow("connection refused");
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,46 @@
|
||||
import type { BackpressureSignalSource, BackpressureVerdict } from "./backpressureMonitor.js";
|
||||
|
||||
// Reads the apiserver's stored-pod-object count from a Prometheus /metrics scrape.
|
||||
const POD_COUNT_RE = /^apiserver_storage_objects\{[^}]*resource="pods"[^}]*\}\s+([0-9.eE+]+)/m;
|
||||
|
||||
export function parsePodCount(metricsText: string): number {
|
||||
const match = metricsText.match(POD_COUNT_RE);
|
||||
if (!match) {
|
||||
throw new Error('apiserver_storage_objects{resource="pods"} not found in metrics');
|
||||
}
|
||||
const value = Number(match[1]);
|
||||
if (!Number.isFinite(value)) {
|
||||
throw new Error(`unparseable pod count: ${match[1]}`);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
export type K8sPodCountSignalSourceOptions = {
|
||||
fetchMetrics: () => Promise<string>;
|
||||
engageThreshold: number;
|
||||
releaseThreshold: number;
|
||||
reportPodCount?: (count: number) => void;
|
||||
};
|
||||
|
||||
// Engage/release with hysteresis so a count hovering near the line doesn't flap.
|
||||
export class K8sPodCountSignalSource implements BackpressureSignalSource {
|
||||
private engaged = false;
|
||||
|
||||
constructor(private readonly opts: K8sPodCountSignalSourceOptions) {}
|
||||
|
||||
async read(): Promise<BackpressureVerdict> {
|
||||
const text = await this.opts.fetchMetrics();
|
||||
const count = parsePodCount(text);
|
||||
this.opts.reportPodCount?.(count);
|
||||
|
||||
if (this.engaged) {
|
||||
if (count < this.opts.releaseThreshold) {
|
||||
this.engaged = false;
|
||||
}
|
||||
} else if (count >= this.opts.engageThreshold) {
|
||||
this.engaged = true;
|
||||
}
|
||||
|
||||
return { engaged: this.engaged, ts: Date.now() };
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
import { redisTest } from "@internal/testcontainers";
|
||||
import { Redis } from "ioredis";
|
||||
import { describe, expect } from "vitest";
|
||||
import { RedisBackpressureSignalSource } from "./redisBackpressureSignalSource.js";
|
||||
|
||||
const KEY = "backpressure:test";
|
||||
|
||||
describe("RedisBackpressureSignalSource", () => {
|
||||
redisTest("returns null when the key is absent", async ({ redisOptions }) => {
|
||||
const redis = new Redis(redisOptions);
|
||||
try {
|
||||
const source = new RedisBackpressureSignalSource(redis, KEY);
|
||||
expect(await source.read()).toBeNull();
|
||||
} finally {
|
||||
await redis.quit();
|
||||
}
|
||||
});
|
||||
|
||||
redisTest("parses a valid engaged verdict", async ({ redisOptions }) => {
|
||||
const redis = new Redis(redisOptions);
|
||||
try {
|
||||
await redis.set(KEY, JSON.stringify({ engaged: true, ts: 1_700_000_000_000 }));
|
||||
const source = new RedisBackpressureSignalSource(redis, KEY);
|
||||
expect(await source.read()).toEqual({ engaged: true, ts: 1_700_000_000_000 });
|
||||
} finally {
|
||||
await redis.quit();
|
||||
}
|
||||
});
|
||||
|
||||
redisTest("parses a clear verdict", async ({ redisOptions }) => {
|
||||
const redis = new Redis(redisOptions);
|
||||
try {
|
||||
await redis.set(KEY, JSON.stringify({ engaged: false }));
|
||||
const source = new RedisBackpressureSignalSource(redis, KEY);
|
||||
expect(await source.read()).toEqual({ engaged: false });
|
||||
} finally {
|
||||
await redis.quit();
|
||||
}
|
||||
});
|
||||
|
||||
redisTest("returns null for malformed JSON (fail-open)", async ({ redisOptions }) => {
|
||||
const redis = new Redis(redisOptions);
|
||||
try {
|
||||
await redis.set(KEY, "not json {");
|
||||
const source = new RedisBackpressureSignalSource(redis, KEY);
|
||||
expect(await source.read()).toBeNull();
|
||||
} finally {
|
||||
await redis.quit();
|
||||
}
|
||||
});
|
||||
|
||||
redisTest(
|
||||
"returns null for valid JSON of the wrong shape (fail-open)",
|
||||
async ({ redisOptions }) => {
|
||||
const redis = new Redis(redisOptions);
|
||||
try {
|
||||
await redis.set(KEY, JSON.stringify({ foo: "bar" }));
|
||||
const source = new RedisBackpressureSignalSource(redis, KEY);
|
||||
expect(await source.read()).toBeNull();
|
||||
} finally {
|
||||
await redis.quit();
|
||||
}
|
||||
}
|
||||
);
|
||||
});
|
||||
@@ -0,0 +1,35 @@
|
||||
import type { Redis } from "ioredis";
|
||||
import { z } from "zod";
|
||||
import type { BackpressureSignalSource, BackpressureVerdict } from "./backpressureMonitor.js";
|
||||
|
||||
const VerdictSchema = z.object({
|
||||
engaged: z.boolean(),
|
||||
ts: z.number().optional(),
|
||||
});
|
||||
|
||||
/** Reads the backpressure verdict from a Redis key written by the cluster-side aggregator. */
|
||||
export class RedisBackpressureSignalSource implements BackpressureSignalSource {
|
||||
constructor(
|
||||
private readonly redis: Redis,
|
||||
private readonly key: string
|
||||
) {}
|
||||
|
||||
async read(): Promise<BackpressureVerdict | null> {
|
||||
const raw = await this.redis.get(this.key);
|
||||
if (raw === null) {
|
||||
return null;
|
||||
}
|
||||
|
||||
// A malformed or wrong-shaped value is treated as unknown (null) so the
|
||||
// monitor fails open rather than acting on garbage.
|
||||
let json: unknown;
|
||||
try {
|
||||
json = JSON.parse(raw);
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
|
||||
const parsed = VerdictSchema.safeParse(json);
|
||||
return parsed.success ? parsed.data : null;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,108 @@
|
||||
import * as k8s from "@kubernetes/client-node";
|
||||
import type { Informer, KubernetesObject, ListPromise } from "@kubernetes/client-node";
|
||||
import { assertExhaustive } from "@trigger.dev/core/utils";
|
||||
import { SimpleStructuredLogger } from "@trigger.dev/core/v3/utils/structuredLogger";
|
||||
import * as https from "node:https";
|
||||
|
||||
export const RUNTIME_ENV = process.env.KUBERNETES_PORT ? "kubernetes" : "local";
|
||||
|
||||
const logger = new SimpleStructuredLogger("kubernetes-client");
|
||||
|
||||
export function createK8sApi() {
|
||||
const kubeConfig = getKubeConfig();
|
||||
|
||||
function makeInformer<T extends KubernetesObject>(
|
||||
path: string,
|
||||
listPromiseFn: ListPromise<T>,
|
||||
labelSelector?: string,
|
||||
fieldSelector?: string
|
||||
): Informer<T> {
|
||||
return k8s.makeInformer(kubeConfig, path, listPromiseFn, labelSelector, fieldSelector);
|
||||
}
|
||||
|
||||
const api = {
|
||||
core: kubeConfig.makeApiClient(k8s.CoreV1Api),
|
||||
batch: kubeConfig.makeApiClient(k8s.BatchV1Api),
|
||||
apps: kubeConfig.makeApiClient(k8s.AppsV1Api),
|
||||
makeInformer,
|
||||
};
|
||||
|
||||
return api;
|
||||
}
|
||||
|
||||
export type K8sApi = ReturnType<typeof createK8sApi>;
|
||||
|
||||
function getKubeConfig() {
|
||||
logger.debug("getKubeConfig()", { RUNTIME_ENV });
|
||||
|
||||
const kubeConfig = new k8s.KubeConfig();
|
||||
|
||||
switch (RUNTIME_ENV) {
|
||||
case "local":
|
||||
kubeConfig.loadFromDefault();
|
||||
break;
|
||||
case "kubernetes":
|
||||
kubeConfig.loadFromCluster();
|
||||
break;
|
||||
default:
|
||||
assertExhaustive(RUNTIME_ENV);
|
||||
}
|
||||
|
||||
return kubeConfig;
|
||||
}
|
||||
|
||||
export { k8s };
|
||||
|
||||
/**
|
||||
* Builds a function that scrapes the apiserver's Prometheus /metrics endpoint.
|
||||
* One lightweight aggregate read - not a pod listing. Requires the service
|
||||
* account to be granted GET on the /metrics non-resource URL.
|
||||
*/
|
||||
export function createApiserverMetricsFetcher(timeoutMs: number): () => Promise<string> {
|
||||
const kubeConfig = getKubeConfig();
|
||||
|
||||
return async () => {
|
||||
const cluster = kubeConfig.getCurrentCluster();
|
||||
if (!cluster) {
|
||||
throw new Error("no current cluster in kubeconfig");
|
||||
}
|
||||
const url = new URL(`${cluster.server}/metrics`);
|
||||
const opts: https.RequestOptions = {
|
||||
method: "GET",
|
||||
protocol: url.protocol,
|
||||
hostname: url.hostname,
|
||||
port: url.port,
|
||||
path: url.pathname,
|
||||
};
|
||||
// applyToHTTPSOptions sets the cluster CA, client cert/key, and auth headers
|
||||
// (incl. exec plugins) on the request - so TLS verifies against the cluster
|
||||
// CA, not the system store. The fetch-options path attaches the CA as an
|
||||
// https.Agent, which global fetch (undici) ignores.
|
||||
await kubeConfig.applyToHTTPSOptions(opts);
|
||||
|
||||
return new Promise<string>((resolve, reject) => {
|
||||
const req = https.request(opts, (res) => {
|
||||
const status = res.statusCode ?? 0;
|
||||
let body = "";
|
||||
res.setEncoding("utf8");
|
||||
res.on("data", (chunk) => {
|
||||
body += chunk;
|
||||
});
|
||||
res.on("end", () => {
|
||||
if (status >= 200 && status < 300) {
|
||||
resolve(body);
|
||||
} else {
|
||||
reject(new Error(`apiserver /metrics scrape failed: ${status}`));
|
||||
}
|
||||
});
|
||||
});
|
||||
// Without this a hung connect/TLS/read never settles, and the monitor's
|
||||
// refreshInFlight guard would freeze the source (silent fail-open).
|
||||
req.setTimeout(timeoutMs, () => {
|
||||
req.destroy(new Error(`apiserver /metrics scrape timed out after ${timeoutMs}ms`));
|
||||
});
|
||||
req.on("error", reject);
|
||||
req.end();
|
||||
});
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
import { describe, it, expect, vi } from "vitest";
|
||||
|
||||
// Mock std-env before importing env.ts so the module-level `Env.parse(stdEnv)`
|
||||
// doesn't fail in a test environment that lacks required vars.
|
||||
vi.mock("std-env", () => ({
|
||||
env: {
|
||||
TRIGGER_API_URL: "http://localhost:3030",
|
||||
TRIGGER_WORKER_TOKEN: "test-token",
|
||||
MANAGED_WORKER_SECRET: "test-secret",
|
||||
OTEL_EXPORTER_OTLP_ENDPOINT: "http://localhost:4318",
|
||||
},
|
||||
}));
|
||||
|
||||
const { Env } = await import("./env.js");
|
||||
|
||||
// Minimal env that satisfies all required fields; everything else has defaults.
|
||||
const base = {
|
||||
TRIGGER_API_URL: "http://localhost:3030",
|
||||
TRIGGER_WORKER_TOKEN: "test-token",
|
||||
MANAGED_WORKER_SECRET: "test-secret",
|
||||
OTEL_EXPORTER_OTLP_ENDPOINT: "http://localhost:4318",
|
||||
};
|
||||
|
||||
describe("Env superRefine - backpressure source awareness", () => {
|
||||
it("pod-count source can be enabled without a Redis host", () => {
|
||||
expect(() =>
|
||||
Env.parse({
|
||||
...base,
|
||||
TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_ENABLED: "true",
|
||||
})
|
||||
).not.toThrow();
|
||||
});
|
||||
|
||||
it("redis source requires a Redis host", () => {
|
||||
expect(() =>
|
||||
Env.parse({
|
||||
...base,
|
||||
TRIGGER_DEQUEUE_BACKPRESSURE_ENABLED: "true",
|
||||
})
|
||||
).toThrow();
|
||||
});
|
||||
|
||||
it("both sources can be enabled together (with a Redis host)", () => {
|
||||
expect(() =>
|
||||
Env.parse({
|
||||
...base,
|
||||
TRIGGER_DEQUEUE_BACKPRESSURE_ENABLED: "true",
|
||||
TRIGGER_DEQUEUE_BACKPRESSURE_REDIS_HOST: "localhost",
|
||||
TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_ENABLED: "true",
|
||||
})
|
||||
).not.toThrow();
|
||||
});
|
||||
|
||||
it("rejects pod-count release >= engage when the source is enabled", () => {
|
||||
expect(() =>
|
||||
Env.parse({
|
||||
...base,
|
||||
TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_ENABLED: "true",
|
||||
TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_ENGAGE: "100",
|
||||
TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_RELEASE: "100",
|
||||
})
|
||||
).toThrow();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,385 @@
|
||||
import { randomUUID } from "crypto";
|
||||
import { env as stdEnv } from "std-env";
|
||||
import { z } from "zod";
|
||||
import { AdditionalEnvVars, BoolEnv } from "./envUtil.js";
|
||||
|
||||
export const Env = z
|
||||
.object({
|
||||
// This will come from `spec.nodeName` in k8s
|
||||
TRIGGER_WORKER_INSTANCE_NAME: z.string().default(randomUUID()),
|
||||
TRIGGER_WORKER_HEARTBEAT_INTERVAL_SECONDS: z.coerce.number().default(30),
|
||||
|
||||
// Opt-in, dev-only: stream this process's logs over a local telnet/TCP socket on this port.
|
||||
SUPERVISOR_TELNET_LOGS_PORT: z.coerce.number().optional(),
|
||||
|
||||
// Required settings
|
||||
TRIGGER_API_URL: z.string().url(),
|
||||
TRIGGER_WORKER_TOKEN: z.string(), // accepts file:// path to read from a file
|
||||
MANAGED_WORKER_SECRET: z.string(),
|
||||
OTEL_EXPORTER_OTLP_ENDPOINT: z.string().url(), // set on the runners
|
||||
|
||||
// Workload API settings (coordinator mode) - the workload API is what the run controller connects to
|
||||
TRIGGER_WORKLOAD_API_ENABLED: BoolEnv.default(true),
|
||||
TRIGGER_WORKLOAD_API_PROTOCOL: z
|
||||
.string()
|
||||
.transform((s) => z.enum(["http", "https"]).parse(s.toLowerCase()))
|
||||
.default("http"),
|
||||
TRIGGER_WORKLOAD_API_DOMAIN: z.string().optional(), // If unset, will use orchestrator-specific default
|
||||
TRIGGER_WORKLOAD_API_HOST_INTERNAL: z.string().default("0.0.0.0"),
|
||||
TRIGGER_WORKLOAD_API_PORT_INTERNAL: z.coerce.number().default(8020), // This is the port the workload API listens on
|
||||
TRIGGER_WORKLOAD_API_PORT_EXTERNAL: z.coerce.number().default(8020), // This is the exposed port passed to the run controller
|
||||
|
||||
// Runner settings
|
||||
RUNNER_HEARTBEAT_INTERVAL_SECONDS: z.coerce.number().optional(),
|
||||
RUNNER_SNAPSHOT_POLL_INTERVAL_SECONDS: z.coerce.number().optional(),
|
||||
RUNNER_ADDITIONAL_ENV_VARS: AdditionalEnvVars, // optional (csv)
|
||||
RUNNER_PRETTY_LOGS: BoolEnv.default(false),
|
||||
|
||||
// Dequeue settings (provider mode)
|
||||
TRIGGER_DEQUEUE_ENABLED: BoolEnv.default(true),
|
||||
// Which worker-queue class this supervisor fleet serves. "default" pulls the
|
||||
// region queue (standard/agent runs); "scheduled" pulls the dedicated
|
||||
// scheduled-lineage queue. Run a separate fleet per class for isolation.
|
||||
TRIGGER_WORKER_QUEUE_CLASS: z.enum(["default", "scheduled"]).default("default"),
|
||||
TRIGGER_DEQUEUE_INTERVAL_MS: z.coerce.number().int().default(250),
|
||||
TRIGGER_DEQUEUE_IDLE_INTERVAL_MS: z.coerce.number().int().default(1000),
|
||||
TRIGGER_DEQUEUE_MAX_RUN_COUNT: z.coerce.number().int().default(1),
|
||||
TRIGGER_DEQUEUE_MIN_CONSUMER_COUNT: z.coerce.number().int().default(1),
|
||||
TRIGGER_DEQUEUE_MAX_CONSUMER_COUNT: z.coerce.number().int().default(10),
|
||||
TRIGGER_DEQUEUE_SCALING_STRATEGY: z.enum(["none", "smooth", "aggressive"]).default("none"),
|
||||
TRIGGER_DEQUEUE_SCALING_UP_COOLDOWN_MS: z.coerce.number().int().default(5000), // 5 seconds
|
||||
TRIGGER_DEQUEUE_SCALING_DOWN_COOLDOWN_MS: z.coerce.number().int().default(30000), // 30 seconds
|
||||
TRIGGER_DEQUEUE_SCALING_TARGET_RATIO: z.coerce.number().default(1.0), // Target ratio of queue items to consumers (1.0 = 1 item per consumer)
|
||||
TRIGGER_DEQUEUE_SCALING_EWMA_ALPHA: z.coerce.number().min(0).max(1).default(0.3), // Smooths queue length measurements (0=historical, 1=current)
|
||||
TRIGGER_DEQUEUE_SCALING_BATCH_WINDOW_MS: z.coerce.number().int().positive().default(1000), // Batch window for metrics processing (ms)
|
||||
TRIGGER_DEQUEUE_SCALING_DAMPING_FACTOR: z.coerce.number().min(0).max(1).default(0.7), // Smooths consumer count changes after EWMA (0=no scaling, 1=immediate)
|
||||
|
||||
// Dequeue backpressure - off by default. When enabled, the supervisor reads a
|
||||
// verdict from Redis (written by the cluster-side aggregator) and pauses dequeues
|
||||
// while the worker cluster can't schedule pods. Disabled = total no-op: no Redis
|
||||
// client is created, no reads happen, and the dequeue loop is unaffected.
|
||||
TRIGGER_DEQUEUE_BACKPRESSURE_ENABLED: BoolEnv.default(false),
|
||||
// Safety default: even when enabled, backpressure only logs what it would do.
|
||||
// Set to false to actually skip dequeues / freeze scale-up.
|
||||
TRIGGER_DEQUEUE_BACKPRESSURE_DRY_RUN: BoolEnv.default(true),
|
||||
TRIGGER_DEQUEUE_BACKPRESSURE_REDIS_KEY: z.string().default("engine:dequeue:backpressure"),
|
||||
TRIGGER_DEQUEUE_BACKPRESSURE_REFRESH_MS: z.coerce.number().int().positive().default(1000),
|
||||
TRIGGER_DEQUEUE_BACKPRESSURE_RAMP_MS: z.coerce.number().int().min(0).default(30_000), // Resume ramp window after release; 0 = instant resume
|
||||
|
||||
TRIGGER_DEQUEUE_BACKPRESSURE_MAX_VERDICT_AGE_MS: z.coerce
|
||||
.number()
|
||||
.int()
|
||||
.positive()
|
||||
.default(15_000), // Stale verdict → fail-open (treat as not engaged)
|
||||
TRIGGER_DEQUEUE_BACKPRESSURE_REDIS_HOST: z.string().optional(),
|
||||
TRIGGER_DEQUEUE_BACKPRESSURE_REDIS_PORT: z.coerce.number().int().optional(),
|
||||
TRIGGER_DEQUEUE_BACKPRESSURE_REDIS_USERNAME: z.string().optional(),
|
||||
TRIGGER_DEQUEUE_BACKPRESSURE_REDIS_PASSWORD: z.string().optional(),
|
||||
TRIGGER_DEQUEUE_BACKPRESSURE_REDIS_TLS_DISABLED: BoolEnv.default(false),
|
||||
TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_ENABLED: BoolEnv.default(false),
|
||||
TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_DRY_RUN: BoolEnv.default(true),
|
||||
TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_ENGAGE: z.coerce
|
||||
.number()
|
||||
.int()
|
||||
.positive()
|
||||
.default(10_000),
|
||||
TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_RELEASE: z.coerce
|
||||
.number()
|
||||
.int()
|
||||
.positive()
|
||||
.default(5_000),
|
||||
TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_REFRESH_MS: z.coerce
|
||||
.number()
|
||||
.int()
|
||||
.positive()
|
||||
.default(5_000),
|
||||
// Hard timeout on the apiserver /metrics scrape. A hung request would otherwise
|
||||
// never settle and freeze the monitor's refresh loop (fail-open silently).
|
||||
TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_SCRAPE_TIMEOUT_MS: z.coerce
|
||||
.number()
|
||||
.int()
|
||||
.positive()
|
||||
.default(10_000),
|
||||
|
||||
// Optional services
|
||||
TRIGGER_WARM_START_URL: z.string().optional(),
|
||||
TRIGGER_CHECKPOINT_URL: z.string().optional(),
|
||||
TRIGGER_METADATA_URL: z.string().optional(),
|
||||
|
||||
// Warm-start delivery verification: after a warm-start hit, probe the
|
||||
// platform and cold-start the run if no runner acted on the dispatch
|
||||
TRIGGER_WARM_START_VERIFY_ENABLED: BoolEnv.default(false),
|
||||
TRIGGER_WARM_START_VERIFY_DELAY_MS: z.coerce
|
||||
.number()
|
||||
.int()
|
||||
.min(1_000)
|
||||
.max(60_000)
|
||||
.default(10_000),
|
||||
|
||||
// Used by the resource monitor
|
||||
RESOURCE_MONITOR_ENABLED: BoolEnv.default(false),
|
||||
RESOURCE_MONITOR_OVERRIDE_CPU_TOTAL: z.coerce.number().optional(),
|
||||
RESOURCE_MONITOR_OVERRIDE_MEMORY_TOTAL_GB: z.coerce.number().optional(),
|
||||
|
||||
// Docker settings
|
||||
DOCKER_API_VERSION: z.string().optional(),
|
||||
DOCKER_PLATFORM: z.string().optional(), // e.g. linux/amd64, linux/arm64
|
||||
DOCKER_STRIP_IMAGE_DIGEST: BoolEnv.default(true),
|
||||
DOCKER_REGISTRY_USERNAME: z.string().optional(),
|
||||
DOCKER_REGISTRY_PASSWORD: z.string().optional(),
|
||||
DOCKER_REGISTRY_URL: z.string().optional(), // e.g. https://index.docker.io/v1
|
||||
DOCKER_ENFORCE_MACHINE_PRESETS: BoolEnv.default(true),
|
||||
DOCKER_AUTOREMOVE_EXITED_CONTAINERS: BoolEnv.default(true),
|
||||
/**
|
||||
* Network mode to use for all runners. Supported standard values are: `bridge`, `host`, `none`, and `container:<name|id>`.
|
||||
* Any other value is taken as a custom network's name to which all runners should connect to.
|
||||
*
|
||||
* Accepts a list of comma-separated values to attach to multiple networks. Additional networks are interpreted as network names and will be attached after container creation.
|
||||
*
|
||||
* **WARNING**: Specifying multiple networks will slightly increase startup times.
|
||||
*
|
||||
* @default "host"
|
||||
*/
|
||||
DOCKER_RUNNER_NETWORKS: z.string().default("host"),
|
||||
|
||||
// Compute settings
|
||||
COMPUTE_GATEWAY_URL: z.string().url().optional(),
|
||||
COMPUTE_GATEWAY_AUTH_TOKEN: z.string().optional(),
|
||||
COMPUTE_GATEWAY_TIMEOUT_MS: z.coerce.number().int().default(30_000),
|
||||
COMPUTE_SNAPSHOTS_ENABLED: BoolEnv.default(false),
|
||||
COMPUTE_TRACE_SPANS_ENABLED: BoolEnv.default(true),
|
||||
COMPUTE_TRACE_OTLP_ENDPOINT: z.string().url().optional(), // Override for span export (derived from TRIGGER_API_URL if unset)
|
||||
COMPUTE_SNAPSHOT_DELAY_MS: z.coerce.number().int().min(0).max(60_000).default(5_000),
|
||||
COMPUTE_SNAPSHOT_DISPATCH_LIMIT: z.coerce.number().int().min(1).max(100).default(10),
|
||||
// Instance create retries for transient placement failures (1 = no retries)
|
||||
COMPUTE_INSTANCE_CREATE_MAX_ATTEMPTS: z.coerce.number().int().min(1).max(10).default(3),
|
||||
COMPUTE_INSTANCE_CREATE_RETRY_BASE_DELAY_MS: z.coerce
|
||||
.number()
|
||||
.int()
|
||||
.min(0)
|
||||
.max(10_000)
|
||||
.default(250),
|
||||
|
||||
// Kubernetes settings
|
||||
KUBERNETES_FORCE_ENABLED: BoolEnv.default(false),
|
||||
KUBERNETES_NAMESPACE: z.string().default("default"),
|
||||
KUBERNETES_WORKER_NODETYPE_LABEL: z.string().default("v4-worker"),
|
||||
KUBERNETES_IMAGE_PULL_SECRETS: z.string().optional(), // csv
|
||||
KUBERNETES_EPHEMERAL_STORAGE_SIZE_LIMIT: z.string().default("10Gi"),
|
||||
KUBERNETES_EPHEMERAL_STORAGE_SIZE_REQUEST: z.string().default("2Gi"),
|
||||
KUBERNETES_STRIP_IMAGE_DIGEST: BoolEnv.default(false),
|
||||
KUBERNETES_CPU_REQUEST_MIN_CORES: z.coerce.number().min(0).default(0),
|
||||
KUBERNETES_CPU_REQUEST_RATIO: z.coerce.number().min(0).max(1).default(0.75), // Ratio of CPU limit, so 0.75 = 75% of CPU limit
|
||||
KUBERNETES_MEMORY_REQUEST_MIN_GB: z.coerce.number().min(0).default(0),
|
||||
KUBERNETES_MEMORY_REQUEST_RATIO: z.coerce.number().min(0).max(1).default(1), // Ratio of memory limit, so 1 = 100% of memory limit
|
||||
|
||||
// Per-preset overrides of the global KUBERNETES_CPU_REQUEST_RATIO
|
||||
KUBERNETES_CPU_REQUEST_RATIO_MICRO: z.coerce.number().min(0).max(1).optional(),
|
||||
KUBERNETES_CPU_REQUEST_RATIO_SMALL_1X: z.coerce.number().min(0).max(1).optional(),
|
||||
KUBERNETES_CPU_REQUEST_RATIO_SMALL_2X: z.coerce.number().min(0).max(1).optional(),
|
||||
KUBERNETES_CPU_REQUEST_RATIO_MEDIUM_1X: z.coerce.number().min(0).max(1).optional(),
|
||||
KUBERNETES_CPU_REQUEST_RATIO_MEDIUM_2X: z.coerce.number().min(0).max(1).optional(),
|
||||
KUBERNETES_CPU_REQUEST_RATIO_LARGE_1X: z.coerce.number().min(0).max(1).optional(),
|
||||
KUBERNETES_CPU_REQUEST_RATIO_LARGE_2X: z.coerce.number().min(0).max(1).optional(),
|
||||
|
||||
// Per-preset overrides of the global KUBERNETES_MEMORY_REQUEST_RATIO
|
||||
KUBERNETES_MEMORY_REQUEST_RATIO_MICRO: z.coerce.number().min(0).max(1).optional(),
|
||||
KUBERNETES_MEMORY_REQUEST_RATIO_SMALL_1X: z.coerce.number().min(0).max(1).optional(),
|
||||
KUBERNETES_MEMORY_REQUEST_RATIO_SMALL_2X: z.coerce.number().min(0).max(1).optional(),
|
||||
KUBERNETES_MEMORY_REQUEST_RATIO_MEDIUM_1X: z.coerce.number().min(0).max(1).optional(),
|
||||
KUBERNETES_MEMORY_REQUEST_RATIO_MEDIUM_2X: z.coerce.number().min(0).max(1).optional(),
|
||||
KUBERNETES_MEMORY_REQUEST_RATIO_LARGE_1X: z.coerce.number().min(0).max(1).optional(),
|
||||
KUBERNETES_MEMORY_REQUEST_RATIO_LARGE_2X: z.coerce.number().min(0).max(1).optional(),
|
||||
|
||||
KUBERNETES_MEMORY_OVERHEAD_GB: z.coerce.number().min(0).optional(), // Optional memory overhead to add to the limit in GB
|
||||
KUBERNETES_SCHEDULER_NAME: z.string().optional(), // Custom scheduler name for pods
|
||||
|
||||
// Pod DNS config — override the cluster default ndots to `KUBERNETES_POD_DNS_NDOTS`.
|
||||
// Default k8s ndots is 5: any name with fewer than 5 dots (e.g. `api.example.com`, 2 dots) is first walked
|
||||
// through every entry in the cluster search list (`<ns>.svc.cluster.local`, `svc.cluster.local`, `cluster.local`)
|
||||
// before being tried as-is, turning one resolution into 4+ CoreDNS queries (×2 with A+AAAA).
|
||||
// Overriding the default can be useful to cut CoreDNS query amplification for external domains.
|
||||
// Note: before enabling, make sure no code path relies on search-list expansion for names with dots ≥ the value
|
||||
// set here — those names will now hit their as-is form first and could resolve externally before falling back.
|
||||
KUBERNETES_POD_DNS_NDOTS_OVERRIDE_ENABLED: BoolEnv.default(false),
|
||||
KUBERNETES_POD_DNS_NDOTS: z.coerce.number().int().min(1).max(15).default(2),
|
||||
// Large machine affinity settings - large-* presets prefer a dedicated pool
|
||||
KUBERNETES_LARGE_MACHINE_AFFINITY_ENABLED: BoolEnv.default(false),
|
||||
KUBERNETES_LARGE_MACHINE_AFFINITY_POOL_LABEL_KEY: z
|
||||
.string()
|
||||
.trim()
|
||||
.min(1)
|
||||
.default("node.cluster.x-k8s.io/machinepool"),
|
||||
KUBERNETES_LARGE_MACHINE_AFFINITY_POOL_LABEL_VALUE: z
|
||||
.string()
|
||||
.trim()
|
||||
.min(1)
|
||||
.default("large-machines"),
|
||||
KUBERNETES_LARGE_MACHINE_AFFINITY_WEIGHT: z.coerce.number().int().min(1).max(100).default(100),
|
||||
|
||||
// Project affinity settings - pods from the same project prefer the same node
|
||||
KUBERNETES_PROJECT_AFFINITY_ENABLED: BoolEnv.default(false),
|
||||
KUBERNETES_PROJECT_AFFINITY_WEIGHT: z.coerce.number().int().min(1).max(100).default(50),
|
||||
KUBERNETES_PROJECT_AFFINITY_TOPOLOGY_KEY: z
|
||||
.string()
|
||||
.trim()
|
||||
.min(1)
|
||||
.default("kubernetes.io/hostname"),
|
||||
|
||||
// Schedule affinity settings - runs from schedule trees prefer a dedicated pool
|
||||
KUBERNETES_SCHEDULED_RUN_AFFINITY_ENABLED: BoolEnv.default(false),
|
||||
KUBERNETES_SCHEDULED_RUN_AFFINITY_POOL_LABEL_KEY: z
|
||||
.string()
|
||||
.trim()
|
||||
.min(1)
|
||||
.default("node.cluster.x-k8s.io/machinepool"),
|
||||
KUBERNETES_SCHEDULED_RUN_AFFINITY_POOL_LABEL_VALUE: z
|
||||
.string()
|
||||
.trim()
|
||||
.min(1)
|
||||
.default("scheduled-runs"),
|
||||
KUBERNETES_SCHEDULED_RUN_AFFINITY_WEIGHT: z.coerce.number().int().min(1).max(100).default(80),
|
||||
KUBERNETES_SCHEDULED_RUN_ANTI_AFFINITY_WEIGHT: z.coerce
|
||||
.number()
|
||||
.int()
|
||||
.min(1)
|
||||
.max(100)
|
||||
.default(20),
|
||||
|
||||
// Schedule toleration settings - scheduled runs tolerate taints on the dedicated pool
|
||||
// Comma-separated list of tolerations in the format: key=value:effect
|
||||
// For Exists operator (no value): key:effect
|
||||
KUBERNETES_SCHEDULED_RUN_TOLERATIONS: z
|
||||
.string()
|
||||
.transform((val, ctx) => {
|
||||
const tolerations = val
|
||||
.split(",")
|
||||
.map((entry) => entry.trim())
|
||||
.filter((entry) => entry.length > 0)
|
||||
.map((entry) => {
|
||||
const colonIdx = entry.lastIndexOf(":");
|
||||
if (colonIdx === -1) {
|
||||
ctx.addIssue({
|
||||
code: z.ZodIssueCode.custom,
|
||||
message: `Invalid toleration format (missing effect): "${entry}"`,
|
||||
});
|
||||
return z.NEVER;
|
||||
}
|
||||
|
||||
const effect = entry.slice(colonIdx + 1);
|
||||
const validEffects = ["NoSchedule", "NoExecute", "PreferNoSchedule"];
|
||||
if (!validEffects.includes(effect)) {
|
||||
ctx.addIssue({
|
||||
code: z.ZodIssueCode.custom,
|
||||
message: `Invalid toleration effect "${effect}" in "${entry}". Must be one of: ${validEffects.join(
|
||||
", "
|
||||
)}`,
|
||||
});
|
||||
return z.NEVER;
|
||||
}
|
||||
|
||||
const keyValue = entry.slice(0, colonIdx);
|
||||
const eqIdx = keyValue.indexOf("=");
|
||||
const key = eqIdx === -1 ? keyValue : keyValue.slice(0, eqIdx);
|
||||
|
||||
if (!key) {
|
||||
ctx.addIssue({
|
||||
code: z.ZodIssueCode.custom,
|
||||
message: `Invalid toleration format (empty key): "${entry}"`,
|
||||
});
|
||||
return z.NEVER;
|
||||
}
|
||||
|
||||
if (eqIdx === -1) {
|
||||
return { key, operator: "Exists" as const, effect };
|
||||
}
|
||||
|
||||
return {
|
||||
key,
|
||||
operator: "Equal" as const,
|
||||
value: keyValue.slice(eqIdx + 1),
|
||||
effect,
|
||||
};
|
||||
});
|
||||
|
||||
return tolerations;
|
||||
})
|
||||
.optional(),
|
||||
|
||||
// Placement tags settings
|
||||
PLACEMENT_TAGS_ENABLED: BoolEnv.default(false),
|
||||
PLACEMENT_TAGS_PREFIX: z.string().default("node.cluster.x-k8s.io"),
|
||||
|
||||
// Metrics
|
||||
METRICS_ENABLED: BoolEnv.default(true),
|
||||
METRICS_COLLECT_DEFAULTS: BoolEnv.default(true),
|
||||
METRICS_HOST: z.string().default("127.0.0.1"),
|
||||
METRICS_PORT: z.coerce.number().int().default(9090),
|
||||
|
||||
// Pod cleaner
|
||||
POD_CLEANER_ENABLED: BoolEnv.default(true),
|
||||
POD_CLEANER_INTERVAL_MS: z.coerce.number().int().default(10000),
|
||||
POD_CLEANER_BATCH_SIZE: z.coerce.number().int().default(500),
|
||||
|
||||
// Failed pod handler
|
||||
FAILED_POD_HANDLER_ENABLED: BoolEnv.default(true),
|
||||
FAILED_POD_HANDLER_RECONNECT_INTERVAL_MS: z.coerce.number().int().default(1000),
|
||||
|
||||
// Debug
|
||||
DEBUG: BoolEnv.default(false),
|
||||
SEND_RUN_DEBUG_LOGS: BoolEnv.default(false),
|
||||
|
||||
// Wide-event observability - off by default. Emits one flat-keyed JSON
|
||||
// line per natural unit of work (dequeue iteration, HTTP request, socket
|
||||
// lifecycle). High-QPS hotpath, so the kill switch must be honoured.
|
||||
TRIGGER_WIDE_EVENTS_ENABLED: BoolEnv.default(false),
|
||||
// When true, also emit wide events for high-frequency HTTP routes
|
||||
// (heartbeat, snapshots-since, logs/debug). Off in prod to keep event
|
||||
// volume manageable; on in test environments for full-fidelity debugging.
|
||||
TRIGGER_WIDE_EVENTS_NOISY_ROUTES: BoolEnv.default(false),
|
||||
})
|
||||
.superRefine((data, ctx) => {
|
||||
if (
|
||||
data.TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_ENABLED &&
|
||||
data.TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_RELEASE >=
|
||||
data.TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_ENGAGE
|
||||
) {
|
||||
ctx.addIssue({
|
||||
code: z.ZodIssueCode.custom,
|
||||
message:
|
||||
"TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_RELEASE must be less than TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_ENGAGE",
|
||||
path: ["TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_RELEASE"],
|
||||
});
|
||||
}
|
||||
if (data.COMPUTE_SNAPSHOTS_ENABLED && !data.TRIGGER_METADATA_URL) {
|
||||
ctx.addIssue({
|
||||
code: z.ZodIssueCode.custom,
|
||||
message: "TRIGGER_METADATA_URL is required when COMPUTE_SNAPSHOTS_ENABLED is true",
|
||||
path: ["TRIGGER_METADATA_URL"],
|
||||
});
|
||||
}
|
||||
if (data.COMPUTE_SNAPSHOTS_ENABLED && !data.TRIGGER_WORKLOAD_API_DOMAIN) {
|
||||
ctx.addIssue({
|
||||
code: z.ZodIssueCode.custom,
|
||||
message: "TRIGGER_WORKLOAD_API_DOMAIN is required when COMPUTE_SNAPSHOTS_ENABLED is true",
|
||||
path: ["TRIGGER_WORKLOAD_API_DOMAIN"],
|
||||
});
|
||||
}
|
||||
if (
|
||||
data.TRIGGER_DEQUEUE_BACKPRESSURE_ENABLED &&
|
||||
!data.TRIGGER_DEQUEUE_BACKPRESSURE_REDIS_HOST
|
||||
) {
|
||||
ctx.addIssue({
|
||||
code: z.ZodIssueCode.custom,
|
||||
message:
|
||||
"TRIGGER_DEQUEUE_BACKPRESSURE_REDIS_HOST is required when TRIGGER_DEQUEUE_BACKPRESSURE_ENABLED is true",
|
||||
path: ["TRIGGER_DEQUEUE_BACKPRESSURE_REDIS_HOST"],
|
||||
});
|
||||
}
|
||||
})
|
||||
.transform((data) => ({
|
||||
...data,
|
||||
COMPUTE_TRACE_OTLP_ENDPOINT: data.COMPUTE_TRACE_OTLP_ENDPOINT ?? `${data.TRIGGER_API_URL}/otel`,
|
||||
}));
|
||||
|
||||
export const env = Env.parse(stdEnv);
|
||||
@@ -0,0 +1,80 @@
|
||||
import { describe, it, expect } from "vitest";
|
||||
import { BoolEnv, AdditionalEnvVars } from "./envUtil.js";
|
||||
|
||||
describe("BoolEnv", () => {
|
||||
it("should parse string 'true' as true", () => {
|
||||
expect(BoolEnv.parse("true")).toBe(true);
|
||||
expect(BoolEnv.parse("TRUE")).toBe(true);
|
||||
expect(BoolEnv.parse("True")).toBe(true);
|
||||
});
|
||||
|
||||
it("should parse string '1' as true", () => {
|
||||
expect(BoolEnv.parse("1")).toBe(true);
|
||||
});
|
||||
|
||||
it("should parse string 'false' as false", () => {
|
||||
expect(BoolEnv.parse("false")).toBe(false);
|
||||
expect(BoolEnv.parse("FALSE")).toBe(false);
|
||||
expect(BoolEnv.parse("False")).toBe(false);
|
||||
});
|
||||
|
||||
it("should handle whitespace", () => {
|
||||
expect(BoolEnv.parse(" true ")).toBe(true);
|
||||
expect(BoolEnv.parse(" 1 ")).toBe(true);
|
||||
});
|
||||
|
||||
it("should pass through boolean values", () => {
|
||||
expect(BoolEnv.parse(true)).toBe(true);
|
||||
expect(BoolEnv.parse(false)).toBe(false);
|
||||
});
|
||||
|
||||
it("should return false for invalid inputs", () => {
|
||||
expect(BoolEnv.parse("invalid")).toBe(false);
|
||||
expect(BoolEnv.parse("")).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("AdditionalEnvVars", () => {
|
||||
it("should parse single key-value pair", () => {
|
||||
expect(AdditionalEnvVars.parse("FOO=bar")).toEqual({ FOO: "bar" });
|
||||
});
|
||||
|
||||
it("should parse multiple key-value pairs", () => {
|
||||
expect(AdditionalEnvVars.parse("FOO=bar,BAZ=qux")).toEqual({
|
||||
FOO: "bar",
|
||||
BAZ: "qux",
|
||||
});
|
||||
});
|
||||
|
||||
it("should handle whitespace", () => {
|
||||
expect(AdditionalEnvVars.parse(" FOO = bar , BAZ = qux ")).toEqual({
|
||||
FOO: "bar",
|
||||
BAZ: "qux",
|
||||
});
|
||||
});
|
||||
|
||||
it("should return undefined for empty string", () => {
|
||||
expect(AdditionalEnvVars.parse("")).toBeUndefined();
|
||||
});
|
||||
|
||||
it("should return undefined for invalid format", () => {
|
||||
expect(AdditionalEnvVars.parse("invalid")).toBeUndefined();
|
||||
});
|
||||
|
||||
it("should skip invalid pairs but include valid ones", () => {
|
||||
expect(AdditionalEnvVars.parse("FOO=bar,INVALID,BAZ=qux")).toEqual({
|
||||
FOO: "bar",
|
||||
BAZ: "qux",
|
||||
});
|
||||
});
|
||||
|
||||
it("should pass through undefined", () => {
|
||||
expect(AdditionalEnvVars.parse(undefined)).toBeUndefined();
|
||||
});
|
||||
|
||||
it("should handle empty values", () => {
|
||||
expect(AdditionalEnvVars.parse("FOO=,BAR=value")).toEqual({
|
||||
BAR: "value",
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,47 @@
|
||||
import { z } from "zod";
|
||||
import { SimpleStructuredLogger } from "@trigger.dev/core/v3/utils/structuredLogger";
|
||||
|
||||
const logger = new SimpleStructuredLogger("env-util");
|
||||
|
||||
const baseBoolEnv = z.preprocess((val) => {
|
||||
if (typeof val !== "string") {
|
||||
return val;
|
||||
}
|
||||
|
||||
return ["true", "1"].includes(val.toLowerCase().trim());
|
||||
}, z.boolean());
|
||||
|
||||
// Create a type-safe version that only accepts boolean defaults
|
||||
export const BoolEnv = baseBoolEnv as Omit<typeof baseBoolEnv, "default"> & {
|
||||
default: (value: boolean) => z.ZodDefault<typeof baseBoolEnv>;
|
||||
};
|
||||
|
||||
export const AdditionalEnvVars = z.preprocess((val) => {
|
||||
if (typeof val !== "string") {
|
||||
return val;
|
||||
}
|
||||
|
||||
if (!val) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
try {
|
||||
const result = val.split(",").reduce(
|
||||
(acc, pair) => {
|
||||
const [key, value] = pair.split("=");
|
||||
if (!key || !value) {
|
||||
return acc;
|
||||
}
|
||||
acc[key.trim()] = value.trim();
|
||||
return acc;
|
||||
},
|
||||
{} as Record<string, string>
|
||||
);
|
||||
|
||||
// Return undefined if no valid key-value pairs were found
|
||||
return Object.keys(result).length === 0 ? undefined : result;
|
||||
} catch (error) {
|
||||
logger.warn("Failed to parse additional env vars", { error, val });
|
||||
return undefined;
|
||||
}
|
||||
}, z.record(z.string(), z.string()).optional());
|
||||
@@ -0,0 +1,763 @@
|
||||
import { SupervisorSession } from "@trigger.dev/core/v3/workers";
|
||||
import { SimpleStructuredLogger } from "@trigger.dev/core/v3/utils/structuredLogger";
|
||||
import { formatLogLine, startTelnetLogServer } from "@trigger.dev/core/v3/telnetLogServer";
|
||||
import { env } from "./env.js";
|
||||
import { WorkloadServer } from "./workloadServer/index.js";
|
||||
import type { WorkloadManagerOptions, WorkloadManager } from "./workloadManager/types.js";
|
||||
import Docker from "dockerode";
|
||||
import { z } from "zod";
|
||||
import { type DequeuedMessage } from "@trigger.dev/core/v3";
|
||||
import {
|
||||
DockerResourceMonitor,
|
||||
KubernetesResourceMonitor,
|
||||
NoopResourceMonitor,
|
||||
type ResourceMonitor,
|
||||
} from "./resourceMonitor.js";
|
||||
import { KubernetesWorkloadManager } from "./workloadManager/kubernetes.js";
|
||||
import { DockerWorkloadManager } from "./workloadManager/docker.js";
|
||||
import { ComputeWorkloadManager } from "./workloadManager/compute.js";
|
||||
import {
|
||||
HttpServer,
|
||||
CheckpointClient,
|
||||
isKubernetesEnvironment,
|
||||
} from "@trigger.dev/core/v3/serverOnly";
|
||||
import { createK8sApi, createApiserverMetricsFetcher } from "./clients/kubernetes.js";
|
||||
import { collectDefaultMetrics, Gauge, Histogram } from "prom-client";
|
||||
import { register } from "./metrics.js";
|
||||
import { PodCleaner } from "./services/podCleaner.js";
|
||||
import { FailedPodHandler } from "./services/failedPodHandler.js";
|
||||
import { getWorkerToken } from "./workerToken.js";
|
||||
import { OtlpTraceService } from "./services/otlpTraceService.js";
|
||||
import {
|
||||
WarmStartVerificationService,
|
||||
type WarmStartTimings,
|
||||
} from "./services/warmStartVerificationService.js";
|
||||
import { extractTraceparent, getRestoreRunnerId } from "./util.js";
|
||||
import { Redis } from "ioredis";
|
||||
import { BackpressureMonitor } from "./backpressure/backpressureMonitor.js";
|
||||
import { RedisBackpressureSignalSource } from "./backpressure/redisBackpressureSignalSource.js";
|
||||
import { BackpressureMetrics } from "./backpressure/backpressureMetrics.js";
|
||||
import { K8sPodCountSignalSource } from "./backpressure/k8sPodCountSignalSource.js";
|
||||
import {
|
||||
fromContext,
|
||||
recordPhaseSince,
|
||||
runWideEvent,
|
||||
setExtra,
|
||||
setMeta,
|
||||
type WideEventOptions,
|
||||
} from "./wideEvents/index.js";
|
||||
|
||||
if (env.METRICS_COLLECT_DEFAULTS) {
|
||||
collectDefaultMetrics({ register });
|
||||
}
|
||||
|
||||
const workloadCreateDuration = new Histogram({
|
||||
name: "workload_create_duration_seconds",
|
||||
help: "Duration of workload manager create calls. A create may include backend-internal retries, so one observation can span multiple attempts.",
|
||||
labelNames: ["backend", "outcome"],
|
||||
buckets: [0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10, 30, 60],
|
||||
registers: [register],
|
||||
});
|
||||
|
||||
class ManagedSupervisor {
|
||||
private readonly workerSession: SupervisorSession;
|
||||
private readonly metricsServer?: HttpServer;
|
||||
private readonly workloadServer: WorkloadServer;
|
||||
private readonly workloadManager: WorkloadManager;
|
||||
private readonly workloadManagerBackend: "compute" | "kubernetes" | "docker";
|
||||
private readonly computeManager?: ComputeWorkloadManager;
|
||||
private readonly logger = new SimpleStructuredLogger("managed-supervisor");
|
||||
private readonly resourceMonitor: ResourceMonitor;
|
||||
private readonly checkpointClient?: CheckpointClient;
|
||||
private readonly warmStartVerifier?: WarmStartVerificationService;
|
||||
|
||||
private readonly podCleaner?: PodCleaner;
|
||||
private readonly failedPodHandler?: FailedPodHandler;
|
||||
private readonly tracing?: OtlpTraceService;
|
||||
private readonly backpressureMonitors: BackpressureMonitor[] = [];
|
||||
private readonly backpressureRedis?: Redis;
|
||||
|
||||
private readonly isKubernetes = isKubernetesEnvironment(env.KUBERNETES_FORCE_ENABLED);
|
||||
private readonly warmStartUrl = env.TRIGGER_WARM_START_URL;
|
||||
|
||||
private readonly wideEventOpts: WideEventOptions = {
|
||||
service: "supervisor",
|
||||
env: { nodeId: env.TRIGGER_WORKER_INSTANCE_NAME },
|
||||
enabled: env.TRIGGER_WIDE_EVENTS_ENABLED,
|
||||
};
|
||||
private readonly wideEventsNoisyRoutes = env.TRIGGER_WIDE_EVENTS_NOISY_ROUTES;
|
||||
|
||||
constructor() {
|
||||
// Strip secret-like env vars before debug-logging the rest. Add any new
|
||||
// secret env var here so it never lands in the DEBUG "Starting up" log.
|
||||
const {
|
||||
TRIGGER_WORKER_TOKEN,
|
||||
MANAGED_WORKER_SECRET,
|
||||
COMPUTE_GATEWAY_AUTH_TOKEN,
|
||||
DOCKER_REGISTRY_PASSWORD,
|
||||
TRIGGER_DEQUEUE_BACKPRESSURE_REDIS_PASSWORD,
|
||||
...envWithoutSecrets
|
||||
} = env;
|
||||
|
||||
if (env.DEBUG) {
|
||||
this.logger.debug("Starting up", { envWithoutSecrets });
|
||||
}
|
||||
|
||||
if (this.warmStartUrl) {
|
||||
this.logger.log("🔥 Warm starts enabled", {
|
||||
warmStartUrl: this.warmStartUrl,
|
||||
});
|
||||
}
|
||||
|
||||
const workloadManagerOptions = {
|
||||
workloadApiProtocol: env.TRIGGER_WORKLOAD_API_PROTOCOL,
|
||||
workloadApiDomain: env.TRIGGER_WORKLOAD_API_DOMAIN,
|
||||
workloadApiPort: env.TRIGGER_WORKLOAD_API_PORT_EXTERNAL,
|
||||
warmStartUrl: this.warmStartUrl,
|
||||
metadataUrl: env.TRIGGER_METADATA_URL,
|
||||
imagePullSecrets: env.KUBERNETES_IMAGE_PULL_SECRETS?.split(","),
|
||||
heartbeatIntervalSeconds: env.RUNNER_HEARTBEAT_INTERVAL_SECONDS,
|
||||
snapshotPollIntervalSeconds: env.RUNNER_SNAPSHOT_POLL_INTERVAL_SECONDS,
|
||||
additionalEnvVars: env.RUNNER_ADDITIONAL_ENV_VARS,
|
||||
dockerAutoremove: env.DOCKER_AUTOREMOVE_EXITED_CONTAINERS,
|
||||
} satisfies WorkloadManagerOptions;
|
||||
|
||||
this.resourceMonitor = env.RESOURCE_MONITOR_ENABLED
|
||||
? this.isKubernetes
|
||||
? new KubernetesResourceMonitor(createK8sApi(), env.TRIGGER_WORKER_INSTANCE_NAME)
|
||||
: new DockerResourceMonitor(new Docker())
|
||||
: new NoopResourceMonitor();
|
||||
|
||||
if (env.COMPUTE_GATEWAY_URL) {
|
||||
if (!env.TRIGGER_WORKLOAD_API_DOMAIN) {
|
||||
throw new Error("TRIGGER_WORKLOAD_API_DOMAIN is not set, cannot create compute manager");
|
||||
}
|
||||
|
||||
const callbackUrl = `${env.TRIGGER_WORKLOAD_API_PROTOCOL}://${env.TRIGGER_WORKLOAD_API_DOMAIN}:${env.TRIGGER_WORKLOAD_API_PORT_EXTERNAL}/api/v1/compute/snapshot-complete`;
|
||||
|
||||
if (env.COMPUTE_TRACE_SPANS_ENABLED) {
|
||||
this.tracing = new OtlpTraceService({
|
||||
endpointUrl: env.COMPUTE_TRACE_OTLP_ENDPOINT,
|
||||
});
|
||||
}
|
||||
|
||||
const computeManager = new ComputeWorkloadManager({
|
||||
...workloadManagerOptions,
|
||||
gateway: {
|
||||
url: env.COMPUTE_GATEWAY_URL,
|
||||
authToken: env.COMPUTE_GATEWAY_AUTH_TOKEN,
|
||||
timeoutMs: env.COMPUTE_GATEWAY_TIMEOUT_MS,
|
||||
},
|
||||
snapshots: {
|
||||
enabled: env.COMPUTE_SNAPSHOTS_ENABLED,
|
||||
delayMs: env.COMPUTE_SNAPSHOT_DELAY_MS,
|
||||
dispatchLimit: env.COMPUTE_SNAPSHOT_DISPATCH_LIMIT,
|
||||
callbackUrl,
|
||||
},
|
||||
tracing: this.tracing,
|
||||
runner: {
|
||||
instanceName: env.TRIGGER_WORKER_INSTANCE_NAME,
|
||||
otelEndpoint: env.OTEL_EXPORTER_OTLP_ENDPOINT,
|
||||
prettyLogs: env.RUNNER_PRETTY_LOGS,
|
||||
sendRunDebugLogs: env.SEND_RUN_DEBUG_LOGS,
|
||||
},
|
||||
createRetry: {
|
||||
maxAttempts: env.COMPUTE_INSTANCE_CREATE_MAX_ATTEMPTS,
|
||||
baseDelayMs: env.COMPUTE_INSTANCE_CREATE_RETRY_BASE_DELAY_MS,
|
||||
},
|
||||
});
|
||||
this.computeManager = computeManager;
|
||||
this.workloadManager = computeManager;
|
||||
this.workloadManagerBackend = "compute";
|
||||
} else if (this.isKubernetes) {
|
||||
this.workloadManager = new KubernetesWorkloadManager(workloadManagerOptions);
|
||||
this.workloadManagerBackend = "kubernetes";
|
||||
} else {
|
||||
this.workloadManager = new DockerWorkloadManager(workloadManagerOptions);
|
||||
this.workloadManagerBackend = "docker";
|
||||
}
|
||||
|
||||
if (this.isKubernetes) {
|
||||
if (env.POD_CLEANER_ENABLED) {
|
||||
this.logger.log("🧹 Pod cleaner enabled", {
|
||||
namespace: env.KUBERNETES_NAMESPACE,
|
||||
batchSize: env.POD_CLEANER_BATCH_SIZE,
|
||||
intervalMs: env.POD_CLEANER_INTERVAL_MS,
|
||||
});
|
||||
this.podCleaner = new PodCleaner({
|
||||
register,
|
||||
namespace: env.KUBERNETES_NAMESPACE,
|
||||
batchSize: env.POD_CLEANER_BATCH_SIZE,
|
||||
intervalMs: env.POD_CLEANER_INTERVAL_MS,
|
||||
});
|
||||
} else {
|
||||
this.logger.warn("Pod cleaner disabled");
|
||||
}
|
||||
|
||||
if (env.FAILED_POD_HANDLER_ENABLED) {
|
||||
this.logger.log("🔁 Failed pod handler enabled", {
|
||||
namespace: env.KUBERNETES_NAMESPACE,
|
||||
reconnectIntervalMs: env.FAILED_POD_HANDLER_RECONNECT_INTERVAL_MS,
|
||||
});
|
||||
this.failedPodHandler = new FailedPodHandler({
|
||||
register,
|
||||
namespace: env.KUBERNETES_NAMESPACE,
|
||||
reconnectIntervalMs: env.FAILED_POD_HANDLER_RECONNECT_INTERVAL_MS,
|
||||
});
|
||||
} else {
|
||||
this.logger.warn("Failed pod handler disabled");
|
||||
}
|
||||
}
|
||||
|
||||
if (env.TRIGGER_DEQUEUE_INTERVAL_MS > env.TRIGGER_DEQUEUE_IDLE_INTERVAL_MS) {
|
||||
this.logger.warn(
|
||||
`⚠️ TRIGGER_DEQUEUE_INTERVAL_MS (${env.TRIGGER_DEQUEUE_INTERVAL_MS}) is greater than TRIGGER_DEQUEUE_IDLE_INTERVAL_MS (${env.TRIGGER_DEQUEUE_IDLE_INTERVAL_MS}) - did you mix them up?`
|
||||
);
|
||||
}
|
||||
|
||||
// Redis-verdict source (external aggregator). Keeps existing metric names.
|
||||
if (env.TRIGGER_DEQUEUE_BACKPRESSURE_ENABLED) {
|
||||
this.backpressureRedis = new Redis({
|
||||
host: env.TRIGGER_DEQUEUE_BACKPRESSURE_REDIS_HOST,
|
||||
port: env.TRIGGER_DEQUEUE_BACKPRESSURE_REDIS_PORT,
|
||||
username: env.TRIGGER_DEQUEUE_BACKPRESSURE_REDIS_USERNAME,
|
||||
password: env.TRIGGER_DEQUEUE_BACKPRESSURE_REDIS_PASSWORD,
|
||||
...(env.TRIGGER_DEQUEUE_BACKPRESSURE_REDIS_TLS_DISABLED ? {} : { tls: {} }),
|
||||
maxRetriesPerRequest: null,
|
||||
});
|
||||
this.backpressureRedis.on("error", (error) =>
|
||||
this.logger.error("Backpressure redis error", { error: error.message })
|
||||
);
|
||||
this.backpressureMonitors.push(
|
||||
new BackpressureMonitor({
|
||||
enabled: true,
|
||||
source: new RedisBackpressureSignalSource(
|
||||
this.backpressureRedis,
|
||||
env.TRIGGER_DEQUEUE_BACKPRESSURE_REDIS_KEY
|
||||
),
|
||||
refreshIntervalMs: env.TRIGGER_DEQUEUE_BACKPRESSURE_REFRESH_MS,
|
||||
maxVerdictAgeMs: env.TRIGGER_DEQUEUE_BACKPRESSURE_MAX_VERDICT_AGE_MS,
|
||||
rampMs: env.TRIGGER_DEQUEUE_BACKPRESSURE_RAMP_MS,
|
||||
dryRun: env.TRIGGER_DEQUEUE_BACKPRESSURE_DRY_RUN,
|
||||
logger: this.logger,
|
||||
metrics: new BackpressureMetrics({ register }),
|
||||
})
|
||||
);
|
||||
this.logger.log("🛑 Dequeue backpressure enabled (redis source)", {
|
||||
key: env.TRIGGER_DEQUEUE_BACKPRESSURE_REDIS_KEY,
|
||||
refreshIntervalMs: env.TRIGGER_DEQUEUE_BACKPRESSURE_REFRESH_MS,
|
||||
dryRun: env.TRIGGER_DEQUEUE_BACKPRESSURE_DRY_RUN,
|
||||
});
|
||||
}
|
||||
|
||||
// Pod-count source (in-process apiserver scrape). Namespaced metrics so the
|
||||
// redis source's metric names are preserved.
|
||||
if (env.TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_ENABLED) {
|
||||
// RELEASE < ENGAGE is enforced in env.ts (superRefine), so it's valid here.
|
||||
const podCountGauge = new Gauge({
|
||||
name: "supervisor_cluster_pod_count",
|
||||
help: "Total pod objects stored in the cluster, scraped for backpressure",
|
||||
registers: [register],
|
||||
});
|
||||
this.backpressureMonitors.push(
|
||||
new BackpressureMonitor({
|
||||
enabled: true,
|
||||
source: new K8sPodCountSignalSource({
|
||||
fetchMetrics: createApiserverMetricsFetcher(
|
||||
env.TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_SCRAPE_TIMEOUT_MS
|
||||
),
|
||||
engageThreshold: env.TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_ENGAGE,
|
||||
releaseThreshold: env.TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_RELEASE,
|
||||
reportPodCount: (count) => podCountGauge.set(count),
|
||||
}),
|
||||
refreshIntervalMs: env.TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_REFRESH_MS,
|
||||
maxVerdictAgeMs: env.TRIGGER_DEQUEUE_BACKPRESSURE_MAX_VERDICT_AGE_MS,
|
||||
rampMs: env.TRIGGER_DEQUEUE_BACKPRESSURE_RAMP_MS,
|
||||
dryRun: env.TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_DRY_RUN,
|
||||
logger: this.logger,
|
||||
metrics: new BackpressureMetrics({
|
||||
register,
|
||||
prefix: "supervisor_backpressure_pod_count",
|
||||
}),
|
||||
})
|
||||
);
|
||||
this.logger.log("🛑 Dequeue backpressure enabled (pod-count source)", {
|
||||
engage: env.TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_ENGAGE,
|
||||
release: env.TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_RELEASE,
|
||||
refreshIntervalMs: env.TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_REFRESH_MS,
|
||||
dryRun: env.TRIGGER_DEQUEUE_BACKPRESSURE_POD_COUNT_DRY_RUN,
|
||||
});
|
||||
}
|
||||
|
||||
this.workerSession = new SupervisorSession({
|
||||
workerToken: getWorkerToken(),
|
||||
apiUrl: env.TRIGGER_API_URL,
|
||||
instanceName: env.TRIGGER_WORKER_INSTANCE_NAME,
|
||||
managedWorkerSecret: env.MANAGED_WORKER_SECRET,
|
||||
dequeueIntervalMs: env.TRIGGER_DEQUEUE_INTERVAL_MS,
|
||||
dequeueIdleIntervalMs: env.TRIGGER_DEQUEUE_IDLE_INTERVAL_MS,
|
||||
queueConsumerEnabled: env.TRIGGER_DEQUEUE_ENABLED,
|
||||
maxRunCount: env.TRIGGER_DEQUEUE_MAX_RUN_COUNT,
|
||||
queueClass: env.TRIGGER_WORKER_QUEUE_CLASS,
|
||||
metricsRegistry: register,
|
||||
scaling: {
|
||||
strategy: env.TRIGGER_DEQUEUE_SCALING_STRATEGY,
|
||||
minConsumerCount: env.TRIGGER_DEQUEUE_MIN_CONSUMER_COUNT,
|
||||
maxConsumerCount: env.TRIGGER_DEQUEUE_MAX_CONSUMER_COUNT,
|
||||
scaleUpCooldownMs: env.TRIGGER_DEQUEUE_SCALING_UP_COOLDOWN_MS,
|
||||
scaleDownCooldownMs: env.TRIGGER_DEQUEUE_SCALING_DOWN_COOLDOWN_MS,
|
||||
targetRatio: env.TRIGGER_DEQUEUE_SCALING_TARGET_RATIO,
|
||||
ewmaAlpha: env.TRIGGER_DEQUEUE_SCALING_EWMA_ALPHA,
|
||||
batchWindowMs: env.TRIGGER_DEQUEUE_SCALING_BATCH_WINDOW_MS,
|
||||
dampingFactor: env.TRIGGER_DEQUEUE_SCALING_DAMPING_FACTOR,
|
||||
// Freeze scale-up while backpressure is hard-engaged (not during the resume
|
||||
// ramp). Undefined when backpressure is disabled → no effect on scaling.
|
||||
shouldPauseScaling: () => this.backpressureMonitors.some((m) => m.isEngaged()),
|
||||
},
|
||||
runNotificationsEnabled: env.TRIGGER_WORKLOAD_API_ENABLED,
|
||||
heartbeatIntervalSeconds: env.TRIGGER_WORKER_HEARTBEAT_INTERVAL_SECONDS,
|
||||
sendRunDebugLogs: env.SEND_RUN_DEBUG_LOGS,
|
||||
preDequeue: async () => {
|
||||
// Synchronous, hot-path-safe cached read; false when no monitors are active.
|
||||
const skipForBackpressure = this.backpressureMonitors.some((m) => m.shouldSkipDequeue());
|
||||
|
||||
if (!env.RESOURCE_MONITOR_ENABLED || this.isKubernetes) {
|
||||
// Resource monitor is not used in k8s; backpressure is the only gate there.
|
||||
return { skipDequeue: skipForBackpressure };
|
||||
}
|
||||
|
||||
const resources = await this.resourceMonitor.getNodeResources();
|
||||
|
||||
return {
|
||||
maxResources: {
|
||||
cpu: resources.cpuAvailable,
|
||||
memory: resources.memoryAvailable,
|
||||
},
|
||||
skipDequeue:
|
||||
skipForBackpressure ||
|
||||
resources.cpuAvailable < 0.25 ||
|
||||
resources.memoryAvailable < 0.25,
|
||||
};
|
||||
},
|
||||
preSkip: async () => {
|
||||
// When the node is full, it should still try to warm start runs
|
||||
// await this.tryWarmStartAllThisNode();
|
||||
},
|
||||
});
|
||||
|
||||
if (env.TRIGGER_CHECKPOINT_URL) {
|
||||
this.logger.log("🥶 Checkpoints enabled", {
|
||||
checkpointUrl: env.TRIGGER_CHECKPOINT_URL,
|
||||
});
|
||||
|
||||
this.checkpointClient = new CheckpointClient({
|
||||
apiUrl: new URL(env.TRIGGER_CHECKPOINT_URL),
|
||||
workerClient: this.workerSession.httpClient,
|
||||
orchestrator: this.isKubernetes ? "KUBERNETES" : "DOCKER",
|
||||
});
|
||||
}
|
||||
|
||||
if (env.TRIGGER_WARM_START_VERIFY_ENABLED && this.warmStartUrl) {
|
||||
this.logger.log("Warm-start delivery verification enabled", {
|
||||
delayMs: env.TRIGGER_WARM_START_VERIFY_DELAY_MS,
|
||||
});
|
||||
|
||||
this.warmStartVerifier = new WarmStartVerificationService({
|
||||
workerClient: this.workerSession.httpClient,
|
||||
delayMs: env.TRIGGER_WARM_START_VERIFY_DELAY_MS,
|
||||
createWorkload: (message, timings) => this.createWorkload(message, timings),
|
||||
wideEventOpts: this.wideEventOpts,
|
||||
});
|
||||
}
|
||||
|
||||
this.workerSession.on("runNotification", async ({ time, run }) => {
|
||||
this.logger.verbose("runNotification", { time, run });
|
||||
|
||||
this.workloadServer.notifyRun({ run });
|
||||
});
|
||||
|
||||
this.workerSession.on(
|
||||
"runQueueMessage",
|
||||
async ({ time, message, dequeueResponseMs, pollingIntervalMs }) => {
|
||||
this.logger.verbose(`Received message with timestamp ${time.toLocaleString()}`, message);
|
||||
|
||||
const traceparent = extractTraceparent(message.run.traceContext);
|
||||
|
||||
await runWideEvent(
|
||||
{
|
||||
...this.wideEventOpts,
|
||||
op: "dequeue",
|
||||
kind: "inbound",
|
||||
traceparent,
|
||||
setup: (state) => {
|
||||
setMeta(state, "run_id", message.run.friendlyId);
|
||||
setMeta(state, "env_id", message.environment.id);
|
||||
setMeta(state, "org_id", message.organization.id);
|
||||
setMeta(state, "project_id", message.project.id);
|
||||
if (message.deployment.friendlyId) {
|
||||
setMeta(state, "deployment_id", message.deployment.friendlyId);
|
||||
}
|
||||
setMeta(state, "machine_preset", message.run.machine.name);
|
||||
state.extras.iteration = "dequeue";
|
||||
state.extras.dequeue_response_ms = dequeueResponseMs;
|
||||
state.extras.polling_interval_ms = pollingIntervalMs;
|
||||
state.extras.completed_waitpoints = message.completedWaitpoints.length;
|
||||
},
|
||||
},
|
||||
async () => {
|
||||
if (message.completedWaitpoints.length > 0) {
|
||||
this.logger.debug("Run has completed waitpoints", {
|
||||
runId: message.run.id,
|
||||
completedWaitpoints: message.completedWaitpoints.length,
|
||||
});
|
||||
}
|
||||
|
||||
if (!message.image) {
|
||||
setExtra(fromContext(), "path_taken", "skipped_no_image");
|
||||
this.logger.error("Run has no image", { runId: message.run.id });
|
||||
return;
|
||||
}
|
||||
|
||||
const { checkpoint, ...rest } = message;
|
||||
|
||||
// Register trace context early so snapshot spans work for all paths
|
||||
// (cold create, restore, warm start). Re-registration on restore is safe
|
||||
// since dequeue always provides fresh context.
|
||||
if (this.computeManager?.traceSpansEnabled && traceparent) {
|
||||
this.workloadServer.registerRunTraceContext(message.run.friendlyId, {
|
||||
traceparent,
|
||||
envId: message.environment.id,
|
||||
orgId: message.organization.id,
|
||||
projectId: message.project.id,
|
||||
});
|
||||
}
|
||||
|
||||
if (checkpoint) {
|
||||
setExtra(fromContext(), "path_taken", "restore");
|
||||
this.logger.debug("Restoring run", { runId: message.run.id });
|
||||
|
||||
if (this.computeManager) {
|
||||
const restoreStart = performance.now();
|
||||
try {
|
||||
const runnerId = getRestoreRunnerId(message.run.friendlyId, checkpoint.id);
|
||||
|
||||
const didRestore = await this.computeManager.restore({
|
||||
snapshotId: checkpoint.location,
|
||||
runnerId,
|
||||
runFriendlyId: message.run.friendlyId,
|
||||
snapshotFriendlyId: message.snapshot.friendlyId,
|
||||
machine: message.run.machine,
|
||||
traceContext: message.run.traceContext,
|
||||
envId: message.environment.id,
|
||||
orgId: message.organization.id,
|
||||
projectId: message.project.id,
|
||||
hasPrivateLink: message.organization.hasPrivateLink,
|
||||
dequeuedAt: message.dequeuedAt,
|
||||
});
|
||||
recordPhaseSince("restore", restoreStart, undefined);
|
||||
setExtra(fromContext(), "did_restore", didRestore);
|
||||
|
||||
if (didRestore) {
|
||||
this.logger.debug("Compute restore successful", {
|
||||
runId: message.run.id,
|
||||
runnerId,
|
||||
});
|
||||
} else {
|
||||
this.logger.error("Compute restore failed", {
|
||||
runId: message.run.id,
|
||||
runnerId,
|
||||
});
|
||||
}
|
||||
} catch (error) {
|
||||
recordPhaseSince(
|
||||
"restore",
|
||||
restoreStart,
|
||||
error instanceof Error ? error : new Error(String(error))
|
||||
);
|
||||
this.logger.error("Failed to restore run (compute)", { error });
|
||||
}
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
if (!this.checkpointClient) {
|
||||
this.logger.error("No checkpoint client", { runId: message.run.id });
|
||||
return;
|
||||
}
|
||||
|
||||
const restoreStart = performance.now();
|
||||
try {
|
||||
const didRestore = await this.checkpointClient.restoreRun({
|
||||
runFriendlyId: message.run.friendlyId,
|
||||
snapshotFriendlyId: message.snapshot.friendlyId,
|
||||
body: {
|
||||
...rest,
|
||||
checkpoint,
|
||||
},
|
||||
});
|
||||
recordPhaseSince("restore", restoreStart, undefined);
|
||||
setExtra(fromContext(), "did_restore", didRestore);
|
||||
|
||||
if (didRestore) {
|
||||
this.logger.debug("Restore successful", { runId: message.run.id });
|
||||
} else {
|
||||
this.logger.error("Restore failed", { runId: message.run.id });
|
||||
}
|
||||
} catch (error) {
|
||||
recordPhaseSince(
|
||||
"restore",
|
||||
restoreStart,
|
||||
error instanceof Error ? error : new Error(String(error))
|
||||
);
|
||||
this.logger.error("Failed to restore run", { error });
|
||||
}
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
this.logger.debug("Scheduling run", { runId: message.run.id });
|
||||
|
||||
const warmStartStart = performance.now();
|
||||
const didWarmStart = await this.tryWarmStart(message, traceparent);
|
||||
const warmStartCheckMs = Math.round(performance.now() - warmStartStart);
|
||||
recordPhaseSince("warm_start", warmStartStart, undefined);
|
||||
setExtra(fromContext(), "did_warm_start", didWarmStart);
|
||||
|
||||
if (didWarmStart) {
|
||||
setExtra(fromContext(), "path_taken", "warm_start");
|
||||
this.logger.debug("Warm start successful", { runId: message.run.id });
|
||||
// A hit only means the response was written to the long-poll
|
||||
// socket, not that the runner received it. Schedule a delivery
|
||||
// verification that cold-starts the run if nobody acts on it.
|
||||
this.warmStartVerifier?.schedule(message, {
|
||||
dequeueResponseMs,
|
||||
pollingIntervalMs,
|
||||
warmStartCheckMs,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
setExtra(fromContext(), "path_taken", "cold_create");
|
||||
|
||||
await this.createWorkload(message, {
|
||||
dequeueResponseMs,
|
||||
pollingIntervalMs,
|
||||
warmStartCheckMs,
|
||||
});
|
||||
}
|
||||
);
|
||||
}
|
||||
);
|
||||
|
||||
if (env.METRICS_ENABLED) {
|
||||
this.metricsServer = new HttpServer({
|
||||
port: env.METRICS_PORT,
|
||||
host: env.METRICS_HOST,
|
||||
metrics: {
|
||||
register,
|
||||
expose: true,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
// Responds to workload requests only
|
||||
this.workloadServer = new WorkloadServer({
|
||||
port: env.TRIGGER_WORKLOAD_API_PORT_INTERNAL,
|
||||
host: env.TRIGGER_WORKLOAD_API_HOST_INTERNAL,
|
||||
workerClient: this.workerSession.httpClient,
|
||||
checkpointClient: this.checkpointClient,
|
||||
computeManager: this.computeManager,
|
||||
tracing: this.tracing,
|
||||
wideEventOpts: this.wideEventOpts,
|
||||
wideEventsNoisyRoutes: this.wideEventsNoisyRoutes,
|
||||
});
|
||||
|
||||
this.workloadServer.on("runConnected", this.onRunConnected.bind(this));
|
||||
this.workloadServer.on("runDisconnected", this.onRunDisconnected.bind(this));
|
||||
}
|
||||
|
||||
async onRunConnected({ run }: { run: { friendlyId: string } }) {
|
||||
this.logger.debug("Run connected", { run });
|
||||
// The dispatched run reached a runner on this node - no fallback needed.
|
||||
this.warmStartVerifier?.cancel(run.friendlyId);
|
||||
this.workerSession.subscribeToRunNotifications([run.friendlyId]);
|
||||
}
|
||||
|
||||
async onRunDisconnected({ run }: { run: { friendlyId: string } }) {
|
||||
this.logger.debug("Run disconnected", { run });
|
||||
this.workerSession.unsubscribeFromRunNotifications([run.friendlyId]);
|
||||
}
|
||||
|
||||
private async createWorkload(message: DequeuedMessage, timings: WarmStartTimings) {
|
||||
const createStart = performance.now();
|
||||
try {
|
||||
if (!message.deployment.friendlyId) {
|
||||
// mostly a type guard, deployments always exists for deployed environments
|
||||
// a proper fix would be to use a discriminated union schema to differentiate between dequeued runs in dev and in deployed environments.
|
||||
throw new Error("Deployment is missing");
|
||||
}
|
||||
|
||||
if (!message.image) {
|
||||
// same type-guard situation as deployment above
|
||||
throw new Error("Image is missing");
|
||||
}
|
||||
|
||||
await this.workloadManager.create({
|
||||
dequeuedAt: message.dequeuedAt,
|
||||
dequeueResponseMs: timings.dequeueResponseMs,
|
||||
pollingIntervalMs: timings.pollingIntervalMs,
|
||||
warmStartCheckMs: timings.warmStartCheckMs,
|
||||
envId: message.environment.id,
|
||||
envType: message.environment.type,
|
||||
image: message.image,
|
||||
machine: message.run.machine,
|
||||
orgId: message.organization.id,
|
||||
projectId: message.project.id,
|
||||
deploymentFriendlyId: message.deployment.friendlyId,
|
||||
deploymentVersion: message.backgroundWorker.version,
|
||||
runId: message.run.id,
|
||||
runFriendlyId: message.run.friendlyId,
|
||||
version: message.version,
|
||||
nextAttemptNumber: message.run.attemptNumber,
|
||||
snapshotId: message.snapshot.id,
|
||||
snapshotFriendlyId: message.snapshot.friendlyId,
|
||||
placementTags: message.placementTags,
|
||||
traceContext: message.run.traceContext,
|
||||
annotations: message.run.annotations,
|
||||
hasPrivateLink: message.organization.hasPrivateLink,
|
||||
});
|
||||
recordPhaseSince("workload_create", createStart, undefined);
|
||||
workloadCreateDuration.observe(
|
||||
{ backend: this.workloadManagerBackend, outcome: "success" },
|
||||
(performance.now() - createStart) / 1000
|
||||
);
|
||||
|
||||
// Disabled for now
|
||||
// this.resourceMonitor.blockResources({
|
||||
// cpu: message.run.machine.cpu,
|
||||
// memory: message.run.machine.memory,
|
||||
// });
|
||||
} catch (error) {
|
||||
recordPhaseSince(
|
||||
"workload_create",
|
||||
createStart,
|
||||
error instanceof Error ? error : new Error(String(error))
|
||||
);
|
||||
workloadCreateDuration.observe(
|
||||
{ backend: this.workloadManagerBackend, outcome: "error" },
|
||||
(performance.now() - createStart) / 1000
|
||||
);
|
||||
this.logger.error("Failed to create workload", {
|
||||
runId: message.run.friendlyId,
|
||||
error,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
private async tryWarmStart(
|
||||
dequeuedMessage: DequeuedMessage,
|
||||
traceparent: string | undefined
|
||||
): Promise<boolean> {
|
||||
if (!this.warmStartUrl) {
|
||||
return false;
|
||||
}
|
||||
|
||||
const warmStartUrlWithPath = new URL("/warm-start", this.warmStartUrl);
|
||||
|
||||
const headers: Record<string, string> = {
|
||||
"Content-Type": "application/json",
|
||||
};
|
||||
// Propagate the inbound W3C traceparent so the upstream warm-start
|
||||
// receiver continues the same trace instead of minting a new one. Gated
|
||||
// by the same kill switch as the wide-event emission so the whole PR is
|
||||
// a no-op on the wire when disabled.
|
||||
if (this.wideEventOpts.enabled && traceparent) {
|
||||
headers.traceparent = traceparent;
|
||||
}
|
||||
|
||||
try {
|
||||
const res = await fetch(warmStartUrlWithPath.href, {
|
||||
method: "POST",
|
||||
headers,
|
||||
body: JSON.stringify({ dequeuedMessage }),
|
||||
});
|
||||
|
||||
if (!res.ok) {
|
||||
this.logger.error("Warm start failed", {
|
||||
runId: dequeuedMessage.run.id,
|
||||
});
|
||||
return false;
|
||||
}
|
||||
|
||||
const data = await res.json();
|
||||
const parsedData = z.object({ didWarmStart: z.boolean() }).safeParse(data);
|
||||
|
||||
if (!parsedData.success) {
|
||||
this.logger.error("Warm start response invalid", {
|
||||
runId: dequeuedMessage.run.id,
|
||||
data,
|
||||
});
|
||||
return false;
|
||||
}
|
||||
|
||||
return parsedData.data.didWarmStart;
|
||||
} catch (error) {
|
||||
this.logger.error("Warm start error", {
|
||||
runId: dequeuedMessage.run.id,
|
||||
error,
|
||||
});
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
async start() {
|
||||
this.logger.log("Starting up");
|
||||
|
||||
// Optional services
|
||||
this.backpressureMonitors.forEach((m) => m.start());
|
||||
await this.podCleaner?.start();
|
||||
await this.failedPodHandler?.start();
|
||||
await this.metricsServer?.start();
|
||||
|
||||
if (env.TRIGGER_WORKLOAD_API_ENABLED) {
|
||||
this.logger.log("Workload API enabled", {
|
||||
protocol: env.TRIGGER_WORKLOAD_API_PROTOCOL,
|
||||
domain: env.TRIGGER_WORKLOAD_API_DOMAIN,
|
||||
port: env.TRIGGER_WORKLOAD_API_PORT_INTERNAL,
|
||||
});
|
||||
await this.workloadServer.start();
|
||||
} else {
|
||||
this.logger.warn("Workload API disabled");
|
||||
}
|
||||
|
||||
await this.workerSession.start();
|
||||
}
|
||||
|
||||
async stop() {
|
||||
this.logger.log("Shutting down");
|
||||
// Stop the verifier first: its timer can otherwise fire mid-shutdown and
|
||||
// cold-create a workload on a node that is going down.
|
||||
this.warmStartVerifier?.stop();
|
||||
await this.workloadServer.stop();
|
||||
await this.workerSession.stop();
|
||||
|
||||
// Optional services
|
||||
this.backpressureMonitors.forEach((m) => m.stop());
|
||||
await this.backpressureRedis?.quit();
|
||||
await this.podCleaner?.stop();
|
||||
await this.failedPodHandler?.stop();
|
||||
await this.metricsServer?.stop();
|
||||
}
|
||||
}
|
||||
|
||||
// Opt-in, dev-only: mirror this process's structured logs to a local telnet/TCP stream.
|
||||
if (env.SUPERVISOR_TELNET_LOGS_PORT && env.SUPERVISOR_TELNET_LOGS_PORT > 0) {
|
||||
const telnetLogServer = startTelnetLogServer({
|
||||
port: env.SUPERVISOR_TELNET_LOGS_PORT,
|
||||
name: "supervisor",
|
||||
});
|
||||
SimpleStructuredLogger.onLog = (log) => telnetLogServer.broadcast(formatLogLine(log));
|
||||
}
|
||||
|
||||
const worker = new ManagedSupervisor();
|
||||
worker.start();
|
||||
@@ -0,0 +1,3 @@
|
||||
import { Registry } from "prom-client";
|
||||
|
||||
export const register = new Registry();
|
||||
@@ -0,0 +1,278 @@
|
||||
import type Docker from "dockerode";
|
||||
import type { MachineResources } from "@trigger.dev/core/v3";
|
||||
import { SimpleStructuredLogger } from "@trigger.dev/core/v3/utils/structuredLogger";
|
||||
import { env } from "./env.js";
|
||||
import type { K8sApi } from "./clients/kubernetes.js";
|
||||
|
||||
const logger = new SimpleStructuredLogger("resource-monitor");
|
||||
|
||||
interface NodeResources {
|
||||
cpuTotal: number; // in cores
|
||||
cpuAvailable: number;
|
||||
memoryTotal: number; // in bytes
|
||||
memoryAvailable: number;
|
||||
}
|
||||
|
||||
interface ResourceRequest {
|
||||
cpu: number; // in cores
|
||||
memory: number; // in bytes
|
||||
}
|
||||
|
||||
export abstract class ResourceMonitor {
|
||||
protected cacheTimeoutMs = 5_000;
|
||||
protected lastUpdateMs = 0;
|
||||
|
||||
protected cachedResources: NodeResources = {
|
||||
cpuTotal: 0,
|
||||
cpuAvailable: 0,
|
||||
memoryTotal: 0,
|
||||
memoryAvailable: 0,
|
||||
};
|
||||
|
||||
protected resourceParser: ResourceParser;
|
||||
|
||||
constructor(Parser: new () => ResourceParser) {
|
||||
this.resourceParser = new Parser();
|
||||
}
|
||||
|
||||
abstract getNodeResources(fromCache?: boolean): Promise<NodeResources>;
|
||||
|
||||
blockResources(resources: MachineResources): void {
|
||||
const { cpu, memory } = this.toResourceRequest(resources);
|
||||
|
||||
logger.debug("[ResourceMonitor] Blocking resources", {
|
||||
raw: resources,
|
||||
converted: { cpu, memory },
|
||||
});
|
||||
|
||||
this.cachedResources.cpuAvailable -= cpu;
|
||||
this.cachedResources.memoryAvailable -= memory;
|
||||
}
|
||||
|
||||
async wouldFit(request: ResourceRequest): Promise<boolean> {
|
||||
const resources = await this.getNodeResources();
|
||||
return resources.cpuAvailable >= request.cpu && resources.memoryAvailable >= request.memory;
|
||||
}
|
||||
|
||||
private toResourceRequest(resources: MachineResources): ResourceRequest {
|
||||
return {
|
||||
cpu: resources.cpu ?? 0,
|
||||
memory: this.gbToBytes(resources.memory ?? 0),
|
||||
};
|
||||
}
|
||||
|
||||
private gbToBytes(gb: number): number {
|
||||
return gb * 1024 * 1024 * 1024;
|
||||
}
|
||||
|
||||
protected isCacheValid(): boolean {
|
||||
return this.cachedResources !== null && Date.now() - this.lastUpdateMs < this.cacheTimeoutMs;
|
||||
}
|
||||
|
||||
protected applyOverrides(resources: NodeResources): NodeResources {
|
||||
if (
|
||||
!env.RESOURCE_MONITOR_OVERRIDE_CPU_TOTAL &&
|
||||
!env.RESOURCE_MONITOR_OVERRIDE_MEMORY_TOTAL_GB
|
||||
) {
|
||||
return resources;
|
||||
}
|
||||
|
||||
logger.debug("[ResourceMonitor] 🛡️ Applying resource overrides", {
|
||||
cpuTotal: env.RESOURCE_MONITOR_OVERRIDE_CPU_TOTAL,
|
||||
memoryTotalGb: env.RESOURCE_MONITOR_OVERRIDE_MEMORY_TOTAL_GB,
|
||||
});
|
||||
|
||||
const cpuTotal = env.RESOURCE_MONITOR_OVERRIDE_CPU_TOTAL ?? resources.cpuTotal;
|
||||
const memoryTotal = env.RESOURCE_MONITOR_OVERRIDE_MEMORY_TOTAL_GB
|
||||
? this.gbToBytes(env.RESOURCE_MONITOR_OVERRIDE_MEMORY_TOTAL_GB)
|
||||
: resources.memoryTotal;
|
||||
|
||||
const cpuDiff = cpuTotal - resources.cpuTotal;
|
||||
const memoryDiff = memoryTotal - resources.memoryTotal;
|
||||
|
||||
const cpuAvailable = Math.max(0, resources.cpuAvailable + cpuDiff);
|
||||
const memoryAvailable = Math.max(0, resources.memoryAvailable + memoryDiff);
|
||||
|
||||
return {
|
||||
cpuTotal,
|
||||
cpuAvailable,
|
||||
memoryTotal,
|
||||
memoryAvailable,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
type SystemInfo = {
|
||||
NCPU: number | undefined;
|
||||
MemTotal: number | undefined;
|
||||
};
|
||||
|
||||
export class DockerResourceMonitor extends ResourceMonitor {
|
||||
private docker: Docker;
|
||||
|
||||
constructor(docker: Docker) {
|
||||
super(DockerResourceParser);
|
||||
this.docker = docker;
|
||||
}
|
||||
|
||||
async getNodeResources(fromCache?: boolean): Promise<NodeResources> {
|
||||
if (this.isCacheValid() || fromCache) {
|
||||
// logger.debug("[ResourceMonitor] Using cached resources");
|
||||
return this.cachedResources;
|
||||
}
|
||||
|
||||
const info: SystemInfo = await this.docker.info();
|
||||
const stats = await this.docker.listContainers({ all: true });
|
||||
|
||||
// Get system-wide resources
|
||||
const cpuTotal = info.NCPU ?? 0;
|
||||
const memoryTotal = info.MemTotal ?? 0;
|
||||
|
||||
// Calculate used resources from running containers
|
||||
let cpuUsed = 0;
|
||||
let memoryUsed = 0;
|
||||
|
||||
for (const container of stats) {
|
||||
if (container.State === "running") {
|
||||
const c = this.docker.getContainer(container.Id);
|
||||
const { HostConfig } = await c.inspect();
|
||||
|
||||
const cpu = this.resourceParser.cpu(HostConfig.NanoCpus ?? 0);
|
||||
const memory = this.resourceParser.memory(HostConfig.Memory ?? 0);
|
||||
|
||||
cpuUsed += cpu;
|
||||
memoryUsed += memory;
|
||||
}
|
||||
}
|
||||
|
||||
this.cachedResources = this.applyOverrides({
|
||||
cpuTotal,
|
||||
cpuAvailable: cpuTotal - cpuUsed,
|
||||
memoryTotal,
|
||||
memoryAvailable: memoryTotal - memoryUsed,
|
||||
});
|
||||
|
||||
this.lastUpdateMs = Date.now();
|
||||
|
||||
return this.cachedResources;
|
||||
}
|
||||
}
|
||||
|
||||
export class KubernetesResourceMonitor extends ResourceMonitor {
|
||||
private k8s: K8sApi;
|
||||
private nodeName: string;
|
||||
|
||||
constructor(k8s: K8sApi, nodeName: string) {
|
||||
super(KubernetesResourceParser);
|
||||
this.k8s = k8s;
|
||||
this.nodeName = nodeName;
|
||||
}
|
||||
|
||||
async getNodeResources(fromCache?: boolean): Promise<NodeResources> {
|
||||
if (this.isCacheValid() || fromCache) {
|
||||
logger.debug("[ResourceMonitor] Using cached resources");
|
||||
return this.cachedResources;
|
||||
}
|
||||
|
||||
const node = await this.k8s.core.readNode({ name: this.nodeName });
|
||||
const pods = await this.k8s.core.listPodForAllNamespaces({
|
||||
// TODO: ensure this includes all pods that consume resources
|
||||
fieldSelector: `spec.nodeName=${this.nodeName},status.phase=Running`,
|
||||
});
|
||||
|
||||
const allocatable = node.status?.allocatable;
|
||||
const cpuTotal = this.resourceParser.cpu(allocatable?.cpu ?? "0");
|
||||
const memoryTotal = this.resourceParser.memory(allocatable?.memory ?? "0");
|
||||
|
||||
// Sum up resources requested by all pods on this node
|
||||
let cpuRequested = 0;
|
||||
let memoryRequested = 0;
|
||||
|
||||
for (const pod of pods.items) {
|
||||
if (pod.status?.phase === "Running") {
|
||||
if (!pod.spec) {
|
||||
continue;
|
||||
}
|
||||
|
||||
for (const container of pod.spec.containers) {
|
||||
const resources = container.resources?.requests ?? {};
|
||||
cpuRequested += this.resourceParser.cpu(resources.cpu ?? "0");
|
||||
memoryRequested += this.resourceParser.memory(resources.memory ?? "0");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
this.cachedResources = this.applyOverrides({
|
||||
cpuTotal,
|
||||
cpuAvailable: cpuTotal - cpuRequested,
|
||||
memoryTotal,
|
||||
memoryAvailable: memoryTotal - memoryRequested,
|
||||
});
|
||||
|
||||
this.lastUpdateMs = Date.now();
|
||||
|
||||
return this.cachedResources;
|
||||
}
|
||||
}
|
||||
|
||||
export class NoopResourceMonitor extends ResourceMonitor {
|
||||
constructor() {
|
||||
super(NoopResourceParser);
|
||||
}
|
||||
|
||||
async getNodeResources(): Promise<NodeResources> {
|
||||
return {
|
||||
cpuTotal: 0,
|
||||
cpuAvailable: Infinity,
|
||||
memoryTotal: 0,
|
||||
memoryAvailable: Infinity,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
abstract class ResourceParser {
|
||||
abstract cpu(cpu: number | string): number;
|
||||
abstract memory(memory: number | string): number;
|
||||
}
|
||||
|
||||
class DockerResourceParser extends ResourceParser {
|
||||
cpu(cpu: number): number {
|
||||
return cpu / 1e9;
|
||||
}
|
||||
|
||||
memory(memory: number): number {
|
||||
return memory;
|
||||
}
|
||||
}
|
||||
|
||||
class KubernetesResourceParser extends ResourceParser {
|
||||
cpu(cpu: string): number {
|
||||
if (cpu.endsWith("m")) {
|
||||
return parseInt(cpu.slice(0, -1)) / 1000;
|
||||
}
|
||||
return parseInt(cpu);
|
||||
}
|
||||
|
||||
memory(memory: string): number {
|
||||
if (memory.endsWith("Ki")) {
|
||||
return parseInt(memory.slice(0, -2)) * 1024;
|
||||
}
|
||||
if (memory.endsWith("Mi")) {
|
||||
return parseInt(memory.slice(0, -2)) * 1024 * 1024;
|
||||
}
|
||||
if (memory.endsWith("Gi")) {
|
||||
return parseInt(memory.slice(0, -2)) * 1024 * 1024 * 1024;
|
||||
}
|
||||
return parseInt(memory);
|
||||
}
|
||||
}
|
||||
|
||||
class NoopResourceParser extends ResourceParser {
|
||||
cpu(cpu: number): number {
|
||||
return cpu;
|
||||
}
|
||||
|
||||
memory(memory: number): number {
|
||||
return memory;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,130 @@
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import { setTimeout as sleep } from "node:timers/promises";
|
||||
import { ComputeSnapshotService } from "./computeSnapshotService.js";
|
||||
import type { ComputeWorkloadManager } from "../workloadManager/compute.js";
|
||||
import type { SupervisorHttpClient } from "@trigger.dev/core/v3/workers";
|
||||
|
||||
// The TimerWheel ticks every 100ms, so a 200ms delay dispatches within ~300ms.
|
||||
const DELAY_MS = 200;
|
||||
// Long enough that a pending snapshot would certainly have dispatched.
|
||||
const SETTLE_MS = 600;
|
||||
|
||||
function createService() {
|
||||
const snapshot = vi.fn(
|
||||
async (_opts: { runnerId: string; metadata: Record<string, string> }) => true
|
||||
);
|
||||
|
||||
const computeManager = {
|
||||
snapshotDelayMs: DELAY_MS,
|
||||
snapshotDispatchLimit: 1,
|
||||
snapshot,
|
||||
} as unknown as ComputeWorkloadManager;
|
||||
|
||||
const service = new ComputeSnapshotService({
|
||||
computeManager,
|
||||
workerClient: {} as SupervisorHttpClient,
|
||||
wideEventOpts: { service: "supervisor-test", env: {}, enabled: false },
|
||||
});
|
||||
|
||||
return { service, snapshot };
|
||||
}
|
||||
|
||||
function delayedSnapshot(runnerId = "runner-1") {
|
||||
return {
|
||||
runnerId,
|
||||
runFriendlyId: "run_1",
|
||||
snapshotFriendlyId: "snapshot_1",
|
||||
};
|
||||
}
|
||||
|
||||
describe("ComputeSnapshotService", () => {
|
||||
it("dispatches a scheduled snapshot after the delay", async () => {
|
||||
const { service, snapshot } = createService();
|
||||
try {
|
||||
service.schedule("run_1", delayedSnapshot());
|
||||
|
||||
await vi.waitFor(() => expect(snapshot).toHaveBeenCalledTimes(1), { timeout: 2_000 });
|
||||
expect(snapshot).toHaveBeenCalledWith({
|
||||
runnerId: "runner-1",
|
||||
metadata: { runId: "run_1", snapshotFriendlyId: "snapshot_1" },
|
||||
});
|
||||
} finally {
|
||||
service.stop();
|
||||
}
|
||||
});
|
||||
|
||||
it("cancel before the delay expires prevents the dispatch", async () => {
|
||||
const { service, snapshot } = createService();
|
||||
try {
|
||||
service.schedule("run_1", delayedSnapshot());
|
||||
|
||||
expect(service.cancel("run_1")).toBe(true);
|
||||
|
||||
await sleep(SETTLE_MS);
|
||||
expect(snapshot).not.toHaveBeenCalled();
|
||||
} finally {
|
||||
service.stop();
|
||||
}
|
||||
});
|
||||
|
||||
it("cancel returns false when nothing is pending", () => {
|
||||
const { service } = createService();
|
||||
try {
|
||||
expect(service.cancel("run_1")).toBe(false);
|
||||
} finally {
|
||||
service.stop();
|
||||
}
|
||||
});
|
||||
|
||||
it("cancel with a matching runnerId cancels the pending snapshot", async () => {
|
||||
const { service, snapshot } = createService();
|
||||
try {
|
||||
service.schedule("run_1", delayedSnapshot("runner-a"));
|
||||
|
||||
expect(service.cancel("run_1", "runner-a")).toBe(true);
|
||||
|
||||
await sleep(SETTLE_MS);
|
||||
expect(snapshot).not.toHaveBeenCalled();
|
||||
} finally {
|
||||
service.stop();
|
||||
}
|
||||
});
|
||||
|
||||
it("cancel with a different runnerId leaves the pending snapshot alone", async () => {
|
||||
const { service, snapshot } = createService();
|
||||
try {
|
||||
service.schedule("run_1", delayedSnapshot("runner-a"));
|
||||
|
||||
// A stale runner for a reassigned run must not cancel the new runner's snapshot.
|
||||
expect(service.cancel("run_1", "runner-b")).toBe(false);
|
||||
|
||||
await vi.waitFor(() => expect(snapshot).toHaveBeenCalledTimes(1), { timeout: 2_000 });
|
||||
expect(snapshot).toHaveBeenCalledWith(expect.objectContaining({ runnerId: "runner-a" }));
|
||||
} finally {
|
||||
service.stop();
|
||||
}
|
||||
});
|
||||
|
||||
it("re-scheduling the same run replaces the pending snapshot", async () => {
|
||||
const { service, snapshot } = createService();
|
||||
try {
|
||||
service.schedule("run_1", delayedSnapshot());
|
||||
service.schedule("run_1", {
|
||||
runnerId: "runner-1",
|
||||
runFriendlyId: "run_1",
|
||||
snapshotFriendlyId: "snapshot_2",
|
||||
});
|
||||
|
||||
await vi.waitFor(() => expect(snapshot).toHaveBeenCalledTimes(1), { timeout: 2_000 });
|
||||
await sleep(SETTLE_MS);
|
||||
|
||||
expect(snapshot).toHaveBeenCalledTimes(1);
|
||||
expect(snapshot).toHaveBeenCalledWith({
|
||||
runnerId: "runner-1",
|
||||
metadata: { runId: "run_1", snapshotFriendlyId: "snapshot_2" },
|
||||
});
|
||||
} finally {
|
||||
service.stop();
|
||||
}
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,323 @@
|
||||
import pLimit from "p-limit";
|
||||
import { SimpleStructuredLogger } from "@trigger.dev/core/v3/utils/structuredLogger";
|
||||
import { parseTraceparent } from "@trigger.dev/core/v3/isomorphic";
|
||||
import type { SupervisorHttpClient } from "@trigger.dev/core/v3/workers";
|
||||
import { type SnapshotCallbackPayload } from "@internal/compute";
|
||||
import type { ComputeWorkloadManager } from "../workloadManager/compute.js";
|
||||
import { TimerWheel } from "./timerWheel.js";
|
||||
import type { OtlpTraceService } from "./otlpTraceService.js";
|
||||
import {
|
||||
emitOneShot,
|
||||
fromContext,
|
||||
recordPhaseSince,
|
||||
runWideEvent,
|
||||
setExtra,
|
||||
setMeta,
|
||||
type WideEventOptions,
|
||||
} from "../wideEvents/index.js";
|
||||
|
||||
type DelayedSnapshot = {
|
||||
runnerId: string;
|
||||
runFriendlyId: string;
|
||||
snapshotFriendlyId: string;
|
||||
};
|
||||
|
||||
export type RunTraceContext = {
|
||||
traceparent: string;
|
||||
envId: string;
|
||||
orgId: string;
|
||||
projectId: string;
|
||||
};
|
||||
|
||||
export type ComputeSnapshotServiceOptions = {
|
||||
computeManager: ComputeWorkloadManager;
|
||||
workerClient: SupervisorHttpClient;
|
||||
tracing?: OtlpTraceService;
|
||||
wideEventOpts: WideEventOptions;
|
||||
};
|
||||
|
||||
export class ComputeSnapshotService {
|
||||
private readonly logger = new SimpleStructuredLogger("compute-snapshot-service");
|
||||
|
||||
private static readonly MAX_TRACE_CONTEXTS = 10_000;
|
||||
private readonly runTraceContexts = new Map<string, RunTraceContext>();
|
||||
private readonly timerWheel: TimerWheel<DelayedSnapshot>;
|
||||
private readonly dispatchLimit: ReturnType<typeof pLimit>;
|
||||
|
||||
private readonly computeManager: ComputeWorkloadManager;
|
||||
private readonly workerClient: SupervisorHttpClient;
|
||||
private readonly tracing?: OtlpTraceService;
|
||||
private readonly wideEventOpts: WideEventOptions;
|
||||
|
||||
constructor(opts: ComputeSnapshotServiceOptions) {
|
||||
this.computeManager = opts.computeManager;
|
||||
this.workerClient = opts.workerClient;
|
||||
this.tracing = opts.tracing;
|
||||
this.wideEventOpts = opts.wideEventOpts;
|
||||
|
||||
this.dispatchLimit = pLimit(this.computeManager.snapshotDispatchLimit);
|
||||
this.timerWheel = new TimerWheel<DelayedSnapshot>({
|
||||
delayMs: this.computeManager.snapshotDelayMs,
|
||||
onExpire: (item) => {
|
||||
this.dispatchLimit(() => this.dispatch(item.data)).catch((error) => {
|
||||
this.logger.error("Snapshot dispatch failed", {
|
||||
runId: item.data.runFriendlyId,
|
||||
runnerId: item.data.runnerId,
|
||||
error,
|
||||
});
|
||||
});
|
||||
},
|
||||
});
|
||||
this.timerWheel.start();
|
||||
}
|
||||
|
||||
/** Schedule a delayed snapshot for a run. Replaces any pending snapshot for the same run. */
|
||||
schedule(runFriendlyId: string, data: DelayedSnapshot) {
|
||||
this.timerWheel.submit(runFriendlyId, data);
|
||||
emitOneShot({
|
||||
...this.wideEventOpts,
|
||||
op: "snapshot.schedule",
|
||||
kind: "event",
|
||||
populate: (state) => {
|
||||
state.meta.run_id = runFriendlyId;
|
||||
state.meta.snapshot_id = data.snapshotFriendlyId;
|
||||
state.extras.runner_id = data.runnerId;
|
||||
state.extras.delay_ms = this.computeManager.snapshotDelayMs;
|
||||
},
|
||||
});
|
||||
this.logger.debug("Snapshot scheduled", {
|
||||
runFriendlyId,
|
||||
snapshotFriendlyId: data.snapshotFriendlyId,
|
||||
delayMs: this.computeManager.snapshotDelayMs,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Cancel a pending delayed snapshot. Returns true if one was cancelled.
|
||||
* When `runnerId` is given, only a snapshot scheduled for that same runner
|
||||
* is cancelled - a stale runner for a run that has since been reassigned
|
||||
* must not cancel the new runner's pending snapshot.
|
||||
*/
|
||||
cancel(runFriendlyId: string, runnerId?: string): boolean {
|
||||
if (runnerId) {
|
||||
const pending = this.timerWheel.peek(runFriendlyId);
|
||||
if (pending && pending.data.runnerId !== runnerId) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
const cancelled = this.timerWheel.cancel(runFriendlyId);
|
||||
if (cancelled) {
|
||||
emitOneShot({
|
||||
...this.wideEventOpts,
|
||||
op: "snapshot.canceled",
|
||||
kind: "event",
|
||||
populate: (state) => {
|
||||
state.meta.run_id = runFriendlyId;
|
||||
},
|
||||
});
|
||||
this.logger.debug("Snapshot cancelled", { runFriendlyId });
|
||||
}
|
||||
return cancelled;
|
||||
}
|
||||
|
||||
/** Handle the callback from the gateway after a snapshot completes or fails. */
|
||||
async handleCallback(body: SnapshotCallbackPayload) {
|
||||
const snapshotId = body.status === "completed" ? body.snapshot_id : undefined;
|
||||
const runId = body.metadata?.runId;
|
||||
const snapshotFriendlyId = body.metadata?.snapshotFriendlyId;
|
||||
|
||||
// Enrich the wrapping route's wide event with snapshot metadata. The
|
||||
// `/api/v1/compute/snapshot-complete` route is registered with `wideRoute`,
|
||||
// so `fromContext()` returns the State of that route and these calls
|
||||
// become extras/meta on the same wide event - no nested emission.
|
||||
const state = fromContext();
|
||||
if (state) {
|
||||
state.extras["snapshot.status"] = body.status;
|
||||
if (body.instance_id) state.extras["snapshot.instance_id"] = body.instance_id;
|
||||
if (body.duration_ms !== undefined) state.extras["snapshot.duration_ms"] = body.duration_ms;
|
||||
if (snapshotId) state.extras["snapshot.id"] = snapshotId;
|
||||
if (body.status === "failed" && body.error) state.extras["snapshot.error"] = body.error;
|
||||
}
|
||||
if (runId) setMeta(state, "run_id", runId);
|
||||
if (snapshotFriendlyId) setMeta(state, "snapshot_id", snapshotFriendlyId);
|
||||
|
||||
this.logger.debug("Snapshot callback", {
|
||||
snapshotId,
|
||||
instanceId: body.instance_id,
|
||||
status: body.status,
|
||||
error: body.status === "failed" ? body.error : undefined,
|
||||
metadata: body.metadata,
|
||||
durationMs: body.duration_ms,
|
||||
});
|
||||
|
||||
if (!runId || !snapshotFriendlyId) {
|
||||
this.logger.error("Snapshot callback missing metadata", { body });
|
||||
return { ok: false as const, status: 400 };
|
||||
}
|
||||
|
||||
this.#emitSnapshotSpan(runId, body.duration_ms, snapshotId);
|
||||
|
||||
if (body.status === "completed") {
|
||||
const submitStart = performance.now();
|
||||
const result = await this.workerClient.submitSuspendCompletion({
|
||||
runId,
|
||||
snapshotId: snapshotFriendlyId,
|
||||
body: {
|
||||
success: true,
|
||||
checkpoint: {
|
||||
type: "COMPUTE",
|
||||
location: body.snapshot_id,
|
||||
},
|
||||
},
|
||||
});
|
||||
recordPhaseSince(
|
||||
"submit_completion",
|
||||
submitStart,
|
||||
result.success ? undefined : new Error(String(result.error))
|
||||
);
|
||||
|
||||
if (result.success) {
|
||||
this.logger.debug("Suspend completion submitted", {
|
||||
runId,
|
||||
instanceId: body.instance_id,
|
||||
snapshotId: body.snapshot_id,
|
||||
});
|
||||
} else {
|
||||
setExtra(state, "submit_completion.error", String(result.error));
|
||||
this.logger.error("Failed to submit suspend completion", {
|
||||
runId,
|
||||
snapshotFriendlyId,
|
||||
error: result.error,
|
||||
});
|
||||
}
|
||||
} else {
|
||||
const submitStart = performance.now();
|
||||
const result = await this.workerClient.submitSuspendCompletion({
|
||||
runId,
|
||||
snapshotId: snapshotFriendlyId,
|
||||
body: {
|
||||
success: false,
|
||||
error: body.error ?? "Snapshot failed",
|
||||
},
|
||||
});
|
||||
recordPhaseSince(
|
||||
"submit_completion",
|
||||
submitStart,
|
||||
result.success ? undefined : new Error(String(result.error))
|
||||
);
|
||||
|
||||
if (!result.success) {
|
||||
setExtra(state, "submit_completion.error", String(result.error));
|
||||
this.logger.error("Failed to submit suspend failure", {
|
||||
runId,
|
||||
snapshotFriendlyId,
|
||||
error: result.error,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
return { ok: true as const, status: 200 };
|
||||
}
|
||||
|
||||
registerTraceContext(runFriendlyId: string, ctx: RunTraceContext) {
|
||||
// Evict oldest entries if we've hit the cap. This is best-effort: on a busy
|
||||
// supervisor, entries for long-lived runs may be evicted before their snapshot
|
||||
// callback arrives, causing those snapshot spans to be silently dropped.
|
||||
// That's acceptable - trace spans are observability sugar, not correctness.
|
||||
if (this.runTraceContexts.size >= ComputeSnapshotService.MAX_TRACE_CONTEXTS) {
|
||||
const firstKey = this.runTraceContexts.keys().next().value;
|
||||
if (firstKey) {
|
||||
this.runTraceContexts.delete(firstKey);
|
||||
}
|
||||
}
|
||||
|
||||
this.runTraceContexts.set(runFriendlyId, ctx);
|
||||
}
|
||||
|
||||
/** Stop the timer wheel, dropping pending snapshots. */
|
||||
stop(): string[] {
|
||||
// Intentionally drop pending snapshots rather than dispatching them. The supervisor
|
||||
// is shutting down, so our callback URL will be dead by the time the gateway responds.
|
||||
// Runners detect the supervisor is gone and reconnect to a new instance, which
|
||||
// re-triggers the snapshot workflow. Snapshots are an optimization, not a correctness
|
||||
// requirement - runs continue fine without them.
|
||||
const remaining = this.timerWheel.stop();
|
||||
const droppedRuns = remaining.map((item) => item.key);
|
||||
|
||||
if (droppedRuns.length > 0) {
|
||||
this.logger.info("Stopped, dropped pending snapshots", { count: droppedRuns.length });
|
||||
this.logger.debug("Dropped snapshot details", { runs: droppedRuns });
|
||||
}
|
||||
|
||||
return droppedRuns;
|
||||
}
|
||||
|
||||
/** Dispatch a snapshot request to the gateway. */
|
||||
private async dispatch(snapshot: DelayedSnapshot): Promise<void> {
|
||||
await runWideEvent(
|
||||
{
|
||||
...this.wideEventOpts,
|
||||
op: "snapshot.dispatch",
|
||||
kind: "scheduled",
|
||||
setup: (state) => {
|
||||
state.meta.run_id = snapshot.runFriendlyId;
|
||||
state.meta.snapshot_id = snapshot.snapshotFriendlyId;
|
||||
state.extras.runner_id = snapshot.runnerId;
|
||||
},
|
||||
},
|
||||
async () => {
|
||||
const result = await this.computeManager.snapshot({
|
||||
runnerId: snapshot.runnerId,
|
||||
metadata: {
|
||||
runId: snapshot.runFriendlyId,
|
||||
snapshotFriendlyId: snapshot.snapshotFriendlyId,
|
||||
},
|
||||
});
|
||||
|
||||
if (!result) {
|
||||
throw new Error("Snapshot dispatch returned no result");
|
||||
}
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#emitSnapshotSpan(runFriendlyId: string, durationMs?: number, snapshotId?: string) {
|
||||
if (!this.tracing) return;
|
||||
|
||||
const ctx = this.runTraceContexts.get(runFriendlyId);
|
||||
if (!ctx) return;
|
||||
|
||||
const parsed = parseTraceparent(ctx.traceparent);
|
||||
if (!parsed) return;
|
||||
|
||||
const endEpochMs = Date.now();
|
||||
const startEpochMs = durationMs ? endEpochMs - durationMs : endEpochMs;
|
||||
|
||||
const spanAttributes: Record<string, string | number | boolean> = {
|
||||
"compute.type": "snapshot",
|
||||
};
|
||||
|
||||
if (durationMs !== undefined) {
|
||||
spanAttributes["compute.total_ms"] = durationMs;
|
||||
}
|
||||
|
||||
if (snapshotId) {
|
||||
spanAttributes["compute.snapshot_id"] = snapshotId;
|
||||
}
|
||||
|
||||
this.tracing.emit({
|
||||
traceId: parsed.traceId,
|
||||
parentSpanId: parsed.spanId,
|
||||
spanName: "compute.snapshot",
|
||||
startTimeMs: startEpochMs,
|
||||
endTimeMs: endEpochMs,
|
||||
resourceAttributes: {
|
||||
"ctx.environment.id": ctx.envId,
|
||||
"ctx.organization.id": ctx.orgId,
|
||||
"ctx.project.id": ctx.projectId,
|
||||
"ctx.run.id": runFriendlyId,
|
||||
},
|
||||
spanAttributes,
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,581 @@
|
||||
import { describe, it, expect, beforeAll, afterEach } from "vitest";
|
||||
import { FailedPodHandler } from "./failedPodHandler.js";
|
||||
import { type K8sApi, createK8sApi } from "../clients/kubernetes.js";
|
||||
import { Registry } from "prom-client";
|
||||
import { setTimeout } from "timers/promises";
|
||||
|
||||
// These tests require live K8s cluster credentials - skip by default
|
||||
describe.skipIf(!process.env.K8S_INTEGRATION_TESTS)("FailedPodHandler Integration Tests", () => {
|
||||
const k8s = createK8sApi();
|
||||
const namespace = "integration-test";
|
||||
const register = new Registry();
|
||||
|
||||
beforeAll(async () => {
|
||||
// Create the test namespace if it doesn't exist
|
||||
try {
|
||||
await k8s.core.readNamespace({ name: namespace });
|
||||
} catch (_error) {
|
||||
await k8s.core.createNamespace({
|
||||
body: {
|
||||
metadata: {
|
||||
name: namespace,
|
||||
},
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
// Clear any existing pods in the namespace
|
||||
await deleteAllPodsInNamespace({ k8sApi: k8s, namespace });
|
||||
});
|
||||
|
||||
afterEach(async () => {
|
||||
// Clear metrics to avoid conflicts
|
||||
register.clear();
|
||||
|
||||
// Delete any remaining pods in the namespace
|
||||
await deleteAllPodsInNamespace({ k8sApi: k8s, namespace });
|
||||
});
|
||||
|
||||
it("should process and delete failed pods with app=task-run label", async () => {
|
||||
const handler = new FailedPodHandler({ namespace, k8s, register });
|
||||
|
||||
try {
|
||||
// Create failed pods with the correct label
|
||||
const podNames = await createTestPods({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
count: 2,
|
||||
shouldFail: true,
|
||||
});
|
||||
|
||||
// Wait for pods to reach Failed state
|
||||
await waitForPodsPhase({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
podNames,
|
||||
phase: "Failed",
|
||||
});
|
||||
|
||||
// Start the handler
|
||||
await handler.start();
|
||||
|
||||
// Wait for pods to be deleted
|
||||
await waitForPodsDeletion({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
podNames,
|
||||
});
|
||||
|
||||
// Verify metrics
|
||||
const metrics = handler.getMetrics();
|
||||
|
||||
// Check informer events were recorded
|
||||
const informerEvents = await metrics.informerEventsTotal.get();
|
||||
expect(informerEvents.values).toContainEqual(
|
||||
expect.objectContaining({
|
||||
labels: expect.objectContaining({
|
||||
namespace,
|
||||
verb: "add",
|
||||
}),
|
||||
value: 2,
|
||||
})
|
||||
);
|
||||
expect(informerEvents.values).toContainEqual(
|
||||
expect.objectContaining({
|
||||
labels: expect.objectContaining({
|
||||
namespace,
|
||||
verb: "connect",
|
||||
}),
|
||||
value: 1,
|
||||
})
|
||||
);
|
||||
expect(informerEvents.values).not.toContainEqual(
|
||||
expect.objectContaining({
|
||||
labels: expect.objectContaining({
|
||||
namespace,
|
||||
verb: "error",
|
||||
}),
|
||||
})
|
||||
);
|
||||
|
||||
// Check pods were processed
|
||||
const processedPods = await metrics.processedPodsTotal.get();
|
||||
expect(processedPods.values).toContainEqual(
|
||||
expect.objectContaining({
|
||||
labels: expect.objectContaining({
|
||||
namespace,
|
||||
status: "Failed",
|
||||
}),
|
||||
value: 2,
|
||||
})
|
||||
);
|
||||
|
||||
// Check pods were deleted
|
||||
const deletedPods = await metrics.deletedPodsTotal.get();
|
||||
expect(deletedPods.values).toContainEqual(
|
||||
expect.objectContaining({
|
||||
labels: expect.objectContaining({
|
||||
namespace,
|
||||
status: "Failed",
|
||||
}),
|
||||
value: 2,
|
||||
})
|
||||
);
|
||||
|
||||
// Check no deletion errors were recorded
|
||||
const deletionErrors = await metrics.deletionErrorsTotal.get();
|
||||
expect(deletionErrors.values).toHaveLength(0);
|
||||
|
||||
// Check processing durations were recorded
|
||||
const durations = await metrics.processingDurationSeconds.get();
|
||||
const failedDurations = durations.values.filter(
|
||||
(v) => v.labels.namespace === namespace && v.labels.status === "Failed"
|
||||
);
|
||||
expect(failedDurations.length).toBeGreaterThan(0);
|
||||
} finally {
|
||||
await handler.stop();
|
||||
}
|
||||
}, 30000);
|
||||
|
||||
it("should ignore pods without app=task-run label", async () => {
|
||||
const handler = new FailedPodHandler({ namespace, k8s, register });
|
||||
|
||||
try {
|
||||
// Create failed pods without the task-run label
|
||||
const podNames = await createTestPods({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
count: 1,
|
||||
shouldFail: true,
|
||||
labels: { app: "not-task-run" },
|
||||
});
|
||||
|
||||
// Wait for pod to reach Failed state
|
||||
await waitForPodsPhase({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
podNames,
|
||||
phase: "Failed",
|
||||
});
|
||||
|
||||
await handler.start();
|
||||
|
||||
// Wait a reasonable time to ensure pod isn't deleted
|
||||
await setTimeout(5000);
|
||||
|
||||
// Verify pod still exists
|
||||
const exists = await podExists({ k8sApi: k8s, namespace, podName: podNames[0]! });
|
||||
expect(exists).toBe(true);
|
||||
|
||||
// Verify no metrics were recorded
|
||||
const metrics = handler.getMetrics();
|
||||
const processedPods = await metrics.processedPodsTotal.get();
|
||||
expect(processedPods.values).toHaveLength(0);
|
||||
} finally {
|
||||
await handler.stop();
|
||||
}
|
||||
}, 30000);
|
||||
|
||||
it("should not process pods that are being deleted", async () => {
|
||||
const handler = new FailedPodHandler({ namespace, k8s, register });
|
||||
|
||||
try {
|
||||
// Create a failed pod that we'll mark for deletion
|
||||
const podNames = await createTestPods({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
count: 1,
|
||||
shouldFail: true,
|
||||
command: ["/bin/sh", "-c", "sleep 30"],
|
||||
});
|
||||
|
||||
// Wait for pod to reach Failed state
|
||||
await waitForPodsPhase({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
podNames,
|
||||
phase: "Running",
|
||||
});
|
||||
|
||||
// Delete the pod but don't wait for deletion
|
||||
await k8s.core.deleteNamespacedPod({
|
||||
namespace,
|
||||
name: podNames[0]!,
|
||||
gracePeriodSeconds: 5,
|
||||
});
|
||||
|
||||
// Start the handler
|
||||
await handler.start();
|
||||
|
||||
// Wait for pod to be fully deleted
|
||||
await waitForPodsDeletion({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
podNames,
|
||||
});
|
||||
|
||||
// Verify metrics show we skipped processing
|
||||
const metrics = handler.getMetrics();
|
||||
const processedPods = await metrics.processedPodsTotal.get();
|
||||
expect(processedPods.values).toHaveLength(0);
|
||||
} finally {
|
||||
await handler.stop();
|
||||
}
|
||||
}, 30000);
|
||||
|
||||
it("should detect and process pods that fail after handler starts", async () => {
|
||||
const handler = new FailedPodHandler({ namespace, k8s, register });
|
||||
|
||||
try {
|
||||
// Start the handler
|
||||
await handler.start();
|
||||
|
||||
// Create failed pods with the correct label
|
||||
const podNames = await createTestPods({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
count: 3,
|
||||
shouldFail: true,
|
||||
});
|
||||
|
||||
// Wait for pods to be deleted
|
||||
await waitForPodsDeletion({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
podNames,
|
||||
});
|
||||
|
||||
// Verify metrics
|
||||
const metrics = handler.getMetrics();
|
||||
|
||||
// Check informer events were recorded
|
||||
const informerEvents = await metrics.informerEventsTotal.get();
|
||||
expect(informerEvents.values).toContainEqual(
|
||||
expect.objectContaining({
|
||||
labels: expect.objectContaining({
|
||||
namespace,
|
||||
verb: "add",
|
||||
}),
|
||||
value: 3,
|
||||
})
|
||||
);
|
||||
expect(informerEvents.values).toContainEqual(
|
||||
expect.objectContaining({
|
||||
labels: expect.objectContaining({
|
||||
namespace,
|
||||
verb: "connect",
|
||||
}),
|
||||
value: 1,
|
||||
})
|
||||
);
|
||||
expect(informerEvents.values).not.toContainEqual(
|
||||
expect.objectContaining({
|
||||
labels: expect.objectContaining({
|
||||
namespace,
|
||||
verb: "error",
|
||||
}),
|
||||
})
|
||||
);
|
||||
|
||||
// Check pods were processed
|
||||
const processedPods = await metrics.processedPodsTotal.get();
|
||||
expect(processedPods.values).toContainEqual(
|
||||
expect.objectContaining({
|
||||
labels: expect.objectContaining({
|
||||
namespace,
|
||||
status: "Failed",
|
||||
}),
|
||||
value: 3,
|
||||
})
|
||||
);
|
||||
|
||||
// Check pods were deleted
|
||||
const deletedPods = await metrics.deletedPodsTotal.get();
|
||||
expect(deletedPods.values).toContainEqual(
|
||||
expect.objectContaining({
|
||||
labels: expect.objectContaining({
|
||||
namespace,
|
||||
status: "Failed",
|
||||
}),
|
||||
value: 3,
|
||||
})
|
||||
);
|
||||
|
||||
// Check no deletion errors were recorded
|
||||
const deletionErrors = await metrics.deletionErrorsTotal.get();
|
||||
expect(deletionErrors.values).toHaveLength(0);
|
||||
|
||||
// Check processing durations were recorded
|
||||
const durations = await metrics.processingDurationSeconds.get();
|
||||
const failedDurations = durations.values.filter(
|
||||
(v) => v.labels.namespace === namespace && v.labels.status === "Failed"
|
||||
);
|
||||
expect(failedDurations.length).toBeGreaterThan(0);
|
||||
} finally {
|
||||
await handler.stop();
|
||||
}
|
||||
}, 60000);
|
||||
|
||||
it("should handle graceful shutdown pods differently", async () => {
|
||||
const handler = new FailedPodHandler({ namespace, k8s, register });
|
||||
|
||||
try {
|
||||
// Create first batch of pods before starting handler
|
||||
const firstBatchPodNames = await createTestPods({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
count: 2,
|
||||
exitCode: FailedPodHandler.GRACEFUL_SHUTDOWN_EXIT_CODE,
|
||||
});
|
||||
|
||||
// Wait for pods to reach Failed state
|
||||
await waitForPodsPhase({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
podNames: firstBatchPodNames,
|
||||
phase: "Failed",
|
||||
});
|
||||
|
||||
// Start the handler
|
||||
await handler.start();
|
||||
|
||||
// Wait for first batch to be deleted
|
||||
await waitForPodsDeletion({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
podNames: firstBatchPodNames,
|
||||
});
|
||||
|
||||
// Create second batch of pods after handler is running
|
||||
const secondBatchPodNames = await createTestPods({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
count: 3,
|
||||
exitCode: FailedPodHandler.GRACEFUL_SHUTDOWN_EXIT_CODE,
|
||||
});
|
||||
|
||||
// Wait for second batch to be deleted
|
||||
await waitForPodsDeletion({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
podNames: secondBatchPodNames,
|
||||
});
|
||||
|
||||
// Verify metrics
|
||||
const metrics = handler.getMetrics();
|
||||
|
||||
// Check informer events were recorded for both batches
|
||||
const informerEvents = await metrics.informerEventsTotal.get();
|
||||
expect(informerEvents.values).toContainEqual(
|
||||
expect.objectContaining({
|
||||
labels: expect.objectContaining({
|
||||
namespace,
|
||||
verb: "add",
|
||||
}),
|
||||
value: 5, // 2 from first batch + 3 from second batch
|
||||
})
|
||||
);
|
||||
|
||||
// Check pods were processed as graceful shutdowns
|
||||
const processedPods = await metrics.processedPodsTotal.get();
|
||||
|
||||
// Should not be marked as Failed
|
||||
const failedPods = processedPods.values.find(
|
||||
(v) => v.labels.namespace === namespace && v.labels.status === "Failed"
|
||||
);
|
||||
expect(failedPods).toBeUndefined();
|
||||
|
||||
// Should be marked as GracefulShutdown
|
||||
const gracefulShutdowns = processedPods.values.find(
|
||||
(v) => v.labels.namespace === namespace && v.labels.status === "GracefulShutdown"
|
||||
);
|
||||
expect(gracefulShutdowns).toBeDefined();
|
||||
expect(gracefulShutdowns?.value).toBe(5); // Total from both batches
|
||||
|
||||
// Check pods were still deleted
|
||||
const deletedPods = await metrics.deletedPodsTotal.get();
|
||||
expect(deletedPods.values).toContainEqual(
|
||||
expect.objectContaining({
|
||||
labels: expect.objectContaining({
|
||||
namespace,
|
||||
status: "Failed",
|
||||
}),
|
||||
value: 5, // Total from both batches
|
||||
})
|
||||
);
|
||||
|
||||
// Check no deletion errors were recorded
|
||||
const deletionErrors = await metrics.deletionErrorsTotal.get();
|
||||
expect(deletionErrors.values).toHaveLength(0);
|
||||
} finally {
|
||||
await handler.stop();
|
||||
}
|
||||
}, 30000);
|
||||
});
|
||||
|
||||
async function createTestPods({
|
||||
k8sApi,
|
||||
namespace,
|
||||
count,
|
||||
labels = { app: "task-run" },
|
||||
shouldFail = false,
|
||||
namePrefix = "test-pod",
|
||||
command = ["/bin/sh", "-c", shouldFail ? "exit 1" : "exit 0"],
|
||||
randomizeName = true,
|
||||
exitCode,
|
||||
}: {
|
||||
k8sApi: K8sApi;
|
||||
namespace: string;
|
||||
count: number;
|
||||
labels?: Record<string, string>;
|
||||
shouldFail?: boolean;
|
||||
namePrefix?: string;
|
||||
command?: string[];
|
||||
randomizeName?: boolean;
|
||||
exitCode?: number;
|
||||
}) {
|
||||
const createdPods: string[] = [];
|
||||
|
||||
// If exitCode is specified, override the command
|
||||
if (exitCode !== undefined) {
|
||||
command = ["/bin/sh", "-c", `exit ${exitCode}`];
|
||||
}
|
||||
|
||||
for (let i = 0; i < count; i++) {
|
||||
const podName = randomizeName
|
||||
? `${namePrefix}-${i}-${Math.random().toString(36).substring(2, 15)}`
|
||||
: `${namePrefix}-${i}`;
|
||||
await k8sApi.core.createNamespacedPod({
|
||||
namespace,
|
||||
body: {
|
||||
metadata: {
|
||||
name: podName,
|
||||
labels,
|
||||
},
|
||||
spec: {
|
||||
restartPolicy: "Never",
|
||||
containers: [
|
||||
{
|
||||
name: "run-controller", // Changed to match the name we check in failedPodHandler
|
||||
image: "busybox:1.37.0",
|
||||
command,
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
});
|
||||
createdPods.push(podName);
|
||||
}
|
||||
|
||||
return createdPods;
|
||||
}
|
||||
|
||||
async function waitForPodsDeletion({
|
||||
k8sApi,
|
||||
namespace,
|
||||
podNames,
|
||||
timeoutMs = 10000,
|
||||
waitMs = 1000,
|
||||
}: {
|
||||
k8sApi: K8sApi;
|
||||
namespace: string;
|
||||
podNames: string[];
|
||||
timeoutMs?: number;
|
||||
waitMs?: number;
|
||||
}) {
|
||||
const startTime = Date.now();
|
||||
const pendingPods = new Set(podNames);
|
||||
|
||||
while (pendingPods.size > 0 && Date.now() - startTime < timeoutMs) {
|
||||
const pods = await k8sApi.core.listNamespacedPod({ namespace });
|
||||
const existingPods = new Set(pods.items.map((pod) => pod.metadata?.name ?? ""));
|
||||
|
||||
for (const podName of pendingPods) {
|
||||
if (!existingPods.has(podName)) {
|
||||
pendingPods.delete(podName);
|
||||
}
|
||||
}
|
||||
|
||||
if (pendingPods.size > 0) {
|
||||
await setTimeout(waitMs);
|
||||
}
|
||||
}
|
||||
|
||||
if (pendingPods.size > 0) {
|
||||
throw new Error(
|
||||
`Pods [${Array.from(pendingPods).join(", ")}] were not deleted within ${timeoutMs}ms`
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
async function podExists({
|
||||
k8sApi,
|
||||
namespace,
|
||||
podName,
|
||||
}: {
|
||||
k8sApi: K8sApi;
|
||||
namespace: string;
|
||||
podName: string;
|
||||
}) {
|
||||
const pods = await k8sApi.core.listNamespacedPod({ namespace });
|
||||
return pods.items.some((p) => p.metadata?.name === podName);
|
||||
}
|
||||
|
||||
async function waitForPodsPhase({
|
||||
k8sApi,
|
||||
namespace,
|
||||
podNames,
|
||||
phase,
|
||||
timeoutMs = 10000,
|
||||
waitMs = 1000,
|
||||
}: {
|
||||
k8sApi: K8sApi;
|
||||
namespace: string;
|
||||
podNames: string[];
|
||||
phase: "Pending" | "Running" | "Succeeded" | "Failed" | "Unknown";
|
||||
timeoutMs?: number;
|
||||
waitMs?: number;
|
||||
}) {
|
||||
const startTime = Date.now();
|
||||
const pendingPods = new Set(podNames);
|
||||
|
||||
while (pendingPods.size > 0 && Date.now() - startTime < timeoutMs) {
|
||||
const pods = await k8sApi.core.listNamespacedPod({ namespace });
|
||||
|
||||
for (const pod of pods.items) {
|
||||
if (pendingPods.has(pod.metadata?.name ?? "") && pod.status?.phase === phase) {
|
||||
pendingPods.delete(pod.metadata?.name ?? "");
|
||||
}
|
||||
}
|
||||
|
||||
if (pendingPods.size > 0) {
|
||||
await setTimeout(waitMs);
|
||||
}
|
||||
}
|
||||
|
||||
if (pendingPods.size > 0) {
|
||||
throw new Error(
|
||||
`Pods [${Array.from(pendingPods).join(
|
||||
", "
|
||||
)}] did not reach phase ${phase} within ${timeoutMs}ms`
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
async function deleteAllPodsInNamespace({
|
||||
k8sApi,
|
||||
namespace,
|
||||
}: {
|
||||
k8sApi: K8sApi;
|
||||
namespace: string;
|
||||
}) {
|
||||
// Get all pods
|
||||
const pods = await k8sApi.core.listNamespacedPod({ namespace });
|
||||
const podNames = pods.items.map((p) => p.metadata?.name ?? "");
|
||||
|
||||
// Delete all pods
|
||||
await k8sApi.core.deleteCollectionNamespacedPod({ namespace, gracePeriodSeconds: 0 });
|
||||
|
||||
// Wait for all pods to be deleted
|
||||
await waitForPodsDeletion({ k8sApi, namespace, podNames });
|
||||
}
|
||||
@@ -0,0 +1,327 @@
|
||||
import type { Informer, V1Pod } from "@kubernetes/client-node";
|
||||
import { LogLevel, SimpleStructuredLogger } from "@trigger.dev/core/v3/utils/structuredLogger";
|
||||
import type { Registry } from "prom-client";
|
||||
import { Counter, Histogram } from "prom-client";
|
||||
import { setTimeout } from "timers/promises";
|
||||
import type { K8sApi } from "../clients/kubernetes.js";
|
||||
import { createK8sApi } from "../clients/kubernetes.js";
|
||||
import { register } from "../metrics.js";
|
||||
|
||||
type PodStatus = "Pending" | "Running" | "Succeeded" | "Failed" | "Unknown" | "GracefulShutdown";
|
||||
|
||||
export type FailedPodHandlerOptions = {
|
||||
namespace: string;
|
||||
reconnectIntervalMs?: number;
|
||||
k8s?: K8sApi;
|
||||
register?: Registry;
|
||||
};
|
||||
|
||||
export class FailedPodHandler {
|
||||
private readonly id: string;
|
||||
private readonly logger: SimpleStructuredLogger;
|
||||
private readonly k8s: K8sApi;
|
||||
private readonly namespace: string;
|
||||
|
||||
private isRunning = false;
|
||||
|
||||
private readonly informer: Informer<V1Pod>;
|
||||
private readonly reconnectIntervalMs: number;
|
||||
private reconnecting = false;
|
||||
|
||||
// Metrics
|
||||
private readonly register: Registry;
|
||||
private readonly processedPodsTotal: Counter;
|
||||
private readonly deletedPodsTotal: Counter;
|
||||
private readonly deletionErrorsTotal: Counter;
|
||||
private readonly processingDurationSeconds: Histogram<string>;
|
||||
private readonly informerEventsTotal: Counter;
|
||||
|
||||
static readonly GRACEFUL_SHUTDOWN_EXIT_CODE = 200;
|
||||
|
||||
constructor(opts: FailedPodHandlerOptions) {
|
||||
this.id = Math.random().toString(36).substring(2, 15);
|
||||
this.logger = new SimpleStructuredLogger("failed-pod-handler", LogLevel.debug, {
|
||||
id: this.id,
|
||||
});
|
||||
|
||||
this.k8s = opts.k8s ?? createK8sApi();
|
||||
|
||||
this.namespace = opts.namespace;
|
||||
this.reconnectIntervalMs = opts.reconnectIntervalMs ?? 1000;
|
||||
|
||||
this.informer = this.k8s.makeInformer(
|
||||
`/api/v1/namespaces/${this.namespace}/pods`,
|
||||
() =>
|
||||
this.k8s.core.listNamespacedPod({
|
||||
namespace: this.namespace,
|
||||
labelSelector: "app=task-run",
|
||||
fieldSelector: "status.phase=Failed",
|
||||
}),
|
||||
"app=task-run",
|
||||
"status.phase=Failed"
|
||||
);
|
||||
|
||||
// Whenever a matching pod is added to the informer cache
|
||||
this.informer.on("add", this.onPodCompleted.bind(this));
|
||||
|
||||
// Informer events
|
||||
this.informer.on("connect", this.makeOnConnect("failed-pod-informer").bind(this));
|
||||
this.informer.on("error", this.makeOnError("failed-pod-informer").bind(this));
|
||||
|
||||
// Initialize metrics
|
||||
this.register = opts.register ?? register;
|
||||
|
||||
this.processedPodsTotal = new Counter({
|
||||
name: "failed_pod_handler_processed_pods_total",
|
||||
help: "Total number of failed pods processed",
|
||||
labelNames: ["namespace", "status"],
|
||||
registers: [this.register],
|
||||
});
|
||||
|
||||
this.deletedPodsTotal = new Counter({
|
||||
name: "failed_pod_handler_deleted_pods_total",
|
||||
help: "Total number of pods deleted",
|
||||
labelNames: ["namespace", "status"],
|
||||
registers: [this.register],
|
||||
});
|
||||
|
||||
this.deletionErrorsTotal = new Counter({
|
||||
name: "failed_pod_handler_deletion_errors_total",
|
||||
help: "Total number of errors encountered while deleting pods",
|
||||
labelNames: ["namespace", "error_type"],
|
||||
registers: [this.register],
|
||||
});
|
||||
|
||||
this.processingDurationSeconds = new Histogram({
|
||||
name: "failed_pod_handler_processing_duration_seconds",
|
||||
help: "The duration of pod processing",
|
||||
labelNames: ["namespace", "status"],
|
||||
registers: [this.register],
|
||||
});
|
||||
|
||||
this.informerEventsTotal = new Counter({
|
||||
name: "failed_pod_handler_informer_events_total",
|
||||
help: "Total number of informer events",
|
||||
labelNames: ["namespace", "verb"],
|
||||
registers: [this.register],
|
||||
});
|
||||
}
|
||||
|
||||
async start() {
|
||||
if (this.isRunning) {
|
||||
this.logger.warn("failed pod handler already running");
|
||||
return;
|
||||
}
|
||||
|
||||
this.isRunning = true;
|
||||
|
||||
this.logger.info("starting failed pod handler");
|
||||
await this.informer.start();
|
||||
}
|
||||
|
||||
async stop() {
|
||||
if (!this.isRunning) {
|
||||
this.logger.warn("failed pod handler not running");
|
||||
return;
|
||||
}
|
||||
|
||||
this.isRunning = false;
|
||||
|
||||
this.logger.info("stopping failed pod handler");
|
||||
await this.informer.stop();
|
||||
}
|
||||
|
||||
private async withHistogram<T>(
|
||||
histogram: Histogram<string>,
|
||||
promise: Promise<T>,
|
||||
labels?: Record<string, string>
|
||||
): Promise<T> {
|
||||
const end = histogram.startTimer({ namespace: this.namespace, ...labels });
|
||||
try {
|
||||
return await promise;
|
||||
} finally {
|
||||
end();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns the non-nullable status of a pod
|
||||
*/
|
||||
private podStatus(pod: V1Pod): PodStatus {
|
||||
return (pod.status?.phase ?? "Unknown") as PodStatus;
|
||||
}
|
||||
|
||||
private async onPodCompleted(pod: V1Pod) {
|
||||
this.logger.debug("pod-completed", this.podSummary(pod));
|
||||
this.informerEventsTotal.inc({ namespace: this.namespace, verb: "add" });
|
||||
|
||||
if (!pod.metadata?.name) {
|
||||
this.logger.error("pod-completed: no name", this.podSummary(pod));
|
||||
return;
|
||||
}
|
||||
|
||||
if (!pod.status) {
|
||||
this.logger.error("pod-completed: no status", this.podSummary(pod));
|
||||
return;
|
||||
}
|
||||
|
||||
if (pod.metadata?.deletionTimestamp) {
|
||||
this.logger.verbose("pod-completed: pod is being deleted", this.podSummary(pod));
|
||||
return;
|
||||
}
|
||||
|
||||
const podStatus = this.podStatus(pod);
|
||||
|
||||
switch (podStatus) {
|
||||
case "Succeeded":
|
||||
await this.withHistogram(this.processingDurationSeconds, this.onPodSucceeded(pod), {
|
||||
status: podStatus,
|
||||
});
|
||||
break;
|
||||
case "Failed":
|
||||
await this.withHistogram(this.processingDurationSeconds, this.onPodFailed(pod), {
|
||||
status: podStatus,
|
||||
});
|
||||
break;
|
||||
default:
|
||||
this.logger.error("pod-completed: unknown phase", this.podSummary(pod));
|
||||
}
|
||||
}
|
||||
|
||||
private async onPodSucceeded(pod: V1Pod) {
|
||||
this.logger.debug("pod-succeeded", this.podSummary(pod));
|
||||
this.processedPodsTotal.inc({
|
||||
namespace: this.namespace,
|
||||
status: this.podStatus(pod),
|
||||
});
|
||||
}
|
||||
|
||||
private async onPodFailed(pod: V1Pod) {
|
||||
this.logger.debug("pod-failed", this.podSummary(pod));
|
||||
|
||||
try {
|
||||
await this.processFailedPod(pod);
|
||||
} catch (error) {
|
||||
this.logger.error("pod-failed: error processing pod", this.podSummary(pod), { error });
|
||||
} finally {
|
||||
await this.deletePod(pod);
|
||||
}
|
||||
}
|
||||
|
||||
private async processFailedPod(pod: V1Pod) {
|
||||
this.logger.verbose("pod-failed: processing pod", this.podSummary(pod));
|
||||
|
||||
const mainContainer = pod.status?.containerStatuses?.find((c) => c.name === "run-controller");
|
||||
|
||||
// If it's our special "graceful shutdown" exit code, don't process it further, just delete it
|
||||
if (
|
||||
mainContainer?.state?.terminated?.exitCode === FailedPodHandler.GRACEFUL_SHUTDOWN_EXIT_CODE
|
||||
) {
|
||||
this.logger.debug("pod-failed: graceful shutdown detected", this.podSummary(pod));
|
||||
this.processedPodsTotal.inc({
|
||||
namespace: this.namespace,
|
||||
status: "GracefulShutdown",
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
this.processedPodsTotal.inc({
|
||||
namespace: this.namespace,
|
||||
status: this.podStatus(pod),
|
||||
});
|
||||
}
|
||||
|
||||
private async deletePod(pod: V1Pod) {
|
||||
this.logger.verbose("pod-failed: deleting pod", this.podSummary(pod));
|
||||
try {
|
||||
await this.k8s.core.deleteNamespacedPod({
|
||||
name: pod.metadata!.name!,
|
||||
namespace: this.namespace,
|
||||
});
|
||||
this.deletedPodsTotal.inc({
|
||||
namespace: this.namespace,
|
||||
status: this.podStatus(pod),
|
||||
});
|
||||
} catch (error) {
|
||||
this.logger.error("pod-failed: error deleting pod", this.podSummary(pod), { error });
|
||||
this.deletionErrorsTotal.inc({
|
||||
namespace: this.namespace,
|
||||
error_type: error instanceof Error ? error.name : "unknown",
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
private makeOnError(informerName: string) {
|
||||
return (err?: unknown) => this.onError(informerName, err);
|
||||
}
|
||||
|
||||
private async onError(informerName: string, err?: unknown) {
|
||||
if (!this.isRunning) {
|
||||
this.logger.warn("onError: informer not running");
|
||||
return;
|
||||
}
|
||||
|
||||
// Guard against multiple simultaneous reconnections
|
||||
if (this.reconnecting) {
|
||||
this.logger.debug("onError: reconnection already in progress, skipping", {
|
||||
informerName,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
this.reconnecting = true;
|
||||
|
||||
try {
|
||||
const error = err instanceof Error ? err : undefined;
|
||||
this.logger.error("error event fired", {
|
||||
informerName,
|
||||
error: error?.message,
|
||||
errorType: error?.name,
|
||||
});
|
||||
this.informerEventsTotal.inc({ namespace: this.namespace, verb: "error" });
|
||||
|
||||
// Reconnect on errors
|
||||
await setTimeout(this.reconnectIntervalMs);
|
||||
await this.informer.start();
|
||||
} catch (handlerError) {
|
||||
const error = handlerError instanceof Error ? handlerError : undefined;
|
||||
this.logger.error("onError: reconnection attempt failed", {
|
||||
informerName,
|
||||
error: error?.message,
|
||||
errorType: error?.name,
|
||||
errorStack: error?.stack,
|
||||
});
|
||||
} finally {
|
||||
this.reconnecting = false;
|
||||
}
|
||||
}
|
||||
|
||||
private makeOnConnect(informerName: string) {
|
||||
return () => this.onConnect(informerName);
|
||||
}
|
||||
|
||||
private async onConnect(informerName: string) {
|
||||
this.logger.info(`informer connected: ${informerName}`);
|
||||
this.informerEventsTotal.inc({ namespace: this.namespace, verb: "connect" });
|
||||
}
|
||||
|
||||
private podSummary(pod: V1Pod) {
|
||||
return {
|
||||
name: pod.metadata?.name,
|
||||
namespace: pod.metadata?.namespace,
|
||||
status: pod.status?.phase,
|
||||
deletionTimestamp: pod.metadata?.deletionTimestamp,
|
||||
};
|
||||
}
|
||||
|
||||
// Method to expose metrics for testing
|
||||
public getMetrics() {
|
||||
return {
|
||||
processedPodsTotal: this.processedPodsTotal,
|
||||
deletedPodsTotal: this.deletedPodsTotal,
|
||||
deletionErrorsTotal: this.deletionErrorsTotal,
|
||||
informerEventsTotal: this.informerEventsTotal,
|
||||
processingDurationSeconds: this.processingDurationSeconds,
|
||||
};
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,177 @@
|
||||
import { describe, it, expect } from "vitest";
|
||||
import { buildPayload } from "./otlpTraceService.js";
|
||||
|
||||
describe("buildPayload", () => {
|
||||
it("builds valid OTLP JSON with timing attributes", () => {
|
||||
const payload = buildPayload({
|
||||
traceId: "abcd1234abcd1234abcd1234abcd1234",
|
||||
parentSpanId: "1234567890abcdef",
|
||||
spanName: "compute.provision",
|
||||
startTimeMs: 1000,
|
||||
endTimeMs: 1250,
|
||||
resourceAttributes: {
|
||||
"ctx.environment.id": "env_123",
|
||||
"ctx.organization.id": "org_456",
|
||||
"ctx.project.id": "proj_789",
|
||||
"ctx.run.id": "run_abc",
|
||||
},
|
||||
spanAttributes: {
|
||||
"compute.total_ms": 250,
|
||||
"compute.gateway.schedule_ms": 1,
|
||||
"compute.cache.image_cached": true,
|
||||
},
|
||||
});
|
||||
|
||||
expect(payload.resourceSpans).toHaveLength(1);
|
||||
|
||||
const resourceSpan = payload.resourceSpans[0]!;
|
||||
|
||||
// $trigger=true so the webapp accepts it
|
||||
const triggerAttr = resourceSpan.resource.attributes.find((a) => a.key === "$trigger");
|
||||
expect(triggerAttr).toEqual({ key: "$trigger", value: { boolValue: true } });
|
||||
|
||||
// Resource attributes
|
||||
const envAttr = resourceSpan.resource.attributes.find((a) => a.key === "ctx.environment.id");
|
||||
expect(envAttr).toEqual({
|
||||
key: "ctx.environment.id",
|
||||
value: { stringValue: "env_123" },
|
||||
});
|
||||
|
||||
// Span basics
|
||||
const span = resourceSpan.scopeSpans[0]!.spans[0]!;
|
||||
expect(span.name).toBe("compute.provision");
|
||||
expect(span.traceId).toBe("abcd1234abcd1234abcd1234abcd1234");
|
||||
expect(span.parentSpanId).toBe("1234567890abcdef");
|
||||
|
||||
// Integer attribute
|
||||
const totalMs = span.attributes.find((a) => a.key === "compute.total_ms");
|
||||
expect(totalMs).toEqual({ key: "compute.total_ms", value: { intValue: 250 } });
|
||||
|
||||
// Boolean attribute
|
||||
const cached = span.attributes.find((a) => a.key === "compute.cache.image_cached");
|
||||
expect(cached).toEqual({ key: "compute.cache.image_cached", value: { boolValue: true } });
|
||||
});
|
||||
|
||||
it("generates a valid 16-char hex span ID", () => {
|
||||
const payload = buildPayload({
|
||||
traceId: "abcd1234abcd1234abcd1234abcd1234",
|
||||
spanName: "test",
|
||||
startTimeMs: 1000,
|
||||
endTimeMs: 1001,
|
||||
resourceAttributes: {},
|
||||
spanAttributes: {},
|
||||
});
|
||||
|
||||
const span = payload.resourceSpans[0]!.scopeSpans[0]!.spans[0]!;
|
||||
expect(span.spanId).toMatch(/^[0-9a-f]{16}$/);
|
||||
});
|
||||
|
||||
it("converts timestamps to nanoseconds", () => {
|
||||
const payload = buildPayload({
|
||||
traceId: "abcd1234abcd1234abcd1234abcd1234",
|
||||
spanName: "test",
|
||||
startTimeMs: 1000,
|
||||
endTimeMs: 1250,
|
||||
resourceAttributes: {},
|
||||
spanAttributes: {},
|
||||
});
|
||||
|
||||
const span = payload.resourceSpans[0]!.scopeSpans[0]!.spans[0]!;
|
||||
expect(span.startTimeUnixNano).toBe("1000000000");
|
||||
expect(span.endTimeUnixNano).toBe("1250000000");
|
||||
});
|
||||
|
||||
it("converts real epoch timestamps without precision loss", () => {
|
||||
// Date.now() values exceed Number.MAX_SAFE_INTEGER when multiplied by 1e6
|
||||
const startMs = 1711929600000; // 2024-04-01T00:00:00Z
|
||||
const endMs = 1711929600250;
|
||||
|
||||
const payload = buildPayload({
|
||||
traceId: "abcd1234abcd1234abcd1234abcd1234",
|
||||
spanName: "test",
|
||||
startTimeMs: startMs,
|
||||
endTimeMs: endMs,
|
||||
resourceAttributes: {},
|
||||
spanAttributes: {},
|
||||
});
|
||||
|
||||
const span = payload.resourceSpans[0]!.scopeSpans[0]!.spans[0]!;
|
||||
expect(span.startTimeUnixNano).toBe("1711929600000000000");
|
||||
expect(span.endTimeUnixNano).toBe("1711929600250000000");
|
||||
});
|
||||
|
||||
it("preserves sub-millisecond precision from performance.now() arithmetic", () => {
|
||||
// provisionStartEpochMs = Date.now() - (performance.now() - startMs) produces fractional ms.
|
||||
// Use small epoch + fraction to avoid IEEE 754 noise in the fractional part.
|
||||
const startMs = 1000.322;
|
||||
const endMs = 1045.789;
|
||||
|
||||
const payload = buildPayload({
|
||||
traceId: "abcd1234abcd1234abcd1234abcd1234",
|
||||
spanName: "test",
|
||||
startTimeMs: startMs,
|
||||
endTimeMs: endMs,
|
||||
resourceAttributes: {},
|
||||
spanAttributes: {},
|
||||
});
|
||||
|
||||
const span = payload.resourceSpans[0]!.scopeSpans[0]!.spans[0]!;
|
||||
expect(span.startTimeUnixNano).toBe("1000322000");
|
||||
expect(span.endTimeUnixNano).toBe("1045789000");
|
||||
});
|
||||
|
||||
it("sub-ms precision affects ordering for real epoch values", () => {
|
||||
// Two spans within the same millisecond should have different nanosecond timestamps
|
||||
const spanA = buildPayload({
|
||||
traceId: "abcd1234abcd1234abcd1234abcd1234",
|
||||
spanName: "a",
|
||||
startTimeMs: 1711929600000.3,
|
||||
endTimeMs: 1711929600001,
|
||||
resourceAttributes: {},
|
||||
spanAttributes: {},
|
||||
});
|
||||
|
||||
const spanB = buildPayload({
|
||||
traceId: "abcd1234abcd1234abcd1234abcd1234",
|
||||
spanName: "b",
|
||||
startTimeMs: 1711929600000.7,
|
||||
endTimeMs: 1711929600001,
|
||||
resourceAttributes: {},
|
||||
spanAttributes: {},
|
||||
});
|
||||
|
||||
const startA = BigInt(spanA.resourceSpans[0]!.scopeSpans[0]!.spans[0]!.startTimeUnixNano);
|
||||
const startB = BigInt(spanB.resourceSpans[0]!.scopeSpans[0]!.spans[0]!.startTimeUnixNano);
|
||||
// A should sort before B (both in the same ms but different sub-ms positions)
|
||||
expect(startA).toBeLessThan(startB);
|
||||
});
|
||||
|
||||
it("omits parentSpanId when not provided", () => {
|
||||
const payload = buildPayload({
|
||||
traceId: "abcd1234abcd1234abcd1234abcd1234",
|
||||
spanName: "test",
|
||||
startTimeMs: 1000,
|
||||
endTimeMs: 1001,
|
||||
resourceAttributes: {},
|
||||
spanAttributes: {},
|
||||
});
|
||||
|
||||
const span = payload.resourceSpans[0]!.scopeSpans[0]!.spans[0]!;
|
||||
expect(span.parentSpanId).toBeUndefined();
|
||||
});
|
||||
|
||||
it("handles double values for non-integer numbers", () => {
|
||||
const payload = buildPayload({
|
||||
traceId: "abcd1234abcd1234abcd1234abcd1234",
|
||||
spanName: "test",
|
||||
startTimeMs: 1000,
|
||||
endTimeMs: 1001,
|
||||
resourceAttributes: {},
|
||||
spanAttributes: { "compute.cpu": 0.25 },
|
||||
});
|
||||
|
||||
const span = payload.resourceSpans[0]!.scopeSpans[0]!.spans[0]!;
|
||||
const cpu = span.attributes.find((a) => a.key === "compute.cpu");
|
||||
expect(cpu).toEqual({ key: "compute.cpu", value: { doubleValue: 0.25 } });
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,104 @@
|
||||
import { randomBytes } from "crypto";
|
||||
import { SimpleStructuredLogger } from "@trigger.dev/core/v3/utils/structuredLogger";
|
||||
|
||||
export type OtlpTraceServiceOptions = {
|
||||
endpointUrl: string;
|
||||
timeoutMs?: number;
|
||||
};
|
||||
|
||||
export type OtlpTraceSpan = {
|
||||
traceId: string;
|
||||
parentSpanId?: string;
|
||||
spanName: string;
|
||||
startTimeMs: number;
|
||||
endTimeMs: number;
|
||||
resourceAttributes: Record<string, string | number | boolean>;
|
||||
spanAttributes: Record<string, string | number | boolean>;
|
||||
};
|
||||
|
||||
export class OtlpTraceService {
|
||||
private readonly logger = new SimpleStructuredLogger("otlp-trace");
|
||||
|
||||
constructor(private opts: OtlpTraceServiceOptions) {}
|
||||
|
||||
/** Fire-and-forget: build payload and send to the configured OTLP endpoint */
|
||||
emit(span: OtlpTraceSpan): void {
|
||||
const payload = buildPayload(span);
|
||||
|
||||
fetch(`${this.opts.endpointUrl}/v1/traces`, {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify(payload),
|
||||
signal: AbortSignal.timeout(this.opts.timeoutMs ?? 5_000),
|
||||
}).catch((err) => {
|
||||
this.logger.warn("failed to send compute trace span", {
|
||||
error: err instanceof Error ? err.message : String(err),
|
||||
});
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// ── Payload builder (internal) ───────────────────────────────────────────────
|
||||
|
||||
/** @internal Exported for tests only */
|
||||
export function buildPayload(span: OtlpTraceSpan) {
|
||||
const spanId = randomBytes(8).toString("hex");
|
||||
|
||||
return {
|
||||
resourceSpans: [
|
||||
{
|
||||
resource: {
|
||||
attributes: [
|
||||
{ key: "$trigger", value: { boolValue: true } },
|
||||
...toOtlpAttributes(span.resourceAttributes),
|
||||
],
|
||||
},
|
||||
scopeSpans: [
|
||||
{
|
||||
scope: { name: "supervisor.compute" },
|
||||
spans: [
|
||||
{
|
||||
traceId: span.traceId,
|
||||
spanId,
|
||||
parentSpanId: span.parentSpanId,
|
||||
name: span.spanName,
|
||||
kind: 3, // SPAN_KIND_CLIENT
|
||||
startTimeUnixNano: msToNano(span.startTimeMs),
|
||||
endTimeUnixNano: msToNano(span.endTimeMs),
|
||||
attributes: toOtlpAttributes(span.spanAttributes),
|
||||
status: { code: 1 }, // STATUS_CODE_OK
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
};
|
||||
}
|
||||
|
||||
function toOtlpAttributes(
|
||||
attrs: Record<string, string | number | boolean>
|
||||
): Array<{ key: string; value: Record<string, unknown> }> {
|
||||
return Object.entries(attrs).map(([key, value]) => ({
|
||||
key,
|
||||
value: toOtlpValue(value),
|
||||
}));
|
||||
}
|
||||
|
||||
function toOtlpValue(value: string | number | boolean): Record<string, unknown> {
|
||||
if (typeof value === "string") return { stringValue: value };
|
||||
if (typeof value === "boolean") return { boolValue: value };
|
||||
if (Number.isInteger(value)) return { intValue: value };
|
||||
return { doubleValue: value };
|
||||
}
|
||||
|
||||
/**
|
||||
* Convert epoch milliseconds to nanosecond string, preserving sub-ms precision.
|
||||
* Fractional ms from performance.now() arithmetic carry meaningful microsecond
|
||||
* data that affects span sort ordering when events happen within the same ms.
|
||||
*/
|
||||
function msToNano(ms: number): string {
|
||||
const wholeMs = Math.trunc(ms);
|
||||
const fracNs = Math.round((ms - wholeMs) * 1_000_000);
|
||||
return String(BigInt(wholeMs) * 1_000_000n + BigInt(fracNs));
|
||||
}
|
||||
@@ -0,0 +1,473 @@
|
||||
import { PodCleaner } from "./podCleaner.js";
|
||||
import { type K8sApi, createK8sApi } from "../clients/kubernetes.js";
|
||||
import { setTimeout } from "timers/promises";
|
||||
import { describe, it, expect, beforeAll, afterEach } from "vitest";
|
||||
import { Registry } from "prom-client";
|
||||
|
||||
// These tests require live K8s cluster credentials - skip by default
|
||||
describe.skipIf(!process.env.K8S_INTEGRATION_TESTS)("PodCleaner Integration Tests", () => {
|
||||
const k8s = createK8sApi();
|
||||
const namespace = "integration-test";
|
||||
const register = new Registry();
|
||||
|
||||
beforeAll(async () => {
|
||||
// Create the test namespace, only if it doesn't exist
|
||||
try {
|
||||
await k8s.core.readNamespace({ name: namespace });
|
||||
} catch (_error) {
|
||||
await k8s.core.createNamespace({
|
||||
body: {
|
||||
metadata: {
|
||||
name: namespace,
|
||||
},
|
||||
},
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
afterEach(async () => {
|
||||
// Clear metrics to avoid conflicts
|
||||
register.clear();
|
||||
|
||||
// Delete all pods in the namespace
|
||||
await k8s.core.deleteCollectionNamespacedPod({ namespace, gracePeriodSeconds: 0 });
|
||||
});
|
||||
|
||||
it("should clean up succeeded pods", async () => {
|
||||
const podCleaner = new PodCleaner({ namespace, k8s, register });
|
||||
|
||||
try {
|
||||
// Create a test pod that's in succeeded state
|
||||
const podNames = await createTestPods({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
count: 1,
|
||||
namePrefix: "test-succeeded-pod",
|
||||
});
|
||||
|
||||
if (!podNames[0]) {
|
||||
throw new Error("Failed to create test pod");
|
||||
}
|
||||
const podName = podNames[0];
|
||||
|
||||
// Wait for pod to complete
|
||||
await waitForPodPhase({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
podName,
|
||||
phase: "Succeeded",
|
||||
});
|
||||
|
||||
// Start the pod cleaner
|
||||
await podCleaner.start();
|
||||
|
||||
// Wait for pod to be deleted
|
||||
await waitForPodDeletion({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
podName,
|
||||
});
|
||||
|
||||
// Verify pod was deleted
|
||||
expect(await podExists({ k8sApi: k8s, namespace, podName })).toBe(false);
|
||||
} finally {
|
||||
await podCleaner.stop();
|
||||
}
|
||||
}, 30000);
|
||||
|
||||
it("should accurately track deletion metrics", async () => {
|
||||
const podCleaner = new PodCleaner({ namespace, k8s, register });
|
||||
try {
|
||||
// Create a test pod that's in succeeded state
|
||||
const podNames = await createTestPods({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
count: 1,
|
||||
namePrefix: "test-succeeded-pod",
|
||||
});
|
||||
|
||||
// Wait for pod to be in succeeded state
|
||||
await waitForPodsPhase({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
podNames,
|
||||
phase: "Succeeded",
|
||||
});
|
||||
|
||||
await podCleaner.start();
|
||||
|
||||
// Wait for pod to be deleted
|
||||
await waitForPodsDeletion({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
podNames,
|
||||
});
|
||||
|
||||
const metrics = podCleaner.getMetrics();
|
||||
const deletionCycles = await metrics.deletionCyclesTotal.get();
|
||||
const deletionTimestamp = await metrics.lastDeletionTimestamp.get();
|
||||
|
||||
expect(deletionCycles?.values[0]?.value).toBeGreaterThan(0);
|
||||
expect(deletionTimestamp?.values[0]?.value).toBeGreaterThan(0);
|
||||
} finally {
|
||||
await podCleaner.stop();
|
||||
}
|
||||
}, 30000);
|
||||
|
||||
it("should handle different batch sizes - small", async () => {
|
||||
const podCleaner = new PodCleaner({
|
||||
namespace,
|
||||
k8s,
|
||||
register,
|
||||
batchSize: 1,
|
||||
});
|
||||
|
||||
try {
|
||||
// Create some pods that will succeed
|
||||
const podNames = await createTestPods({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
count: 2,
|
||||
});
|
||||
|
||||
await waitForPodsPhase({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
podNames,
|
||||
phase: "Succeeded",
|
||||
});
|
||||
|
||||
await podCleaner.start();
|
||||
|
||||
await waitForPodsDeletion({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
podNames,
|
||||
});
|
||||
|
||||
const metrics = podCleaner.getMetrics();
|
||||
const cycles = await metrics.deletionCyclesTotal.get();
|
||||
|
||||
expect(cycles?.values[0]?.value).toBe(2);
|
||||
} finally {
|
||||
await podCleaner.stop();
|
||||
}
|
||||
}, 30000);
|
||||
|
||||
it("should handle different batch sizes - large", async () => {
|
||||
const podCleaner = new PodCleaner({
|
||||
namespace,
|
||||
k8s,
|
||||
register,
|
||||
batchSize: 5000,
|
||||
});
|
||||
|
||||
try {
|
||||
// Create some pods that will succeed
|
||||
const podNames = await createTestPods({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
count: 10,
|
||||
});
|
||||
|
||||
await waitForPodsPhase({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
podNames,
|
||||
phase: "Succeeded",
|
||||
});
|
||||
|
||||
await podCleaner.start();
|
||||
|
||||
await waitForPodsDeletion({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
podNames,
|
||||
});
|
||||
|
||||
const metrics = podCleaner.getMetrics();
|
||||
const cycles = await metrics.deletionCyclesTotal.get();
|
||||
|
||||
expect(cycles?.values[0]?.value).toBe(1);
|
||||
} finally {
|
||||
await podCleaner.stop();
|
||||
}
|
||||
}, 30000);
|
||||
|
||||
it("should not delete pods without app=task-run label", async () => {
|
||||
const podCleaner = new PodCleaner({ namespace, k8s, register });
|
||||
|
||||
try {
|
||||
// Create a test pod without the task-run label
|
||||
const podNames = await createTestPods({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
count: 1,
|
||||
labels: { app: "different-label" },
|
||||
namePrefix: "non-task-run-pod",
|
||||
});
|
||||
|
||||
if (!podNames[0]) {
|
||||
throw new Error("Failed to create test pod");
|
||||
}
|
||||
const podName = podNames[0];
|
||||
|
||||
// Wait for pod to complete
|
||||
await waitForPodPhase({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
podName,
|
||||
phase: "Succeeded",
|
||||
});
|
||||
|
||||
await podCleaner.start();
|
||||
|
||||
// Wait a reasonable time to ensure pod isn't deleted
|
||||
await setTimeout(5000);
|
||||
|
||||
// Verify pod still exists
|
||||
expect(await podExists({ k8sApi: k8s, namespace, podName })).toBe(true);
|
||||
} finally {
|
||||
await podCleaner.stop();
|
||||
}
|
||||
}, 30000);
|
||||
|
||||
it("should not delete pods that are still running", async () => {
|
||||
const podCleaner = new PodCleaner({ namespace, k8s, register });
|
||||
|
||||
try {
|
||||
// Create a test pod with a long-running command
|
||||
const podNames = await createTestPods({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
count: 1,
|
||||
namePrefix: "running-pod",
|
||||
command: ["sleep", "30"], // Will keep pod running
|
||||
});
|
||||
|
||||
if (!podNames[0]) {
|
||||
throw new Error("Failed to create test pod");
|
||||
}
|
||||
const podName = podNames[0];
|
||||
|
||||
// Wait for pod to be running
|
||||
await waitForPodPhase({
|
||||
k8sApi: k8s,
|
||||
namespace,
|
||||
podName,
|
||||
phase: "Running",
|
||||
});
|
||||
|
||||
await podCleaner.start();
|
||||
|
||||
// Wait a reasonable time to ensure pod isn't deleted
|
||||
await setTimeout(5000);
|
||||
|
||||
// Verify pod still exists
|
||||
expect(await podExists({ k8sApi: k8s, namespace, podName })).toBe(true);
|
||||
} finally {
|
||||
await podCleaner.stop();
|
||||
}
|
||||
}, 30000);
|
||||
});
|
||||
|
||||
// Helper functions
|
||||
async function waitForPodPhase({
|
||||
k8sApi,
|
||||
namespace,
|
||||
podName,
|
||||
phase,
|
||||
timeoutMs = 10000,
|
||||
waitMs = 1000,
|
||||
}: {
|
||||
k8sApi: K8sApi;
|
||||
namespace: string;
|
||||
podName: string;
|
||||
phase: string;
|
||||
timeoutMs?: number;
|
||||
waitMs?: number;
|
||||
}) {
|
||||
const startTime = Date.now();
|
||||
|
||||
while (Date.now() - startTime < timeoutMs) {
|
||||
const pod = await k8sApi.core.readNamespacedPod({
|
||||
namespace,
|
||||
name: podName,
|
||||
});
|
||||
if (pod.status?.phase === phase) {
|
||||
return;
|
||||
}
|
||||
await setTimeout(waitMs);
|
||||
}
|
||||
|
||||
throw new Error(`Pod ${podName} did not reach phase ${phase} within ${timeoutMs}ms`);
|
||||
}
|
||||
|
||||
async function waitForPodDeletion({
|
||||
k8sApi,
|
||||
namespace,
|
||||
podName,
|
||||
timeoutMs = 10000,
|
||||
waitMs = 1000,
|
||||
}: {
|
||||
k8sApi: K8sApi;
|
||||
namespace: string;
|
||||
podName: string;
|
||||
timeoutMs?: number;
|
||||
waitMs?: number;
|
||||
}) {
|
||||
const startTime = Date.now();
|
||||
|
||||
while (Date.now() - startTime < timeoutMs) {
|
||||
try {
|
||||
await k8sApi.core.readNamespacedPod({
|
||||
namespace,
|
||||
name: podName,
|
||||
});
|
||||
await setTimeout(waitMs);
|
||||
} catch (_error) {
|
||||
// Pod was deleted
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
throw new Error(`Pod ${podName} was not deleted within ${timeoutMs}ms`);
|
||||
}
|
||||
|
||||
async function createTestPods({
|
||||
k8sApi,
|
||||
namespace,
|
||||
count,
|
||||
labels = { app: "task-run" },
|
||||
shouldFail = false,
|
||||
namePrefix = "test-pod",
|
||||
command = ["/bin/sh", "-c", shouldFail ? "exit 1" : "exit 0"],
|
||||
}: {
|
||||
k8sApi: K8sApi;
|
||||
namespace: string;
|
||||
count: number;
|
||||
labels?: Record<string, string>;
|
||||
shouldFail?: boolean;
|
||||
namePrefix?: string;
|
||||
command?: string[];
|
||||
}) {
|
||||
const createdPods: string[] = [];
|
||||
|
||||
for (let i = 0; i < count; i++) {
|
||||
const podName = `${namePrefix}-${i}`;
|
||||
await k8sApi.core.createNamespacedPod({
|
||||
namespace,
|
||||
body: {
|
||||
metadata: {
|
||||
name: podName,
|
||||
labels,
|
||||
},
|
||||
spec: {
|
||||
restartPolicy: "Never",
|
||||
containers: [
|
||||
{
|
||||
name: "test",
|
||||
image: "busybox:1.37.0",
|
||||
command,
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
});
|
||||
createdPods.push(podName);
|
||||
}
|
||||
|
||||
return createdPods;
|
||||
}
|
||||
|
||||
async function waitForPodsPhase({
|
||||
k8sApi,
|
||||
namespace,
|
||||
podNames,
|
||||
phase,
|
||||
timeoutMs = 10000,
|
||||
waitMs = 1000,
|
||||
}: {
|
||||
k8sApi: K8sApi;
|
||||
namespace: string;
|
||||
podNames: string[];
|
||||
phase: "Pending" | "Running" | "Succeeded" | "Failed" | "Unknown";
|
||||
timeoutMs?: number;
|
||||
waitMs?: number;
|
||||
}) {
|
||||
const startTime = Date.now();
|
||||
const pendingPods = new Set(podNames);
|
||||
|
||||
while (pendingPods.size > 0 && Date.now() - startTime < timeoutMs) {
|
||||
const pods = await k8sApi.core.listNamespacedPod({ namespace });
|
||||
|
||||
for (const pod of pods.items) {
|
||||
if (pendingPods.has(pod.metadata?.name ?? "") && pod.status?.phase === phase) {
|
||||
pendingPods.delete(pod.metadata?.name ?? "");
|
||||
}
|
||||
}
|
||||
|
||||
if (pendingPods.size > 0) {
|
||||
await setTimeout(waitMs);
|
||||
}
|
||||
}
|
||||
|
||||
if (pendingPods.size > 0) {
|
||||
throw new Error(
|
||||
`Pods [${Array.from(pendingPods).join(
|
||||
", "
|
||||
)}] did not reach phase ${phase} within ${timeoutMs}ms`
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
async function waitForPodsDeletion({
|
||||
k8sApi,
|
||||
namespace,
|
||||
podNames,
|
||||
timeoutMs = 10000,
|
||||
waitMs = 1000,
|
||||
}: {
|
||||
k8sApi: K8sApi;
|
||||
namespace: string;
|
||||
podNames: string[];
|
||||
timeoutMs?: number;
|
||||
waitMs?: number;
|
||||
}) {
|
||||
const startTime = Date.now();
|
||||
const pendingPods = new Set(podNames);
|
||||
|
||||
while (pendingPods.size > 0 && Date.now() - startTime < timeoutMs) {
|
||||
const pods = await k8sApi.core.listNamespacedPod({ namespace });
|
||||
const existingPods = new Set(pods.items.map((pod) => pod.metadata?.name ?? ""));
|
||||
|
||||
for (const podName of pendingPods) {
|
||||
if (!existingPods.has(podName)) {
|
||||
pendingPods.delete(podName);
|
||||
}
|
||||
}
|
||||
|
||||
if (pendingPods.size > 0) {
|
||||
await setTimeout(waitMs);
|
||||
}
|
||||
}
|
||||
|
||||
if (pendingPods.size > 0) {
|
||||
throw new Error(
|
||||
`Pods [${Array.from(pendingPods).join(", ")}] were not deleted within ${timeoutMs}ms`
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
async function podExists({
|
||||
k8sApi,
|
||||
namespace,
|
||||
podName,
|
||||
}: {
|
||||
k8sApi: K8sApi;
|
||||
namespace: string;
|
||||
podName: string;
|
||||
}) {
|
||||
const pods = await k8sApi.core.listNamespacedPod({ namespace });
|
||||
return pods.items.some((p) => p.metadata?.name === podName);
|
||||
}
|
||||
@@ -0,0 +1,119 @@
|
||||
import { IntervalService } from "@trigger.dev/core/v3";
|
||||
import { SimpleStructuredLogger } from "@trigger.dev/core/v3/utils/structuredLogger";
|
||||
import type { Registry } from "prom-client";
|
||||
import { Counter, Gauge } from "prom-client";
|
||||
import type { K8sApi } from "../clients/kubernetes.js";
|
||||
import { createK8sApi } from "../clients/kubernetes.js";
|
||||
import { register } from "../metrics.js";
|
||||
|
||||
export type PodCleanerOptions = {
|
||||
namespace: string;
|
||||
k8s?: K8sApi;
|
||||
register?: Registry;
|
||||
batchSize?: number;
|
||||
intervalMs?: number;
|
||||
};
|
||||
|
||||
export class PodCleaner {
|
||||
private readonly logger = new SimpleStructuredLogger("pod-cleaner");
|
||||
private readonly k8s: K8sApi;
|
||||
private readonly namespace: string;
|
||||
|
||||
private readonly batchSize: number;
|
||||
private readonly deletionInterval: IntervalService;
|
||||
|
||||
// Metrics
|
||||
private readonly register: Registry;
|
||||
private readonly deletionCyclesTotal: Counter;
|
||||
private readonly lastDeletionTimestamp: Gauge;
|
||||
|
||||
constructor(opts: PodCleanerOptions) {
|
||||
this.k8s = opts.k8s ?? createK8sApi();
|
||||
|
||||
this.namespace = opts.namespace;
|
||||
this.batchSize = opts.batchSize ?? 500;
|
||||
|
||||
this.deletionInterval = new IntervalService({
|
||||
intervalMs: opts.intervalMs ?? 10000,
|
||||
leadingEdge: true,
|
||||
onInterval: this.deleteCompletedPods.bind(this),
|
||||
});
|
||||
|
||||
// Initialize metrics
|
||||
this.register = opts.register ?? register;
|
||||
|
||||
this.deletionCyclesTotal = new Counter({
|
||||
name: "pod_cleaner_deletion_cycles_total",
|
||||
help: "Total number of pod deletion cycles run",
|
||||
labelNames: ["namespace", "status", "batch_size"],
|
||||
registers: [this.register],
|
||||
});
|
||||
|
||||
this.lastDeletionTimestamp = new Gauge({
|
||||
name: "pod_cleaner_last_deletion_timestamp",
|
||||
help: "Timestamp of the last deletion cycle",
|
||||
labelNames: ["namespace"],
|
||||
registers: [this.register],
|
||||
});
|
||||
}
|
||||
|
||||
async start() {
|
||||
this.deletionInterval.start();
|
||||
}
|
||||
|
||||
async stop() {
|
||||
this.deletionInterval.stop();
|
||||
}
|
||||
|
||||
private async deleteCompletedPods() {
|
||||
let continuationToken: string | undefined;
|
||||
|
||||
do {
|
||||
try {
|
||||
const result = await this.k8s.core.deleteCollectionNamespacedPod({
|
||||
namespace: this.namespace,
|
||||
labelSelector: "app=task-run",
|
||||
fieldSelector: "status.phase=Succeeded",
|
||||
limit: this.batchSize,
|
||||
_continue: continuationToken,
|
||||
gracePeriodSeconds: 0,
|
||||
propagationPolicy: "Background",
|
||||
timeoutSeconds: 30,
|
||||
});
|
||||
|
||||
// Update continuation token for next batch
|
||||
continuationToken = result.metadata?._continue;
|
||||
|
||||
// Increment the deletion cycles counter
|
||||
this.deletionCyclesTotal.inc({
|
||||
namespace: this.namespace,
|
||||
batch_size: this.batchSize,
|
||||
status: "succeeded",
|
||||
});
|
||||
|
||||
this.logger.debug("Deleted batch of pods", { continuationToken });
|
||||
} catch (err) {
|
||||
this.logger.error("Failed to delete batch of pods", {
|
||||
err: err instanceof Error ? err.message : String(err),
|
||||
});
|
||||
|
||||
this.deletionCyclesTotal.inc({
|
||||
namespace: this.namespace,
|
||||
batch_size: this.batchSize,
|
||||
status: "failed",
|
||||
});
|
||||
break;
|
||||
}
|
||||
} while (continuationToken);
|
||||
|
||||
this.lastDeletionTimestamp.set({ namespace: this.namespace }, Date.now());
|
||||
}
|
||||
|
||||
// Method to expose metrics for testing
|
||||
public getMetrics() {
|
||||
return {
|
||||
deletionCyclesTotal: this.deletionCyclesTotal,
|
||||
lastDeletionTimestamp: this.lastDeletionTimestamp,
|
||||
};
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,271 @@
|
||||
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
|
||||
import { TimerWheel } from "./timerWheel.js";
|
||||
|
||||
describe("TimerWheel", () => {
|
||||
beforeEach(() => {
|
||||
vi.useFakeTimers();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.useRealTimers();
|
||||
});
|
||||
|
||||
it("dispatches item after delay", () => {
|
||||
const dispatched: string[] = [];
|
||||
const wheel = new TimerWheel<string>({
|
||||
delayMs: 3000,
|
||||
onExpire: (item) => dispatched.push(item.key),
|
||||
});
|
||||
|
||||
wheel.start();
|
||||
wheel.submit("run-1", "snapshot-data");
|
||||
|
||||
// Not yet
|
||||
vi.advanceTimersByTime(2900);
|
||||
expect(dispatched).toEqual([]);
|
||||
|
||||
// After delay
|
||||
vi.advanceTimersByTime(200);
|
||||
expect(dispatched).toEqual(["run-1"]);
|
||||
|
||||
wheel.stop();
|
||||
});
|
||||
|
||||
it("cancels item before it fires", () => {
|
||||
const dispatched: string[] = [];
|
||||
const wheel = new TimerWheel<string>({
|
||||
delayMs: 3000,
|
||||
onExpire: (item) => dispatched.push(item.key),
|
||||
});
|
||||
|
||||
wheel.start();
|
||||
wheel.submit("run-1", "data");
|
||||
|
||||
vi.advanceTimersByTime(1000);
|
||||
expect(wheel.cancel("run-1")).toBe(true);
|
||||
|
||||
vi.advanceTimersByTime(5000);
|
||||
expect(dispatched).toEqual([]);
|
||||
expect(wheel.size).toBe(0);
|
||||
|
||||
wheel.stop();
|
||||
});
|
||||
|
||||
it("peek returns the pending item without removing it", () => {
|
||||
const wheel = new TimerWheel<string>({ delayMs: 3000, onExpire: () => {} });
|
||||
|
||||
wheel.start();
|
||||
wheel.submit("run-1", "data");
|
||||
|
||||
expect(wheel.peek("run-1")).toEqual({ key: "run-1", data: "data" });
|
||||
expect(wheel.size).toBe(1);
|
||||
expect(wheel.peek("run-2")).toBeUndefined();
|
||||
|
||||
// Dispatched items are no longer peekable
|
||||
vi.advanceTimersByTime(3100);
|
||||
expect(wheel.peek("run-1")).toBeUndefined();
|
||||
|
||||
wheel.stop();
|
||||
});
|
||||
|
||||
it("cancel returns false for unknown key", () => {
|
||||
const wheel = new TimerWheel<string>({
|
||||
delayMs: 3000,
|
||||
onExpire: () => {},
|
||||
});
|
||||
expect(wheel.cancel("nonexistent")).toBe(false);
|
||||
});
|
||||
|
||||
it("deduplicates: resubmitting same key replaces the entry", () => {
|
||||
const dispatched: { key: string; data: string }[] = [];
|
||||
const wheel = new TimerWheel<string>({
|
||||
delayMs: 3000,
|
||||
onExpire: (item) => dispatched.push({ key: item.key, data: item.data }),
|
||||
});
|
||||
|
||||
wheel.start();
|
||||
wheel.submit("run-1", "old-data");
|
||||
|
||||
vi.advanceTimersByTime(1000);
|
||||
wheel.submit("run-1", "new-data");
|
||||
|
||||
// Original would have fired at t=3000, but was replaced
|
||||
// New one fires at t=1000+3000=4000
|
||||
vi.advanceTimersByTime(2100);
|
||||
expect(dispatched).toEqual([]);
|
||||
|
||||
vi.advanceTimersByTime(1000);
|
||||
expect(dispatched).toEqual([{ key: "run-1", data: "new-data" }]);
|
||||
|
||||
wheel.stop();
|
||||
});
|
||||
|
||||
it("handles many concurrent items", () => {
|
||||
const dispatched: string[] = [];
|
||||
const wheel = new TimerWheel<string>({
|
||||
delayMs: 3000,
|
||||
onExpire: (item) => dispatched.push(item.key),
|
||||
});
|
||||
|
||||
wheel.start();
|
||||
|
||||
for (let i = 0; i < 1000; i++) {
|
||||
wheel.submit(`run-${i}`, `data-${i}`);
|
||||
}
|
||||
expect(wheel.size).toBe(1000);
|
||||
|
||||
vi.advanceTimersByTime(3100);
|
||||
expect(dispatched.length).toBe(1000);
|
||||
expect(wheel.size).toBe(0);
|
||||
|
||||
wheel.stop();
|
||||
});
|
||||
|
||||
it("handles items submitted at different times", () => {
|
||||
const dispatched: string[] = [];
|
||||
const wheel = new TimerWheel<string>({
|
||||
delayMs: 3000,
|
||||
onExpire: (item) => dispatched.push(item.key),
|
||||
});
|
||||
|
||||
wheel.start();
|
||||
|
||||
wheel.submit("run-1", "data");
|
||||
vi.advanceTimersByTime(1000);
|
||||
wheel.submit("run-2", "data");
|
||||
vi.advanceTimersByTime(1000);
|
||||
wheel.submit("run-3", "data");
|
||||
|
||||
// t=2000: nothing yet
|
||||
expect(dispatched).toEqual([]);
|
||||
|
||||
// t=3100: run-1 fires
|
||||
vi.advanceTimersByTime(1100);
|
||||
expect(dispatched).toEqual(["run-1"]);
|
||||
|
||||
// t=4100: run-2 fires
|
||||
vi.advanceTimersByTime(1000);
|
||||
expect(dispatched).toEqual(["run-1", "run-2"]);
|
||||
|
||||
// t=5100: run-3 fires
|
||||
vi.advanceTimersByTime(1000);
|
||||
expect(dispatched).toEqual(["run-1", "run-2", "run-3"]);
|
||||
|
||||
wheel.stop();
|
||||
});
|
||||
|
||||
it("setDelay changes delay for new items only", () => {
|
||||
const dispatched: string[] = [];
|
||||
const wheel = new TimerWheel<string>({
|
||||
delayMs: 3000,
|
||||
onExpire: (item) => dispatched.push(item.key),
|
||||
});
|
||||
|
||||
wheel.start();
|
||||
|
||||
wheel.submit("run-1", "data"); // 3s delay
|
||||
|
||||
vi.advanceTimersByTime(500);
|
||||
wheel.setDelay(1000);
|
||||
wheel.submit("run-2", "data"); // 1s delay
|
||||
|
||||
// t=1500: run-2 should have fired (submitted at t=500 with 1s delay)
|
||||
vi.advanceTimersByTime(1100);
|
||||
expect(dispatched).toEqual(["run-2"]);
|
||||
|
||||
// t=3100: run-1 fires at its original 3s delay
|
||||
vi.advanceTimersByTime(1500);
|
||||
expect(dispatched).toEqual(["run-2", "run-1"]);
|
||||
|
||||
wheel.stop();
|
||||
});
|
||||
|
||||
it("stop returns unprocessed items", () => {
|
||||
const dispatched: string[] = [];
|
||||
const wheel = new TimerWheel<string>({
|
||||
delayMs: 3000,
|
||||
onExpire: (item) => dispatched.push(item.key),
|
||||
});
|
||||
|
||||
wheel.start();
|
||||
wheel.submit("run-1", "data-1");
|
||||
wheel.submit("run-2", "data-2");
|
||||
wheel.submit("run-3", "data-3");
|
||||
|
||||
const remaining = wheel.stop();
|
||||
expect(dispatched).toEqual([]);
|
||||
expect(wheel.size).toBe(0);
|
||||
expect(remaining.length).toBe(3);
|
||||
expect(remaining.map((r) => r.key).sort()).toEqual(["run-1", "run-2", "run-3"]);
|
||||
expect(remaining.find((r) => r.key === "run-1")?.data).toBe("data-1");
|
||||
});
|
||||
|
||||
it("after stop, new submissions are silently dropped", () => {
|
||||
const dispatched: string[] = [];
|
||||
const wheel = new TimerWheel<string>({
|
||||
delayMs: 3000,
|
||||
onExpire: (item) => dispatched.push(item.key),
|
||||
});
|
||||
|
||||
wheel.start();
|
||||
wheel.stop();
|
||||
|
||||
wheel.submit("run-late", "data");
|
||||
expect(dispatched).toEqual([]);
|
||||
expect(wheel.size).toBe(0);
|
||||
});
|
||||
|
||||
it("tracks size correctly through submit/cancel/dispatch", () => {
|
||||
const wheel = new TimerWheel<string>({
|
||||
delayMs: 3000,
|
||||
onExpire: () => {},
|
||||
});
|
||||
|
||||
wheel.start();
|
||||
|
||||
wheel.submit("a", "data");
|
||||
wheel.submit("b", "data");
|
||||
expect(wheel.size).toBe(2);
|
||||
|
||||
wheel.cancel("a");
|
||||
expect(wheel.size).toBe(1);
|
||||
|
||||
vi.advanceTimersByTime(3100);
|
||||
expect(wheel.size).toBe(0);
|
||||
|
||||
wheel.stop();
|
||||
});
|
||||
|
||||
it("clamps delay to valid range", () => {
|
||||
const dispatched: string[] = [];
|
||||
|
||||
// Very small delay (should be at least 1 tick = 100ms)
|
||||
const wheel = new TimerWheel<string>({
|
||||
delayMs: 0,
|
||||
onExpire: (item) => dispatched.push(item.key),
|
||||
});
|
||||
|
||||
wheel.start();
|
||||
wheel.submit("run-1", "data");
|
||||
|
||||
vi.advanceTimersByTime(200);
|
||||
expect(dispatched).toEqual(["run-1"]);
|
||||
|
||||
wheel.stop();
|
||||
});
|
||||
|
||||
it("multiple cancel calls are safe", () => {
|
||||
const wheel = new TimerWheel<string>({
|
||||
delayMs: 3000,
|
||||
onExpire: () => {},
|
||||
});
|
||||
|
||||
wheel.start();
|
||||
wheel.submit("run-1", "data");
|
||||
|
||||
expect(wheel.cancel("run-1")).toBe(true);
|
||||
expect(wheel.cancel("run-1")).toBe(false);
|
||||
|
||||
wheel.stop();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,166 @@
|
||||
/**
|
||||
* TimerWheel implements a hashed timer wheel for efficiently managing large numbers
|
||||
* of delayed operations with O(1) submit, cancel, and per-item dispatch.
|
||||
*
|
||||
* Used by the supervisor to delay snapshot requests so that short-lived waitpoints
|
||||
* (e.g. triggerAndWait that resolves in <3s) skip the snapshot entirely.
|
||||
*
|
||||
* The wheel is a ring buffer of slots. A single setInterval advances a cursor.
|
||||
* When the cursor reaches a slot, all items in that slot are dispatched.
|
||||
*
|
||||
* Fixed capacity: 600 slots at 100ms tick = 60s max delay.
|
||||
*/
|
||||
|
||||
const TICK_MS = 100;
|
||||
const NUM_SLOTS = 600; // 60s max delay at 100ms tick
|
||||
|
||||
export type TimerWheelItem<T> = {
|
||||
key: string;
|
||||
data: T;
|
||||
};
|
||||
|
||||
export type TimerWheelOptions<T> = {
|
||||
/** Called when an item's delay expires. */
|
||||
onExpire: (item: TimerWheelItem<T>) => void;
|
||||
/** Delay in milliseconds before items fire. Clamped to [100, 60000]. */
|
||||
delayMs: number;
|
||||
};
|
||||
|
||||
type Entry<T> = {
|
||||
key: string;
|
||||
data: T;
|
||||
slotIndex: number;
|
||||
};
|
||||
|
||||
export class TimerWheel<T> {
|
||||
private slots: Set<string>[];
|
||||
private entries: Map<string, Entry<T>>;
|
||||
private cursor: number;
|
||||
private intervalId: ReturnType<typeof setInterval> | null;
|
||||
private onExpire: (item: TimerWheelItem<T>) => void;
|
||||
private delaySlots: number;
|
||||
|
||||
constructor(opts: TimerWheelOptions<T>) {
|
||||
this.slots = Array.from({ length: NUM_SLOTS }, () => new Set());
|
||||
this.entries = new Map();
|
||||
this.cursor = 0;
|
||||
this.intervalId = null;
|
||||
this.onExpire = opts.onExpire;
|
||||
this.delaySlots = Math.max(1, Math.min(NUM_SLOTS, Math.ceil(opts.delayMs / TICK_MS)));
|
||||
}
|
||||
|
||||
/** Start the timer wheel. Must be called before submitting items. */
|
||||
start(): void {
|
||||
if (this.intervalId) return;
|
||||
this.intervalId = setInterval(() => this.tick(), TICK_MS);
|
||||
// Don't hold the process open just for the timer wheel
|
||||
if (this.intervalId && typeof this.intervalId === "object" && "unref" in this.intervalId) {
|
||||
this.intervalId.unref();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Stop the timer wheel and return all unprocessed items.
|
||||
* The wheel keeps running normally during graceful shutdown - call stop()
|
||||
* only when you're ready to tear down. Caller decides what to do with leftovers.
|
||||
*/
|
||||
stop(): TimerWheelItem<T>[] {
|
||||
if (this.intervalId) {
|
||||
clearInterval(this.intervalId);
|
||||
this.intervalId = null;
|
||||
}
|
||||
|
||||
const remaining: TimerWheelItem<T>[] = [];
|
||||
for (const [key, entry] of this.entries) {
|
||||
remaining.push({ key, data: entry.data });
|
||||
}
|
||||
|
||||
for (const slot of this.slots) {
|
||||
slot.clear();
|
||||
}
|
||||
this.entries.clear();
|
||||
|
||||
return remaining;
|
||||
}
|
||||
|
||||
/**
|
||||
* Update the delay for future submissions. Already-queued items keep their original timing.
|
||||
* Clamped to [TICK_MS, 60000ms].
|
||||
*/
|
||||
setDelay(delayMs: number): void {
|
||||
this.delaySlots = Math.max(1, Math.min(NUM_SLOTS, Math.ceil(delayMs / TICK_MS)));
|
||||
}
|
||||
|
||||
/**
|
||||
* Submit an item to be dispatched after the configured delay.
|
||||
* If an item with the same key already exists, it is replaced (dedup).
|
||||
* No-op if the wheel is stopped.
|
||||
*/
|
||||
submit(key: string, data: T): void {
|
||||
if (!this.intervalId) return;
|
||||
|
||||
// Dedup: remove existing entry for this key
|
||||
this.cancel(key);
|
||||
|
||||
const slotIndex = (this.cursor + this.delaySlots) % NUM_SLOTS;
|
||||
const entry: Entry<T> = { key, data, slotIndex };
|
||||
|
||||
this.entries.set(key, entry);
|
||||
this.slot(slotIndex).add(key);
|
||||
}
|
||||
|
||||
/**
|
||||
* Cancel a pending item. Returns true if the item was found and removed.
|
||||
*/
|
||||
cancel(key: string): boolean {
|
||||
const entry = this.entries.get(key);
|
||||
if (!entry) return false;
|
||||
|
||||
this.slot(entry.slotIndex).delete(key);
|
||||
this.entries.delete(key);
|
||||
return true;
|
||||
}
|
||||
|
||||
/** Look up a pending item without removing it. */
|
||||
peek(key: string): TimerWheelItem<T> | undefined {
|
||||
const entry = this.entries.get(key);
|
||||
return entry ? { key, data: entry.data } : undefined;
|
||||
}
|
||||
|
||||
/** Number of pending items in the wheel. */
|
||||
get size(): number {
|
||||
return this.entries.size;
|
||||
}
|
||||
|
||||
/** Whether the wheel is running. */
|
||||
get running(): boolean {
|
||||
return this.intervalId !== null;
|
||||
}
|
||||
|
||||
/** Get a slot by index. The array is fully initialized so this always returns a Set. */
|
||||
private slot(index: number): Set<string> {
|
||||
const s = this.slots[index];
|
||||
if (!s) throw new Error(`TimerWheel: invalid slot index ${index}`);
|
||||
return s;
|
||||
}
|
||||
|
||||
/** Advance the cursor and dispatch all items in the current slot. */
|
||||
private tick(): void {
|
||||
this.cursor = (this.cursor + 1) % NUM_SLOTS;
|
||||
const slot = this.slot(this.cursor);
|
||||
|
||||
if (slot.size === 0) return;
|
||||
|
||||
// Collect items to dispatch (copy keys since we mutate during iteration)
|
||||
const keys = [...slot];
|
||||
slot.clear();
|
||||
|
||||
for (const key of keys) {
|
||||
const entry = this.entries.get(key);
|
||||
if (!entry) continue;
|
||||
|
||||
this.entries.delete(key);
|
||||
this.onExpire({ key, data: entry.data });
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,129 @@
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import { setTimeout as sleep } from "node:timers/promises";
|
||||
import { WarmStartVerificationService } from "./warmStartVerificationService.js";
|
||||
import type { DequeuedMessage } from "@trigger.dev/core/v3";
|
||||
import type { SupervisorHttpClient } from "@trigger.dev/core/v3/workers";
|
||||
|
||||
// The TimerWheel ticks every 100ms, so a 1000ms delay (the env minimum)
|
||||
// fires within ~1.1s.
|
||||
const DELAY_MS = 1_000;
|
||||
// Long enough that a pending verification would certainly have fired.
|
||||
const SETTLE_MS = 1_600;
|
||||
|
||||
const DEQUEUED_SNAPSHOT_ID = "snapshot_dequeued";
|
||||
|
||||
function makeMessage(runFriendlyId = "run_1"): DequeuedMessage {
|
||||
return {
|
||||
run: { friendlyId: runFriendlyId },
|
||||
snapshot: { friendlyId: DEQUEUED_SNAPSHOT_ID },
|
||||
} as unknown as DequeuedMessage;
|
||||
}
|
||||
|
||||
function createService(opts: { latestSnapshotId?: string; probeError?: boolean }) {
|
||||
const getLatestSnapshot = vi.fn(async (_runId: string) =>
|
||||
opts.probeError
|
||||
? { success: false as const, error: "connection refused" }
|
||||
: {
|
||||
success: true as const,
|
||||
data: { execution: { snapshot: { friendlyId: opts.latestSnapshotId } } },
|
||||
}
|
||||
);
|
||||
|
||||
const createWorkload = vi.fn(async () => {});
|
||||
|
||||
const service = new WarmStartVerificationService({
|
||||
workerClient: { getLatestSnapshot } as unknown as SupervisorHttpClient,
|
||||
delayMs: DELAY_MS,
|
||||
createWorkload,
|
||||
wideEventOpts: { service: "supervisor-test", env: {}, enabled: false },
|
||||
});
|
||||
|
||||
return { service, getLatestSnapshot, createWorkload };
|
||||
}
|
||||
|
||||
describe("WarmStartVerificationService", () => {
|
||||
it("falls back to a cold create when the snapshot is unchanged", async () => {
|
||||
const { service, createWorkload } = createService({
|
||||
latestSnapshotId: DEQUEUED_SNAPSHOT_ID,
|
||||
});
|
||||
try {
|
||||
const message = makeMessage();
|
||||
const timings = { warmStartCheckMs: 12 };
|
||||
service.schedule(message, timings);
|
||||
|
||||
await vi.waitFor(() => expect(createWorkload).toHaveBeenCalledTimes(1), {
|
||||
timeout: 3_000,
|
||||
});
|
||||
expect(createWorkload).toHaveBeenCalledWith(message, timings);
|
||||
} finally {
|
||||
service.stop();
|
||||
}
|
||||
});
|
||||
|
||||
it("does nothing when the snapshot has moved on (delivered)", async () => {
|
||||
const { service, getLatestSnapshot, createWorkload } = createService({
|
||||
latestSnapshotId: "snapshot_executing",
|
||||
});
|
||||
try {
|
||||
service.schedule(makeMessage(), { warmStartCheckMs: 12 });
|
||||
|
||||
await vi.waitFor(() => expect(getLatestSnapshot).toHaveBeenCalledTimes(1), {
|
||||
timeout: 3_000,
|
||||
});
|
||||
await sleep(100);
|
||||
expect(createWorkload).not.toHaveBeenCalled();
|
||||
} finally {
|
||||
service.stop();
|
||||
}
|
||||
});
|
||||
|
||||
it("never falls back when the probe errors", async () => {
|
||||
const { service, getLatestSnapshot, createWorkload } = createService({ probeError: true });
|
||||
try {
|
||||
service.schedule(makeMessage(), { warmStartCheckMs: 12 });
|
||||
|
||||
await vi.waitFor(() => expect(getLatestSnapshot).toHaveBeenCalledTimes(1), {
|
||||
timeout: 3_000,
|
||||
});
|
||||
await sleep(100);
|
||||
expect(createWorkload).not.toHaveBeenCalled();
|
||||
} finally {
|
||||
service.stop();
|
||||
}
|
||||
});
|
||||
|
||||
it("cancel before the delay prevents the probe entirely", async () => {
|
||||
const { service, getLatestSnapshot, createWorkload } = createService({
|
||||
latestSnapshotId: DEQUEUED_SNAPSHOT_ID,
|
||||
});
|
||||
try {
|
||||
service.schedule(makeMessage(), { warmStartCheckMs: 12 });
|
||||
|
||||
expect(service.cancel("run_1")).toBe(true);
|
||||
|
||||
await sleep(SETTLE_MS);
|
||||
expect(getLatestSnapshot).not.toHaveBeenCalled();
|
||||
expect(createWorkload).not.toHaveBeenCalled();
|
||||
} finally {
|
||||
service.stop();
|
||||
}
|
||||
});
|
||||
|
||||
it("re-scheduling the same run replaces the pending verification", async () => {
|
||||
const { service, getLatestSnapshot } = createService({
|
||||
latestSnapshotId: "snapshot_executing",
|
||||
});
|
||||
try {
|
||||
service.schedule(makeMessage(), { warmStartCheckMs: 1 });
|
||||
service.schedule(makeMessage(), { warmStartCheckMs: 2 });
|
||||
|
||||
await vi.waitFor(() => expect(getLatestSnapshot).toHaveBeenCalledTimes(1), {
|
||||
timeout: 3_000,
|
||||
});
|
||||
await sleep(SETTLE_MS);
|
||||
expect(getLatestSnapshot).toHaveBeenCalledTimes(1);
|
||||
} finally {
|
||||
service.stop();
|
||||
}
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,173 @@
|
||||
import pLimit from "p-limit";
|
||||
import { SimpleStructuredLogger } from "@trigger.dev/core/v3/utils/structuredLogger";
|
||||
import type { DequeuedMessage } from "@trigger.dev/core/v3";
|
||||
import type { SupervisorHttpClient } from "@trigger.dev/core/v3/workers";
|
||||
import { tryCatch } from "@trigger.dev/core";
|
||||
import { TimerWheel } from "./timerWheel.js";
|
||||
import { emitOneShot, type WideEventOptions } from "../wideEvents/index.js";
|
||||
|
||||
const PROBE_CONCURRENCY_LIMIT = 10;
|
||||
|
||||
export type WarmStartTimings = {
|
||||
dequeueResponseMs?: number;
|
||||
pollingIntervalMs?: number;
|
||||
warmStartCheckMs: number;
|
||||
};
|
||||
|
||||
type PendingVerification = {
|
||||
message: DequeuedMessage;
|
||||
timings: WarmStartTimings;
|
||||
};
|
||||
|
||||
export type WarmStartVerificationServiceOptions = {
|
||||
workerClient: SupervisorHttpClient;
|
||||
/** How long after a warm-start hit to verify the runner acted on it. */
|
||||
delayMs: number;
|
||||
/** Cold-creates the workload for a dispatched-but-lost run. */
|
||||
createWorkload: (message: DequeuedMessage, timings: WarmStartTimings) => Promise<void>;
|
||||
wideEventOpts: WideEventOptions;
|
||||
};
|
||||
|
||||
/**
|
||||
* Verifies that warm-start dispatches were actually acted on.
|
||||
*
|
||||
* Firestarter's `didWarmStart: true` means "response written to a socket",
|
||||
* not "runner received it". A silently dead poller (no FIN - e.g. a VM torn
|
||||
* down mid-poll) leaves the dispatched run stuck in PENDING_EXECUTING until
|
||||
* the run engine's heartbeat redrive minutes later, burning a queue
|
||||
* redelivery each time (TRI-10659).
|
||||
*
|
||||
* After a hit, the dequeued message is retained for `delayMs`, then the
|
||||
* platform is asked for the run's latest snapshot. If it is still the exact
|
||||
* snapshot we dequeued, no runner ever started the attempt - fall through to
|
||||
* the regular cold-create path with the original message. Double-starts are
|
||||
* impossible: `startRunAttempt` runs under a per-run lock and rejects stale
|
||||
* snapshot ids, so if the original runner revives and races the fallback
|
||||
* workload, exactly one wins and the loser exits before executing anything.
|
||||
*
|
||||
* On a probe ERROR we deliberately do nothing: the runner's attempt-start
|
||||
* goes through nested retries, so during platform brownouts a healthy runner
|
||||
* can legitimately act late - falling back on uncertainty would stampede
|
||||
* duplicate workloads exactly when the platform is degraded. The heartbeat
|
||||
* redrive remains the backstop for that case (and for supervisor restarts,
|
||||
* which drop the in-memory timers).
|
||||
*/
|
||||
export class WarmStartVerificationService {
|
||||
private readonly logger = new SimpleStructuredLogger("warm-start-verification");
|
||||
|
||||
private readonly timerWheel: TimerWheel<PendingVerification>;
|
||||
private readonly probeLimit: ReturnType<typeof pLimit>;
|
||||
|
||||
private readonly workerClient: SupervisorHttpClient;
|
||||
private readonly delayMs: number;
|
||||
private readonly createWorkload: WarmStartVerificationServiceOptions["createWorkload"];
|
||||
private readonly wideEventOpts: WideEventOptions;
|
||||
|
||||
constructor(opts: WarmStartVerificationServiceOptions) {
|
||||
this.workerClient = opts.workerClient;
|
||||
this.delayMs = opts.delayMs;
|
||||
this.createWorkload = opts.createWorkload;
|
||||
this.wideEventOpts = opts.wideEventOpts;
|
||||
|
||||
this.probeLimit = pLimit(PROBE_CONCURRENCY_LIMIT);
|
||||
this.timerWheel = new TimerWheel<PendingVerification>({
|
||||
delayMs: opts.delayMs,
|
||||
onExpire: (item) => {
|
||||
this.probeLimit(() => this.verify(item.data)).catch((error) => {
|
||||
this.logger.error("Verification failed", {
|
||||
runId: item.data.message.run.friendlyId,
|
||||
error,
|
||||
});
|
||||
});
|
||||
},
|
||||
});
|
||||
this.timerWheel.start();
|
||||
}
|
||||
|
||||
/** Schedule delivery verification for a warm-start hit. */
|
||||
schedule(message: DequeuedMessage, timings: WarmStartTimings) {
|
||||
this.timerWheel.submit(message.run.friendlyId, { message, timings });
|
||||
this.logger.debug("Verification scheduled", {
|
||||
runId: message.run.friendlyId,
|
||||
snapshotId: message.snapshot.friendlyId,
|
||||
delayMs: this.delayMs,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Cancel a pending verification, e.g. when the runner connects to this
|
||||
* supervisor. Purely an optimization: the matched runner often lives on a
|
||||
* different node and connects to that node's supervisor, so most healthy
|
||||
* deliveries are confirmed by the probe, not by this.
|
||||
*/
|
||||
cancel(runFriendlyId: string): boolean {
|
||||
return this.timerWheel.cancel(runFriendlyId);
|
||||
}
|
||||
|
||||
/** Stop the timer wheel, dropping pending verifications. The run engine's
|
||||
* heartbeat redrive covers anything dropped here. */
|
||||
stop() {
|
||||
const remaining = this.timerWheel.stop();
|
||||
if (remaining.length > 0) {
|
||||
this.logger.info("Stopped, dropped pending verifications", { count: remaining.length });
|
||||
}
|
||||
}
|
||||
|
||||
private async verify({ message, timings }: PendingVerification) {
|
||||
const runFriendlyId = message.run.friendlyId;
|
||||
|
||||
const result = await this.workerClient.getLatestSnapshot(runFriendlyId);
|
||||
|
||||
if (!result.success) {
|
||||
// Never fall back on uncertainty - see class docs.
|
||||
this.emitOutcome(message, "probe_error", String(result.error));
|
||||
this.logger.warn("Verification probe failed, skipping", {
|
||||
runId: runFriendlyId,
|
||||
error: result.error,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
const latestSnapshotId = result.data.execution.snapshot.friendlyId;
|
||||
|
||||
if (latestSnapshotId !== message.snapshot.friendlyId) {
|
||||
// Something acted on the run (attempt started, or it was cancelled or
|
||||
// requeued) - the dispatch is no longer ours to worry about.
|
||||
this.emitOutcome(message, "delivered");
|
||||
return;
|
||||
}
|
||||
|
||||
this.emitOutcome(message, "fallback");
|
||||
this.logger.warn("Warm start dispatch was never acted on, cold starting", {
|
||||
runId: runFriendlyId,
|
||||
snapshotId: message.snapshot.friendlyId,
|
||||
});
|
||||
|
||||
const [error] = await tryCatch(this.createWorkload(message, timings));
|
||||
if (error) {
|
||||
this.logger.error("Fallback workload create failed", {
|
||||
runId: runFriendlyId,
|
||||
error: error instanceof Error ? error.message : String(error),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
private emitOutcome(
|
||||
message: DequeuedMessage,
|
||||
outcome: "delivered" | "fallback" | "probe_error",
|
||||
error?: string
|
||||
) {
|
||||
emitOneShot({
|
||||
...this.wideEventOpts,
|
||||
op: "warmstart.verify",
|
||||
kind: "event",
|
||||
populate: (state) => {
|
||||
state.meta.run_id = message.run.friendlyId;
|
||||
state.meta.snapshot_id = message.snapshot.friendlyId;
|
||||
state.extras.outcome = outcome;
|
||||
state.extras.delay_ms = this.delayMs;
|
||||
if (error) state.extras.error = error;
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
import { isMacOS, isWindows } from "std-env";
|
||||
|
||||
export function normalizeDockerHostUrl(url: string) {
|
||||
const $url = new URL(url);
|
||||
|
||||
if ($url.hostname === "localhost") {
|
||||
$url.hostname = getDockerHostDomain();
|
||||
}
|
||||
|
||||
return $url.toString();
|
||||
}
|
||||
|
||||
export function getDockerHostDomain() {
|
||||
return isMacOS || isWindows ? "host.docker.internal" : "localhost";
|
||||
}
|
||||
|
||||
/** Extract the W3C traceparent string from an untyped trace context record */
|
||||
export function extractTraceparent(traceContext?: Record<string, unknown>): string | undefined {
|
||||
if (
|
||||
traceContext &&
|
||||
"traceparent" in traceContext &&
|
||||
typeof traceContext.traceparent === "string"
|
||||
) {
|
||||
return traceContext.traceparent;
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
export function getRunnerId(runId: string, attemptNumber?: number) {
|
||||
const parts = ["runner", runId.replace("run_", "")];
|
||||
|
||||
if (attemptNumber && attemptNumber > 1) {
|
||||
parts.push(`attempt-${attemptNumber}`);
|
||||
}
|
||||
|
||||
return parts.join("-");
|
||||
}
|
||||
|
||||
/** Derive a unique runnerId for a restore cycle using the checkpoint suffix */
|
||||
export function getRestoreRunnerId(runFriendlyId: string, checkpointId: string) {
|
||||
const runIdShort = runFriendlyId.replace("run_", "");
|
||||
const checkpointSuffix = checkpointId.slice(-8);
|
||||
return `runner-${runIdShort}-${checkpointSuffix}`;
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
import { describe, it, expect } from "vitest";
|
||||
import { encodeBaggage } from "./baggage.js";
|
||||
|
||||
describe("encodeBaggage", () => {
|
||||
it("returns empty string for an empty map", () => {
|
||||
expect(encodeBaggage({})).toBe("");
|
||||
});
|
||||
|
||||
it("encodes a single entry as k=v", () => {
|
||||
expect(encodeBaggage({ run_id: "run-1" })).toBe("run_id=run-1");
|
||||
});
|
||||
|
||||
it("sorts keys for stable output across hops", () => {
|
||||
expect(encodeBaggage({ b: "2", a: "1", c: "3" })).toBe("a=1,b=2,c=3");
|
||||
});
|
||||
|
||||
it("skips empty keys and empty values", () => {
|
||||
expect(encodeBaggage({ "": "v", k: "", real: "x" })).toBe("real=x");
|
||||
});
|
||||
|
||||
it("truncates values longer than the cap", () => {
|
||||
const long = "x".repeat(1024);
|
||||
const got = encodeBaggage({ k: long });
|
||||
const value = got.slice("k=".length);
|
||||
expect(value.length).toBe(256);
|
||||
});
|
||||
|
||||
it("caps multibyte values by UTF-8 bytes, not code units", () => {
|
||||
const long = "あ".repeat(512); // 3 UTF-8 bytes each
|
||||
const got = encodeBaggage({ k: long });
|
||||
const value = got.slice("k=".length);
|
||||
expect(Buffer.byteLength(value, "utf8")).toBeLessThanOrEqual(256);
|
||||
});
|
||||
|
||||
it("caps the number of entries", () => {
|
||||
const meta: Record<string, string> = {};
|
||||
for (let i = 0; i < 50; i++) {
|
||||
// Sortable two-digit keys so we know which 32 survive.
|
||||
meta[`k${String(i).padStart(2, "0")}`] = "v";
|
||||
}
|
||||
const got = encodeBaggage(meta);
|
||||
expect(got.split(",").length).toBe(32);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,45 @@
|
||||
/**
|
||||
* W3C Baggage (https://www.w3.org/TR/baggage/) encoding for outbound peer
|
||||
* calls. Serialises a State's `meta` map into a `Baggage` header value so
|
||||
* the downstream service auto-stamps the same labels onto its own wide
|
||||
* events - even on early-error paths that bail before parsing the request
|
||||
* body.
|
||||
*
|
||||
* Outbound discipline: only call this on peer-to-peer hops within the trust
|
||||
* boundary. External-endpoint calls (image registries, cloud-provider
|
||||
* APIs, third-party webhooks) must not include the Baggage header.
|
||||
*/
|
||||
|
||||
import { truncateUtf8 } from "./truncate.js";
|
||||
|
||||
/**
|
||||
* Cap the number of entries serialised onto the header. A misbehaving
|
||||
* caller's `meta` map shouldn't blow up downstream event width.
|
||||
*/
|
||||
const MAX_BAGGAGE_ENTRIES = 32;
|
||||
|
||||
/**
|
||||
* Cap each value's length. Defense against an upstream that stuffs
|
||||
* unbounded payloads into a meta value.
|
||||
*/
|
||||
const MAX_BAGGAGE_VALUE_BYTES = 256;
|
||||
|
||||
/**
|
||||
* Encode a `meta` map as a Baggage header value (`k1=v1,k2=v2`). Keys are
|
||||
* sorted for stable output across hops; an empty input yields the empty
|
||||
* string so the caller can skip emitting the header entirely.
|
||||
*/
|
||||
export function encodeBaggage(meta: Record<string, string>): string {
|
||||
const entries = Object.entries(meta).filter(([k, v]) => k && v);
|
||||
if (entries.length === 0) return "";
|
||||
|
||||
entries.sort(([a], [b]) => (a < b ? -1 : a > b ? 1 : 0));
|
||||
|
||||
const out: string[] = [];
|
||||
for (const [k, raw] of entries) {
|
||||
if (out.length >= MAX_BAGGAGE_ENTRIES) break;
|
||||
const v = truncateUtf8(raw, MAX_BAGGAGE_VALUE_BYTES);
|
||||
out.push(`${k}=${v}`);
|
||||
}
|
||||
return out.join(",");
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
import { AsyncLocalStorage } from "node:async_hooks";
|
||||
import type { State } from "./state.js";
|
||||
|
||||
/**
|
||||
* AsyncLocalStorage threading per-operation `State` through the call stack.
|
||||
* Wrappers enter a state via `wideEventStorage.run(state, () => fn())` and
|
||||
* any code in the async call tree retrieves it via `fromContext()`.
|
||||
*/
|
||||
export const wideEventStorage = new AsyncLocalStorage<State>();
|
||||
|
||||
/** Returns the State attached to the current async context, or null. */
|
||||
export function fromContext(): State | null {
|
||||
return wideEventStorage.getStore() ?? null;
|
||||
}
|
||||
@@ -0,0 +1,134 @@
|
||||
import { describe, it, expect } from "vitest";
|
||||
import { emit, EmitMessage } from "./emit.js";
|
||||
import { newState } from "./new.js";
|
||||
|
||||
function captureEmit(state: Parameters<typeof emit>[0]): Record<string, unknown> {
|
||||
const captured: string[] = [];
|
||||
const origWrite = process.stdout.write;
|
||||
process.stdout.write = ((chunk: unknown) => {
|
||||
captured.push(String(chunk));
|
||||
return true;
|
||||
}) as typeof process.stdout.write;
|
||||
try {
|
||||
emit(state);
|
||||
} finally {
|
||||
process.stdout.write = origWrite;
|
||||
}
|
||||
expect(captured).toHaveLength(1);
|
||||
const line = captured[0];
|
||||
if (!line) throw new Error("no captured line");
|
||||
return JSON.parse(line) as Record<string, unknown>;
|
||||
}
|
||||
|
||||
describe("emit", () => {
|
||||
it("emits a single line with the stable message + request_id", () => {
|
||||
const s = newState({ service: "supervisor", env: {} });
|
||||
s.statusCode = 200;
|
||||
s.ok = true;
|
||||
s.durationMs = 5;
|
||||
const out = captureEmit(s);
|
||||
expect(out.msg).toBe(EmitMessage);
|
||||
expect(out.request_id).toBe(s.requestId);
|
||||
expect(out.service).toBe("supervisor");
|
||||
expect(out.ok).toBe(true);
|
||||
expect(out.status).toBe(200);
|
||||
expect(out.duration_ms).toBe(5);
|
||||
});
|
||||
|
||||
it("emits start_time as an ISO timestamp set by newState", () => {
|
||||
const s = newState({ service: "supervisor", env: {} });
|
||||
s.statusCode = 200;
|
||||
s.ok = true;
|
||||
const out = captureEmit(s);
|
||||
expect(typeof out.start_time).toBe("string");
|
||||
// Microsecond-precision RFC3339 (6 fractional digits), parseable as a date.
|
||||
expect(out.start_time).toMatch(/^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d{6}Z$/);
|
||||
expect(Number.isNaN(new Date(out.start_time as string).getTime())).toBe(false);
|
||||
});
|
||||
|
||||
it("omits start_time when unset", () => {
|
||||
const s = newState({ service: "supervisor", env: {} });
|
||||
delete s.startTime;
|
||||
s.statusCode = 200;
|
||||
s.ok = true;
|
||||
const out = captureEmit(s);
|
||||
expect(out).not.toHaveProperty("start_time");
|
||||
});
|
||||
|
||||
it("omits empty optional fields", () => {
|
||||
const s = newState({ service: "supervisor", env: {} });
|
||||
s.statusCode = 200;
|
||||
s.ok = true;
|
||||
const out = captureEmit(s);
|
||||
expect(out).not.toHaveProperty("trace_id");
|
||||
expect(out).not.toHaveProperty("version");
|
||||
expect(out).not.toHaveProperty("commit_sha");
|
||||
expect(out).not.toHaveProperty("error.code");
|
||||
});
|
||||
|
||||
it("flattens meta keys as meta.<key>", () => {
|
||||
const s = newState({ service: "supervisor", env: {} });
|
||||
s.statusCode = 200;
|
||||
s.ok = true;
|
||||
s.meta.run_id = "run_abc";
|
||||
s.meta.deployment_id = "dep_xyz";
|
||||
const out = captureEmit(s);
|
||||
expect(out["meta.run_id"]).toBe("run_abc");
|
||||
expect(out["meta.deployment_id"]).toBe("dep_xyz");
|
||||
expect(out).not.toHaveProperty("meta");
|
||||
});
|
||||
|
||||
it("flattens phases as phase.<name>.<field>", () => {
|
||||
const s = newState({ service: "supervisor", env: {} });
|
||||
s.statusCode = 200;
|
||||
s.ok = true;
|
||||
s.phases.push({ name: "warm_start", durationMs: 12, ok: true, attempts: 1 });
|
||||
s.phases.push({
|
||||
name: "workload_create",
|
||||
durationMs: 3,
|
||||
ok: false,
|
||||
attempts: 2,
|
||||
errorCode: "Error",
|
||||
errorMsg: "boom",
|
||||
sub: { create_ms: 1 },
|
||||
});
|
||||
const out = captureEmit(s);
|
||||
expect(out["phase.warm_start.duration_ms"]).toBe(12);
|
||||
expect(out["phase.warm_start.ok"]).toBe(true);
|
||||
expect(out["phase.warm_start.attempts"]).toBe(1);
|
||||
expect(out["phase.workload_create.duration_ms"]).toBe(3);
|
||||
expect(out["phase.workload_create.ok"]).toBe(false);
|
||||
expect(out["phase.workload_create.attempts"]).toBe(2);
|
||||
expect(out["phase.workload_create.error_code"]).toBe("Error");
|
||||
expect(out["phase.workload_create.error_message"]).toBe("boom");
|
||||
expect(out["phase.workload_create.create_ms"]).toBe(1);
|
||||
});
|
||||
|
||||
it("includes error.code/message/kind when state.error is set", () => {
|
||||
const s = newState({ service: "supervisor", env: {} });
|
||||
s.statusCode = 500;
|
||||
s.error = { code: "InternalError", message: "kaboom", kind: "internal" };
|
||||
const out = captureEmit(s);
|
||||
expect(out["error.code"]).toBe("InternalError");
|
||||
expect(out["error.message"]).toBe("kaboom");
|
||||
expect(out["error.kind"]).toBe("internal");
|
||||
});
|
||||
|
||||
it("truncates very long error messages", () => {
|
||||
const s = newState({ service: "supervisor", env: {} });
|
||||
s.error = { code: "Big", message: "x".repeat(2000), kind: "internal" };
|
||||
const out = captureEmit(s);
|
||||
expect((out["error.message"] as string).length).toBe(512);
|
||||
});
|
||||
|
||||
it("flattens extras at the top level", () => {
|
||||
const s = newState({ service: "supervisor", env: {} });
|
||||
s.statusCode = 200;
|
||||
s.ok = true;
|
||||
s.extras.route = "/health";
|
||||
s.extras["dispatch.result"] = "hit";
|
||||
const out = captureEmit(s);
|
||||
expect(out.route).toBe("/health");
|
||||
expect(out["dispatch.result"]).toBe("hit");
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,84 @@
|
||||
import type { State } from "./state.js";
|
||||
import { truncateUtf8 } from "./truncate.js";
|
||||
|
||||
/**
|
||||
* Stable slog message string for every wide event. Downstream filters (jq,
|
||||
* Axiom queries, Vector pipelines) pin to this constant. The `service` field
|
||||
* disambiguates which service emitted it.
|
||||
*/
|
||||
export const EmitMessage = "wide_event";
|
||||
|
||||
const MAX_ERROR_MSG_BYTES = 512;
|
||||
|
||||
/**
|
||||
* Serializes a State as a single flat-keyed JSON line on stdout. Keys are
|
||||
* flat (no nested objects) to keep jq filtering and Axiom indexing cheap.
|
||||
* Empty optional fields are omitted.
|
||||
*/
|
||||
export function emit(state: State): void {
|
||||
// Best-effort: an observability failure (serialization, stdout write) must
|
||||
// never break or mask the caller's operation. Every call site relies on this.
|
||||
try {
|
||||
const out: Record<string, unknown> = {
|
||||
msg: EmitMessage,
|
||||
request_id: state.requestId,
|
||||
};
|
||||
|
||||
if (state.traceId) out.trace_id = state.traceId;
|
||||
appendIfSet(out, "start_time", state.startTime);
|
||||
appendIfSet(out, "service", state.service);
|
||||
appendIfSet(out, "version", state.version);
|
||||
appendIfSet(out, "commit_sha", state.commitSha);
|
||||
appendIfSet(out, "region", state.region);
|
||||
appendIfSet(out, "node_id", state.nodeId);
|
||||
|
||||
appendIfSet(out, "op", state.op);
|
||||
appendIfSet(out, "kind", state.kind);
|
||||
|
||||
out.ok = state.ok;
|
||||
if (state.statusCode !== 0) out.status = state.statusCode;
|
||||
out.duration_ms = state.durationMs;
|
||||
|
||||
if (state.error) {
|
||||
appendIfSet(out, "error.code", state.error.code);
|
||||
appendIfSet(out, "error.message", truncateUtf8(state.error.message, MAX_ERROR_MSG_BYTES));
|
||||
appendIfSet(out, "error.kind", state.error.kind);
|
||||
}
|
||||
|
||||
for (const [k, v] of Object.entries(state.meta)) {
|
||||
out["meta." + k] = v;
|
||||
}
|
||||
|
||||
for (const p of state.phases) {
|
||||
const prefix = "phase." + p.name + ".";
|
||||
out[prefix + "duration_ms"] = p.durationMs;
|
||||
out[prefix + "ok"] = p.ok;
|
||||
out[prefix + "attempts"] = p.attempts;
|
||||
if (p.errorCode) out[prefix + "error_code"] = p.errorCode;
|
||||
if (p.errorMsg) out[prefix + "error_message"] = p.errorMsg;
|
||||
if (p.sub) {
|
||||
for (const [sk, sv] of Object.entries(p.sub)) {
|
||||
out[prefix + sk] = sv;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (const [k, v] of Object.entries(state.extras)) {
|
||||
out[k] = v;
|
||||
}
|
||||
|
||||
process.stdout.write(JSON.stringify(out) + "\n");
|
||||
} catch (err) {
|
||||
try {
|
||||
process.stderr.write(
|
||||
`wide_event_emit_failed: ${err instanceof Error ? err.message : String(err)}\n`
|
||||
);
|
||||
} catch {
|
||||
// last resort - drop the event rather than throw
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function appendIfSet(out: Record<string, unknown>, key: string, value: string | undefined): void {
|
||||
if (value) out[key] = value;
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
/**
|
||||
* Wide-event observability surface for the supervisor. One flat-keyed JSON
|
||||
* line per natural unit of work (HTTP request, dequeue iteration, socket
|
||||
* lifecycle event). Events join across services via `trace_id` (parsed from
|
||||
* the inbound W3C `traceparent`) and `meta.run_id`.
|
||||
*
|
||||
* Off by default behind a kill switch - the dispatch hotpath runs at high
|
||||
* QPS, so logging pressure must be cleanly removable.
|
||||
*/
|
||||
export { type Env, isValidRequestId, newState, type NewStateOptions } from "./new.js";
|
||||
export { emit, EmitMessage } from "./emit.js";
|
||||
export { parseTraceId } from "./traceparent.js";
|
||||
export { fromContext, wideEventStorage } from "./context.js";
|
||||
export { type PhaseOpt, recordPhase, recordPhaseSince, timePhase } from "./record.js";
|
||||
export {
|
||||
emitOneShot,
|
||||
runWideEvent,
|
||||
setExtra,
|
||||
setMeta,
|
||||
type WideEventLifecycleOptions,
|
||||
type WideEventOptions,
|
||||
} from "./middleware.js";
|
||||
export type { ErrorInfo, PhaseRecord, State } from "./state.js";
|
||||
export { encodeBaggage } from "./baggage.js";
|
||||
@@ -0,0 +1,206 @@
|
||||
import { describe, it, expect } from "vitest";
|
||||
import { fromContext } from "./context.js";
|
||||
import { emitOneShot, runWideEvent, setMeta } from "./middleware.js";
|
||||
|
||||
function captureStdout(fn: () => Promise<unknown> | unknown): Promise<string[]> {
|
||||
const captured: string[] = [];
|
||||
const orig = process.stdout.write;
|
||||
process.stdout.write = ((chunk: unknown) => {
|
||||
captured.push(String(chunk));
|
||||
return true;
|
||||
}) as typeof process.stdout.write;
|
||||
return Promise.resolve(fn())
|
||||
.finally(() => {
|
||||
process.stdout.write = orig;
|
||||
})
|
||||
.then(() => captured);
|
||||
}
|
||||
|
||||
describe("runWideEvent", () => {
|
||||
it("emits one event with ok=true when no statusCode is set", async () => {
|
||||
const lines = await captureStdout(async () => {
|
||||
await runWideEvent(
|
||||
{ service: "supervisor", env: {}, enabled: true, op: "test", route: "/x", method: "POST" },
|
||||
async () => undefined
|
||||
);
|
||||
});
|
||||
expect(lines).toHaveLength(1);
|
||||
const line = lines[0];
|
||||
if (!line) throw new Error("no line");
|
||||
const ev = JSON.parse(line) as Record<string, unknown>;
|
||||
expect(ev.ok).toBe(true);
|
||||
expect(ev.service).toBe("supervisor");
|
||||
expect(ev.route).toBe("/x");
|
||||
expect(ev.method).toBe("POST");
|
||||
expect(typeof ev.duration_ms).toBe("number");
|
||||
expect(typeof ev.request_id).toBe("string");
|
||||
});
|
||||
|
||||
it("derives ok from statusCode set via finalize", async () => {
|
||||
const lines = await captureStdout(async () => {
|
||||
await runWideEvent(
|
||||
{ service: "supervisor", env: {}, enabled: true, op: "test" },
|
||||
async () => undefined,
|
||||
(state) => {
|
||||
state.statusCode = 200;
|
||||
}
|
||||
);
|
||||
});
|
||||
const line = lines[0];
|
||||
if (!line) throw new Error("no line");
|
||||
const ev = JSON.parse(line) as Record<string, unknown>;
|
||||
expect(ev.ok).toBe(true);
|
||||
expect(ev.status).toBe(200);
|
||||
});
|
||||
|
||||
it("treats 4xx as ok=false", async () => {
|
||||
const lines = await captureStdout(async () => {
|
||||
await runWideEvent(
|
||||
{ service: "supervisor", env: {}, enabled: true, op: "test" },
|
||||
async () => undefined,
|
||||
(state) => {
|
||||
state.statusCode = 400;
|
||||
}
|
||||
);
|
||||
});
|
||||
const line = lines[0];
|
||||
if (!line) throw new Error("no line");
|
||||
const ev = JSON.parse(line) as Record<string, unknown>;
|
||||
expect(ev.ok).toBe(false);
|
||||
expect(ev.status).toBe(400);
|
||||
});
|
||||
|
||||
it("emits ok=false with error.kind=internal on throw", async () => {
|
||||
const lines = await captureStdout(async () => {
|
||||
await runWideEvent(
|
||||
{ service: "supervisor", env: {}, enabled: true, op: "test" },
|
||||
async () => {
|
||||
throw new Error("boom");
|
||||
}
|
||||
).catch(() => undefined);
|
||||
});
|
||||
const line = lines[0];
|
||||
if (!line) throw new Error("no line");
|
||||
const ev = JSON.parse(line) as Record<string, unknown>;
|
||||
expect(ev.ok).toBe(false);
|
||||
expect(ev.status).toBe(500);
|
||||
expect(ev["error.kind"]).toBe("internal");
|
||||
expect(ev["error.message"]).toBe("boom");
|
||||
});
|
||||
|
||||
it("threads state through AsyncLocalStorage", async () => {
|
||||
const lines = await captureStdout(async () => {
|
||||
await runWideEvent(
|
||||
{ service: "supervisor", env: {}, enabled: true, op: "test" },
|
||||
async () => {
|
||||
setMeta(fromContext(), "run_id", "run_abc");
|
||||
}
|
||||
);
|
||||
});
|
||||
const line = lines[0];
|
||||
if (!line) throw new Error("no line");
|
||||
const ev = JSON.parse(line) as Record<string, unknown>;
|
||||
expect(ev["meta.run_id"]).toBe("run_abc");
|
||||
expect(ev.ok).toBe(true);
|
||||
});
|
||||
|
||||
it("picks up inbound traceparent for trace_id", async () => {
|
||||
const tp = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01";
|
||||
const lines = await captureStdout(async () => {
|
||||
await runWideEvent(
|
||||
{ service: "supervisor", env: {}, enabled: true, op: "test", traceparent: tp },
|
||||
async () => undefined
|
||||
);
|
||||
});
|
||||
const line = lines[0];
|
||||
if (!line) throw new Error("no line");
|
||||
const ev = JSON.parse(line) as Record<string, unknown>;
|
||||
expect(ev.trace_id).toBe("4bf92f3577b34da6a3ce929d0e0e4736");
|
||||
});
|
||||
|
||||
it("honours setup() to attach meta and extras before fn runs", async () => {
|
||||
const lines = await captureStdout(async () => {
|
||||
await runWideEvent(
|
||||
{
|
||||
service: "supervisor",
|
||||
env: {},
|
||||
enabled: true,
|
||||
op: "test",
|
||||
setup: (state) => {
|
||||
state.meta.run_id = "run_abc";
|
||||
state.extras.iteration = "dequeue";
|
||||
},
|
||||
},
|
||||
async () => undefined
|
||||
);
|
||||
});
|
||||
const line = lines[0];
|
||||
if (!line) throw new Error("no line");
|
||||
const ev = JSON.parse(line) as Record<string, unknown>;
|
||||
expect(ev["meta.run_id"]).toBe("run_abc");
|
||||
expect(ev.iteration).toBe("dequeue");
|
||||
});
|
||||
|
||||
it("short-circuits to pass-through when enabled=false", async () => {
|
||||
let seenState: ReturnType<typeof fromContext> = null;
|
||||
const lines = await captureStdout(async () => {
|
||||
await runWideEvent(
|
||||
{ service: "supervisor", env: {}, enabled: false, op: "test" },
|
||||
async () => {
|
||||
seenState = fromContext();
|
||||
}
|
||||
);
|
||||
});
|
||||
expect(lines).toHaveLength(0);
|
||||
expect(seenState).toBe(null);
|
||||
});
|
||||
|
||||
it("isolates state across concurrent invocations", async () => {
|
||||
const lines = await captureStdout(async () => {
|
||||
await Promise.all(
|
||||
["a", "b", "c"].map((tag) =>
|
||||
runWideEvent({ service: "supervisor", env: {}, enabled: true, op: "test" }, async () => {
|
||||
const s = fromContext();
|
||||
if (!s) throw new Error("no state");
|
||||
s.meta.tag = tag;
|
||||
await new Promise((r) => setTimeout(r, 5));
|
||||
expect(s.meta.tag).toBe(tag);
|
||||
})
|
||||
)
|
||||
);
|
||||
});
|
||||
const tags = lines.map((l) => (JSON.parse(l) as Record<string, unknown>)["meta.tag"]);
|
||||
expect(tags.sort()).toEqual(["a", "b", "c"]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("emitOneShot", () => {
|
||||
it("emits a single event with populated meta when enabled", async () => {
|
||||
const lines = await captureStdout(() => {
|
||||
emitOneShot({
|
||||
service: "supervisor",
|
||||
env: {},
|
||||
enabled: true,
|
||||
op: "test",
|
||||
populate: (s) => {
|
||||
s.meta.run_id = "run_abc";
|
||||
s.extras.event = "run:start";
|
||||
},
|
||||
});
|
||||
});
|
||||
expect(lines).toHaveLength(1);
|
||||
const line = lines[0];
|
||||
if (!line) throw new Error("no line");
|
||||
const ev = JSON.parse(line) as Record<string, unknown>;
|
||||
expect(ev.ok).toBe(true);
|
||||
expect(ev["meta.run_id"]).toBe("run_abc");
|
||||
expect(ev.event).toBe("run:start");
|
||||
});
|
||||
|
||||
it("emits nothing when disabled", async () => {
|
||||
const lines = await captureStdout(() => {
|
||||
emitOneShot({ service: "supervisor", env: {}, enabled: false, op: "test" });
|
||||
});
|
||||
expect(lines).toHaveLength(0);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,132 @@
|
||||
import { emit } from "./emit.js";
|
||||
import { newState, type Env } from "./new.js";
|
||||
import { wideEventStorage } from "./context.js";
|
||||
import type { State } from "./state.js";
|
||||
|
||||
/** Options common to every wide-event lifecycle. */
|
||||
export type WideEventOptions = {
|
||||
service: string;
|
||||
env: Env;
|
||||
/**
|
||||
* Kill switch. When false, lifecycles degenerate into transparent
|
||||
* pass-through - no State allocation, no AsyncLocalStorage run, no emit.
|
||||
* Important for the dispatch hotpath where logging pressure must be
|
||||
* cleanly removable.
|
||||
*/
|
||||
enabled: boolean;
|
||||
};
|
||||
|
||||
/** Per-invocation options layered on top of `WideEventOptions`. */
|
||||
export type WideEventLifecycleOptions = WideEventOptions & {
|
||||
/** Operation discriminator (`instance.create`, `dequeue`, ...). Required. */
|
||||
op: string;
|
||||
/** Event shape: `inbound` | `outbound` | `event` | `scheduled`. Optional. */
|
||||
kind?: string;
|
||||
/** Route template (HTTP only) captured into `extras.route`. */
|
||||
route?: string;
|
||||
/** HTTP method captured into `extras.method`. */
|
||||
method?: string;
|
||||
/** Inbound W3C traceparent (HTTP header, queue message field). */
|
||||
traceparent?: string;
|
||||
/** Inbound request id (e.g. `x-request-id` header). */
|
||||
inboundRequestId?: string;
|
||||
/** Runs after the state is built, before the wrapped fn. Use to attach meta. */
|
||||
setup?: (state: State) => void;
|
||||
};
|
||||
|
||||
/**
|
||||
* Runs `fn` inside an AsyncLocalStorage state and emits one wide event on
|
||||
* completion or error. `finalize` runs after `fn` returns but before emit -
|
||||
* use it to read out-of-band outcome info (e.g. `res.statusCode` for an HTTP
|
||||
* route) and assign to `state.statusCode`. The wrapper computes `ok` from
|
||||
* `statusCode` if it's set; otherwise it defaults to true on success.
|
||||
*
|
||||
* Returns the original `fn` result. When `enabled=false`, `fn` runs unchanged
|
||||
* with no event emitted.
|
||||
*/
|
||||
export async function runWideEvent<T>(
|
||||
opts: WideEventLifecycleOptions,
|
||||
fn: () => Promise<T> | T,
|
||||
finalize?: (state: State) => void
|
||||
): Promise<T> {
|
||||
if (!opts.enabled) {
|
||||
return fn();
|
||||
}
|
||||
|
||||
const state = newState({
|
||||
service: opts.service,
|
||||
env: opts.env,
|
||||
inboundRequestId: opts.inboundRequestId,
|
||||
traceparent: opts.traceparent,
|
||||
op: opts.op,
|
||||
kind: opts.kind,
|
||||
});
|
||||
if (opts.route) state.extras.route = opts.route;
|
||||
if (opts.method) state.extras.method = opts.method;
|
||||
|
||||
const start = performance.now();
|
||||
try {
|
||||
if (opts.setup) opts.setup(state);
|
||||
const result = await wideEventStorage.run(state, () => Promise.resolve(fn()));
|
||||
state.durationMs = Math.round(performance.now() - start);
|
||||
if (finalize) finalize(state);
|
||||
if (state.statusCode !== 0) {
|
||||
state.ok = state.statusCode >= 200 && state.statusCode < 300;
|
||||
} else {
|
||||
state.ok = true;
|
||||
}
|
||||
emit(state);
|
||||
return result;
|
||||
} catch (err) {
|
||||
state.durationMs = Math.round(performance.now() - start);
|
||||
const e = err instanceof Error ? err : new Error(String(err));
|
||||
if (state.statusCode === 0) state.statusCode = 500;
|
||||
state.ok = false;
|
||||
state.error = {
|
||||
code: e.name || "Error",
|
||||
message: e.message,
|
||||
kind: "internal",
|
||||
};
|
||||
emit(state);
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* One-shot wide event with no wrapped operation. Use for socket lifecycle
|
||||
* events (`run:start`, `run:stop`) where there is no surrounding async unit
|
||||
* of work to time. `populate` runs synchronously to attach meta/extras
|
||||
* before emit.
|
||||
*/
|
||||
export function emitOneShot(
|
||||
opts: WideEventOptions & {
|
||||
op: string;
|
||||
kind?: string;
|
||||
traceparent?: string;
|
||||
populate?: (state: State) => void;
|
||||
}
|
||||
): void {
|
||||
if (!opts.enabled) return;
|
||||
const state = newState({
|
||||
service: opts.service,
|
||||
env: opts.env,
|
||||
traceparent: opts.traceparent,
|
||||
op: opts.op,
|
||||
kind: opts.kind,
|
||||
});
|
||||
if (opts.populate) opts.populate(state);
|
||||
state.ok = true;
|
||||
emit(state);
|
||||
}
|
||||
|
||||
/** Convenience accessor for in-handler meta mutation. */
|
||||
export function setMeta(state: State | null, key: string, value: string): void {
|
||||
if (!state) return;
|
||||
state.meta[key] = value;
|
||||
}
|
||||
|
||||
/** Convenience for free-form fields (did_warm_start, dispatch.result, ...). */
|
||||
export function setExtra(state: State | null, key: string, value: unknown): void {
|
||||
if (!state) return;
|
||||
state.extras[key] = value;
|
||||
}
|
||||
@@ -0,0 +1,81 @@
|
||||
import { describe, it, expect } from "vitest";
|
||||
import { isValidRequestId, newState } from "./new.js";
|
||||
|
||||
describe("isValidRequestId", () => {
|
||||
it("accepts visible ASCII", () => {
|
||||
expect(isValidRequestId("req-abc-123_456.7")).toBe(true);
|
||||
expect(isValidRequestId("a")).toBe(true);
|
||||
});
|
||||
|
||||
it("rejects empty string", () => {
|
||||
expect(isValidRequestId("")).toBe(false);
|
||||
});
|
||||
|
||||
it("rejects overlong strings (>128 bytes)", () => {
|
||||
expect(isValidRequestId("a".repeat(128))).toBe(true);
|
||||
expect(isValidRequestId("a".repeat(129))).toBe(false);
|
||||
});
|
||||
|
||||
it("rejects whitespace, newlines, control chars", () => {
|
||||
expect(isValidRequestId("has space")).toBe(false);
|
||||
expect(isValidRequestId("has\ttab")).toBe(false);
|
||||
expect(isValidRequestId("has\nnewline")).toBe(false);
|
||||
expect(isValidRequestId("\x00null")).toBe(false);
|
||||
});
|
||||
|
||||
it("rejects high-bit / non-ASCII", () => {
|
||||
expect(isValidRequestId("café")).toBe(false);
|
||||
expect(isValidRequestId("a\x7f")).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("newState", () => {
|
||||
const env = { version: "1.0.0", commitSha: "abc123", region: "us-east-1", nodeId: "node-1" };
|
||||
|
||||
it("populates service identity from env", () => {
|
||||
const s = newState({ service: "supervisor", env });
|
||||
expect(s.service).toBe("supervisor");
|
||||
expect(s.version).toBe("1.0.0");
|
||||
expect(s.commitSha).toBe("abc123");
|
||||
expect(s.region).toBe("us-east-1");
|
||||
expect(s.nodeId).toBe("node-1");
|
||||
});
|
||||
|
||||
it("mints a fresh request id when none provided", () => {
|
||||
const s = newState({ service: "test", env: {} });
|
||||
expect(s.requestId).toMatch(/^req-[0-9a-f]{32}$/);
|
||||
});
|
||||
|
||||
it("honours a valid inbound request id", () => {
|
||||
const s = newState({ service: "test", env: {}, inboundRequestId: "trace-abc-123" });
|
||||
expect(s.requestId).toBe("trace-abc-123");
|
||||
});
|
||||
|
||||
it("rejects unsafe inbound request id and mints a fresh one", () => {
|
||||
const s = newState({ service: "test", env: {}, inboundRequestId: "has space" });
|
||||
expect(s.requestId).toMatch(/^req-[0-9a-f]{32}$/);
|
||||
});
|
||||
|
||||
it("parses traceparent into traceId and preserves the raw header", () => {
|
||||
const tp = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01";
|
||||
const s = newState({ service: "test", env: {}, traceparent: tp });
|
||||
expect(s.traceId).toBe("4bf92f3577b34da6a3ce929d0e0e4736");
|
||||
expect(s.traceparent).toBe(tp);
|
||||
});
|
||||
|
||||
it("leaves traceId empty when no traceparent provided", () => {
|
||||
const s = newState({ service: "test", env: {} });
|
||||
expect(s.traceId).toBe("");
|
||||
expect(s.traceparent).toBe("");
|
||||
});
|
||||
|
||||
it("initialises empty meta/extras/phases", () => {
|
||||
const s = newState({ service: "test", env: {} });
|
||||
expect(s.meta).toEqual({});
|
||||
expect(s.extras).toEqual({});
|
||||
expect(s.phases).toEqual([]);
|
||||
expect(s.ok).toBe(false);
|
||||
expect(s.statusCode).toBe(0);
|
||||
expect(s.durationMs).toBe(0);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,96 @@
|
||||
import { randomBytes } from "node:crypto";
|
||||
import { parseTraceId } from "./traceparent.js";
|
||||
import type { State } from "./state.js";
|
||||
|
||||
const MAX_REQUEST_ID_LEN = 128;
|
||||
|
||||
/**
|
||||
* Validates an inbound request id. Non-empty, no longer than 128 bytes,
|
||||
* composed entirely of visible ASCII (0x21..0x7E). Rejects newlines, control
|
||||
* characters, whitespace, DEL, high-bit bytes - any of which could poison the
|
||||
* log pipeline if echoed back verbatim.
|
||||
*/
|
||||
export function isValidRequestId(s: string): boolean {
|
||||
if (s.length === 0 || s.length > MAX_REQUEST_ID_LEN) return false;
|
||||
for (let i = 0; i < s.length; i++) {
|
||||
const c = s.charCodeAt(i);
|
||||
if (c < 0x21 || c > 0x7e) return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
/**
|
||||
* Service-level identity that's constant for the lifetime of the process.
|
||||
* Populated once at startup, copied into every State.
|
||||
*/
|
||||
export type Env = {
|
||||
version?: string;
|
||||
commitSha?: string;
|
||||
region?: string;
|
||||
nodeId?: string;
|
||||
};
|
||||
|
||||
export type NewStateOptions = {
|
||||
service: string;
|
||||
env: Env;
|
||||
/** Optional inbound request id (e.g. from `x-request-id`). If unsafe or absent, a fresh `req-<hex>` is minted. */
|
||||
inboundRequestId?: string;
|
||||
/** Optional inbound W3C traceparent (HTTP header, queue message field). */
|
||||
traceparent?: string;
|
||||
/** Operation discriminator. Dotted `noun.verb`. Defaults to empty (set later). */
|
||||
op?: string;
|
||||
/** Event shape: `inbound` | `outbound` | `event` | `scheduled`. Defaults to empty. */
|
||||
kind?: string;
|
||||
};
|
||||
|
||||
/**
|
||||
* Builds a State for a wide-event lifecycle.
|
||||
*
|
||||
* - requestId: honours `inboundRequestId` if present and safe; otherwise
|
||||
* mints a fresh `req-<hex>` id.
|
||||
* - traceId: parsed from the provided traceparent (graceful empty if
|
||||
* absent or malformed).
|
||||
* - traceparent: preserved verbatim for downstream propagation.
|
||||
*/
|
||||
export function newState(opts: NewStateOptions): State {
|
||||
const traceparent = opts.traceparent ?? "";
|
||||
const inbound = opts.inboundRequestId ?? "";
|
||||
const requestId = isValidRequestId(inbound) ? inbound : newRequestId();
|
||||
|
||||
return {
|
||||
startTime: nowRfc3339(),
|
||||
requestId,
|
||||
traceId: parseTraceId(traceparent),
|
||||
traceparent,
|
||||
service: opts.service,
|
||||
version: opts.env.version,
|
||||
commitSha: opts.env.commitSha,
|
||||
region: opts.env.region,
|
||||
nodeId: opts.env.nodeId,
|
||||
op: opts.op ?? "",
|
||||
kind: opts.kind ?? "",
|
||||
meta: {},
|
||||
phases: [],
|
||||
ok: false,
|
||||
statusCode: 0,
|
||||
durationMs: 0,
|
||||
extras: {},
|
||||
};
|
||||
}
|
||||
|
||||
function newRequestId(): string {
|
||||
return "req-" + randomBytes(16).toString("hex");
|
||||
}
|
||||
|
||||
/**
|
||||
* Current wall-clock time as an RFC3339 string with microsecond precision.
|
||||
* `Date.toISOString()` only has millisecond resolution, which is too coarse to
|
||||
* order multiple wide events emitted within the same millisecond.
|
||||
* `performance.timeOrigin + performance.now()` gives a sub-millisecond wall-clock
|
||||
* reading; we append the microsecond digits to the millisecond ISO string.
|
||||
*/
|
||||
function nowRfc3339(): string {
|
||||
const ms = performance.timeOrigin + performance.now();
|
||||
const micros = Math.floor((ms % 1) * 1000); // microseconds within the millisecond (0..999)
|
||||
return new Date(ms).toISOString().slice(0, -1) + String(micros).padStart(3, "0") + "Z";
|
||||
}
|
||||
@@ -0,0 +1,112 @@
|
||||
import { describe, it, expect } from "vitest";
|
||||
import { fromContext, wideEventStorage } from "./context.js";
|
||||
import { recordPhase, recordPhaseSince, timePhase } from "./record.js";
|
||||
import { newState } from "./new.js";
|
||||
import type { State } from "./state.js";
|
||||
|
||||
function makeState(): State {
|
||||
return newState({ service: "test", env: {} });
|
||||
}
|
||||
|
||||
describe("recordPhase", () => {
|
||||
it("appends a successful phase", () => {
|
||||
const s = makeState();
|
||||
recordPhase(s, "lookup", performance.now() - 50, undefined);
|
||||
expect(s.phases).toHaveLength(1);
|
||||
const phase = s.phases[0];
|
||||
if (!phase) throw new Error("missing phase");
|
||||
expect(phase.name).toBe("lookup");
|
||||
expect(phase.ok).toBe(true);
|
||||
expect(phase.attempts).toBe(1);
|
||||
expect(phase.durationMs).toBeGreaterThanOrEqual(45);
|
||||
});
|
||||
|
||||
it("appends a failed phase with error code/message", () => {
|
||||
const s = makeState();
|
||||
recordPhase(s, "dispatch", performance.now(), new Error("nope"));
|
||||
const phase = s.phases[0];
|
||||
if (!phase) throw new Error("missing phase");
|
||||
expect(phase.ok).toBe(false);
|
||||
expect(phase.errorCode).toBe("Error");
|
||||
expect(phase.errorMsg).toBe("nope");
|
||||
});
|
||||
|
||||
it("truncates very long error messages", () => {
|
||||
const s = makeState();
|
||||
recordPhase(s, "x", performance.now(), new Error("y".repeat(2000)));
|
||||
const phase = s.phases[0];
|
||||
if (!phase) throw new Error("missing phase");
|
||||
expect(phase.errorMsg?.length).toBe(512);
|
||||
});
|
||||
|
||||
it("honours opts.attempts", () => {
|
||||
const s = makeState();
|
||||
recordPhase(s, "retry", performance.now(), undefined, { attempts: 3 });
|
||||
expect(s.phases[0]?.attempts).toBe(3);
|
||||
});
|
||||
|
||||
it("attaches sub-timings", () => {
|
||||
const s = makeState();
|
||||
recordPhase(s, "complex", performance.now(), undefined, { sub: { setup_ms: 10, work_ms: 5 } });
|
||||
expect(s.phases[0]?.sub).toEqual({ setup_ms: 10, work_ms: 5 });
|
||||
});
|
||||
|
||||
it("is a no-op when state is null", () => {
|
||||
expect(() => recordPhase(null, "x", performance.now(), undefined)).not.toThrow();
|
||||
});
|
||||
});
|
||||
|
||||
describe("timePhase + AsyncLocalStorage threading", () => {
|
||||
it("records via fromContext on success", async () => {
|
||||
const s = makeState();
|
||||
const value = await wideEventStorage.run(s, () => timePhase("work", async () => 42));
|
||||
expect(value).toBe(42);
|
||||
expect(s.phases).toHaveLength(1);
|
||||
expect(s.phases[0]?.ok).toBe(true);
|
||||
});
|
||||
|
||||
it("records via fromContext on error and rethrows", async () => {
|
||||
const s = makeState();
|
||||
await expect(
|
||||
wideEventStorage.run(s, () =>
|
||||
timePhase("work", async () => {
|
||||
throw new Error("boom");
|
||||
})
|
||||
)
|
||||
).rejects.toThrow("boom");
|
||||
expect(s.phases).toHaveLength(1);
|
||||
expect(s.phases[0]?.ok).toBe(false);
|
||||
expect(s.phases[0]?.errorMsg).toBe("boom");
|
||||
});
|
||||
|
||||
it("runs fn unchanged when no state on context", async () => {
|
||||
const value = await timePhase("work", async () => "ok");
|
||||
expect(value).toBe("ok");
|
||||
});
|
||||
});
|
||||
|
||||
describe("recordPhaseSince", () => {
|
||||
it("records using a caller-captured start time", async () => {
|
||||
const s = makeState();
|
||||
await wideEventStorage.run(s, async () => {
|
||||
const start = performance.now();
|
||||
await new Promise((r) => setTimeout(r, 10));
|
||||
recordPhaseSince("spanning", start, undefined);
|
||||
});
|
||||
expect(s.phases).toHaveLength(1);
|
||||
expect(s.phases[0]?.durationMs).toBeGreaterThanOrEqual(8);
|
||||
});
|
||||
});
|
||||
|
||||
describe("fromContext", () => {
|
||||
it("returns null when no state attached", () => {
|
||||
expect(fromContext()).toBe(null);
|
||||
});
|
||||
|
||||
it("returns the state when inside wideEventStorage.run", () => {
|
||||
const s = makeState();
|
||||
wideEventStorage.run(s, () => {
|
||||
expect(fromContext()).toBe(s);
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,82 @@
|
||||
import { fromContext } from "./context.js";
|
||||
import type { PhaseRecord, State } from "./state.js";
|
||||
import { truncateUtf8 } from "./truncate.js";
|
||||
|
||||
const MAX_ERROR_MSG_BYTES = 512;
|
||||
|
||||
/** Optional knobs for a phase record. */
|
||||
export type PhaseOpt = {
|
||||
/** Attempt count for the phase (default 1). */
|
||||
attempts?: number;
|
||||
/** Sub-timings to fold into `phase.<name>.<key>`. */
|
||||
sub?: Record<string, number>;
|
||||
};
|
||||
|
||||
/**
|
||||
* Appends a phase outcome to `state.phases`. Safe to call on success
|
||||
* (`err === undefined`) and error paths. `errorMsg` is truncated to 512 bytes
|
||||
* to keep the wide event compact. No-op if state is null.
|
||||
*/
|
||||
export function recordPhase(
|
||||
state: State | null,
|
||||
name: string,
|
||||
startMs: number,
|
||||
err: Error | undefined,
|
||||
opts: PhaseOpt = {}
|
||||
): void {
|
||||
if (!state) return;
|
||||
const p: PhaseRecord = {
|
||||
name,
|
||||
durationMs: Math.round(performance.now() - startMs),
|
||||
ok: err === undefined,
|
||||
attempts: opts.attempts ?? 1,
|
||||
};
|
||||
if (err) {
|
||||
p.errorCode = err.name || "Error";
|
||||
p.errorMsg = truncateUtf8(err.message, MAX_ERROR_MSG_BYTES);
|
||||
}
|
||||
if (opts.sub) p.sub = opts.sub;
|
||||
state.phases.push(p);
|
||||
}
|
||||
|
||||
/**
|
||||
* Runs `fn` and appends a phase outcome to the State attached to the current
|
||||
* async context. If no state is on context (test paths, background work),
|
||||
* `fn` runs unchanged. The phase is recorded on both success and error paths
|
||||
* so failed phases still appear in the wide event with duration_ms +
|
||||
* error_code.
|
||||
*/
|
||||
export async function timePhase<T>(
|
||||
name: string,
|
||||
fn: () => Promise<T> | T,
|
||||
opts: PhaseOpt = {}
|
||||
): Promise<T> {
|
||||
const start = performance.now();
|
||||
try {
|
||||
const result = await fn();
|
||||
recordPhase(fromContext(), name, start, undefined, opts);
|
||||
return result;
|
||||
} catch (err) {
|
||||
recordPhase(fromContext(), name, start, asError(err), opts);
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Appends a phase outcome to the State attached to the current async context
|
||||
* using a `startMs` captured by the caller. Use when the phase boundary spans
|
||||
* multiple calls with intermediate error handling that can't fit inside a
|
||||
* single `timePhase` closure. Nil-state safe.
|
||||
*/
|
||||
export function recordPhaseSince(
|
||||
name: string,
|
||||
startMs: number,
|
||||
err: Error | undefined,
|
||||
opts: PhaseOpt = {}
|
||||
): void {
|
||||
recordPhase(fromContext(), name, startMs, err, opts);
|
||||
}
|
||||
|
||||
function asError(e: unknown): Error {
|
||||
return e instanceof Error ? e : new Error(String(e));
|
||||
}
|
||||
@@ -0,0 +1,84 @@
|
||||
/**
|
||||
* Per-event accumulator backing a single wide event. The supervisor emits one
|
||||
* flat-keyed JSON line per natural unit of work (dequeue iteration, HTTP
|
||||
* request, socket lifecycle event). Optional fields are omitted on emit so
|
||||
* events stay compact.
|
||||
*/
|
||||
export type State = {
|
||||
/**
|
||||
* Wall-clock time the event began, as an ISO-8601 string. Emitted as
|
||||
* `start_time` so log collection orders events by when work started rather
|
||||
* than by the collector's ingestion time.
|
||||
*/
|
||||
startTime?: string;
|
||||
|
||||
// Cross-stack correlation.
|
||||
requestId: string;
|
||||
traceId: string;
|
||||
/**
|
||||
* Raw inbound W3C `traceparent`, preserved verbatim so outbound calls can
|
||||
* propagate the same trace context without losing the parent span-id.
|
||||
* Empty when no inbound traceparent was set.
|
||||
*/
|
||||
traceparent: string;
|
||||
|
||||
// Service identity (set by `newState` from Env).
|
||||
service: string;
|
||||
version?: string;
|
||||
commitSha?: string;
|
||||
region?: string;
|
||||
nodeId?: string;
|
||||
|
||||
/**
|
||||
* Operation discriminator. Dotted `noun.verb` (e.g. `instance.create`,
|
||||
* `snapshot.dispatch`). Low cardinality - bounded set per service, not
|
||||
* unbounded. Empty allowed during construction but expected to be set
|
||||
* before emit.
|
||||
*/
|
||||
op: string;
|
||||
|
||||
/**
|
||||
* Event shape. `inbound` for received requests, `outbound` for outgoing
|
||||
* calls, `event` for ambient occurrences with no meaningful duration,
|
||||
* `scheduled` for timer-driven work. Empty allowed; omitted from emit
|
||||
* when empty.
|
||||
*/
|
||||
kind: string;
|
||||
|
||||
// Caller-attached opaque metadata, flattened to `meta.<key>` on emit.
|
||||
meta: Record<string, string>;
|
||||
|
||||
// Per-phase outcomes, in completion order.
|
||||
phases: PhaseRecord[];
|
||||
|
||||
// Top-level outcome (set after the wrapped operation returns).
|
||||
ok: boolean;
|
||||
statusCode: number;
|
||||
durationMs: number;
|
||||
error?: ErrorInfo;
|
||||
|
||||
// Free-form ad-hoc additions (route, method, did_warm_start, ...).
|
||||
extras: Record<string, unknown>;
|
||||
};
|
||||
|
||||
/**
|
||||
* Single named phase outcome. Retries collapse into `attempts > 1` with the
|
||||
* last error reflected in errorCode/errorMsg.
|
||||
*/
|
||||
export type PhaseRecord = {
|
||||
name: string;
|
||||
durationMs: number;
|
||||
ok: boolean;
|
||||
attempts: number;
|
||||
errorCode?: string;
|
||||
errorMsg?: string;
|
||||
sub?: Record<string, number>;
|
||||
};
|
||||
|
||||
/** Top-level error summary for a failed operation. */
|
||||
export type ErrorInfo = {
|
||||
code: string;
|
||||
message: string;
|
||||
/** Coarse classification - "client" | "upstream" | "internal" | "timeout". */
|
||||
kind: string;
|
||||
};
|
||||
@@ -0,0 +1,43 @@
|
||||
import { describe, it, expect } from "vitest";
|
||||
import { parseTraceId } from "./traceparent.js";
|
||||
|
||||
describe("parseTraceId", () => {
|
||||
const validTraceId = "4bf92f3577b34da6a3ce929d0e0e4736";
|
||||
const validHeader = `00-${validTraceId}-00f067aa0ba902b7-01`;
|
||||
|
||||
it("extracts the trace-id from a valid W3C traceparent", () => {
|
||||
expect(parseTraceId(validHeader)).toBe(validTraceId);
|
||||
});
|
||||
|
||||
it("returns empty string for empty/null/undefined input", () => {
|
||||
expect(parseTraceId("")).toBe("");
|
||||
expect(parseTraceId(null)).toBe("");
|
||||
expect(parseTraceId(undefined)).toBe("");
|
||||
});
|
||||
|
||||
it("returns empty for wrong segment count", () => {
|
||||
expect(parseTraceId("00-abc-def")).toBe("");
|
||||
expect(parseTraceId("00-abc-def-01-extra")).toBe("");
|
||||
});
|
||||
|
||||
it("returns empty for non-zero version byte", () => {
|
||||
expect(parseTraceId(`01-${validTraceId}-00f067aa0ba902b7-01`)).toBe("");
|
||||
});
|
||||
|
||||
it("returns empty for wrong-length trace-id", () => {
|
||||
expect(parseTraceId("00-abc-00f067aa0ba902b7-01")).toBe("");
|
||||
});
|
||||
|
||||
it("returns empty for non-hex trace-id", () => {
|
||||
expect(parseTraceId("00-zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz-00f067aa0ba902b7-01")).toBe("");
|
||||
});
|
||||
|
||||
it("returns empty for all-zero trace-id", () => {
|
||||
expect(parseTraceId("00-00000000000000000000000000000000-00f067aa0ba902b7-01")).toBe("");
|
||||
});
|
||||
|
||||
it("accepts uppercase hex", () => {
|
||||
const tid = "4BF92F3577B34DA6A3CE929D0E0E4736";
|
||||
expect(parseTraceId(`00-${tid}-00f067aa0ba902b7-01`)).toBe(tid);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,39 @@
|
||||
/**
|
||||
* Extracts the trace-id from a W3C `traceparent` header. Returns "" when the
|
||||
* header is absent, malformed, or carries an all-zero trace-id.
|
||||
*
|
||||
* Format: `<version>-<trace-id>-<span-id>-<flags>`
|
||||
* version : 2 hex chars, must be "00"
|
||||
* trace-id: 32 hex chars, non-zero
|
||||
* span-id : 16 hex chars (not validated - we only need trace-id)
|
||||
* flags : 2 hex chars (not validated)
|
||||
*/
|
||||
export function parseTraceId(header: string | null | undefined): string {
|
||||
if (!header) return "";
|
||||
const parts = header.split("-");
|
||||
if (parts.length !== 4) return "";
|
||||
if (parts[0] !== "00") return "";
|
||||
const tid = parts[1];
|
||||
if (!tid || tid.length !== 32) return "";
|
||||
if (!isHex(tid)) return "";
|
||||
if (isAllZero(tid)) return "";
|
||||
return tid;
|
||||
}
|
||||
|
||||
function isHex(s: string): boolean {
|
||||
for (let i = 0; i < s.length; i++) {
|
||||
const c = s.charCodeAt(i);
|
||||
const isDigit = c >= 0x30 && c <= 0x39;
|
||||
const isLower = c >= 0x61 && c <= 0x66;
|
||||
const isUpper = c >= 0x41 && c <= 0x46;
|
||||
if (!isDigit && !isLower && !isUpper) return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
function isAllZero(s: string): boolean {
|
||||
for (let i = 0; i < s.length; i++) {
|
||||
if (s.charCodeAt(i) !== 0x30) return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
import { describe, it, expect } from "vitest";
|
||||
import { truncateUtf8 } from "./truncate.js";
|
||||
|
||||
describe("truncateUtf8", () => {
|
||||
it("returns short ASCII unchanged", () => {
|
||||
expect(truncateUtf8("hello", 512)).toBe("hello");
|
||||
});
|
||||
|
||||
it("truncates ASCII to the byte cap", () => {
|
||||
expect(truncateUtf8("x".repeat(1024), 256)).toBe("x".repeat(256));
|
||||
});
|
||||
|
||||
it("never exceeds the byte cap for multibyte input", () => {
|
||||
// "あ" is 3 UTF-8 bytes; 200 of them = 600 bytes.
|
||||
const got = truncateUtf8("あ".repeat(200), 256);
|
||||
expect(Buffer.byteLength(got, "utf8")).toBeLessThanOrEqual(256);
|
||||
});
|
||||
|
||||
it("does not split a multibyte sequence", () => {
|
||||
// 256 / 3 bytes = 85 whole chars (255 bytes), the 86th would overflow.
|
||||
expect(truncateUtf8("あ".repeat(200), 256)).toBe("あ".repeat(85));
|
||||
});
|
||||
|
||||
it("does not split a surrogate pair", () => {
|
||||
// "😀" is 2 UTF-16 units / 4 UTF-8 bytes; only one fits under a 5-byte cap.
|
||||
const got = truncateUtf8("😀😀", 5);
|
||||
expect(got).toBe("😀");
|
||||
expect(Buffer.byteLength(got, "utf8")).toBe(4);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,20 @@
|
||||
/**
|
||||
* Truncate `value` to at most `maxBytes` UTF-8 bytes without splitting a
|
||||
* multi-byte sequence or surrogate pair. Plain `.slice()` counts UTF-16 code
|
||||
* units, so multibyte text can blow past a byte cap and cutting mid-pair
|
||||
* leaves a lone surrogate that downstream JSON / Postgres consumers reject.
|
||||
*/
|
||||
export function truncateUtf8(value: string, maxBytes: number): string {
|
||||
if (Buffer.byteLength(value, "utf8") <= maxBytes) return value;
|
||||
|
||||
let bytes = 0;
|
||||
let end = 0;
|
||||
// `for..of` yields whole code points, so a surrogate pair is never split.
|
||||
for (const ch of value) {
|
||||
const size = Buffer.byteLength(ch, "utf8");
|
||||
if (bytes + size > maxBytes) break;
|
||||
bytes += size;
|
||||
end += ch.length;
|
||||
}
|
||||
return value.slice(0, end);
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
import { readFileSync } from "fs";
|
||||
import { env } from "./env.js";
|
||||
|
||||
export function getWorkerToken() {
|
||||
if (!env.TRIGGER_WORKER_TOKEN.startsWith("file://")) {
|
||||
return env.TRIGGER_WORKER_TOKEN;
|
||||
}
|
||||
|
||||
const tokenPath = env.TRIGGER_WORKER_TOKEN.replace("file://", "");
|
||||
|
||||
console.debug(
|
||||
JSON.stringify({
|
||||
message: "🔑 Reading worker token from file",
|
||||
tokenPath,
|
||||
})
|
||||
);
|
||||
|
||||
try {
|
||||
const token = readFileSync(tokenPath, "utf8").trim();
|
||||
return token;
|
||||
} catch (error) {
|
||||
console.error(`Failed to read worker token from file: ${tokenPath}`, error);
|
||||
throw new Error(
|
||||
`Unable to read worker token from file: ${
|
||||
error instanceof Error ? error.message : "Unknown error"
|
||||
}`
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { ComputeClientError } from "@internal/compute";
|
||||
import { isRetryableCreateError, runnerNameForAttempt } from "./compute.js";
|
||||
|
||||
describe("runnerNameForAttempt", () => {
|
||||
it("keeps the unsuffixed name for the first attempt", () => {
|
||||
expect(runnerNameForAttempt("runner-abc123", 1)).toBe("runner-abc123");
|
||||
});
|
||||
|
||||
it("suffixes retry attempts deterministically", () => {
|
||||
expect(runnerNameForAttempt("runner-abc123", 2)).toBe("runner-abc123-r2");
|
||||
expect(runnerNameForAttempt("runner-abc123", 3)).toBe("runner-abc123-r3");
|
||||
});
|
||||
});
|
||||
|
||||
describe("isRetryableCreateError", () => {
|
||||
it("retries statuses where the create definitely did not commit", () => {
|
||||
expect(isRetryableCreateError(new ComputeClientError(500, "tap busy", "http://gw"))).toBe(true);
|
||||
expect(isRetryableCreateError(new ComputeClientError(503, "no placement", "http://gw"))).toBe(
|
||||
true
|
||||
);
|
||||
});
|
||||
|
||||
it("does not retry lost-response statuses (create may have committed)", () => {
|
||||
expect(isRetryableCreateError(new ComputeClientError(502, "bad gateway", "http://gw"))).toBe(
|
||||
false
|
||||
);
|
||||
expect(
|
||||
isRetryableCreateError(new ComputeClientError(504, "gateway timeout", "http://gw"))
|
||||
).toBe(false);
|
||||
});
|
||||
|
||||
it("does not retry 4xx responses", () => {
|
||||
expect(isRetryableCreateError(new ComputeClientError(400, "bad request", "http://gw"))).toBe(
|
||||
false
|
||||
);
|
||||
expect(isRetryableCreateError(new ComputeClientError(409, "conflict", "http://gw"))).toBe(
|
||||
false
|
||||
);
|
||||
});
|
||||
|
||||
it("does not retry timeouts (instance may still be provisioning)", () => {
|
||||
expect(isRetryableCreateError(new DOMException("timed out", "TimeoutError"))).toBe(false);
|
||||
});
|
||||
|
||||
it("retries network-level fetch failures", () => {
|
||||
expect(isRetryableCreateError(new TypeError("fetch failed"))).toBe(true);
|
||||
});
|
||||
|
||||
it("does not retry unknown errors", () => {
|
||||
expect(isRetryableCreateError(new Error("something else"))).toBe(false);
|
||||
expect(isRetryableCreateError("string error")).toBe(false);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,511 @@
|
||||
import { SimpleStructuredLogger } from "@trigger.dev/core/v3/utils/structuredLogger";
|
||||
import { parseTraceparent } from "@trigger.dev/core/v3/isomorphic";
|
||||
import { flattenAttributes } from "@trigger.dev/core/v3/utils/flattenAttributes";
|
||||
import {
|
||||
type WorkloadManager,
|
||||
type WorkloadManagerCreateOptions,
|
||||
type WorkloadManagerOptions,
|
||||
} from "./types.js";
|
||||
import { ComputeClient, ComputeClientError, stripImageDigest } from "@internal/compute";
|
||||
import { setTimeout as sleep } from "node:timers/promises";
|
||||
import { extractTraceparent, getRunnerId } from "../util.js";
|
||||
import type { OtlpTraceService } from "../services/otlpTraceService.js";
|
||||
import { tryCatch } from "@trigger.dev/core";
|
||||
import { encodeBaggage, fromContext } from "../wideEvents/index.js";
|
||||
|
||||
const DEFAULT_CREATE_MAX_ATTEMPTS = 3;
|
||||
const DEFAULT_CREATE_RETRY_BASE_DELAY_MS = 250;
|
||||
|
||||
/**
|
||||
* TEMPORARY (TRI-10293): a failed create can leave its instance name
|
||||
* registered gateway/fcrun-side until async cleanup runs, so a same-name
|
||||
* retry can 409 against our own residue. Until the gateway cleans up
|
||||
* failed-create registrations properly, retry attempts get a deterministic
|
||||
* suffix. Attempt 1 keeps the unsuffixed name so the non-retry path is
|
||||
* unchanged; the suffixed name flows into both the instance name and
|
||||
* TRIGGER_RUNNER_ID, which downstream flows treat as one opaque
|
||||
* self-reported token. Only attempts following a ComputeClientError are
|
||||
* suffixed - network-failure retries keep the same name on purpose, because
|
||||
* the gateway's name-collision 409 is their safety net against
|
||||
* double-creating an instance whose create response was lost.
|
||||
*/
|
||||
export function runnerNameForAttempt(runnerId: string, attempt: number): string {
|
||||
return attempt === 1 ? runnerId : `${runnerId}-r${attempt}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Whether a failed instance create is worth retrying. Only statuses where
|
||||
* the create definitely did NOT commit are retried: 500 means the agent or
|
||||
* fcrun returned a create error (e.g. a netns slot holding the tap busy, a
|
||||
* full node disk - placement may differ on retry), 503 means the gateway
|
||||
* had nowhere to place it. 502/504 are excluded: the gateway emits those
|
||||
* when it fails to reach the node or read its response, which can happen
|
||||
* AFTER the agent committed the create - and the gateway only records the
|
||||
* instance name on a clean 201, so a same-name retry would miss the
|
||||
* collision check and could double-create the VM on another node. 4xx won't
|
||||
* heal on retry, and timeouts may still be provisioning. Network-level
|
||||
* fetch failures are safe: if the gateway processed the create, its name
|
||||
* index is populated and the retry 409s harmlessly.
|
||||
*/
|
||||
export function isRetryableCreateError(error: unknown): boolean {
|
||||
if (error instanceof ComputeClientError) {
|
||||
return error.status === 500 || error.status === 503;
|
||||
}
|
||||
if (error instanceof DOMException && error.name === "TimeoutError") {
|
||||
return false;
|
||||
}
|
||||
// Network-level fetch failures (gateway briefly unreachable)
|
||||
return error instanceof TypeError;
|
||||
}
|
||||
|
||||
type ComputeWorkloadManagerOptions = WorkloadManagerOptions & {
|
||||
gateway: {
|
||||
url: string;
|
||||
authToken?: string;
|
||||
timeoutMs: number;
|
||||
};
|
||||
snapshots: {
|
||||
enabled: boolean;
|
||||
delayMs: number;
|
||||
dispatchLimit: number;
|
||||
callbackUrl: string;
|
||||
};
|
||||
tracing?: OtlpTraceService;
|
||||
runner: {
|
||||
instanceName: string;
|
||||
otelEndpoint: string;
|
||||
prettyLogs: boolean;
|
||||
sendRunDebugLogs: boolean;
|
||||
};
|
||||
createRetry?: {
|
||||
maxAttempts: number;
|
||||
baseDelayMs: number;
|
||||
};
|
||||
};
|
||||
|
||||
export class ComputeWorkloadManager implements WorkloadManager {
|
||||
private readonly logger = new SimpleStructuredLogger("compute-workload-manager");
|
||||
private readonly compute: ComputeClient;
|
||||
private readonly createMaxAttempts: number;
|
||||
private readonly createRetryBaseDelayMs: number;
|
||||
|
||||
constructor(private opts: ComputeWorkloadManagerOptions) {
|
||||
this.createMaxAttempts = opts.createRetry?.maxAttempts ?? DEFAULT_CREATE_MAX_ATTEMPTS;
|
||||
this.createRetryBaseDelayMs =
|
||||
opts.createRetry?.baseDelayMs ?? DEFAULT_CREATE_RETRY_BASE_DELAY_MS;
|
||||
|
||||
if (opts.workloadApiDomain) {
|
||||
this.logger.warn("⚠️ Custom workload API domain", {
|
||||
domain: opts.workloadApiDomain,
|
||||
});
|
||||
}
|
||||
|
||||
this.compute = new ComputeClient({
|
||||
gatewayUrl: opts.gateway.url,
|
||||
authToken: opts.gateway.authToken,
|
||||
timeoutMs: opts.gateway.timeoutMs,
|
||||
// Forward the current wide-event scope's traceparent + request_id so the
|
||||
// downstream service continues the same trace and joins its own wide
|
||||
// events to ours. Additionally serialize caller-supplied meta labels
|
||||
// into the W3C Baggage header so the downstream service auto-stamps
|
||||
// them even on early-error paths that bail before parsing the body.
|
||||
// When called outside a wide-event scope (or when wide events are
|
||||
// disabled), `fromContext` returns undefined and propagation is skipped.
|
||||
getPropagationHeaders: () => {
|
||||
const state = fromContext();
|
||||
if (!state) return {};
|
||||
const headers: Record<string, string> = { "x-request-id": state.requestId };
|
||||
if (state.traceparent) {
|
||||
headers.traceparent = state.traceparent;
|
||||
}
|
||||
const baggage = encodeBaggage(state.meta);
|
||||
if (baggage) {
|
||||
headers.baggage = baggage;
|
||||
}
|
||||
return headers;
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
get snapshotsEnabled(): boolean {
|
||||
return this.opts.snapshots.enabled;
|
||||
}
|
||||
|
||||
get snapshotDelayMs(): number {
|
||||
return this.opts.snapshots.delayMs;
|
||||
}
|
||||
|
||||
get snapshotDispatchLimit(): number {
|
||||
return this.opts.snapshots.dispatchLimit;
|
||||
}
|
||||
|
||||
get traceSpansEnabled(): boolean {
|
||||
return !!this.opts.tracing;
|
||||
}
|
||||
|
||||
async create(opts: WorkloadManagerCreateOptions) {
|
||||
const runnerId = getRunnerId(opts.runFriendlyId, opts.nextAttemptNumber);
|
||||
|
||||
const envVars: Record<string, string> = {
|
||||
OTEL_EXPORTER_OTLP_ENDPOINT: this.opts.runner.otelEndpoint,
|
||||
TRIGGER_DEQUEUED_AT_MS: String(opts.dequeuedAt.getTime()),
|
||||
TRIGGER_POD_SCHEDULED_AT_MS: String(Date.now()),
|
||||
TRIGGER_ENV_ID: opts.envId,
|
||||
TRIGGER_DEPLOYMENT_ID: opts.deploymentFriendlyId,
|
||||
TRIGGER_DEPLOYMENT_VERSION: opts.deploymentVersion,
|
||||
TRIGGER_RUN_ID: opts.runFriendlyId,
|
||||
TRIGGER_SNAPSHOT_ID: opts.snapshotFriendlyId,
|
||||
TRIGGER_SUPERVISOR_API_PROTOCOL: this.opts.workloadApiProtocol,
|
||||
TRIGGER_SUPERVISOR_API_PORT: String(this.opts.workloadApiPort),
|
||||
TRIGGER_SUPERVISOR_API_DOMAIN: this.opts.workloadApiDomain ?? "",
|
||||
TRIGGER_WORKER_INSTANCE_NAME: this.opts.runner.instanceName,
|
||||
TRIGGER_RUNNER_ID: runnerId,
|
||||
TRIGGER_MACHINE_CPU: String(opts.machine.cpu),
|
||||
TRIGGER_MACHINE_MEMORY: String(opts.machine.memory),
|
||||
PRETTY_LOGS: String(this.opts.runner.prettyLogs),
|
||||
TRIGGER_SEND_RUN_DEBUG_LOGS: String(this.opts.runner.sendRunDebugLogs),
|
||||
};
|
||||
|
||||
if (this.opts.warmStartUrl) {
|
||||
envVars.TRIGGER_WARM_START_URL = this.opts.warmStartUrl;
|
||||
}
|
||||
|
||||
if (this.snapshotsEnabled && this.opts.metadataUrl) {
|
||||
envVars.TRIGGER_METADATA_URL = this.opts.metadataUrl;
|
||||
}
|
||||
|
||||
if (this.opts.heartbeatIntervalSeconds) {
|
||||
envVars.TRIGGER_HEARTBEAT_INTERVAL_SECONDS = String(this.opts.heartbeatIntervalSeconds);
|
||||
}
|
||||
|
||||
if (this.opts.snapshotPollIntervalSeconds) {
|
||||
envVars.TRIGGER_SNAPSHOT_POLL_INTERVAL_SECONDS = String(
|
||||
this.opts.snapshotPollIntervalSeconds
|
||||
);
|
||||
}
|
||||
|
||||
if (this.opts.additionalEnvVars) {
|
||||
Object.assign(envVars, this.opts.additionalEnvVars);
|
||||
}
|
||||
|
||||
// Strip image digest - resolve by tag, not digest
|
||||
const imageRef = stripImageDigest(opts.image);
|
||||
|
||||
// Labels forwarded to the compute provider for network-policy selection.
|
||||
// `org` is always set so every run carries its org identity.
|
||||
const labels: Record<string, string> = {
|
||||
org: opts.orgId,
|
||||
};
|
||||
|
||||
// Wide event: single canonical log line emitted in finally
|
||||
const event: Record<string, unknown> = {
|
||||
// High-cardinality identifiers
|
||||
runId: opts.runFriendlyId,
|
||||
runnerId,
|
||||
envId: opts.envId,
|
||||
envType: opts.envType,
|
||||
orgId: opts.orgId,
|
||||
projectId: opts.projectId,
|
||||
deploymentVersion: opts.deploymentVersion,
|
||||
machine: opts.machine.name,
|
||||
// Environment
|
||||
instanceName: this.opts.runner.instanceName,
|
||||
// Supervisor timing
|
||||
dequeueResponseMs: opts.dequeueResponseMs,
|
||||
pollingIntervalMs: opts.pollingIntervalMs,
|
||||
warmStartCheckMs: opts.warmStartCheckMs,
|
||||
// Request
|
||||
image: imageRef,
|
||||
};
|
||||
|
||||
const startMs = performance.now();
|
||||
|
||||
try {
|
||||
const createRequest = {
|
||||
name: runnerId,
|
||||
image: imageRef,
|
||||
env: envVars,
|
||||
cpu: opts.machine.cpu,
|
||||
memory_gb: opts.machine.memory,
|
||||
metadata: {
|
||||
runId: opts.runFriendlyId,
|
||||
envId: opts.envId,
|
||||
envType: opts.envType,
|
||||
orgId: opts.orgId,
|
||||
projectId: opts.projectId,
|
||||
deploymentVersion: opts.deploymentVersion,
|
||||
machine: opts.machine.name,
|
||||
},
|
||||
...(Object.keys(labels).length > 0 ? { labels } : {}),
|
||||
};
|
||||
|
||||
// Retry transient placement failures instead of abandoning the run: a
|
||||
// swallowed create error leaves the run waiting for the run engine's
|
||||
// PENDING_EXECUTING timeout (minutes) before it is redriven, while a
|
||||
// retried create typically succeeds in under a second (TRI-10293).
|
||||
let error: unknown;
|
||||
let data: Awaited<ReturnType<typeof this.compute.instances.create>> | null | undefined;
|
||||
let attempt = 1;
|
||||
// Set after a ComputeClientError: the failed create may have left its
|
||||
// name registered, so subsequent attempts use a suffixed name.
|
||||
let suffixAttempts = false;
|
||||
for (; attempt <= this.createMaxAttempts; attempt++) {
|
||||
const attemptRunnerId = suffixAttempts ? runnerNameForAttempt(runnerId, attempt) : runnerId;
|
||||
[error, data] = await tryCatch(
|
||||
this.compute.instances.create(
|
||||
attemptRunnerId === runnerId
|
||||
? createRequest
|
||||
: {
|
||||
...createRequest,
|
||||
name: attemptRunnerId,
|
||||
env: { ...envVars, TRIGGER_RUNNER_ID: attemptRunnerId },
|
||||
}
|
||||
)
|
||||
);
|
||||
|
||||
if (!error) {
|
||||
event.runnerId = attemptRunnerId;
|
||||
break;
|
||||
}
|
||||
|
||||
if (error instanceof ComputeClientError) {
|
||||
suffixAttempts = true;
|
||||
}
|
||||
|
||||
this.logger.warn("create instance attempt failed", {
|
||||
runnerId: attemptRunnerId,
|
||||
attempt,
|
||||
error: error instanceof Error ? error.message : String(error),
|
||||
});
|
||||
|
||||
if (!isRetryableCreateError(error) || attempt === this.createMaxAttempts) break;
|
||||
await sleep(this.createRetryBaseDelayMs * attempt);
|
||||
}
|
||||
event.createAttempts = attempt;
|
||||
|
||||
if (error || !data) {
|
||||
event.error = error instanceof Error ? error.message : String(error);
|
||||
event.errorType =
|
||||
error instanceof DOMException && error.name === "TimeoutError" ? "timeout" : "fetch";
|
||||
// Intentional: errors are captured in the wide event, not thrown. This matches
|
||||
// the Docker/K8s managers. The run will eventually time out if scheduling fails.
|
||||
return;
|
||||
}
|
||||
|
||||
event.instanceId = data.id;
|
||||
event.ok = true;
|
||||
|
||||
// Parse timing data from compute response (optional - requires gateway timing flag)
|
||||
if (data._timing) {
|
||||
event.timing = data._timing;
|
||||
}
|
||||
|
||||
this.#emitProvisionSpan(opts, startMs, data._timing);
|
||||
} finally {
|
||||
event.durationMs = Math.round(performance.now() - startMs);
|
||||
event.ok ??= false;
|
||||
this.logger.debug("create instance", event);
|
||||
}
|
||||
}
|
||||
|
||||
async snapshot(opts: { runnerId: string; metadata: Record<string, string> }): Promise<boolean> {
|
||||
const [error] = await tryCatch(
|
||||
this.compute.instances.snapshot(opts.runnerId, {
|
||||
callback: {
|
||||
url: this.opts.snapshots.callbackUrl,
|
||||
metadata: opts.metadata,
|
||||
},
|
||||
})
|
||||
);
|
||||
|
||||
if (error) {
|
||||
this.logger.error("snapshot request failed", {
|
||||
runnerId: opts.runnerId,
|
||||
error: error instanceof Error ? error.message : String(error),
|
||||
});
|
||||
return false;
|
||||
}
|
||||
|
||||
this.logger.debug("snapshot request accepted", { runnerId: opts.runnerId });
|
||||
return true;
|
||||
}
|
||||
|
||||
async deleteInstance(runnerId: string): Promise<boolean> {
|
||||
const [error] = await tryCatch(this.compute.instances.delete(runnerId));
|
||||
|
||||
if (error) {
|
||||
this.logger.error("delete instance failed", {
|
||||
runnerId,
|
||||
error: error instanceof Error ? error.message : String(error),
|
||||
});
|
||||
return false;
|
||||
}
|
||||
|
||||
this.logger.debug("delete instance success", { runnerId });
|
||||
return true;
|
||||
}
|
||||
|
||||
#emitProvisionSpan(opts: WorkloadManagerCreateOptions, startMs: number, timing?: unknown) {
|
||||
if (!this.traceSpansEnabled) return;
|
||||
|
||||
const parsed = parseTraceparent(extractTraceparent(opts.traceContext));
|
||||
if (!parsed) return;
|
||||
|
||||
const endMs = performance.now();
|
||||
const now = Date.now();
|
||||
const provisionStartEpochMs = now - (endMs - startMs);
|
||||
const endEpochMs = now;
|
||||
|
||||
// Span starts at dequeue time so events (dequeue) render in the thin-line section
|
||||
// before "Started". The actual provision call time is in provisionStartEpochMs.
|
||||
// Subtract 1ms so compute span always sorts before the attempt span (same dequeue time)
|
||||
const startEpochMs = opts.dequeuedAt.getTime() - 1;
|
||||
|
||||
const spanAttributes: Record<string, string | number | boolean> = {
|
||||
"compute.type": "create",
|
||||
"compute.provision_start_ms": provisionStartEpochMs,
|
||||
...(timing
|
||||
? (flattenAttributes(timing, "compute") as Record<string, string | number | boolean>)
|
||||
: {}),
|
||||
};
|
||||
|
||||
if (opts.dequeueResponseMs !== undefined) {
|
||||
spanAttributes["supervisor.dequeue_response_ms"] = opts.dequeueResponseMs;
|
||||
}
|
||||
if (opts.warmStartCheckMs !== undefined) {
|
||||
spanAttributes["supervisor.warm_start_check_ms"] = opts.warmStartCheckMs;
|
||||
}
|
||||
|
||||
// Use the platform API URL, not the runner OTLP endpoint (which may be a VM gateway IP)
|
||||
this.opts.tracing?.emit({
|
||||
traceId: parsed.traceId,
|
||||
parentSpanId: parsed.spanId,
|
||||
spanName: "compute.provision",
|
||||
startTimeMs: startEpochMs,
|
||||
endTimeMs: endEpochMs,
|
||||
resourceAttributes: {
|
||||
"ctx.environment.id": opts.envId,
|
||||
"ctx.organization.id": opts.orgId,
|
||||
"ctx.project.id": opts.projectId,
|
||||
"ctx.run.id": opts.runFriendlyId,
|
||||
},
|
||||
spanAttributes,
|
||||
});
|
||||
}
|
||||
|
||||
async restore(opts: {
|
||||
snapshotId: string;
|
||||
runnerId: string;
|
||||
runFriendlyId: string;
|
||||
snapshotFriendlyId: string;
|
||||
machine: { cpu: number; memory: number };
|
||||
// Trace context for OTel span emission
|
||||
traceContext?: Record<string, unknown>;
|
||||
envId?: string;
|
||||
orgId?: string;
|
||||
projectId?: string;
|
||||
hasPrivateLink?: boolean;
|
||||
dequeuedAt?: Date;
|
||||
}): Promise<boolean> {
|
||||
const metadata: Record<string, string> = {
|
||||
TRIGGER_RUNNER_ID: opts.runnerId,
|
||||
TRIGGER_RUN_ID: opts.runFriendlyId,
|
||||
TRIGGER_SNAPSHOT_ID: opts.snapshotFriendlyId,
|
||||
TRIGGER_SUPERVISOR_API_PROTOCOL: this.opts.workloadApiProtocol,
|
||||
TRIGGER_SUPERVISOR_API_PORT: String(this.opts.workloadApiPort),
|
||||
TRIGGER_SUPERVISOR_API_DOMAIN: this.opts.workloadApiDomain ?? "",
|
||||
TRIGGER_WORKER_INSTANCE_NAME: this.opts.runner.instanceName,
|
||||
};
|
||||
|
||||
// Resupply labels on restore (the provider doesn't persist them across a
|
||||
// snapshot). orgId is optional on the restore opts type, so guard it.
|
||||
const labels: Record<string, string> = {};
|
||||
if (opts.orgId) {
|
||||
labels.org = opts.orgId;
|
||||
}
|
||||
|
||||
this.logger.verbose("restore request body", {
|
||||
snapshotId: opts.snapshotId,
|
||||
runnerId: opts.runnerId,
|
||||
});
|
||||
|
||||
const startMs = performance.now();
|
||||
|
||||
const [error] = await tryCatch(
|
||||
this.compute.snapshots.restore(opts.snapshotId, {
|
||||
name: opts.runnerId,
|
||||
metadata,
|
||||
cpu: opts.machine.cpu,
|
||||
memory_gb: opts.machine.memory,
|
||||
...(Object.keys(labels).length > 0 ? { labels } : {}),
|
||||
})
|
||||
);
|
||||
|
||||
const durationMs = Math.round(performance.now() - startMs);
|
||||
|
||||
if (error) {
|
||||
this.logger.error("restore request failed", {
|
||||
snapshotId: opts.snapshotId,
|
||||
runnerId: opts.runnerId,
|
||||
error: error instanceof Error ? error.message : String(error),
|
||||
durationMs,
|
||||
});
|
||||
return false;
|
||||
}
|
||||
|
||||
this.logger.debug("restore request success", {
|
||||
snapshotId: opts.snapshotId,
|
||||
runnerId: opts.runnerId,
|
||||
durationMs,
|
||||
});
|
||||
|
||||
this.#emitRestoreSpan(opts, startMs);
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
#emitRestoreSpan(
|
||||
opts: {
|
||||
snapshotId: string;
|
||||
runnerId: string;
|
||||
runFriendlyId: string;
|
||||
traceContext?: Record<string, unknown>;
|
||||
envId?: string;
|
||||
orgId?: string;
|
||||
projectId?: string;
|
||||
dequeuedAt?: Date;
|
||||
},
|
||||
startMs: number
|
||||
) {
|
||||
if (!this.traceSpansEnabled) return;
|
||||
|
||||
const parsed = parseTraceparent(extractTraceparent(opts.traceContext));
|
||||
if (!parsed || !opts.envId || !opts.orgId || !opts.projectId) return;
|
||||
|
||||
const endMs = performance.now();
|
||||
const now = Date.now();
|
||||
const restoreStartEpochMs = now - (endMs - startMs);
|
||||
const endEpochMs = now;
|
||||
|
||||
// Subtract 1ms so restore span always sorts before the attempt span
|
||||
const startEpochMs = (opts.dequeuedAt?.getTime() ?? restoreStartEpochMs) - 1;
|
||||
|
||||
this.opts.tracing?.emit({
|
||||
traceId: parsed.traceId,
|
||||
parentSpanId: parsed.spanId,
|
||||
spanName: "compute.restore",
|
||||
startTimeMs: startEpochMs,
|
||||
endTimeMs: endEpochMs,
|
||||
resourceAttributes: {
|
||||
"ctx.environment.id": opts.envId,
|
||||
"ctx.organization.id": opts.orgId,
|
||||
"ctx.project.id": opts.projectId,
|
||||
"ctx.run.id": opts.runFriendlyId,
|
||||
},
|
||||
spanAttributes: {
|
||||
"compute.type": "restore",
|
||||
"compute.snapshot_id": opts.snapshotId,
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,305 @@
|
||||
import { SimpleStructuredLogger } from "@trigger.dev/core/v3/utils/structuredLogger";
|
||||
import {
|
||||
type WorkloadManager,
|
||||
type WorkloadManagerCreateOptions,
|
||||
type WorkloadManagerOptions,
|
||||
} from "./types.js";
|
||||
import { env } from "../env.js";
|
||||
import { getDockerHostDomain, getRunnerId, normalizeDockerHostUrl } from "../util.js";
|
||||
import Docker from "dockerode";
|
||||
import { tryCatch } from "@trigger.dev/core";
|
||||
import { ECRAuthService } from "./ecrAuth.js";
|
||||
|
||||
export class DockerWorkloadManager implements WorkloadManager {
|
||||
private readonly logger = new SimpleStructuredLogger("docker-workload-manager");
|
||||
private readonly docker: Docker;
|
||||
|
||||
private readonly runnerNetworks: string[];
|
||||
private readonly staticAuth?: Docker.AuthConfig;
|
||||
private readonly platformOverride?: string;
|
||||
private readonly ecrAuthService?: ECRAuthService;
|
||||
|
||||
constructor(private opts: WorkloadManagerOptions) {
|
||||
this.docker = new Docker({
|
||||
version: env.DOCKER_API_VERSION,
|
||||
});
|
||||
|
||||
if (opts.workloadApiDomain) {
|
||||
this.logger.warn("⚠️ Custom workload API domain", {
|
||||
domain: opts.workloadApiDomain,
|
||||
});
|
||||
}
|
||||
|
||||
this.runnerNetworks = env.DOCKER_RUNNER_NETWORKS.split(",");
|
||||
|
||||
this.platformOverride = env.DOCKER_PLATFORM;
|
||||
if (this.platformOverride) {
|
||||
this.logger.info("🖥️ Platform override", {
|
||||
targetPlatform: this.platformOverride,
|
||||
hostPlatform: process.arch,
|
||||
});
|
||||
}
|
||||
|
||||
if (env.DOCKER_REGISTRY_USERNAME && env.DOCKER_REGISTRY_PASSWORD && env.DOCKER_REGISTRY_URL) {
|
||||
this.logger.info("🐋 Using Docker registry credentials", {
|
||||
username: env.DOCKER_REGISTRY_USERNAME,
|
||||
url: env.DOCKER_REGISTRY_URL,
|
||||
});
|
||||
|
||||
this.staticAuth = {
|
||||
username: env.DOCKER_REGISTRY_USERNAME,
|
||||
password: env.DOCKER_REGISTRY_PASSWORD,
|
||||
serveraddress: env.DOCKER_REGISTRY_URL,
|
||||
};
|
||||
} else if (ECRAuthService.hasAWSCredentials()) {
|
||||
this.logger.info("🐋 AWS credentials found, initializing ECR auth service");
|
||||
this.ecrAuthService = new ECRAuthService();
|
||||
} else {
|
||||
this.logger.warn(
|
||||
"🐋 No Docker registry credentials or AWS credentials provided, skipping auth"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
async create(opts: WorkloadManagerCreateOptions) {
|
||||
this.logger.verbose("create()", { opts });
|
||||
|
||||
const runnerId = getRunnerId(opts.runFriendlyId, opts.nextAttemptNumber);
|
||||
|
||||
// Build environment variables
|
||||
const envVars: string[] = [
|
||||
`OTEL_EXPORTER_OTLP_ENDPOINT=${env.OTEL_EXPORTER_OTLP_ENDPOINT}`,
|
||||
`TRIGGER_DEQUEUED_AT_MS=${opts.dequeuedAt.getTime()}`,
|
||||
`TRIGGER_POD_SCHEDULED_AT_MS=${Date.now()}`,
|
||||
`TRIGGER_ENV_ID=${opts.envId}`,
|
||||
`TRIGGER_DEPLOYMENT_ID=${opts.deploymentFriendlyId}`,
|
||||
`TRIGGER_DEPLOYMENT_VERSION=${opts.deploymentVersion}`,
|
||||
`TRIGGER_RUN_ID=${opts.runFriendlyId}`,
|
||||
`TRIGGER_SNAPSHOT_ID=${opts.snapshotFriendlyId}`,
|
||||
`TRIGGER_SUPERVISOR_API_PROTOCOL=${this.opts.workloadApiProtocol}`,
|
||||
`TRIGGER_SUPERVISOR_API_PORT=${this.opts.workloadApiPort}`,
|
||||
`TRIGGER_SUPERVISOR_API_DOMAIN=${this.opts.workloadApiDomain ?? getDockerHostDomain()}`,
|
||||
`TRIGGER_WORKER_INSTANCE_NAME=${env.TRIGGER_WORKER_INSTANCE_NAME}`,
|
||||
`TRIGGER_RUNNER_ID=${runnerId}`,
|
||||
`TRIGGER_MACHINE_CPU=${opts.machine.cpu}`,
|
||||
`TRIGGER_MACHINE_MEMORY=${opts.machine.memory}`,
|
||||
`PRETTY_LOGS=${env.RUNNER_PRETTY_LOGS}`,
|
||||
`TRIGGER_SEND_RUN_DEBUG_LOGS=${env.SEND_RUN_DEBUG_LOGS}`,
|
||||
];
|
||||
|
||||
if (this.opts.warmStartUrl) {
|
||||
envVars.push(`TRIGGER_WARM_START_URL=${normalizeDockerHostUrl(this.opts.warmStartUrl)}`);
|
||||
}
|
||||
|
||||
if (this.opts.metadataUrl) {
|
||||
envVars.push(`TRIGGER_METADATA_URL=${this.opts.metadataUrl}`);
|
||||
}
|
||||
|
||||
if (this.opts.heartbeatIntervalSeconds) {
|
||||
envVars.push(`TRIGGER_HEARTBEAT_INTERVAL_SECONDS=${this.opts.heartbeatIntervalSeconds}`);
|
||||
}
|
||||
|
||||
if (this.opts.snapshotPollIntervalSeconds) {
|
||||
envVars.push(
|
||||
`TRIGGER_SNAPSHOT_POLL_INTERVAL_SECONDS=${this.opts.snapshotPollIntervalSeconds}`
|
||||
);
|
||||
}
|
||||
|
||||
if (this.opts.additionalEnvVars) {
|
||||
Object.entries(this.opts.additionalEnvVars).forEach(([key, value]) => {
|
||||
envVars.push(`${key}=${value}`);
|
||||
});
|
||||
}
|
||||
|
||||
const hostConfig: Docker.HostConfig = {
|
||||
AutoRemove: !!this.opts.dockerAutoremove,
|
||||
};
|
||||
|
||||
const [firstNetwork, ...remainingNetworks] = this.runnerNetworks;
|
||||
|
||||
// Always attach the first network at container creation time. This has the following benefits:
|
||||
// - If there is only a single network to attach, this will prevent having to make a separate request.
|
||||
// - If there are multiple networks to attach, this will ensure the runner won't also be connected to the bridge network
|
||||
hostConfig.NetworkMode = firstNetwork;
|
||||
|
||||
if (env.DOCKER_ENFORCE_MACHINE_PRESETS) {
|
||||
hostConfig.NanoCpus = opts.machine.cpu * 1e9;
|
||||
hostConfig.Memory = opts.machine.memory * 1024 * 1024 * 1024;
|
||||
}
|
||||
|
||||
let imageRef = opts.image;
|
||||
|
||||
if (env.DOCKER_STRIP_IMAGE_DIGEST) {
|
||||
imageRef = opts.image.split("@")[0]!;
|
||||
}
|
||||
|
||||
const containerCreateOpts: Docker.ContainerCreateOptions = {
|
||||
name: runnerId,
|
||||
Hostname: runnerId,
|
||||
HostConfig: hostConfig,
|
||||
Image: imageRef,
|
||||
AttachStdout: false,
|
||||
AttachStderr: false,
|
||||
AttachStdin: false,
|
||||
};
|
||||
|
||||
if (this.platformOverride) {
|
||||
containerCreateOpts.platform = this.platformOverride;
|
||||
}
|
||||
|
||||
const logger = this.logger.child({ opts, containerCreateOpts });
|
||||
|
||||
const [inspectError, inspectResult] = await tryCatch(this.docker.getImage(imageRef).inspect());
|
||||
|
||||
let shouldPull = !!inspectError;
|
||||
if (this.platformOverride) {
|
||||
const imageArchitecture = inspectResult?.Architecture;
|
||||
|
||||
// When the image architecture doesn't match the platform, we need to pull the image
|
||||
if (imageArchitecture && !this.platformOverride.includes(imageArchitecture)) {
|
||||
shouldPull = true;
|
||||
}
|
||||
}
|
||||
|
||||
// If the image is not present, try to pull it
|
||||
if (shouldPull) {
|
||||
logger.info("Pulling image", {
|
||||
error: inspectError,
|
||||
image: opts.image,
|
||||
targetPlatform: this.platformOverride,
|
||||
imageArchitecture: inspectResult?.Architecture,
|
||||
});
|
||||
|
||||
// Get auth config (static or ECR)
|
||||
const authConfig = await this.getAuthConfig();
|
||||
|
||||
// Ensure the image is present
|
||||
const [createImageError, imageResponseReader] = await tryCatch(
|
||||
this.docker.createImage(authConfig, {
|
||||
fromImage: imageRef,
|
||||
...(this.platformOverride ? { platform: this.platformOverride } : {}),
|
||||
})
|
||||
);
|
||||
if (createImageError) {
|
||||
logger.error("Failed to pull image", { error: createImageError });
|
||||
return;
|
||||
}
|
||||
|
||||
const [imageReadError, imageResponse] = await tryCatch(readAllChunks(imageResponseReader));
|
||||
if (imageReadError) {
|
||||
logger.error("failed to read image response", { error: imageReadError });
|
||||
return;
|
||||
}
|
||||
|
||||
logger.debug("pulled image", { image: opts.image, imageResponse });
|
||||
} else {
|
||||
// Image is present, so we can use it to create the container
|
||||
}
|
||||
|
||||
// Create container
|
||||
const [createContainerError, container] = await tryCatch(
|
||||
this.docker.createContainer({
|
||||
...containerCreateOpts,
|
||||
// Add env vars here so they're not logged
|
||||
Env: envVars,
|
||||
})
|
||||
);
|
||||
|
||||
if (createContainerError) {
|
||||
logger.error("Failed to create container", { error: createContainerError });
|
||||
return;
|
||||
}
|
||||
|
||||
// If there are multiple networks to attach to we need to attach the remaining ones after creation
|
||||
if (remainingNetworks.length > 0) {
|
||||
await this.attachContainerToNetworks({
|
||||
containerId: container.id,
|
||||
networkNames: remainingNetworks,
|
||||
});
|
||||
}
|
||||
|
||||
// Start container
|
||||
const [startError, startResult] = await tryCatch(container.start());
|
||||
|
||||
if (startError) {
|
||||
logger.error("Failed to start container", { error: startError, containerId: container.id });
|
||||
return;
|
||||
}
|
||||
|
||||
logger.debug("create succeeded", { startResult, containerId: container.id });
|
||||
}
|
||||
|
||||
/**
|
||||
* Get authentication config for Docker operations
|
||||
* Uses static credentials if available, otherwise attempts ECR auth
|
||||
*/
|
||||
private async getAuthConfig(): Promise<Docker.AuthConfig | undefined> {
|
||||
// Use static credentials if available
|
||||
if (this.staticAuth) {
|
||||
return this.staticAuth;
|
||||
}
|
||||
|
||||
// Use ECR auth if service is available
|
||||
if (this.ecrAuthService) {
|
||||
const ecrAuth = await this.ecrAuthService.getAuthConfig();
|
||||
return ecrAuth || undefined;
|
||||
}
|
||||
|
||||
// No auth available
|
||||
return undefined;
|
||||
}
|
||||
|
||||
private async attachContainerToNetworks({
|
||||
containerId,
|
||||
networkNames,
|
||||
}: {
|
||||
containerId: string;
|
||||
networkNames: string[];
|
||||
}) {
|
||||
this.logger.debug("Attaching container to networks", { containerId, networkNames });
|
||||
|
||||
const [error, networkResults] = await tryCatch(
|
||||
this.docker.listNetworks({
|
||||
filters: {
|
||||
// Full name matches only to prevent unexpected results
|
||||
name: networkNames.map((name) => `^${name}$`),
|
||||
},
|
||||
})
|
||||
);
|
||||
|
||||
if (error) {
|
||||
this.logger.error("Failed to list networks", { networkNames });
|
||||
return;
|
||||
}
|
||||
|
||||
const results = await Promise.allSettled(
|
||||
networkResults.map((networkInfo) => {
|
||||
const network = this.docker.getNetwork(networkInfo.Id);
|
||||
return network.connect({ Container: containerId });
|
||||
})
|
||||
);
|
||||
|
||||
if (results.some((r) => r.status === "rejected")) {
|
||||
this.logger.error("Failed to attach container to some networks", {
|
||||
containerId,
|
||||
networkNames,
|
||||
results,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
this.logger.debug("Attached container to networks", {
|
||||
containerId,
|
||||
networkNames,
|
||||
results,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
async function readAllChunks(reader: NodeJS.ReadableStream) {
|
||||
const chunks = [];
|
||||
for await (const chunk of reader) {
|
||||
chunks.push(chunk.toString());
|
||||
}
|
||||
return chunks;
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
import { ECRClient, GetAuthorizationTokenCommand } from "@aws-sdk/client-ecr";
|
||||
import { SimpleStructuredLogger } from "@trigger.dev/core/v3/utils/structuredLogger";
|
||||
import { tryCatch } from "@trigger.dev/core";
|
||||
import type Docker from "dockerode";
|
||||
|
||||
interface ECRTokenCache {
|
||||
token: string;
|
||||
username: string;
|
||||
serverAddress: string;
|
||||
expiresAt: Date;
|
||||
}
|
||||
|
||||
export class ECRAuthService {
|
||||
private readonly logger = new SimpleStructuredLogger("ecr-auth-service");
|
||||
private readonly ecrClient: ECRClient;
|
||||
private tokenCache: ECRTokenCache | null = null;
|
||||
|
||||
constructor() {
|
||||
this.ecrClient = new ECRClient();
|
||||
|
||||
this.logger.info("🔐 ECR Auth Service initialized", {
|
||||
region: this.ecrClient.config.region,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if we have AWS credentials configured
|
||||
*/
|
||||
static hasAWSCredentials(): boolean {
|
||||
if (process.env.AWS_ACCESS_KEY_ID && process.env.AWS_SECRET_ACCESS_KEY) {
|
||||
return true;
|
||||
}
|
||||
|
||||
if (
|
||||
process.env.AWS_PROFILE ||
|
||||
process.env.AWS_ROLE_ARN ||
|
||||
process.env.AWS_WEB_IDENTITY_TOKEN_FILE
|
||||
) {
|
||||
return true;
|
||||
}
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if the current token is still valid with a 10-minute buffer
|
||||
*/
|
||||
private isTokenValid(): boolean {
|
||||
if (!this.tokenCache) {
|
||||
return false;
|
||||
}
|
||||
|
||||
const now = new Date();
|
||||
const bufferMs = 10 * 60 * 1000; // 10 minute buffer before expiration
|
||||
return now < new Date(this.tokenCache.expiresAt.getTime() - bufferMs);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get a fresh ECR authorization token from AWS
|
||||
*/
|
||||
private async fetchNewToken(): Promise<ECRTokenCache | null> {
|
||||
const [error, response] = await tryCatch(
|
||||
this.ecrClient.send(new GetAuthorizationTokenCommand({}))
|
||||
);
|
||||
|
||||
if (error) {
|
||||
this.logger.error("Failed to get ECR authorization token", { error });
|
||||
return null;
|
||||
}
|
||||
|
||||
const authData = response.authorizationData?.[0];
|
||||
if (!authData?.authorizationToken || !authData.proxyEndpoint) {
|
||||
this.logger.error("Invalid ECR authorization response", { authData });
|
||||
return null;
|
||||
}
|
||||
|
||||
// Decode the base64 token to get username:password
|
||||
const decoded = Buffer.from(authData.authorizationToken, "base64").toString("utf-8");
|
||||
const [username, password] = decoded.split(":", 2);
|
||||
|
||||
if (!username || !password) {
|
||||
this.logger.error("Failed to parse ECR authorization token");
|
||||
return null;
|
||||
}
|
||||
|
||||
const expiresAt = authData.expiresAt || new Date(Date.now() + 12 * 60 * 60 * 1000); // Default 12 hours
|
||||
|
||||
const tokenCache: ECRTokenCache = {
|
||||
token: password,
|
||||
username,
|
||||
serverAddress: authData.proxyEndpoint,
|
||||
expiresAt,
|
||||
};
|
||||
|
||||
this.logger.info("🔐 Successfully fetched ECR token", {
|
||||
username,
|
||||
serverAddress: authData.proxyEndpoint,
|
||||
expiresAt: expiresAt.toISOString(),
|
||||
});
|
||||
|
||||
return tokenCache;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get ECR auth config for Docker operations
|
||||
* Returns cached token if valid, otherwise fetches a new one
|
||||
*/
|
||||
async getAuthConfig(): Promise<Docker.AuthConfig | null> {
|
||||
// Check if cached token is still valid
|
||||
if (this.isTokenValid()) {
|
||||
this.logger.debug("Using cached ECR token");
|
||||
return {
|
||||
username: this.tokenCache!.username,
|
||||
password: this.tokenCache!.token,
|
||||
serveraddress: this.tokenCache!.serverAddress,
|
||||
};
|
||||
}
|
||||
|
||||
// Fetch new token
|
||||
this.logger.info("Fetching new ECR authorization token");
|
||||
const newToken = await this.fetchNewToken();
|
||||
|
||||
if (!newToken) {
|
||||
return null;
|
||||
}
|
||||
|
||||
// Cache the new token
|
||||
this.tokenCache = newToken;
|
||||
|
||||
return {
|
||||
username: newToken.username,
|
||||
password: newToken.token,
|
||||
serveraddress: newToken.serverAddress,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Clear the cached token (useful for testing or forcing refresh)
|
||||
*/
|
||||
clearCache(): void {
|
||||
this.tokenCache = null;
|
||||
this.logger.debug("ECR token cache cleared");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,582 @@
|
||||
import { SimpleStructuredLogger } from "@trigger.dev/core/v3/utils/structuredLogger";
|
||||
import {
|
||||
type WorkloadManager,
|
||||
type WorkloadManagerCreateOptions,
|
||||
type WorkloadManagerOptions,
|
||||
} from "./types.js";
|
||||
import type {
|
||||
EnvironmentType,
|
||||
MachinePreset,
|
||||
MachinePresetName,
|
||||
PlacementTag,
|
||||
} from "@trigger.dev/core/v3";
|
||||
import { PlacementTagProcessor } from "@trigger.dev/core/v3/serverOnly";
|
||||
import { env } from "../env.js";
|
||||
import { type K8sApi, createK8sApi, type k8s } from "../clients/kubernetes.js";
|
||||
import { getRunnerId } from "../util.js";
|
||||
|
||||
type ResourceQuantities = {
|
||||
[K in "cpu" | "memory" | "ephemeral-storage"]?: string;
|
||||
};
|
||||
|
||||
const cpuRequestRatioByMachinePreset: Record<MachinePresetName, number | undefined> = {
|
||||
micro: env.KUBERNETES_CPU_REQUEST_RATIO_MICRO,
|
||||
"small-1x": env.KUBERNETES_CPU_REQUEST_RATIO_SMALL_1X,
|
||||
"small-2x": env.KUBERNETES_CPU_REQUEST_RATIO_SMALL_2X,
|
||||
"medium-1x": env.KUBERNETES_CPU_REQUEST_RATIO_MEDIUM_1X,
|
||||
"medium-2x": env.KUBERNETES_CPU_REQUEST_RATIO_MEDIUM_2X,
|
||||
"large-1x": env.KUBERNETES_CPU_REQUEST_RATIO_LARGE_1X,
|
||||
"large-2x": env.KUBERNETES_CPU_REQUEST_RATIO_LARGE_2X,
|
||||
};
|
||||
|
||||
const memoryRequestRatioByMachinePreset: Record<MachinePresetName, number | undefined> = {
|
||||
micro: env.KUBERNETES_MEMORY_REQUEST_RATIO_MICRO,
|
||||
"small-1x": env.KUBERNETES_MEMORY_REQUEST_RATIO_SMALL_1X,
|
||||
"small-2x": env.KUBERNETES_MEMORY_REQUEST_RATIO_SMALL_2X,
|
||||
"medium-1x": env.KUBERNETES_MEMORY_REQUEST_RATIO_MEDIUM_1X,
|
||||
"medium-2x": env.KUBERNETES_MEMORY_REQUEST_RATIO_MEDIUM_2X,
|
||||
"large-1x": env.KUBERNETES_MEMORY_REQUEST_RATIO_LARGE_1X,
|
||||
"large-2x": env.KUBERNETES_MEMORY_REQUEST_RATIO_LARGE_2X,
|
||||
};
|
||||
|
||||
export class KubernetesWorkloadManager implements WorkloadManager {
|
||||
private readonly logger = new SimpleStructuredLogger("kubernetes-workload-provider");
|
||||
private k8s: K8sApi;
|
||||
private namespace = env.KUBERNETES_NAMESPACE;
|
||||
private placementTagProcessor: PlacementTagProcessor;
|
||||
|
||||
// Resource settings
|
||||
private readonly cpuRequestMinCores = env.KUBERNETES_CPU_REQUEST_MIN_CORES;
|
||||
private readonly cpuRequestRatio = env.KUBERNETES_CPU_REQUEST_RATIO;
|
||||
private readonly memoryRequestMinGb = env.KUBERNETES_MEMORY_REQUEST_MIN_GB;
|
||||
private readonly memoryRequestRatio = env.KUBERNETES_MEMORY_REQUEST_RATIO;
|
||||
private readonly memoryOverheadGb = env.KUBERNETES_MEMORY_OVERHEAD_GB;
|
||||
|
||||
constructor(private opts: WorkloadManagerOptions) {
|
||||
this.k8s = createK8sApi();
|
||||
this.placementTagProcessor = new PlacementTagProcessor({
|
||||
enabled: env.PLACEMENT_TAGS_ENABLED,
|
||||
prefix: env.PLACEMENT_TAGS_PREFIX,
|
||||
});
|
||||
|
||||
if (opts.workloadApiDomain) {
|
||||
this.logger.warn("[KubernetesWorkloadManager] ⚠️ Custom workload API domain", {
|
||||
domain: opts.workloadApiDomain,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
private addPlacementTags(
|
||||
podSpec: Omit<k8s.V1PodSpec, "containers">,
|
||||
placementTags?: PlacementTag[]
|
||||
): Omit<k8s.V1PodSpec, "containers"> {
|
||||
const nodeSelector = this.placementTagProcessor.convertToNodeSelector(
|
||||
placementTags,
|
||||
podSpec.nodeSelector
|
||||
);
|
||||
|
||||
return {
|
||||
...podSpec,
|
||||
nodeSelector,
|
||||
};
|
||||
}
|
||||
|
||||
private stripImageDigest(imageRef: string): string {
|
||||
if (!env.KUBERNETES_STRIP_IMAGE_DIGEST) {
|
||||
return imageRef;
|
||||
}
|
||||
|
||||
const atIndex = imageRef.lastIndexOf("@");
|
||||
|
||||
if (atIndex === -1) {
|
||||
return imageRef;
|
||||
}
|
||||
|
||||
return imageRef.substring(0, atIndex);
|
||||
}
|
||||
|
||||
private clamp(value: number, min: number, max: number): number {
|
||||
return Math.min(Math.max(value, min), max);
|
||||
}
|
||||
|
||||
async create(opts: WorkloadManagerCreateOptions) {
|
||||
this.logger.verbose("[KubernetesWorkloadManager] Creating container", { opts });
|
||||
|
||||
const runnerId = getRunnerId(opts.runFriendlyId, opts.nextAttemptNumber);
|
||||
|
||||
try {
|
||||
await this.k8s.core.createNamespacedPod({
|
||||
namespace: this.namespace,
|
||||
body: {
|
||||
metadata: {
|
||||
name: runnerId,
|
||||
namespace: this.namespace,
|
||||
labels: {
|
||||
...this.#getSharedLabels(opts),
|
||||
app: "task-run",
|
||||
"app.kubernetes.io/part-of": "trigger-worker",
|
||||
"app.kubernetes.io/component": "create",
|
||||
},
|
||||
},
|
||||
spec: {
|
||||
...this.addPlacementTags(this.#defaultPodSpec, opts.placementTags),
|
||||
affinity: this.#getAffinity(opts),
|
||||
tolerations: this.#getScheduleTolerations(this.#isScheduledRun(opts)),
|
||||
terminationGracePeriodSeconds: 60 * 60,
|
||||
containers: [
|
||||
{
|
||||
name: "run-controller",
|
||||
image: this.stripImageDigest(opts.image),
|
||||
ports: [
|
||||
{
|
||||
containerPort: 8000,
|
||||
},
|
||||
],
|
||||
resources: this.#getResourcesForMachine(opts.machine),
|
||||
env: [
|
||||
{
|
||||
name: "TRIGGER_DEQUEUED_AT_MS",
|
||||
value: opts.dequeuedAt.getTime().toString(),
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_POD_SCHEDULED_AT_MS",
|
||||
value: Date.now().toString(),
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_RUN_ID",
|
||||
value: opts.runFriendlyId,
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_ENV_ID",
|
||||
value: opts.envId,
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_DEPLOYMENT_ID",
|
||||
value: opts.deploymentFriendlyId,
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_DEPLOYMENT_VERSION",
|
||||
value: opts.deploymentVersion,
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_SNAPSHOT_ID",
|
||||
value: opts.snapshotFriendlyId,
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_SUPERVISOR_API_PROTOCOL",
|
||||
value: this.opts.workloadApiProtocol,
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_SUPERVISOR_API_PORT",
|
||||
value: `${this.opts.workloadApiPort}`,
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_SUPERVISOR_API_DOMAIN",
|
||||
...(this.opts.workloadApiDomain
|
||||
? {
|
||||
value: this.opts.workloadApiDomain,
|
||||
}
|
||||
: {
|
||||
valueFrom: {
|
||||
fieldRef: {
|
||||
fieldPath: "status.hostIP",
|
||||
},
|
||||
},
|
||||
}),
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_WORKER_INSTANCE_NAME",
|
||||
valueFrom: {
|
||||
fieldRef: {
|
||||
fieldPath: "spec.nodeName",
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "OTEL_EXPORTER_OTLP_ENDPOINT",
|
||||
value: env.OTEL_EXPORTER_OTLP_ENDPOINT,
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_RUNNER_ID",
|
||||
value: runnerId,
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_MACHINE_CPU",
|
||||
value: `${opts.machine.cpu}`,
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_MACHINE_MEMORY",
|
||||
value: `${opts.machine.memory}`,
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_SEND_RUN_DEBUG_LOGS",
|
||||
value: `${env.SEND_RUN_DEBUG_LOGS}`,
|
||||
},
|
||||
{
|
||||
name: "LIMITS_CPU",
|
||||
valueFrom: {
|
||||
resourceFieldRef: {
|
||||
resource: "limits.cpu",
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "LIMITS_MEMORY",
|
||||
valueFrom: {
|
||||
resourceFieldRef: {
|
||||
resource: "limits.memory",
|
||||
},
|
||||
},
|
||||
},
|
||||
...(this.opts.warmStartUrl
|
||||
? [{ name: "TRIGGER_WARM_START_URL", value: this.opts.warmStartUrl }]
|
||||
: []),
|
||||
...(this.opts.metadataUrl
|
||||
? [{ name: "TRIGGER_METADATA_URL", value: this.opts.metadataUrl }]
|
||||
: []),
|
||||
...(this.opts.heartbeatIntervalSeconds
|
||||
? [
|
||||
{
|
||||
name: "TRIGGER_HEARTBEAT_INTERVAL_SECONDS",
|
||||
value: `${this.opts.heartbeatIntervalSeconds}`,
|
||||
},
|
||||
]
|
||||
: []),
|
||||
...(this.opts.snapshotPollIntervalSeconds
|
||||
? [
|
||||
{
|
||||
name: "TRIGGER_SNAPSHOT_POLL_INTERVAL_SECONDS",
|
||||
value: `${this.opts.snapshotPollIntervalSeconds}`,
|
||||
},
|
||||
]
|
||||
: []),
|
||||
...(this.opts.additionalEnvVars
|
||||
? Object.entries(this.opts.additionalEnvVars).map(([key, value]) => ({
|
||||
name: key,
|
||||
value: value,
|
||||
}))
|
||||
: []),
|
||||
],
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
});
|
||||
} catch (err: unknown) {
|
||||
this.#handleK8sError(err);
|
||||
}
|
||||
}
|
||||
|
||||
#throwUnlessRecord(candidate: unknown): asserts candidate is Record<string, unknown> {
|
||||
if (typeof candidate !== "object" || candidate === null) {
|
||||
throw candidate;
|
||||
}
|
||||
}
|
||||
|
||||
#handleK8sError(err: unknown) {
|
||||
this.#throwUnlessRecord(err);
|
||||
|
||||
if ("body" in err && err.body) {
|
||||
this.logger.error("[KubernetesWorkloadManager] Create failed", { rawError: err.body });
|
||||
this.#throwUnlessRecord(err.body);
|
||||
|
||||
if (typeof err.body.message === "string") {
|
||||
throw new Error(err.body?.message);
|
||||
} else {
|
||||
throw err.body;
|
||||
}
|
||||
} else {
|
||||
this.logger.error("[KubernetesWorkloadManager] Create failed", { rawError: err });
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
|
||||
#envTypeToLabelValue(type: EnvironmentType) {
|
||||
switch (type) {
|
||||
case "PRODUCTION":
|
||||
return "prod";
|
||||
case "STAGING":
|
||||
return "stg";
|
||||
case "DEVELOPMENT":
|
||||
return "dev";
|
||||
case "PREVIEW":
|
||||
return "preview";
|
||||
}
|
||||
}
|
||||
|
||||
private getImagePullSecrets(): k8s.V1LocalObjectReference[] | undefined {
|
||||
return this.opts.imagePullSecrets?.map((name) => ({ name }));
|
||||
}
|
||||
|
||||
get #defaultPodSpec(): Omit<k8s.V1PodSpec, "containers"> {
|
||||
return {
|
||||
restartPolicy: "Never",
|
||||
automountServiceAccountToken: false,
|
||||
imagePullSecrets: this.getImagePullSecrets(),
|
||||
...(env.KUBERNETES_SCHEDULER_NAME
|
||||
? {
|
||||
schedulerName: env.KUBERNETES_SCHEDULER_NAME,
|
||||
}
|
||||
: {}),
|
||||
...(env.KUBERNETES_WORKER_NODETYPE_LABEL
|
||||
? {
|
||||
nodeSelector: {
|
||||
nodetype: env.KUBERNETES_WORKER_NODETYPE_LABEL,
|
||||
},
|
||||
}
|
||||
: {}),
|
||||
...(env.KUBERNETES_POD_DNS_NDOTS_OVERRIDE_ENABLED
|
||||
? {
|
||||
dnsConfig: {
|
||||
options: [{ name: "ndots", value: `${env.KUBERNETES_POD_DNS_NDOTS}` }],
|
||||
},
|
||||
}
|
||||
: {}),
|
||||
};
|
||||
}
|
||||
|
||||
get #defaultResourceRequests(): ResourceQuantities {
|
||||
return {
|
||||
"ephemeral-storage": env.KUBERNETES_EPHEMERAL_STORAGE_SIZE_REQUEST,
|
||||
};
|
||||
}
|
||||
|
||||
get #defaultResourceLimits(): ResourceQuantities {
|
||||
return {
|
||||
"ephemeral-storage": env.KUBERNETES_EPHEMERAL_STORAGE_SIZE_LIMIT,
|
||||
};
|
||||
}
|
||||
|
||||
#isScheduledRun(opts: WorkloadManagerCreateOptions): boolean {
|
||||
return opts.annotations?.rootTriggerSource === "schedule";
|
||||
}
|
||||
|
||||
#getSharedLabels(opts: WorkloadManagerCreateOptions): Record<string, string> {
|
||||
const labels: Record<string, string> = {
|
||||
env: opts.envId,
|
||||
envtype: this.#envTypeToLabelValue(opts.envType),
|
||||
org: opts.orgId,
|
||||
project: opts.projectId,
|
||||
machine: opts.machine.name,
|
||||
// We intentionally use a boolean label rather than exposing the full trigger source
|
||||
// (e.g. sdk, api, cli, mcp, schedule) to keep label cardinality low in metrics.
|
||||
// The schedule vs non-schedule distinction is all we need for the current metrics
|
||||
// and pool-level scheduling decisions; finer-grained source breakdowns live in run annotations.
|
||||
scheduled: String(this.#isScheduledRun(opts)),
|
||||
};
|
||||
|
||||
// Add privatelink label for CiliumNetworkPolicy matching
|
||||
if (opts.hasPrivateLink) {
|
||||
labels.privatelink = opts.orgId;
|
||||
}
|
||||
|
||||
return labels;
|
||||
}
|
||||
|
||||
#getResourceRequestsForMachine(preset: MachinePreset): ResourceQuantities {
|
||||
const cpuRatio = cpuRequestRatioByMachinePreset[preset.name] ?? this.cpuRequestRatio;
|
||||
const memoryRatio = memoryRequestRatioByMachinePreset[preset.name] ?? this.memoryRequestRatio;
|
||||
|
||||
const cpuRequest = preset.cpu * cpuRatio;
|
||||
const memoryRequest = preset.memory * memoryRatio;
|
||||
|
||||
// Clamp between min and max
|
||||
const clampedCpu = this.clamp(cpuRequest, this.cpuRequestMinCores, preset.cpu);
|
||||
const clampedMemory = this.clamp(memoryRequest, this.memoryRequestMinGb, preset.memory);
|
||||
|
||||
return {
|
||||
cpu: `${clampedCpu}`,
|
||||
memory: `${clampedMemory}G`,
|
||||
};
|
||||
}
|
||||
|
||||
#getResourceLimitsForMachine(preset: MachinePreset): ResourceQuantities {
|
||||
const memoryLimit = this.memoryOverheadGb
|
||||
? preset.memory + this.memoryOverheadGb
|
||||
: preset.memory;
|
||||
|
||||
return {
|
||||
cpu: `${preset.cpu}`,
|
||||
memory: `${memoryLimit}G`,
|
||||
};
|
||||
}
|
||||
|
||||
#getResourcesForMachine(preset: MachinePreset): k8s.V1ResourceRequirements {
|
||||
return {
|
||||
requests: {
|
||||
...this.#defaultResourceRequests,
|
||||
...this.#getResourceRequestsForMachine(preset),
|
||||
},
|
||||
limits: {
|
||||
...this.#defaultResourceLimits,
|
||||
...this.#getResourceLimitsForMachine(preset),
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
#isLargeMachine(preset: MachinePreset): boolean {
|
||||
return preset.name.startsWith("large-");
|
||||
}
|
||||
|
||||
#getAffinity(opts: WorkloadManagerCreateOptions): k8s.V1Affinity | undefined {
|
||||
const largeNodeAffinity = this.#getNodeAffinityRules(opts.machine);
|
||||
const scheduleNodeAffinity = this.#getScheduleNodeAffinityRules(this.#isScheduledRun(opts));
|
||||
const podAffinity = this.#getProjectPodAffinity(opts.projectId);
|
||||
|
||||
// Merge node affinity rules from multiple sources
|
||||
const preferred = [
|
||||
...(largeNodeAffinity?.preferredDuringSchedulingIgnoredDuringExecution ?? []),
|
||||
...(scheduleNodeAffinity?.preferredDuringSchedulingIgnoredDuringExecution ?? []),
|
||||
];
|
||||
// Only large machine affinity produces hard requirements (non-large runs must stay off the large pool).
|
||||
// Schedule affinity is soft both ways.
|
||||
const required = [
|
||||
...(largeNodeAffinity?.requiredDuringSchedulingIgnoredDuringExecution?.nodeSelectorTerms ??
|
||||
[]),
|
||||
];
|
||||
|
||||
const hasNodeAffinity = preferred.length > 0 || required.length > 0;
|
||||
|
||||
if (!hasNodeAffinity && !podAffinity) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
return {
|
||||
...(hasNodeAffinity && {
|
||||
nodeAffinity: {
|
||||
...(preferred.length > 0 && {
|
||||
preferredDuringSchedulingIgnoredDuringExecution: preferred,
|
||||
}),
|
||||
...(required.length > 0 && {
|
||||
requiredDuringSchedulingIgnoredDuringExecution: { nodeSelectorTerms: required },
|
||||
}),
|
||||
},
|
||||
}),
|
||||
...(podAffinity && { podAffinity }),
|
||||
};
|
||||
}
|
||||
|
||||
#getNodeAffinityRules(preset: MachinePreset): k8s.V1NodeAffinity | undefined {
|
||||
if (!env.KUBERNETES_LARGE_MACHINE_AFFINITY_ENABLED) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
if (this.#isLargeMachine(preset)) {
|
||||
// soft preference for the large-machine pool, falls back to standard if unavailable
|
||||
return {
|
||||
preferredDuringSchedulingIgnoredDuringExecution: [
|
||||
{
|
||||
weight: env.KUBERNETES_LARGE_MACHINE_AFFINITY_WEIGHT,
|
||||
preference: {
|
||||
matchExpressions: [
|
||||
{
|
||||
key: env.KUBERNETES_LARGE_MACHINE_AFFINITY_POOL_LABEL_KEY,
|
||||
operator: "In",
|
||||
values: [env.KUBERNETES_LARGE_MACHINE_AFFINITY_POOL_LABEL_VALUE],
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
],
|
||||
};
|
||||
}
|
||||
|
||||
// not schedulable in the large-machine pool
|
||||
return {
|
||||
requiredDuringSchedulingIgnoredDuringExecution: {
|
||||
nodeSelectorTerms: [
|
||||
{
|
||||
matchExpressions: [
|
||||
{
|
||||
key: env.KUBERNETES_LARGE_MACHINE_AFFINITY_POOL_LABEL_KEY,
|
||||
operator: "NotIn",
|
||||
values: [env.KUBERNETES_LARGE_MACHINE_AFFINITY_POOL_LABEL_VALUE],
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
#getScheduleNodeAffinityRules(isScheduledRun: boolean): k8s.V1NodeAffinity | undefined {
|
||||
if (
|
||||
!env.KUBERNETES_SCHEDULED_RUN_AFFINITY_ENABLED ||
|
||||
!env.KUBERNETES_SCHEDULED_RUN_AFFINITY_POOL_LABEL_VALUE
|
||||
) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
if (isScheduledRun) {
|
||||
// soft preference for the schedule pool
|
||||
return {
|
||||
preferredDuringSchedulingIgnoredDuringExecution: [
|
||||
{
|
||||
weight: env.KUBERNETES_SCHEDULED_RUN_AFFINITY_WEIGHT,
|
||||
preference: {
|
||||
matchExpressions: [
|
||||
{
|
||||
key: env.KUBERNETES_SCHEDULED_RUN_AFFINITY_POOL_LABEL_KEY,
|
||||
operator: "In",
|
||||
values: [env.KUBERNETES_SCHEDULED_RUN_AFFINITY_POOL_LABEL_VALUE],
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
],
|
||||
};
|
||||
}
|
||||
|
||||
// soft anti-affinity: non-schedule runs prefer to avoid the schedule pool
|
||||
return {
|
||||
preferredDuringSchedulingIgnoredDuringExecution: [
|
||||
{
|
||||
weight: env.KUBERNETES_SCHEDULED_RUN_ANTI_AFFINITY_WEIGHT,
|
||||
preference: {
|
||||
matchExpressions: [
|
||||
{
|
||||
key: env.KUBERNETES_SCHEDULED_RUN_AFFINITY_POOL_LABEL_KEY,
|
||||
operator: "NotIn",
|
||||
values: [env.KUBERNETES_SCHEDULED_RUN_AFFINITY_POOL_LABEL_VALUE],
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
],
|
||||
};
|
||||
}
|
||||
|
||||
#getScheduleTolerations(isScheduledRun: boolean): k8s.V1Toleration[] | undefined {
|
||||
if (!isScheduledRun || !env.KUBERNETES_SCHEDULED_RUN_TOLERATIONS?.length) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
return env.KUBERNETES_SCHEDULED_RUN_TOLERATIONS;
|
||||
}
|
||||
|
||||
#getProjectPodAffinity(projectId: string): k8s.V1PodAffinity | undefined {
|
||||
if (!env.KUBERNETES_PROJECT_AFFINITY_ENABLED) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
return {
|
||||
preferredDuringSchedulingIgnoredDuringExecution: [
|
||||
{
|
||||
weight: env.KUBERNETES_PROJECT_AFFINITY_WEIGHT,
|
||||
podAffinityTerm: {
|
||||
labelSelector: {
|
||||
matchExpressions: [
|
||||
{
|
||||
key: "project",
|
||||
operator: "In",
|
||||
values: [projectId],
|
||||
},
|
||||
],
|
||||
},
|
||||
topologyKey: env.KUBERNETES_PROJECT_AFFINITY_TOPOLOGY_KEY,
|
||||
},
|
||||
},
|
||||
],
|
||||
};
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
import type {
|
||||
EnvironmentType,
|
||||
MachinePreset,
|
||||
PlacementTag,
|
||||
RunAnnotations,
|
||||
} from "@trigger.dev/core/v3";
|
||||
|
||||
export interface WorkloadManagerOptions {
|
||||
workloadApiProtocol: "http" | "https";
|
||||
workloadApiDomain?: string; // If unset, will use orchestrator-specific default
|
||||
workloadApiPort: number;
|
||||
warmStartUrl?: string;
|
||||
metadataUrl?: string;
|
||||
imagePullSecrets?: string[];
|
||||
heartbeatIntervalSeconds?: number;
|
||||
snapshotPollIntervalSeconds?: number;
|
||||
additionalEnvVars?: Record<string, string>;
|
||||
dockerAutoremove?: boolean;
|
||||
}
|
||||
|
||||
export interface WorkloadManager {
|
||||
create: (opts: WorkloadManagerCreateOptions) => Promise<unknown>;
|
||||
}
|
||||
|
||||
export interface WorkloadManagerCreateOptions {
|
||||
image: string;
|
||||
machine: MachinePreset;
|
||||
version: string;
|
||||
nextAttemptNumber?: number;
|
||||
dequeuedAt: Date;
|
||||
placementTags?: PlacementTag[];
|
||||
// Timing context (populated by supervisor handler, included in wide event)
|
||||
dequeueResponseMs?: number;
|
||||
pollingIntervalMs?: number;
|
||||
warmStartCheckMs?: number;
|
||||
// identifiers
|
||||
envId: string;
|
||||
envType: EnvironmentType;
|
||||
orgId: string;
|
||||
projectId: string;
|
||||
deploymentFriendlyId: string;
|
||||
deploymentVersion: string;
|
||||
runId: string;
|
||||
runFriendlyId: string;
|
||||
snapshotId: string;
|
||||
snapshotFriendlyId: string;
|
||||
// Trace context for OTel span emission (W3C format: { traceparent: "00-...", tracestate?: "..." })
|
||||
traceContext?: Record<string, unknown>;
|
||||
annotations?: RunAnnotations;
|
||||
// private networking
|
||||
hasPrivateLink?: boolean;
|
||||
}
|
||||
@@ -0,0 +1,883 @@
|
||||
import { SnapshotCallbackPayloadSchema } from "@internal/compute";
|
||||
import { type CheckpointClient, HttpServer } from "@trigger.dev/core/v3/serverOnly";
|
||||
import { SimpleStructuredLogger } from "@trigger.dev/core/v3/utils/structuredLogger";
|
||||
import {
|
||||
type WorkloadRunSnapshotsSinceResponseBody,
|
||||
type SupervisorHttpClient,
|
||||
WORKLOAD_HEADERS,
|
||||
type WorkloadClientSocketData,
|
||||
type WorkloadClientToServerEvents,
|
||||
type WorkloadContinueRunExecutionResponseBody,
|
||||
WorkloadDebugLogRequestBody,
|
||||
type WorkloadDequeueFromVersionResponseBody,
|
||||
WorkloadHeartbeatRequestBody,
|
||||
type WorkloadHeartbeatResponseBody,
|
||||
WorkloadRunAttemptCompleteRequestBody,
|
||||
type WorkloadRunAttemptCompleteResponseBody,
|
||||
WorkloadRunAttemptStartRequestBody,
|
||||
type WorkloadRunAttemptStartResponseBody,
|
||||
type WorkloadServerToClientEvents,
|
||||
type WorkloadSuspendRunResponseBody,
|
||||
} from "@trigger.dev/core/v3/workers";
|
||||
import EventEmitter from "node:events";
|
||||
import type { IncomingMessage, ServerResponse } from "node:http";
|
||||
import { type Namespace, Server, type Socket } from "socket.io";
|
||||
import { z } from "zod";
|
||||
import { env } from "../env.js";
|
||||
import { register } from "../metrics.js";
|
||||
import {
|
||||
ComputeSnapshotService,
|
||||
type RunTraceContext,
|
||||
} from "../services/computeSnapshotService.js";
|
||||
import type { OtlpTraceService } from "../services/otlpTraceService.js";
|
||||
import {
|
||||
emitOneShot,
|
||||
runWideEvent,
|
||||
setMeta,
|
||||
type State,
|
||||
type WideEventOptions,
|
||||
} from "../wideEvents/index.js";
|
||||
import type { ComputeWorkloadManager } from "../workloadManager/compute.js";
|
||||
|
||||
// Use the official export when upgrading to socket.io@4.8.0
|
||||
interface DefaultEventsMap {
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||||
[event: string]: (...args: any[]) => void;
|
||||
}
|
||||
|
||||
const WorkloadActionParams = z.object({
|
||||
runFriendlyId: z.string(),
|
||||
snapshotFriendlyId: z.string(),
|
||||
});
|
||||
|
||||
// Workloads bundled into customer task images before CLI v4.4.4 use a strict
|
||||
// zod enum for checkpoint type that only allows DOCKER and KUBERNETES. The
|
||||
// workload never reads this field - it only validates the response shape - so
|
||||
// rewriting it to a known value keeps older runners working without affecting
|
||||
// the value stored in the database or seen by internal services.
|
||||
function legacifyCheckpointType<T extends { checkpoint?: { type: string } | null }>(item: T): T {
|
||||
if (item.checkpoint?.type === "COMPUTE") {
|
||||
return { ...item, checkpoint: { ...item.checkpoint, type: "KUBERNETES" } } as T;
|
||||
}
|
||||
return item;
|
||||
}
|
||||
|
||||
type WorkloadServerEvents = {
|
||||
runConnected: [
|
||||
{
|
||||
run: {
|
||||
friendlyId: string;
|
||||
};
|
||||
},
|
||||
];
|
||||
runDisconnected: [
|
||||
{
|
||||
run: {
|
||||
friendlyId: string;
|
||||
};
|
||||
},
|
||||
];
|
||||
};
|
||||
|
||||
type WorkloadServerOptions = {
|
||||
port: number;
|
||||
host?: string;
|
||||
workerClient: SupervisorHttpClient;
|
||||
checkpointClient?: CheckpointClient;
|
||||
computeManager?: ComputeWorkloadManager;
|
||||
tracing?: OtlpTraceService;
|
||||
wideEventOpts: WideEventOptions;
|
||||
/** When true, high-frequency HTTP routes also emit wide events. */
|
||||
wideEventsNoisyRoutes: boolean;
|
||||
};
|
||||
|
||||
export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
|
||||
private checkpointClient?: CheckpointClient;
|
||||
private readonly snapshotService?: ComputeSnapshotService;
|
||||
|
||||
private readonly logger = new SimpleStructuredLogger("workload-server");
|
||||
private readonly wideEventOpts: WideEventOptions;
|
||||
private readonly wideEventsNoisyRoutes: boolean;
|
||||
|
||||
private readonly httpServer: HttpServer;
|
||||
private readonly websocketServer: Namespace<
|
||||
WorkloadClientToServerEvents,
|
||||
WorkloadServerToClientEvents,
|
||||
DefaultEventsMap,
|
||||
WorkloadClientSocketData
|
||||
>;
|
||||
|
||||
private readonly runSockets = new Map<
|
||||
string,
|
||||
Socket<
|
||||
WorkloadClientToServerEvents,
|
||||
WorkloadServerToClientEvents,
|
||||
DefaultEventsMap,
|
||||
WorkloadClientSocketData
|
||||
>
|
||||
>();
|
||||
|
||||
private readonly workerClient: SupervisorHttpClient;
|
||||
|
||||
constructor(opts: WorkloadServerOptions) {
|
||||
super();
|
||||
|
||||
const host = opts.host ?? "0.0.0.0";
|
||||
const port = opts.port;
|
||||
|
||||
this.workerClient = opts.workerClient;
|
||||
this.checkpointClient = opts.checkpointClient;
|
||||
this.wideEventOpts = opts.wideEventOpts;
|
||||
this.wideEventsNoisyRoutes = opts.wideEventsNoisyRoutes;
|
||||
|
||||
if (opts.computeManager?.snapshotsEnabled) {
|
||||
this.snapshotService = new ComputeSnapshotService({
|
||||
computeManager: opts.computeManager,
|
||||
workerClient: opts.workerClient,
|
||||
tracing: opts.tracing,
|
||||
wideEventOpts: this.wideEventOpts,
|
||||
});
|
||||
}
|
||||
|
||||
this.httpServer = this.createHttpServer({ host, port });
|
||||
this.websocketServer = this.createWebsocketServer();
|
||||
}
|
||||
|
||||
private headerValueFromRequest(req: IncomingMessage, headerName: string): string | undefined {
|
||||
const value = req.headers[headerName];
|
||||
|
||||
if (Array.isArray(value)) {
|
||||
return value[0];
|
||||
}
|
||||
|
||||
return value;
|
||||
}
|
||||
|
||||
private runnerIdFromRequest(req: IncomingMessage): string | undefined {
|
||||
return this.headerValueFromRequest(req, WORKLOAD_HEADERS.RUNNER_ID);
|
||||
}
|
||||
|
||||
private deploymentIdFromRequest(req: IncomingMessage): string | undefined {
|
||||
return this.headerValueFromRequest(req, WORKLOAD_HEADERS.DEPLOYMENT_ID);
|
||||
}
|
||||
|
||||
private deploymentVersionFromRequest(req: IncomingMessage): string | undefined {
|
||||
return this.headerValueFromRequest(req, WORKLOAD_HEADERS.DEPLOYMENT_VERSION);
|
||||
}
|
||||
|
||||
private projectRefFromRequest(req: IncomingMessage): string | undefined {
|
||||
return this.headerValueFromRequest(req, WORKLOAD_HEADERS.PROJECT_REF);
|
||||
}
|
||||
|
||||
/**
|
||||
* Sets common route meta on the wide-event state from URL params.
|
||||
*/
|
||||
private attachRouteMeta(state: State, params: unknown): void {
|
||||
if (!params || typeof params !== "object") return;
|
||||
const p = params as Record<string, unknown>;
|
||||
if (typeof p.runFriendlyId === "string") setMeta(state, "run_id", p.runFriendlyId);
|
||||
if (typeof p.snapshotFriendlyId === "string") {
|
||||
setMeta(state, "snapshot_id", p.snapshotFriendlyId);
|
||||
}
|
||||
if (typeof p.deploymentId === "string") setMeta(state, "deployment_id", p.deploymentId);
|
||||
}
|
||||
|
||||
/**
|
||||
* Wraps an HTTP route handler body with the wide-event lifecycle. Reads
|
||||
* `traceparent` and `x-request-id` from `req.headers`, attaches `run_id` /
|
||||
* `snapshot_id` / `deployment_id` meta from `params` when present, and
|
||||
* captures the response status from `res.statusCode` after `fn` returns.
|
||||
*
|
||||
* Pass `highFrequency: true` for noisy routes (heartbeat, polling). Those
|
||||
* still go through the wrapper but only emit when
|
||||
* `TRIGGER_WIDE_EVENTS_NOISY_ROUTES` is on, so prod can keep them dark
|
||||
* while test envs capture full-fidelity traffic for debugging.
|
||||
*/
|
||||
private wideRoute<T>(
|
||||
ctx: { req: IncomingMessage; res: ServerResponse; params?: unknown },
|
||||
op: string,
|
||||
route: string,
|
||||
method: string,
|
||||
fn: () => Promise<T> | T,
|
||||
routeOpts: { highFrequency?: boolean } = {}
|
||||
): Promise<T> {
|
||||
const enabled =
|
||||
this.wideEventOpts.enabled && (!routeOpts.highFrequency || this.wideEventsNoisyRoutes);
|
||||
return runWideEvent(
|
||||
{
|
||||
...this.wideEventOpts,
|
||||
enabled,
|
||||
op,
|
||||
kind: "inbound",
|
||||
route,
|
||||
method,
|
||||
traceparent: this.headerValueFromRequest(ctx.req, "traceparent"),
|
||||
inboundRequestId: this.headerValueFromRequest(ctx.req, "x-request-id"),
|
||||
setup: (state) => this.attachRouteMeta(state, ctx.params),
|
||||
},
|
||||
fn,
|
||||
(state) => {
|
||||
state.statusCode = ctx.res.statusCode;
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
private createHttpServer({ host, port }: { host: string; port: number }) {
|
||||
const httpServer = new HttpServer({
|
||||
port,
|
||||
host,
|
||||
metrics: {
|
||||
register,
|
||||
expose: false,
|
||||
},
|
||||
})
|
||||
.route("/health", "GET", {
|
||||
handler: async ({ reply }) => {
|
||||
reply.text("OK");
|
||||
},
|
||||
})
|
||||
.route(
|
||||
"/api/v1/workload-actions/runs/:runFriendlyId/snapshots/:snapshotFriendlyId/attempts/start",
|
||||
"POST",
|
||||
{
|
||||
paramsSchema: WorkloadActionParams,
|
||||
bodySchema: WorkloadRunAttemptStartRequestBody,
|
||||
handler: async (ctx) =>
|
||||
this.wideRoute(
|
||||
ctx,
|
||||
"attempt.start",
|
||||
"/api/v1/workload-actions/runs/:runFriendlyId/snapshots/:snapshotFriendlyId/attempts/start",
|
||||
"POST",
|
||||
async () => {
|
||||
const { req, reply, params, body } = ctx;
|
||||
const startResponse = await this.workerClient.startRunAttempt(
|
||||
params.runFriendlyId,
|
||||
params.snapshotFriendlyId,
|
||||
body,
|
||||
this.runnerIdFromRequest(req)
|
||||
);
|
||||
|
||||
if (!startResponse.success) {
|
||||
this.logger.error("Failed to start run", {
|
||||
params,
|
||||
error: startResponse.error,
|
||||
});
|
||||
reply.empty(500);
|
||||
return;
|
||||
}
|
||||
|
||||
reply.json(startResponse.data satisfies WorkloadRunAttemptStartResponseBody);
|
||||
return;
|
||||
}
|
||||
),
|
||||
}
|
||||
)
|
||||
.route(
|
||||
"/api/v1/workload-actions/runs/:runFriendlyId/snapshots/:snapshotFriendlyId/attempts/complete",
|
||||
"POST",
|
||||
{
|
||||
paramsSchema: WorkloadActionParams,
|
||||
bodySchema: WorkloadRunAttemptCompleteRequestBody,
|
||||
handler: async (ctx) =>
|
||||
this.wideRoute(
|
||||
ctx,
|
||||
"attempt.complete",
|
||||
"/api/v1/workload-actions/runs/:runFriendlyId/snapshots/:snapshotFriendlyId/attempts/complete",
|
||||
"POST",
|
||||
async () => {
|
||||
const { req, reply, params, body } = ctx;
|
||||
const runnerId = this.runnerIdFromRequest(req);
|
||||
|
||||
// A completion attempt invalidates any pending delayed snapshot
|
||||
// regardless of outcome: the runner has finished executing, so the
|
||||
// suspended state the snapshot was scheduled to capture no longer
|
||||
// exists. Cancel BEFORE the async completion call - the timer
|
||||
// wheel can tick during the await, so cancelling after it leaves
|
||||
// a real window for a due snapshot to dispatch and pause a VM
|
||||
// that has moved on. The runnerId guard keeps a stale duplicate
|
||||
// runner's completion from cancelling a fresh runner's snapshot,
|
||||
// and the runner can't schedule a new suspend until it receives
|
||||
// this route's reply, so nothing legitimate can be cancelled here.
|
||||
this.snapshotService?.cancel(params.runFriendlyId, runnerId);
|
||||
|
||||
const completeResponse = await this.workerClient.completeRunAttempt(
|
||||
params.runFriendlyId,
|
||||
params.snapshotFriendlyId,
|
||||
body,
|
||||
runnerId
|
||||
);
|
||||
|
||||
if (!completeResponse.success) {
|
||||
this.logger.error("Failed to complete run", {
|
||||
params,
|
||||
error: completeResponse.error,
|
||||
});
|
||||
reply.empty(500);
|
||||
return;
|
||||
}
|
||||
|
||||
reply.json(completeResponse.data satisfies WorkloadRunAttemptCompleteResponseBody);
|
||||
return;
|
||||
}
|
||||
),
|
||||
}
|
||||
)
|
||||
.route(
|
||||
"/api/v1/workload-actions/runs/:runFriendlyId/snapshots/:snapshotFriendlyId/heartbeat",
|
||||
"POST",
|
||||
{
|
||||
paramsSchema: WorkloadActionParams,
|
||||
bodySchema: WorkloadHeartbeatRequestBody,
|
||||
handler: async (ctx) =>
|
||||
this.wideRoute(
|
||||
ctx,
|
||||
"heartbeat",
|
||||
"/api/v1/workload-actions/runs/:runFriendlyId/snapshots/:snapshotFriendlyId/heartbeat",
|
||||
"POST",
|
||||
async () => {
|
||||
const { req, reply, params, body } = ctx;
|
||||
const heartbeatResponse = await this.workerClient.heartbeatRun(
|
||||
params.runFriendlyId,
|
||||
params.snapshotFriendlyId,
|
||||
body,
|
||||
this.runnerIdFromRequest(req)
|
||||
);
|
||||
|
||||
if (!heartbeatResponse.success) {
|
||||
this.logger.error("Failed to heartbeat run", {
|
||||
params,
|
||||
error: heartbeatResponse.error,
|
||||
});
|
||||
reply.empty(500);
|
||||
return;
|
||||
}
|
||||
|
||||
reply.json({
|
||||
ok: true,
|
||||
} satisfies WorkloadHeartbeatResponseBody);
|
||||
},
|
||||
{ highFrequency: true }
|
||||
),
|
||||
}
|
||||
)
|
||||
.route(
|
||||
"/api/v1/workload-actions/runs/:runFriendlyId/snapshots/:snapshotFriendlyId/suspend",
|
||||
"GET",
|
||||
{
|
||||
paramsSchema: WorkloadActionParams,
|
||||
handler: async (ctx) =>
|
||||
this.wideRoute(
|
||||
ctx,
|
||||
"suspend",
|
||||
"/api/v1/workload-actions/runs/:runFriendlyId/snapshots/:snapshotFriendlyId/suspend",
|
||||
"GET",
|
||||
async () => {
|
||||
const { reply, params, req } = ctx;
|
||||
const runnerId = this.runnerIdFromRequest(req);
|
||||
const deploymentVersion = this.deploymentVersionFromRequest(req);
|
||||
const projectRef = this.projectRefFromRequest(req);
|
||||
|
||||
this.logger.debug("Suspend request", {
|
||||
params,
|
||||
runnerId,
|
||||
deploymentVersion,
|
||||
projectRef,
|
||||
});
|
||||
|
||||
if (!runnerId || !deploymentVersion || !projectRef) {
|
||||
this.logger.error("Invalid headers for suspend request", {
|
||||
...params,
|
||||
runnerId,
|
||||
deploymentVersion,
|
||||
projectRef,
|
||||
});
|
||||
reply.json(
|
||||
{
|
||||
ok: false,
|
||||
error: "Invalid headers",
|
||||
} satisfies WorkloadSuspendRunResponseBody,
|
||||
false,
|
||||
400
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
if (this.snapshotService) {
|
||||
// Compute mode: delay snapshot to avoid wasted work on short-lived waitpoints.
|
||||
// If the run continues before the delay expires, the snapshot is cancelled.
|
||||
reply.json({ ok: true } satisfies WorkloadSuspendRunResponseBody, false, 202);
|
||||
|
||||
this.snapshotService.schedule(params.runFriendlyId, {
|
||||
runnerId,
|
||||
runFriendlyId: params.runFriendlyId,
|
||||
snapshotFriendlyId: params.snapshotFriendlyId,
|
||||
});
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
if (!this.checkpointClient) {
|
||||
reply.json(
|
||||
{
|
||||
ok: false,
|
||||
error: "Checkpoints disabled",
|
||||
} satisfies WorkloadSuspendRunResponseBody,
|
||||
false,
|
||||
400
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
reply.json(
|
||||
{
|
||||
ok: true,
|
||||
} satisfies WorkloadSuspendRunResponseBody,
|
||||
false,
|
||||
202
|
||||
);
|
||||
|
||||
const suspendResult = await this.checkpointClient.suspendRun({
|
||||
runFriendlyId: params.runFriendlyId,
|
||||
snapshotFriendlyId: params.snapshotFriendlyId,
|
||||
body: {
|
||||
runnerId,
|
||||
runId: params.runFriendlyId,
|
||||
snapshotId: params.snapshotFriendlyId,
|
||||
projectRef,
|
||||
deploymentVersion,
|
||||
},
|
||||
});
|
||||
|
||||
if (!suspendResult) {
|
||||
this.logger.error("Failed to suspend run", { params });
|
||||
return;
|
||||
}
|
||||
}
|
||||
),
|
||||
}
|
||||
)
|
||||
.route(
|
||||
"/api/v1/workload-actions/runs/:runFriendlyId/snapshots/:snapshotFriendlyId/continue",
|
||||
"GET",
|
||||
{
|
||||
paramsSchema: WorkloadActionParams,
|
||||
handler: async (ctx) =>
|
||||
this.wideRoute(
|
||||
ctx,
|
||||
"continue",
|
||||
"/api/v1/workload-actions/runs/:runFriendlyId/snapshots/:snapshotFriendlyId/continue",
|
||||
"GET",
|
||||
async () => {
|
||||
const { req, reply, params } = ctx;
|
||||
this.logger.debug("Run continuation request", { params });
|
||||
|
||||
// Cancel any pending delayed snapshot for this run
|
||||
this.snapshotService?.cancel(params.runFriendlyId);
|
||||
|
||||
const continuationResult = await this.workerClient.continueRunExecution(
|
||||
params.runFriendlyId,
|
||||
params.snapshotFriendlyId,
|
||||
this.runnerIdFromRequest(req)
|
||||
);
|
||||
|
||||
if (!continuationResult.success) {
|
||||
this.logger.error("Failed to continue run execution", { params });
|
||||
reply.json(
|
||||
{
|
||||
ok: false,
|
||||
error: "Failed to continue run execution",
|
||||
},
|
||||
false,
|
||||
400
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
reply.json(continuationResult.data as WorkloadContinueRunExecutionResponseBody);
|
||||
}
|
||||
),
|
||||
}
|
||||
)
|
||||
.route(
|
||||
"/api/v1/workload-actions/runs/:runFriendlyId/snapshots/since/:snapshotFriendlyId",
|
||||
"GET",
|
||||
{
|
||||
paramsSchema: WorkloadActionParams,
|
||||
handler: async (ctx) =>
|
||||
this.wideRoute(
|
||||
ctx,
|
||||
"snapshots.since",
|
||||
"/api/v1/workload-actions/runs/:runFriendlyId/snapshots/since/:snapshotFriendlyId",
|
||||
"GET",
|
||||
async () => {
|
||||
const { req, reply, params } = ctx;
|
||||
const sinceSnapshotResponse = await this.workerClient.getSnapshotsSince(
|
||||
params.runFriendlyId,
|
||||
params.snapshotFriendlyId,
|
||||
this.runnerIdFromRequest(req)
|
||||
);
|
||||
|
||||
if (!sinceSnapshotResponse.success) {
|
||||
this.logger.error("Failed to get snapshots since", {
|
||||
runId: params.runFriendlyId,
|
||||
error: sinceSnapshotResponse.error,
|
||||
});
|
||||
reply.empty(500);
|
||||
return;
|
||||
}
|
||||
|
||||
reply.json({
|
||||
snapshots: sinceSnapshotResponse.data.snapshots.map(legacifyCheckpointType),
|
||||
} satisfies WorkloadRunSnapshotsSinceResponseBody);
|
||||
},
|
||||
{ highFrequency: true }
|
||||
),
|
||||
}
|
||||
)
|
||||
.route("/api/v1/workload-actions/deployments/:deploymentId/dequeue", "GET", {
|
||||
paramsSchema: z.object({
|
||||
deploymentId: z.string(),
|
||||
}),
|
||||
|
||||
handler: async (ctx) =>
|
||||
this.wideRoute(
|
||||
ctx,
|
||||
"deployment.dequeue",
|
||||
"/api/v1/workload-actions/deployments/:deploymentId/dequeue",
|
||||
"GET",
|
||||
async () => {
|
||||
const { req, reply, params } = ctx;
|
||||
const dequeueResponse = await this.workerClient.dequeueFromVersion(
|
||||
params.deploymentId,
|
||||
1,
|
||||
this.runnerIdFromRequest(req)
|
||||
);
|
||||
|
||||
if (!dequeueResponse.success) {
|
||||
this.logger.error("Failed to get latest snapshot", {
|
||||
deploymentId: params.deploymentId,
|
||||
error: dequeueResponse.error,
|
||||
});
|
||||
reply.empty(500);
|
||||
return;
|
||||
}
|
||||
|
||||
reply.json(
|
||||
dequeueResponse.data.map(
|
||||
legacifyCheckpointType
|
||||
) satisfies WorkloadDequeueFromVersionResponseBody
|
||||
);
|
||||
}
|
||||
),
|
||||
});
|
||||
|
||||
if (env.SEND_RUN_DEBUG_LOGS) {
|
||||
httpServer.route("/api/v1/workload-actions/runs/:runFriendlyId/logs/debug", "POST", {
|
||||
paramsSchema: WorkloadActionParams.pick({ runFriendlyId: true }),
|
||||
bodySchema: WorkloadDebugLogRequestBody,
|
||||
handler: async (ctx) =>
|
||||
this.wideRoute(
|
||||
ctx,
|
||||
"logs.debug",
|
||||
"/api/v1/workload-actions/runs/:runFriendlyId/logs/debug",
|
||||
"POST",
|
||||
async () => {
|
||||
const { req, reply, params, body } = ctx;
|
||||
reply.empty(204);
|
||||
|
||||
await this.workerClient.sendDebugLog(
|
||||
params.runFriendlyId,
|
||||
body,
|
||||
this.runnerIdFromRequest(req)
|
||||
);
|
||||
},
|
||||
{ highFrequency: true }
|
||||
),
|
||||
});
|
||||
} else {
|
||||
// Disabled: drop immediately without reading/parsing the body and without
|
||||
// any log we can't switch off. Older runners still POST per log line; the
|
||||
// route stays registered (an unregistered route would log "No route match"
|
||||
// per request) but sheds the request at minimal cost. Request metrics still
|
||||
// count it. 204 is non-retryable on the runner client, so no retry storm.
|
||||
httpServer.route("/api/v1/workload-actions/runs/:runFriendlyId/logs/debug", "POST", {
|
||||
skipBodyParsing: true,
|
||||
handler: async (ctx) => {
|
||||
ctx.reply.empty(204);
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
// Snapshot callback endpoint (inbound from compute path)
|
||||
httpServer.route("/api/v1/compute/snapshot-complete", "POST", {
|
||||
bodySchema: SnapshotCallbackPayloadSchema,
|
||||
handler: async (ctx) =>
|
||||
this.wideRoute(
|
||||
ctx,
|
||||
"snapshot.callback",
|
||||
"/api/v1/compute/snapshot-complete",
|
||||
"POST",
|
||||
async () => {
|
||||
const { reply, body } = ctx;
|
||||
if (!this.snapshotService) {
|
||||
reply.empty(404);
|
||||
return;
|
||||
}
|
||||
|
||||
const result = await this.snapshotService.handleCallback(body);
|
||||
reply.empty(result.status);
|
||||
}
|
||||
),
|
||||
});
|
||||
|
||||
return httpServer;
|
||||
}
|
||||
|
||||
private createWebsocketServer() {
|
||||
const io = new Server(this.httpServer.server);
|
||||
|
||||
const websocketServer: Namespace<
|
||||
WorkloadClientToServerEvents,
|
||||
WorkloadServerToClientEvents,
|
||||
DefaultEventsMap,
|
||||
WorkloadClientSocketData
|
||||
> = io.of("/workload");
|
||||
|
||||
websocketServer.on("disconnect", (socket) => {
|
||||
this.logger.verbose("[WS] disconnect", socket.id);
|
||||
});
|
||||
websocketServer.use(async (socket, next) => {
|
||||
const setSocketDataFromHeader = (
|
||||
dataKey: keyof typeof socket.data,
|
||||
headerName: string,
|
||||
required: boolean = true
|
||||
) => {
|
||||
const value = socket.handshake.headers[headerName];
|
||||
|
||||
if (value) {
|
||||
if (Array.isArray(value)) {
|
||||
if (value[0]) {
|
||||
socket.data[dataKey] = value[0];
|
||||
return;
|
||||
}
|
||||
} else {
|
||||
socket.data[dataKey] = value;
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
if (required) {
|
||||
this.logger.error("[WS] missing required header", { headerName });
|
||||
throw new Error("missing header");
|
||||
}
|
||||
};
|
||||
|
||||
try {
|
||||
setSocketDataFromHeader("deploymentId", WORKLOAD_HEADERS.DEPLOYMENT_ID);
|
||||
setSocketDataFromHeader("runnerId", WORKLOAD_HEADERS.RUNNER_ID);
|
||||
} catch (error) {
|
||||
this.logger.error("[WS] setSocketDataFromHeader error", { error });
|
||||
socket.disconnect(true);
|
||||
return;
|
||||
}
|
||||
|
||||
this.logger.debug("[WS] auth success", socket.data);
|
||||
|
||||
next();
|
||||
});
|
||||
websocketServer.on("connection", (socket) => {
|
||||
const socketLogger = this.logger.child({
|
||||
socketId: socket.id,
|
||||
socketData: socket.data,
|
||||
});
|
||||
|
||||
const getSocketMetadata = () => {
|
||||
return {
|
||||
deploymentId: socket.data.deploymentId,
|
||||
runId: socket.data.runFriendlyId,
|
||||
snapshotId: socket.data.snapshotId,
|
||||
runnerId: socket.data.runnerId,
|
||||
};
|
||||
};
|
||||
|
||||
const emitSocketLifecycle = (
|
||||
event: "run_connected" | "run_disconnected",
|
||||
friendlyId: string,
|
||||
disconnectReason?: string
|
||||
) => {
|
||||
emitOneShot({
|
||||
...this.wideEventOpts,
|
||||
op: event === "run_connected" ? "socket.run.connected" : "socket.run.disconnected",
|
||||
kind: "event",
|
||||
populate: (state) => {
|
||||
state.extras.event = event;
|
||||
setMeta(state, "run_id", friendlyId);
|
||||
if (socket.data.deploymentId) {
|
||||
setMeta(state, "deployment_id", socket.data.deploymentId);
|
||||
}
|
||||
if (socket.data.runnerId) setMeta(state, "runner_id", socket.data.runnerId);
|
||||
state.extras.socket_id = socket.id;
|
||||
if (disconnectReason) state.extras.disconnect_reason = disconnectReason;
|
||||
},
|
||||
});
|
||||
};
|
||||
|
||||
const runConnected = (friendlyId: string) => {
|
||||
socketLogger.debug("runConnected", { ...getSocketMetadata() });
|
||||
|
||||
// If there's already a run ID set, we should "disconnect" it from this socket
|
||||
if (socket.data.runFriendlyId && socket.data.runFriendlyId !== friendlyId) {
|
||||
socketLogger.debug("runConnected: disconnecting existing run", {
|
||||
...getSocketMetadata(),
|
||||
newRunId: friendlyId,
|
||||
oldRunId: socket.data.runFriendlyId,
|
||||
});
|
||||
runDisconnected(socket.data.runFriendlyId, "socket_run_replaced");
|
||||
}
|
||||
|
||||
this.runSockets.set(friendlyId, socket);
|
||||
this.emit("runConnected", { run: { friendlyId } });
|
||||
socket.data.runFriendlyId = friendlyId;
|
||||
emitSocketLifecycle("run_connected", friendlyId);
|
||||
};
|
||||
|
||||
const runDisconnected = (friendlyId: string, reason: string) => {
|
||||
socketLogger.debug("runDisconnected", { ...getSocketMetadata() });
|
||||
|
||||
// The run is gone from this runner (crash, exit, or replaced by a new
|
||||
// run), so a pending delayed snapshot for it is stale. Genuine
|
||||
// waitpoint suspensions keep the socket connected, so this doesn't
|
||||
// cancel a snapshot that's still wanted; the runnerId match guards
|
||||
// against a stale duplicate runner cancelling a fresh runner's
|
||||
// snapshot after the run was reassigned. Caveat: socket.data.runnerId
|
||||
// is frozen at the websocket handshake, so after a same-supervisor
|
||||
// restore (new runner id, socket not recreated) this guard refuses
|
||||
// the cancel - a missed cancel, never a wrong one. The
|
||||
// attempt.complete cancel uses the runner's current HTTP header id
|
||||
// and is unaffected.
|
||||
this.snapshotService?.cancel(friendlyId, socket.data.runnerId);
|
||||
|
||||
this.runSockets.delete(friendlyId);
|
||||
this.emit("runDisconnected", { run: { friendlyId } });
|
||||
socket.data.runFriendlyId = undefined;
|
||||
emitSocketLifecycle("run_disconnected", friendlyId, reason);
|
||||
};
|
||||
|
||||
socketLogger.debug("wsServer socket connected", { ...getSocketMetadata() });
|
||||
|
||||
// FIXME: where does this get set?
|
||||
if (socket.data.runFriendlyId) {
|
||||
runConnected(socket.data.runFriendlyId);
|
||||
}
|
||||
|
||||
socket.on("disconnecting", (reason, description) => {
|
||||
socketLogger.verbose("Socket disconnecting", {
|
||||
...getSocketMetadata(),
|
||||
reason,
|
||||
description,
|
||||
});
|
||||
|
||||
if (socket.data.runFriendlyId) {
|
||||
runDisconnected(socket.data.runFriendlyId, `socket_disconnecting:${reason}`);
|
||||
}
|
||||
});
|
||||
|
||||
socket.on("disconnect", (reason, description) => {
|
||||
socketLogger.debug("Socket disconnected", { ...getSocketMetadata(), reason, description });
|
||||
});
|
||||
|
||||
socket.on("error", (error) => {
|
||||
socketLogger.error("Socket error", {
|
||||
...getSocketMetadata(),
|
||||
error: {
|
||||
name: error.name,
|
||||
message: error.message,
|
||||
stack: error.stack,
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
socket.on("run:start", async (message) => {
|
||||
const log = socketLogger.child({
|
||||
eventName: "run:start",
|
||||
...getSocketMetadata(),
|
||||
...message,
|
||||
});
|
||||
|
||||
log.debug("Handling run:start");
|
||||
|
||||
try {
|
||||
runConnected(message.run.friendlyId);
|
||||
} catch (error) {
|
||||
log.error("run:start error", { error });
|
||||
}
|
||||
});
|
||||
|
||||
socket.on("run:stop", async (message) => {
|
||||
const log = socketLogger.child({
|
||||
eventName: "run:stop",
|
||||
...getSocketMetadata(),
|
||||
...message,
|
||||
});
|
||||
|
||||
log.debug("Handling run:stop");
|
||||
|
||||
try {
|
||||
runDisconnected(message.run.friendlyId, "run_stop_message");
|
||||
// Don't delete trace context here - run:stop fires after each snapshot/shutdown
|
||||
// but the run may be restored on a new VM and snapshot again. Trace context is
|
||||
// re-populated on dequeue, and entries are small (4 strings per run).
|
||||
} catch (error) {
|
||||
log.error("run:stop error", { error });
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
return websocketServer;
|
||||
}
|
||||
|
||||
notifyRun({ run }: { run: { friendlyId: string } }) {
|
||||
try {
|
||||
const runSocket = this.runSockets.get(run.friendlyId);
|
||||
|
||||
if (!runSocket) {
|
||||
this.logger.debug("notifyRun: Run socket not found", { run });
|
||||
|
||||
this.workerClient.sendDebugLog(run.friendlyId, {
|
||||
time: new Date(),
|
||||
message: "run:notify socket not found on supervisor",
|
||||
});
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
runSocket.emit("run:notify", { version: "1", run });
|
||||
this.logger.debug("run:notify sent", { run });
|
||||
|
||||
this.workerClient.sendDebugLog(run.friendlyId, {
|
||||
time: new Date(),
|
||||
message: "run:notify supervisor -> runner",
|
||||
});
|
||||
} catch (error) {
|
||||
this.logger.error("Error in notifyRun", { run, error });
|
||||
|
||||
this.workerClient.sendDebugLog(run.friendlyId, {
|
||||
time: new Date(),
|
||||
message: "run:notify error on supervisor",
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
registerRunTraceContext(runFriendlyId: string, ctx: RunTraceContext) {
|
||||
this.snapshotService?.registerTraceContext(runFriendlyId, ctx);
|
||||
}
|
||||
|
||||
async start() {
|
||||
await this.httpServer.start();
|
||||
}
|
||||
|
||||
async stop() {
|
||||
this.snapshotService?.stop();
|
||||
await this.httpServer.stop();
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user