fix(core): reaper skips own in-flight jobs, heartbeat, no attempt burned on shutdown

Co-Authored-By: Claude Sonnet 5.5 <[email protected]>

packages/core/src/jobs.ts

@@ -6,6 +6,7 @@ import { sql } from "kysely";
 const MAX_ATTEMPTS = 5;
 
 const STALE_LOCK_MS = 60 * 60 * 1000;
+const HEARTBEAT_MS = 5 * 60 * 1000;
 
 export interface JobContext {
   /** Aborted when the worker is stopped. */
@@ -78,7 +79,11 @@ async function claim(db: Db, id: string): Promise<boolean> {
   return r.numUpdatedRows === 1n;
 }
 
-async function runJob(db: Db, handlers: Record<string, JobHandler>, job: { id: string; type: string; payload: string; attempts: number }, signal: AbortSignal): Promise<void> {
+async function runJob(db: Db, handlers: Record<string, JobHandler>, job: { id: string; type: string; payload: string; attempts: number }, signal: AbortSignal, heartbeatMs: number): Promise<void> {
+  // Keep `locked_at` fresh so another process's reaper does not steal a long-running job.
+  const beat = setInterval(() => {
+    db.k.updateTable("jobs").set({ locked_at: nowIso() }).where("id", "=", job.id).where("status", "=", "running").execute().catch(logErr);
+  }, heartbeatMs);
   try {
     try {
       const h = handlers[job.type];
@@ -88,7 +93,10 @@ async function runJob(db: Db, handlers: Record<string, JobHandler>, job: { id: s
     } catch (e) {
       const attempts = job.attempts + 1;
       const err = e instanceof Error ? e.message : String(e);
-      if (e instanceof DropJobError) {
+      if (signal.aborted && !(e instanceof DropJobError) && !(e instanceof NonRetryableError)) {
+        // Interrupted by shutdown: back to the queue untouched, no attempt burned, no backoff.
+        await db.k.updateTable("jobs").set({ status: "queued", locked_at: null }).where("id", "=", job.id).execute();
+      } else if (e instanceof DropJobError) {
         await db.k.updateTable("jobs").set({ status: "done", locked_at: null, last_error: err }).where("id", "=", job.id).execute();
       } else if (e instanceof NonRetryableError || attempts >= MAX_ATTEMPTS) {
         await db.k.updateTable("jobs").set({ status: "failed", attempts, last_error: err, locked_at: null }).where("id", "=", job.id).execute();
@@ -103,31 +111,38 @@ async function runJob(db: Db, handlers: Record<string, JobHandler>, job: { id: s
     }
   } catch (e) {
     logErr(e);
+  } finally {
+    clearInterval(beat);
   }
 }
 
 /** Re-queue jobs a crashed process left `running`, so their lockKey cannot be held forever. */
-async function reapStale(db: Db): Promise<void> {
-  const cutoff = new Date(Date.now() - STALE_LOCK_MS).toISOString();
-  await db.k.updateTable("jobs").set({ status: "queued", locked_at: null }).where("status", "=", "running").where("locked_at", "<", cutoff).execute();
+async function reapStale(db: Db, staleMs: number, mine: Set<string>): Promise<void> {
+  const cutoff = new Date(Date.now() - staleMs).toISOString();
+  let q = db.k.updateTable("jobs").set({ status: "queued", locked_at: null }).where("status", "=", "running").where("locked_at", "<", cutoff);
+  if (mine.size > 0) q = q.where("id", "not in", [...mine]);
+  await q.execute();
 }
 
 /** Poll for due jobs and run them with bounded concurrency. Failures retry with backoff `2^attempts` minutes, then `failed`. */
 export function startWorker(
   db: Db,
   handlers: Record<string, JobHandler>,
-  opts: { pollMs?: number; concurrency?: number } = {},
+  opts: { pollMs?: number; concurrency?: number; staleMs?: number; heartbeatMs?: number } = {},
 ): { stop(): Promise<void> } {
   const pollMs = opts.pollMs ?? 1000;
   const concurrency = opts.concurrency ?? 2;
+  const staleMs = opts.staleMs ?? STALE_LOCK_MS;
+  const heartbeatMs = opts.heartbeatMs ?? HEARTBEAT_MS;
   const inflight = new Set<Promise<void>>();
+  const mine = new Set<string>();
   const abort = new AbortController();
   let stopped = false;
   let timer: ReturnType<typeof setTimeout> | undefined;
   let tickP: Promise<void> = Promise.resolve();
 
   async function tick(): Promise<void> {
-    await reapStale(db);
+    await reapStale(db, staleMs, mine);
     while (!stopped && inflight.size < concurrency) {
       const due = await db.k
         .selectFrom("jobs")
@@ -142,9 +157,13 @@ export function startWorker(
         if (stopped || inflight.size >= concurrency) break;
         if (!(await claim(db, job.id))) continue;
         started++;
-        const p: Promise<void> = runJob(db, handlers, job, abort.signal)
+        mine.add(job.id);
+        const p: Promise<void> = runJob(db, handlers, job, abort.signal, heartbeatMs)
           .catch(logErr)
-          .finally(() => inflight.delete(p));
+          .finally(() => {
+            inflight.delete(p);
+            mine.delete(job.id);
+          });
         inflight.add(p);
       }
       if (started === 0) break;

packages/core/test/jobs.test.ts

@@ -145,6 +145,41 @@ forEachDialect("jobs", (getDb) => {
     err.mockRestore();
   });
 
+  test("a job running past staleMs is not reaped by its own worker and runs once", async () => {
+    const db = getDb();
+    const id = (await enqueue(db, "t", {}, { lockKey: "slow" }))!;
+    let runs = 0;
+    const w = startWorker(db, { t: async () => { runs++; await sleep(300); } }, { pollMs: 10, staleMs: 50, heartbeatMs: 20 });
+    await until(async () => (await row(db, id)).status === "done");
+    await w.stop();
+    expect(runs).toBe(1);
+  });
+
+  test("heartbeat refreshes locked_at", async () => {
+    const db = getDb();
+    const id = (await enqueue(db, "t", {}))!;
+    let release!: () => void;
+    const gate = new Promise<void>((r) => (release = r));
+    const w = startWorker(db, { t: () => gate }, { pollMs: 10, heartbeatMs: 20 });
+    await until(async () => (await row(db, id)).status === "running");
+    const first = (await row(db, id)).locked_at!;
+    await until(async () => (await row(db, id)).locked_at !== first);
+    expect((await row(db, id)).locked_at).not.toBe(first);
+    release();
+    await w.stop();
+  });
+
+  test("stop() re-queues an aborted job without burning an attempt or backoff", async () => {
+    const db = getDb();
+    const id = (await enqueue(db, "t", {}))!;
+    const before = (await row(db, id)).run_at;
+    let started = false;
+    const w = startWorker(db, { t: (_p, { signal }) => new Promise<void>((_, rej) => { started = true; signal.addEventListener("abort", () => rej(new Error("aborted"))); }) }, { pollMs: 10 });
+    await until(async () => started);
+    await w.stop();
+    expect(await row(db, id)).toMatchObject({ status: "queued", attempts: 0, run_at: before, locked_at: null });
+  });
+
   test("stop() aborts the signal handed to handlers", async () => {
     const db = getDb();
     await enqueue(db, "t", {});