cairnq 0.3.0 → 0.5.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

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,
@@ -28,7 +29,7 @@ const BUSY_RETRY_MAX_DELAY_MS = 50;
  * Bounds how long the planner can work from a stale table shape; a minute is
  * arbitrary but small next to the days a worker holds its connection. It does not
  * set how often an ANALYZE actually runs — SQLite decides that itself, and only
- * when a table has outgrown its statistics by roughly 24x, so a shorter interval
+ * once the table has diverged from its statistics by 10x, so a shorter interval
  * costs more no-ops (a few microseconds each) rather than more analyzing.
  */
 const STATS_REFRESH_INTERVAL_MS = 60_000;
@@ -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
@@ -422,10 +460,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> {

package/src/worker.ts

@@ -1,3 +1,4 @@
+import type { BackpressureOptions } from "./backpressure.js";
 import { TaskContext } from "./context.js";
 import { errorEnvelope, LostLease, SerializationError, TaskError } from "./errors.js";
 import { newId } from "./ids.js";
@@ -14,7 +15,12 @@ export type TypedHandler<P, R> = (ctx: TaskContext, payload: P) => R | Promise<R
 /** Where an error the worker recovered from came from. */
 export type ErrorPhase = "claim" | "execute";
 
-export interface WorkerOptions {
+/**
+ * Backpressure is accepted here too, not only on CairnQ: a handler spawning
+ * children through TaskContext.submit is a producer, and in a worker process
+ * there is usually no CairnQ handle to have configured the store.
+ */
+export interface WorkerOptions extends Partial<BackpressureOptions> {
   concurrency?: number;
   leaseMs?: number;
   heartbeatIntervalMs?: number;
@@ -33,6 +39,25 @@ export interface WorkerOptions {
    * a retryable `handler_timeout` failure. Unset disables the ceiling.
    */
   maxRunMs?: number;
+  /**
+   * Resident payload bytes allowed across running handlers, independent of
+   * their count.
+   *
+   * `concurrency` bounds tasks, not memory, so a worker sized for small payloads
+   * holds concurrency * largest-payload bytes the moment a batch of big ones
+   * arrives — for payloads that carry media inline, that is the difference
+   * between megabytes and gigabytes resident. Once the budget is spent the
+   * worker stops claiming until running handlers give it back.
+   *
+   * The bound is on tasks already executing. A claim commits to a whole batch
+   * before any size is known, so one batch can overshoot by up to `claimBatch`
+   * payloads; lower `claimBatch` to tighten that. A single payload larger than
+   * the entire budget still runs — alone, rather than deadlocking the worker.
+   *
+   * Costs one JSON serialization per task to measure, so it is only computed
+   * when set. Unset disables the budget.
+   */
+  maxInFlightBytes?: number;
   /**
    * Called for errors the worker survived — a claim that threw, a store write
    * that failed while finalizing a task. Without it these are silent: the run
@@ -82,9 +107,35 @@ function timeoutEnvelope(name: string, maxRunMs: number): Record<string, unknown
 
 const TIMED_OUT = Symbol("cairnq.timedOut");
 
+/**
+ * Resident size of a task's payload, for the maxInFlightBytes budget.
+ *
+ * Re-serializes because by this point the wire form is gone: `pg` parses a jsonb
+ * column with JSON.parse and discards the text, so on Postgres there is nothing
+ * cheaper to read. On SQLite the column does arrive as a string that rowToTask
+ * sees before parsing — capturing its length there would make this free, at the
+ * cost of carrying a non-protocol field on Task in both SDKs. Left for when the
+ * measurement shows up in a profile.
+ *
+ * What the budget is really after is the memory a payload pins while its handler
+ * runs, and its JSON length tracks that closely enough to size one by.
+ */
+function payloadBytes(task: Task): number {
+  try {
+    return Buffer.byteLength(JSON.stringify(task.payload) ?? "");
+  } catch {
+    // Unmeasurable, and it came out of the store, so it is already resident:
+    // charging nothing under-counts, but failing the claim over an accounting
+    // detail would drop a task the worker can otherwise run.
+    return 0;
+  }
+}
+
 export class Worker {
   private readonly handlers = new Map<string, Handler>();
   private readonly workerId = newId("worker");
+  /** Payload bytes charged to running handlers — see maxInFlightBytes. */
+  private inFlightBytes = 0;
   private stopped = false;
   private stopWake!: () => void;
   // Resolved once by stop(); every sleep races against it. A stopped worker
@@ -102,6 +153,14 @@ export class Worker {
     if (opts.maxRunMs != null && opts.maxRunMs <= 0) {
       throw new Error(`maxRunMs must be > 0, got ${opts.maxRunMs}`);
     }
+    // 0 would make the budget permanently spent, so the worker would claim
+    // nothing and look hung. Rejected here, as the Python SDK does.
+    if (opts.maxInFlightBytes != null && opts.maxInFlightBytes <= 0) {
+      throw new Error(`maxInFlightBytes must be > 0, got ${opts.maxInFlightBytes}`);
+    }
+    if (opts.maxQueueDepth != null) {
+      store.useBackpressure(opts as BackpressureOptions);
+    }
   }
 
   static sqlite(
@@ -208,9 +267,16 @@ export class Worker {
     running: Set<Promise<void>>,
   ): Promise<void> {
     const pollMs = this.opts.pollIntervalMs ?? 500;
+    const byteBudget = this.opts.maxInFlightBytes;
     while (!this.stopped) {
       const free = concurrency - running.size;
-      if (free <= 0) {
+      // Two ceilings, either of which stops the claim: task count and resident
+      // payload bytes. The byte arm is guarded on running.size because it must
+      // never be the reason we race an empty set — Promise.race([]) is pending
+      // forever, past even stop(). With nothing running, nothing is resident, so
+      // the budget cannot be the thing holding us back anyway.
+      const overBudget = byteBudget != null && this.inFlightBytes >= byteBudget;
+      if (running.size > 0 && (free <= 0 || overBudget)) {
         // Wait for a slot rather than spinning. execute() never rejects, so
         // racing these is safe.
         await Promise.race([...running]);
@@ -241,7 +307,14 @@ export class Worker {
         continue;
       }
       for (const task of claimed) {
-        const p = this.execute(task, leaseMs).finally(() => running.delete(p));
+        // Charged before the handler starts and refunded when it settles, so the
+        // budget covers exactly the span the payload is pinned in memory.
+        const bytes = byteBudget == null ? 0 : payloadBytes(task);
+        this.inFlightBytes += bytes;
+        const p = this.execute(task, leaseMs).finally(() => {
+          this.inFlightBytes -= bytes;
+          running.delete(p);
+        });
         running.add(p);
       }
     }