Repository Public
1/** Worker loop: claims runs up to `concurrency`, executes them, heartbeats, stops gracefully. */2import os from 'node:os';3import type pg from 'pg';4import { AosSettings, ENGINE_VERSION } from '@aos/shared';5import { getKernel, query, WAKE_CHANNEL, type Kernel } from './db.js';6import { executeRun, type Execution } from './executor.js';7import { claimRun, DEFAULT_HEARTBEAT_MS, DEFAULT_LEASE_MS } from './lease.js';8import type { DailyBudget, KernelHooks, RunLogic } from './types.js';9import { expireWaits } from './waits.js';1011export interface WorkerLoopOptions {12 logic: RunLogic;13 concurrency?: number; // default env AOS_MAX_CONCURRENT_RUNS, else 414 /** Abort to stop: no new claim, ctx.signal of every run aborted, runs left 'running' with a released lease. */15 signal?: AbortSignal;16 kernel?: Kernel;17 workerId?: string; // default env AOS_WORKER_ID, else hostname-pid18 leaseMs?: number; // 90 s19 heartbeatMs?: number; // 30 s20 /** Poll fallback of the claim loop (the NOTIFY doorbell wakes it earlier; due timers rely on it). */21 pollMs?: number; // 1 s22 /** Timers at most this long are waited in-process instead of suspending the run. */23 inlineSleepMs?: number; // 1.5 s24 /**25 * Daily cap in USD across all missions. Default: the mission owner's `dailyBudgetUsd` in26 * application_settings (falling back to the AosSettings default). `null` disables the cap.27 */28 dailyBudgetUsd?: DailyBudget;29 /** Run expireWaits() every `janitorMs` (30 s). false to disable. */30 janitorMs?: number | false;31 /** How long a graceful stop waits for in-flight steps before returning anyway (leases then expire). */32 stopGraceMs?: number; // 15 s33 version?: string;34 hooks?: KernelHooks;35}3637const DEFAULT_DAILY = AosSettings.parse({}).dailyBudgetUsd;3839async function ownerDailyCap(k: Kernel, missionId: string): Promise<number | null> {40 const rows = await query<{ cap: unknown }>(41 k,42 `SELECT s.settings->>'dailyBudgetUsd' AS cap FROM missions m LEFT JOIN application_settings s ON s.user_id = m.owner_id WHERE m.id = $1`,43 [missionId],44 );45 const v = Number(rows[0]?.cap);46 return Number.isFinite(v) && v > 0 ? v : DEFAULT_DAILY;47}4849/** Resolves when the loop has stopped (signal aborted) and every execution has let go of its run. */50export async function runWorkerLoop(opts: WorkerLoopOptions): Promise<void> {51 const k = opts.kernel ?? getKernel();52 const workerId = opts.workerId ?? process.env.AOS_WORKER_ID ?? `${os.hostname()}-${process.pid}`;53 const concurrency = Math.max(1, opts.concurrency ?? (Number(process.env.AOS_MAX_CONCURRENT_RUNS) || 4));54 const leaseMs = opts.leaseMs ?? DEFAULT_LEASE_MS;55 const heartbeatMs = opts.heartbeatMs ?? Math.min(DEFAULT_HEARTBEAT_MS, Math.floor(leaseMs / 3));56 const pollMs = opts.pollMs ?? 1000;57 const log = k.log.child({ workerId });58 const active = new Map<string, Execution>();59 let stopping = opts.signal?.aborted ?? false;60 let wakeUp: (() => void) | null = null;61 const ring = (): void => {62 const w = wakeUp;63 wakeUp = null;64 w?.();65 };66 const dailyCapUsd = async (missionId: string): Promise<number | null> => {67 const d = opts.dailyBudgetUsd;68 if (d === undefined) return ownerDailyCap(k, missionId);69 return typeof d === 'function' ? d() : d;70 };7172 // Doorbell connection (best effort; the poll is the safety net).73 let listener: pg.PoolClient | null = null;74 const onNote = (msg: pg.Notification): void => {75 if (msg.channel !== WAKE_CHANNEL) return;76 if (msg.payload?.startsWith('c:')) active.get(msg.payload.slice(2))?.poke();77 else ring();78 };79 const dropListener = (err?: Error): void => {80 const c = listener;81 listener = null;82 if (c) {83 c.removeListener('notification', onNote);84 c.release(err ?? true);85 }86 };87 const ensureListener = async (): Promise<void> => {88 if (listener || stopping) return;89 try {90 const c = await k.pool.connect();91 c.on('notification', onNote);92 c.once('error', (err) => {93 if (listener === c) dropListener(err);94 });95 listener = c;96 await c.query(`LISTEN ${WAKE_CHANNEL}`);97 } catch (err) {98 dropListener(err as Error);99 }100 };101102 const beat = (): void => {103 query(k, `INSERT INTO workers (id, last_seen, version) VALUES ($1, now(), $2) ON CONFLICT (id) DO UPDATE SET last_seen = now(), version = EXCLUDED.version`, [104 workerId,105 opts.version ?? ENGINE_VERSION,106 ]).catch((err: Error) => log.debug({ err: err.message }, 'worker heartbeat failed'));107 };108 const beatTimer = setInterval(beat, 10_000);109 beatTimer.unref();110 beat();111 const janitorMs = opts.janitorMs ?? 30_000;112 const janitor =113 janitorMs === false114 ? null115 : setInterval(() => {116 expireWaits(k).then((r) => r.expired > 0 && ring(), (err: Error) => log.debug({ err: err.message }, 'expireWaits failed'));117 }, janitorMs);118 janitor?.unref();119120 const onAbort = (): void => {121 stopping = true;122 for (const e of active.values()) e.stop();123 ring();124 };125 opts.signal?.addEventListener('abort', onAbort, { once: true });126 log.info({ concurrency, leaseMs, heartbeatMs }, 'worker loop started');127128 try {129 while (!stopping) {130 await ensureListener();131 while (!stopping && active.size < concurrency) {132 let claimed;133 try {134 claimed = await claimRun(workerId, { leaseMs, kernel: k });135 } catch (err) {136 log.warn({ err: (err as Error).message }, 'claim failed');137 break;138 }139 if (!claimed) break;140 const exec = executeRun(k, claimed, {141 logic: opts.logic, leaseMs, heartbeatMs, inlineSleepMs: opts.inlineSleepMs ?? 1500, dailyCapUsd, hooks: opts.hooks,142 });143 active.set(exec.runId, exec);144 if (stopping) exec.stop();145 void exec.done.then(() => {146 active.delete(exec.runId);147 ring();148 });149 }150 if (stopping) break;151 await new Promise<void>((resolve) => {152 const t = setTimeout(resolve, pollMs);153 wakeUp = () => {154 clearTimeout(t);155 resolve();156 };157 });158 }159 } finally {160 opts.signal?.removeEventListener('abort', onAbort);161 clearInterval(beatTimer);162 if (janitor) clearInterval(janitor);163 dropListener();164 const pending = [...active.values()].map((e) => e.done);165 if (pending.length > 0) {166 let t: NodeJS.Timeout | undefined;167 const grace = new Promise<'timeout'>((resolve) => {168 t = setTimeout(() => resolve('timeout'), opts.stopGraceMs ?? 15_000);169 });170 const r = await Promise.race([Promise.all(pending), grace]);171 clearTimeout(t);172 if (r === 'timeout') log.warn({ runs: [...active.keys()] }, 'stop grace elapsed; runs left to lease expiry');173 }174 log.info('worker loop stopped');175 }176}177