STTracker/packages/core/src/jobs.ts

import type { Db } from "@sttracker/db/db";
import { newId, nowIso } from "@sttracker/db/ids";
import { fromJson, toJson } from "@sttracker/db/json";
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. */
  signal: AbortSignal;
}
export type JobHandler = (payload: any, ctx: JobContext) => Promise<unknown>;

/** Throw from a handler to fail the job at once, with no retry. */
export class NonRetryableError extends Error {
  override readonly name = "NonRetryableError";
}
/** Throw from a handler to drop the job: it is marked `done` and the message is kept in `last_error`. */
export class DropJobError extends Error {
  override readonly name = "DropJobError";
}

const logErr = (err: unknown) => console.error("[jobs]", err);

/**
 * Queue a job. Returns its id, or null when a `queued|running` job already holds `lockKey`.
 * On Postgres the check and insert run under an advisory transaction lock so two callers cannot both pass.
 */
export async function enqueue(
  db: Db,
  type: string,
  payload: unknown,
  opts: { runAt?: string; lockKey?: string } = {},
): Promise<string | null> {
  const lockKey = opts.lockKey ?? null;
  return db.k.transaction().execute(async (tx) => {
    if (lockKey !== null) {
      if (db.dialect === "postgres") await sql`select pg_advisory_xact_lock(hashtext(${lockKey}))`.execute(tx);
      const held = await tx
        .selectFrom("jobs")
        .select("id")
        .where("lock_key", "=", lockKey)
        .where("status", "in", ["queued", "running"])
        .executeTakeFirst();
      if (held) return null;
    }
    const id = newId();
    const now = nowIso();
    await tx
      .insertInto("jobs")
      .values({
        id,
        type,
        payload: toJson(payload ?? null),
        run_at: opts.runAt ?? now,
        status: "queued",
        attempts: 0,
        last_error: null,
        locked_at: null,
        lock_key: lockKey,
        created_at: now,
      })
      .execute();
    return id;
  });
}

/** Claim one queued job atomically; true when this caller won it. */
async function claim(db: Db, id: string): Promise<boolean> {
  const r = await db.k
    .updateTable("jobs")
    .set({ status: "running", locked_at: nowIso() })
    .where("id", "=", id)
    .where("status", "=", "queued")
    .executeTakeFirst();
  return r.numUpdatedRows === 1n;
}

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];
      if (!h) throw new NonRetryableError(`no handler for job type "${job.type}"`);
      await h(fromJson<unknown>(job.payload, null), { signal });
      await db.k.updateTable("jobs").set({ status: "done", locked_at: null, last_error: null }).where("id", "=", job.id).execute();
    } catch (e) {
      const attempts = job.attempts + 1;
      const err = e instanceof Error ? e.message : String(e);
      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();
      } else {
        const runAt = new Date(Date.now() + 2 ** attempts * 60_000).toISOString();
        await db.k
          .updateTable("jobs")
          .set({ status: "queued", attempts, last_error: err, run_at: runAt, locked_at: null })
          .where("id", "=", job.id)
          .execute();
      }
    }
  } 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, 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; 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, staleMs, mine);
    while (!stopped && inflight.size < concurrency) {
      const due = await db.k
        .selectFrom("jobs")
        .select(["id", "type", "payload", "attempts"])
        .where("status", "=", "queued")
        .where("run_at", "<=", nowIso())
        .orderBy("run_at")
        .limit(concurrency - inflight.size)
        .execute();
      let started = 0;
      for (const job of due) {
        if (stopped || inflight.size >= concurrency) break;
        if (!(await claim(db, job.id))) continue;
        started++;
        mine.add(job.id);
        const p: Promise<void> = runJob(db, handlers, job, abort.signal, heartbeatMs)
          .catch(logErr)
          .finally(() => {
            inflight.delete(p);
            mine.delete(job.id);
          });
        inflight.add(p);
      }
      if (started === 0) break;
    }
  }

  function schedule() {
    if (stopped) return;
    timer = setTimeout(() => {
      tickP = tick().catch(logErr).finally(schedule);
    }, pollMs);
  }
  tickP = tick().catch(logErr).finally(schedule);

  return {
    async stop() {
      stopped = true;
      abort.abort();
      if (timer) clearTimeout(timer);
      await tickP.catch(logErr);
      await Promise.allSettled([...inflight]);
    },
  };
}

/** Every `tickMs`, enqueue a `sync` job for each non-paused (active or errored) connection whose schedule has come due. */
export function startScheduler(db: Db, opts: { tickMs?: number } = {}): { stop(): void; tickNow(): Promise<void> } {
  const tickMs = opts.tickMs ?? 60_000;
  let stopped = false;
  async function tick() {
    const rows = await db.k
      .selectFrom("source_connections")
      .select(["id", "schedule_minutes", "last_run_at"])
      .where("status", "!=", "paused")
      .where("schedule_minutes", "is not", null)
      .execute();
    const now = Date.now();
    for (const r of rows) {
      if (stopped) return;
      const due = r.last_run_at === null || Date.parse(r.last_run_at) + r.schedule_minutes! * 60_000 <= now;
      if (due) await enqueue(db, "sync", { connectionId: r.id }, { lockKey: `sync:${r.id}` });
    }
  }
  const run = () => tick().catch(logErr);
  void run();
  const timer = setInterval(() => void run(), tickMs);
  return {
    /** Run one scheduling pass now (tests, manual triggers). */
    tickNow: run,
    stop() {
      stopped = true;
      clearInterval(timer);
    },
  };
}