cairnq 0.4.0 → 0.6.0

package/src/store/base.ts

@@ -7,6 +7,7 @@ import {
   SerializationError,
 } from "../errors.js";
 import { rowToTask, STATUSES, type Task, type TaskStatus } from "../models.js";
+import { type BackpressureOptions, QueueDepthGate } from "../backpressure.js";
 
 const rejectMangled = function (this: unknown, _key: string, v: unknown): unknown {
   if (typeof v === "number" && !Number.isFinite(v)) {
@@ -64,6 +65,10 @@ export function checkProtocolVersion(version: number): void {
 const CONFLICTS = ["reuse", "reject", "replace"] as const;
 export type Conflict = (typeof CONFLICTS)[number];
 
+/** The queue a submit lands on when it names none. Owned here, where the
+ * default is applied, so nothing above has to re-derive it. */
+export const DEFAULT_QUEUE = "default";
+
 export interface SubmitInput {
   name: string;
   payload: unknown;
@@ -149,6 +154,9 @@ export function statementParams(sql: string): readonly string[] {
  * behavior; the shared SQL already stops them from drifting in wording.
  */
 export abstract class TaskStore {
+  /** Set by useBackpressure; null means submit is ungated. */
+  private gate: QueueDepthGate | null = null;
+
   // ------------------------------------------------------------ dialect seam
   abstract connect(): Promise<void>;
   abstract close(): Promise<void>;
@@ -211,12 +219,24 @@ export abstract class TaskStore {
   }
 
   // ------------------------------------------------------------- client side
+  /**
+   * Bound how deep a queue may get before `submit` blocks. Off unless set.
+   *
+   * It hangs here rather than on `CairnQ` because the store is the one choke
+   * point every submit passes through — a handler spawning children via
+   * `TaskContext.submit` is the shape most likely to outrun its workers, and
+   * gating only the client would leave exactly that path unbounded.
+   */
+  useBackpressure(opts: BackpressureOptions): void {
+    this.gate = new QueueDepthGate(this, opts);
+  }
+
   async submit(input: SubmitInput): Promise<Task> {
     const id = newId("task");
     const ins: Params = {
       id,
       name: input.name,
-      queue: input.queue ?? "default",
+      queue: input.queue ?? DEFAULT_QUEUE,
       payload: dumpJson(input.payload ?? {}),
       metadata: dumpJson(input.metadata ?? {}),
       max_attempts: input.maxAttempts ?? 3,
@@ -243,6 +263,10 @@ export abstract class TaskStore {
     if (input.runAtDelayMs != null && input.runAtDelayMs < 0) {
       throw new Error(`runAtDelayMs must be >= 0, got ${input.runAtDelayMs}`);
     }
+    // After validation and before the first write: bad arguments should fail
+    // now, not after waiting out a full queue. Reads the resolved queue, so the
+    // gate cannot throttle one queue while the row lands on another.
+    if (this.gate) await this.gate.acquire(ins.queue as string);
     if (key === null) return rowToTask((await this.fetch("insert_task", ins))[0]);
 
     // A key makes submit a read-then-write, so it has to be one transaction —
@@ -370,6 +394,22 @@ export abstract class TaskStore {
     return out;
   }
 
+  /**
+   * How many more tasks fit on `queue` under `maxDepth` — 0 once it is full.
+   *
+   * The cheap half of backpressure: bounded at `maxDepth` index entries, unlike
+   * `stats()`, which aggregates the whole table (terminal rows included) and so
+   * costs more the longer a database has been running. Use it directly to shed
+   * load or shape a producer; `QueueDepthGate` builds the blocking form on top.
+   */
+  async queueDepth(queue: string, maxDepth: number): Promise<number> {
+    if (!Number.isInteger(maxDepth) || maxDepth < 0) {
+      throw new Error(`maxDepth must be a non-negative integer, got ${maxDepth}`);
+    }
+    const rows = await this.fetch("queue_depth", { queue, max_depth: maxDepth });
+    return Number(rows[0]?.headroom ?? 0);
+  }
+
   // ------------------------------------------------------------- worker side
   /**
    * Take up to `limit` claimable tasks. `names` restricts the claim to task names
@@ -385,26 +425,82 @@ export abstract class TaskStore {
     limit?: number;
     names?: string[];
   }): Promise<Task[]> {
-    // One queue is the common case and gets its own statement: a list-valued queue
-    // filter cannot be read in claim order, so the planner sorts every claimable
-    // row to take LIMIT of them, and claim's cost grows with the queued backlog
-    // while it holds the claim transaction. See claim_one_queue.sql.
+    const names = input.names ?? null;
+    const claimed = await this.claimSession(
+      { queues: input.queues, workerId: input.workerId, leaseMs: input.leaseMs, names },
+      (claim) => claim(names, input.limit ?? 1),
+    );
+    return claimed ?? [];
+  }
+
+  /**
+   * Open one claim transaction and let the caller draw from it repeatedly.
+   *
+   * The transaction is what has to live here: the read-only probe that keeps an
+   * idle worker off SQLite's single write lock, the `recover_leases` whose
+   * reclaimed leases must be visible to the claims that follow and to nobody in
+   * between, and the write lock itself. *What* gets claimed under it is the
+   * caller's business — a worker drawing a separate quota per task name is
+   * scheduling policy, and this layer has no vocabulary for the "handler call"
+   * that policy is denominated in. It knows queues, names, limits and rows.
+   *
+   * `plan` is handed a `claim(names, limit)` it may call any number of times,
+   * each a separate statement under the same lock and the same recovery, and
+   * each free to size itself from what the previous one returned. That feedback
+   * is the reason this is a callback rather than a list of quotas: a caller
+   * dividing a budget up front has to guess, and every share handed to a name
+   * with nothing queued is a slot left idle until the next poll.
+   *
+   * `plan` runs with the write lock held, so it must await nothing but that
+   * callback.
+   *
+   * `names` is the union `plan` might ask for — the probe and the recovery are
+   * filtered by it. Returns undefined when the probe finds nothing claimable, in
+   * which case `plan` never runs and no transaction is opened.
+   */
+  async claimSession<T>(
+    input: { queues: string[]; workerId: string; leaseMs?: number; names: string[] | null },
+    plan: (claim: (names: string[] | null, limit: number) => Promise<Task[]>) => Promise<T>,
+  ): Promise<T | undefined> {
+    // A list-valued filter cannot be read in claim order, so the planner sorts
+    // every claimable row to take LIMIT of them and the claim's cost grows with
+    // the backlog while it holds the transaction. Both filters therefore have an
+    // equality form, picked per draw: one queue is the common deployment, and one
+    // name is every per-name quota. See claim_one_queue.sql and claim_one_name.sql.
     const oneQueue = input.queues.length === 1;
-    const params: Params = {
+    const base: Params = {
       queues: input.queues,
       queue: oneQueue ? input.queues[0] : null,
-      names: input.names ?? null,
+      names: input.names,
+      name: null,
       worker_id: input.workerId,
       lease_ms: input.leaseMs ?? 30_000,
-      limit: input.limit ?? 1,
+      limit: 1,
       lease_expired_error: LEASE_EXPIRED_ERROR_JSON,
     };
-    if (!(await this.hasClaimableWork(params))) return [];
-    // Recovery must share the claim's transaction: a lease reclaimed here has to
-    // be visible to the claim that follows, and to nobody in between.
+    if (!(await this.hasClaimableWork(base))) return undefined;
     return this.tx(async (fetch) => {
-      await fetch("recover_leases", params);
-      return (await fetch(oneQueue ? "claim_one_queue" : "claim", params)).map(rowToTask);
+      await fetch("recover_leases", base);
+      return plan(async (names, limit) => {
+        // A draw asking for nothing, or filtered to no names, claims nothing —
+        // answer it here rather than spending a statement to learn that.
+        if (limit <= 0 || names?.length === 0) return [];
+        const oneName = names?.length === 1;
+        const statement = oneName
+          ? oneQueue
+            ? "claim_one_queue_one_name"
+            : "claim_one_name"
+          : oneQueue
+            ? "claim_one_queue"
+            : "claim";
+        const rows = await fetch(statement, {
+          ...base,
+          names,
+          name: oneName ? names![0] : null,
+          limit,
+        });
+        return rows.map(rowToTask);
+      });
     });
   }
 
@@ -416,6 +512,33 @@ export abstract class TaskStore {
     });
   }
 
+  /**
+   * Renew several leases in one statement. Returns `taskId -> cancel requested`
+   * for the tasks this worker still holds.
+   *
+   * Deliberately not an ownedWrite: ownership is per task here, so there is no
+   * single answer to "did it work". A task **absent** from the result lost its
+   * lease, and the caller decides what that means for that one task rather than
+   * failing the whole beat.
+   *
+   * It returns flags rather than Tasks because nothing downstream needs a task:
+   * the caller renews leases and observes cancellation, and whole rows would drag
+   * every payload back on every beat for the life of the call.
+   */
+  async heartbeatBatch(input: {
+    taskIds: string[];
+    workerId: string;
+    leaseMs?: number;
+  }): Promise<Map<string, boolean>> {
+    if (!input.taskIds.length) return new Map();
+    const rows = await this.fetch("heartbeat_batch", {
+      ids: input.taskIds,
+      worker_id: input.workerId,
+      lease_ms: input.leaseMs ?? 30_000,
+    });
+    return new Map(rows.map((r) => [r.id as string, r.cancel_requested_at_ms != null]));
+  }
+
   async progress(input: {
     taskId: string;
     workerId: string;

package/src/store/sqlite.ts

@@ -7,6 +7,7 @@ import { nowMs } from "../ids.js";
 import { loadMigrations, loadStatements } from "../sql.js";
 import {
   checkProtocolVersion,
+  COMMENT,
   type Fetch,
   type Params,
   statementParams,
@@ -61,6 +62,34 @@ function isMemory(path: string): boolean {
   return path === ":memory:" || path.includes("mode=memory");
 }
 
+/**
+ * Whether this statement writes, and so belongs in a group commit.
+ *
+ * Read from the SQL rather than from a list of statement names, which would be a
+ * second place to remember when the protocol gains a statement. Every protocol
+ * statement is a single top-level `select`, `insert`, `update` or `delete`.
+ *
+ * Reads must stay out of the batch: `claimable_probe` exists precisely so an idle
+ * worker never takes SQLite's write lock, and a BEGIN IMMEDIATE around it would
+ * hand that back.
+ */
+function isWriteStatement(sql: string): boolean {
+  return !/^\s*select/i.test(sql.replace(COMMENT, ""));
+}
+
+/**
+ * One write waiting for its turn on the shared connection.
+ *
+ * The rows go back to the caller that asked for them, so a batch resolves each
+ * member with its own result rather than a merged one.
+ */
+interface Pending {
+  name: string;
+  params: Params;
+  resolve(rows: any[]): void;
+  reject(err: unknown): void;
+}
+
 /**
  * Whether cairnq_tasks has been analyzed at all.
  *
@@ -189,6 +218,12 @@ export class SQLiteStore extends TaskStore {
   private readonly busyBudgetMs: number;
   /** When this connection may next revisit its planner statistics. */
   private nextStatsRefreshAt = 0;
+  /** Which statements are writes — see isWriteStatement. */
+  private readonly writes: Record<string, boolean>;
+  /** Writes waiting to be group-committed — see flush(). */
+  private pending: Pending[] = [];
+  /** Whether a flusher is already queued to drain `pending`. */
+  private flushing = false;
 
   constructor(
     private readonly path: string,
@@ -197,6 +232,9 @@ export class SQLiteStore extends TaskStore {
     super();
     this.busyBudgetMs = opts.busyTimeoutMs ?? 5000;
     this.statements = loadStatements("sqlite");
+    this.writes = Object.fromEntries(
+      Object.entries(this.statements).map(([name, sql]) => [name, isWriteStatement(sql)]),
+    );
     // Only a bare ":memory:" is guaranteed private to its connection, so only
     // it gets a lock of its own. A "mode=memory" URI stays path-keyed: with
     // cache=shared it names ONE shared database, and on a build without URI
@@ -319,7 +357,10 @@ export class SQLiteStore extends TaskStore {
           bound[name] = now - (params.older_than_ms as number);
           break;
         case "queues":
-          bound[name] = JSON.stringify(params.queues);
+        case "ids":
+          // json_each needs a JSON array. Postgres binds the array itself as
+          // text[], so only this dialect encodes.
+          bound[name] = JSON.stringify(params[name]);
           break;
         case "names":
           // json_each needs a JSON array; null stays null so the SQL's
@@ -422,10 +463,135 @@ export class SQLiteStore extends TaskStore {
     }
   }
 
+  /**
+   * Group commit: one transaction for every write already waiting on the lock.
+   *
+   * A write costs microseconds to execute and a transaction costs a WAL commit, so
+   * N concurrent writes spend nearly all their time on N commits they could have
+   * shared. Measured at 200 finalizes: 80µs each one-transaction-apiece against
+   * 10µs each in one transaction (`bench/sweep` sweep B).
+   *
+   * Nothing waits to form a batch — a flusher takes whatever arrived while the
+   * previous one held the lock, so this trades no latency for the throughput. What
+   * it does trade is atomicity: two callers' writes now land together or not at
+   * all. Under at-least-once that is not observable (a lost batch is a
+   * redelivery), and it is why every member is resolved only after COMMIT.
+   */
+  private flush(db: DB): void {
+    // One writer waiting is the uncontended case, and it stays exactly as cheap as
+    // before: wrapping a single statement in BEGIN/COMMIT would add two statements
+    // to every write on an idle store.
+    if (this.pending.length === 1) {
+      const only = this.pending[0];
+      let rows: any[];
+      try {
+        rows = this.runNow(only.name, only.params);
+      } catch (err) {
+        // Leave it pending on a lost write lock: withLock re-runs this flusher.
+        if (isBusy(err)) throw err;
+        this.pending.shift();
+        only.reject(err);
+        return;
+      }
+      this.pending.shift();
+      only.resolve(rows);
+      return;
+    }
+
+    // BEGIN before consuming, so a lost write lock leaves the batch where the
+    // retry will find it — with anything that arrived meanwhile.
+    db.exec("BEGIN IMMEDIATE");
+    const batch = this.pending;
+    this.pending = [];
+    const out: { rows?: any[]; err?: unknown }[] = [];
+    try {
+      for (const w of batch) {
+        try {
+          out.push({ rows: this.runNow(w.name, w.params) });
+        } catch (err) {
+          // A statement error aborts that statement, not the transaction, so the
+          // rest of the batch is still good and this one waiter carries the error.
+          // If SQLite tore the transaction down instead, nothing in it survived
+          // and every member has to hear about it.
+          if (!db.inTransaction) throw err;
+          out.push({ err });
+        }
+      }
+      db.exec("COMMIT");
+    } catch (err) {
+      if (db.inTransaction) {
+        try {
+          db.exec("ROLLBACK");
+        } catch {
+          // Raced with SQLite's own rollback; the transaction is gone either way.
+        }
+      }
+      if (isBusy(err)) {
+        // Back to the head of the queue, ahead of later arrivals, so the retry
+        // preserves the order the writes were issued in.
+        this.pending = batch.concat(this.pending);
+        throw err;
+      }
+      for (const w of batch) w.reject(err);
+      return;
+    }
+    // Only now: before COMMIT a rollback could still take the write back, and a
+    // caller holding its row would have observed a write that never happened.
+    for (let i = 0; i < batch.length; i++) {
+      // Presence, not truthiness — a thrown value is not guaranteed to be one.
+      if ("err" in out[i]) batch[i].reject(out[i].err);
+      else batch[i].resolve(out[i].rows!);
+    }
+  }
+
+  /**
+   * Make sure some flusher is draining `pending`, without ever running two.
+   *
+   * The flusher loops instead of re-arming itself per batch. A caller that awaits
+   * its writes one at a time resumes and issues the next one *before* the flusher
+   * gets its turn back, so re-arming would cost that write an extra trip through
+   * the lock queue — measured as ~2x on sequential writes, which is most of them.
+   * Looping picks it up in the same session for free.
+   *
+   * The exit is safe because the last `pending` check and clearing the flag happen
+   * in one synchronous step: a write that arrives before it keeps the loop going,
+   * and one that arrives after sees the flag down and starts a new flusher.
+   */
+  private scheduleFlush(db: DB): void {
+    if (this.flushing) return;
+    this.flushing = true;
+    void (async () => {
+      try {
+        while (this.pending.length) {
+          try {
+            await this.withLock(() => this.flush(db));
+          } catch (err) {
+            // flush only throws on a lost write lock, and only after putting its
+            // batch back — so reaching here means withLock spent the whole budget
+            // and those writes are still queued with nobody else coming for them.
+            // Anything that arrived behind them is failed with the same error
+            // rather than left hanging: this store cannot write at all right now,
+            // which is what a lone write would have been told too.
+            const stranded = this.pending;
+            this.pending = [];
+            for (const w of stranded) w.reject(err);
+          }
+        }
+      } finally {
+        this.flushing = false;
+      }
+    })();
+  }
+
   protected async fetch(name: string, params: Params): Promise<any[]> {
     const db = this.ensure();
     await this.maybeRefreshStatistics(db);
-    return this.withLock(() => this.runNow(name, params));
+    // Reads keep their own turn on the lock — see isWriteStatement.
+    if (!this.writes[name]) return this.withLock(() => this.runNow(name, params));
+    return new Promise<any[]>((resolve, reject) => {
+      this.pending.push({ name, params, resolve, reject });
+      this.scheduleFlush(db);
+    });
   }
 
   protected async tx<T>(fn: (fetch: Fetch) => Promise<T>): Promise<T> {