cairnq 0.1.0 → 0.2.0

package/src/store/postgres.ts

@@ -1,21 +1,15 @@
 import type * as PG from "pg";
 
-import { newId } from "../ids.js";
-import { AlreadyExists, errorEnvelope, LostLease, ProtocolVersionMismatch } from "../errors.js";
-import { rowToTask, type Task } from "../models.js";
 import { loadMigrations, loadStatements } from "../sql.js";
-import type { ListInput, SubmitInput, TaskStore } from "./base.js";
-
-const SUPPORTED_PROTOCOL_MAJOR = 1;
-
-const LEASE_EXPIRED_ERROR_JSON = JSON.stringify(
-  errorEnvelope({
-    type: "LeaseExpired",
-    code: "lease_expired",
-    message: "task lease expired and max attempts reached",
-    retryable: false,
-  }),
-);
+import {
+  checkProtocolVersion,
+  COMMENT,
+  type Fetch,
+  NAMED,
+  type Params,
+  statementParams,
+  TaskStore,
+} from "./base.js";
 
 // `pg` is an optional dependency: the SDK is SQLite-first, so it's loaded lazily
 // the first time a PostgresStore connects. Absent -> a clear install hint.
@@ -39,47 +33,111 @@ async function loadPg(): Promise<typeof import("pg")> {
 }
 
 /**
- * Translate the protocol's named-parameter SQL (`:name`) into Postgres positional
+ * Roll back on the way out of a failed transaction, without letting the rollback
+ * become the error the caller sees. A dropped connection fails both the statement
+ * and the rollback, and it is the first one that says what went wrong.
+ */
+async function rollbackQuietly(client: PG.PoolClient): Promise<void> {
+  try {
+    await client.query("rollback");
+  } catch {
+    // Already rolled back, or the connection is gone. Either way the original
+    // error is the one worth propagating.
+  }
+}
+
+// Notification channels, emitted by the 0003_notify trigger.
+const QUEUED_CHANNEL = "cairnq_queued";
+const DONE_CHANNEL = "cairnq_done";
+
+// Backoff between attempts to (re)connect the LISTEN connection after a
+// transient failure. Doubles per failure up to the cap; polling covers the gap.
+const LISTENER_RETRY_MS = 1_000;
+const LISTENER_RETRY_MAX_MS = 30_000;
+
+const translated = new Map<string, { text: string; order: readonly string[] }>();
+
+/**
+ * The rewritten SQL and the order its `$n` slots must be filled in.
+ *
+ * Translates the protocol's named-parameter SQL (`:name`) into Postgres positional
  * placeholders (`$1`), collapsing each DISTINCT name to ONE slot — statements reuse
- * a name across CASE branches / IS NULL guards (e.g. list.sql). SQL comments are
- * stripped first so a `:name` mentioned in a header comment (e.g. "now + :lease_ms")
- * never leaks into the parameter list. Exported for unit testing.
+ * a name across CASE branches / IS NULL guards (e.g. list.sql). Which names count
+ * as parameters is `statementParams`' decision, shared with the SQLite binding path
+ * so the two can't disagree about, say, a `::type` cast.
+ *
+ * Memoized on the statement text, which is loaded once and never varies: this runs
+ * on every query, including the worker's poll loop, for a result that cannot have
+ * changed. Exported for unit testing.
+ */
+export function positionalStatement(sql: string): { text: string; order: readonly string[] } {
+  let entry = translated.get(sql);
+  if (!entry) {
+    const order = statementParams(sql);
+    const slot = new Map(order.map((name, i) => [name, i + 1])); // 1-based $n
+    const text = sql
+      .replace(COMMENT, "")
+      .replace(NAMED, (_m, name: string) => `$${slot.get(name)}`);
+    entry = { text, order };
+    translated.set(sql, entry);
+  }
+  return entry;
+}
+
+/**
+ * The statement's rewritten text plus this call's values, in slot order. Names the
+ * statement does not use are simply not bound, so callers may pass a superset.
  */
 export function toPositional(
   sql: string,
-  params: Record<string, unknown>,
+  params: Params,
 ): { text: string; values: unknown[] } {
-  const body = sql.replace(/--[^\n]*/g, "");
-  const order: string[] = [];
-  const slot = new Map<string, number>();
-  const text = body.replace(/(?<!:):(\w+)/g, (_m, name: string) => {
-    let i = slot.get(name);
-    if (i === undefined) {
-      order.push(name);
-      i = order.length; // 1-based $n
-      slot.set(name, i);
-    }
-    return `$${i}`;
-  });
+  const { text, order } = positionalStatement(sql);
   return { text, values: order.map((n) => params[n]) };
 }
 
 /**
- * PostgresStore — asyncpg's TS counterpart: a `pg` Pool executing the shared
- * cairnq-protocol SQL (postgres dialect). Multi-host capable — unlike SQLite this
- * coordinates API and worker processes across machines through one database. Time
- * comes from the DB clock (now()), claim uses FOR UPDATE SKIP LOCKED, JSON columns
- * are jsonb (bound as JSON text, read back as objects by rowToTask's adaptive parse).
+ * PostgresStore — the Postgres dialect of the shared cairnq-protocol SQL.
+ *
+ * Everything protocol-shaped lives in TaskStore; this file is only what Postgres
+ * does differently: a `pg` Pool, `:name` -> `$n` translation, and time taken from
+ * the DB clock (`now()`) instead of from the SDK, which is what makes this backend
+ * multi-host — unlike SQLite it coordinates API and worker processes across
+ * machines, with no shared clock to agree on. claim uses FOR UPDATE SKIP LOCKED
+ * and needs no claimable_probe, because PG readers don't block writers. JSON
+ * columns are jsonb (bound as JSON text, read back as objects by rowToTask).
  */
-export class PostgresStore implements TaskStore {
+export class PostgresStore extends TaskStore {
   private pool: PG.Pool | null = null;
   private connecting: Promise<void> | null = null;
   private readonly statements: Record<string, string>;
 
+  // ------------------------------------------------------- LISTEN/NOTIFY state
+  // One dedicated connection LISTENs on both channels (see 0003_notify.sql and
+  // the claimWake/taskDoneWake contract on TaskStore). Failure to establish or
+  // keep it silently degrades the store to the base class's plain polling.
+  private listener: PG.Client | null = null;
+  private listenerConnecting: Promise<void> | null = null;
+  /** LISTEN is off for good: the store was closed, or the server accepted a
+   * connection but refused LISTEN (e.g. a transaction-mode pooler) —
+   * deterministic, so retrying would fail the same way every time. */
+  private listenerUnavailable = false;
+  /** A failure to even connect is transient (network blip, server restarting):
+   * retry, but not before this time, backing off so a down server is not
+   * hammered from the poll loop. */
+  private listenerRetryAt = 0;
+  private listenerBackoffMs = LISTENER_RETRY_MS;
+  /** Queues notified while nobody was waiting; consumed by the next claimWake
+   * so a wake that lands between polls is not lost. */
+  private readonly pendingQueues = new Set<string>();
+  /** Wake callbacks by key: "queued" (broadcast) or "done:<task id>". */
+  private readonly waiters = new Map<string, Set<() => void>>();
+
   constructor(
     private readonly dsn: string,
     private readonly opts: { max?: number } = {},
   ) {
+    super();
     this.statements = loadStatements("postgres");
   }
 
@@ -88,6 +146,8 @@ export class PostgresStore implements TaskStore {
   }
 
   async close(): Promise<void> {
+    this.listenerUnavailable = true; // no revival after close
+    this.dropListener();
     if (this.pool) {
       const p = this.pool;
       this.pool = null;
@@ -116,12 +176,7 @@ export class PostgresStore implements TaskStore {
       const client = await pool.connect();
       try {
         await this.applyMigrations(client);
-        const version = await this.readProtocolVersion(client);
-        if (version !== SUPPORTED_PROTOCOL_MAJOR) {
-          throw new ProtocolVersionMismatch(
-            `storage protocol_version=${version}, SDK supports ${SUPPORTED_PROTOCOL_MAJOR}`,
-          );
-        }
+        checkProtocolVersion(await this.readProtocolVersion(client));
       } finally {
         client.release();
       }
@@ -130,6 +185,9 @@ export class PostgresStore implements TaskStore {
       throw e;
     }
     this.pool = pool; // publish only a fully-migrated, version-checked pool
+    // Warm the LISTEN connection in the background so the first idle sleep is
+    // already wakeable. Fire-and-forget: failure just means polling.
+    this.listenerReady();
   }
 
   private async applyMigrations(client: PG.PoolClient): Promise<void> {
@@ -137,33 +195,36 @@ export class PostgresStore implements TaskStore {
       "create table if not exists cairnq_migrations " +
         "(name text primary key, applied_at_ms bigint not null)",
     );
-    const applied = new Set(
-      (await client.query("select name from cairnq_migrations")).rows.map(
-        (r: { name: string }) => r.name,
-      ),
-    );
     for (const { name, sql } of loadMigrations("postgres")) {
-      if (applied.has(name)) continue;
+      // Check and apply inside one transaction, with the table lock taken up
+      // front: two processes cold-starting together would otherwise both see a
+      // migration as unapplied and both run it.
       try {
         await client.query("begin");
-        await client.query(sql); // multi-statement DDL (simple-query, no params)
-        await client.query(
-          "insert into cairnq_migrations (name, applied_at_ms) values " +
-            "($1, (extract(epoch from now()) * 1000)::bigint) on conflict (name) do nothing",
-          [name],
-        );
+        await client.query("lock table cairnq_migrations in exclusive mode");
+        const applied = await client.query("select 1 from cairnq_migrations where name = $1", [
+          name,
+        ]);
+        if (applied.rowCount === 0) {
+          await client.query(sql); // multi-statement DDL (simple-query, no params)
+          await client.query(
+            "insert into cairnq_migrations (name, applied_at_ms) values " +
+              "($1, (extract(epoch from now()) * 1000)::bigint)",
+            [name],
+          );
+        }
         await client.query("commit");
       } catch (e) {
-        await client.query("rollback");
+        await rollbackQuietly(client);
         throw e;
       }
     }
   }
 
+  // Takes an explicit client: during doConnect the pool is not published yet, so
+  // this cannot go through fetch(). The statement binds nothing, so it runs as-is.
   private async readProtocolVersion(client: PG.PoolClient): Promise<number> {
-    const res = await client.query(
-      "select value from cairnq_meta where key = 'protocol_version'",
-    );
+    const res = await client.query(this.statements.protocol_version);
     return res.rows.length ? Number(res.rows[0].value) : 0;
   }
 
@@ -177,233 +238,148 @@ export class PostgresStore implements TaskStore {
     }
   }
 
-  // ------------------------------------------------------------------ helpers
-  private async run(name: string, params: Record<string, unknown>): Promise<any[]> {
-    const { text, values } = toPositional(this.statements[name], params);
-    return (await this.pool!.query(text, values)).rows;
-  }
-
-  private async runOn(
-    client: PG.PoolClient,
-    name: string,
-    params: Record<string, unknown>,
-  ): Promise<any[]> {
-    const { text, values } = toPositional(this.statements[name], params);
-    return (await client.query(text, values)).rows;
-  }
-
-  // BEGIN … COMMIT on a dedicated pool client, rolling back on any error.
-  private async tx<T>(fn: (client: PG.PoolClient) => Promise<T>): Promise<T> {
-    const client = await this.pool!.connect();
-    try {
-      await client.query("begin");
-      const out = await fn(client);
-      await client.query("commit");
-      return out;
-    } catch (e) {
-      await client.query("rollback");
-      throw e;
-    } finally {
-      client.release();
+  // ------------------------------------------------------------ wake channel
+  async claimWake(queues: string[], timeoutMs: number): Promise<void> {
+    if (!this.listenerReady()) return super.claimWake(queues, timeoutMs);
+    const deadline = Date.now() + timeoutMs;
+    for (;;) {
+      // Consume every pending notification for our queues; any hit wakes us.
+      let hit = false;
+      for (const q of queues) if (this.pendingQueues.delete(q)) hit = true;
+      if (hit) return;
+      const remaining = deadline - Date.now();
+      if (remaining <= 0) return;
+      // Broadcast wake: another queue's notification loops us back to waiting
+      // with the remaining budget instead of waking the worker for nothing.
+      await this.wakeOn("queued", remaining);
     }
   }
 
-  // An ownership-checked worker write: each statement's WHERE pins worker_id + a
-  // live lease, so 0 rows back means the lease was lost.
-  private async ownedWrite(
-    name: string,
-    taskId: string,
-    params: Record<string, unknown>,
-  ): Promise<Task> {
-    const rows = await this.run(name, params);
-    if (!rows.length) throw new LostLease(taskId);
-    return rowToTask(rows[0]);
+  taskDoneWake(taskId: string, timeoutMs: number): Promise<void> {
+    if (!this.listenerReady()) return super.taskDoneWake(taskId, timeoutMs);
+    return this.wakeOn(`done:${taskId}`, timeoutMs);
   }
 
-  // ------------------------------------------------------------- client side
-  async submit(input: SubmitInput): Promise<Task> {
-    await this.ensure();
-    const id = newId("task");
-    // No now_ms / run_at_ms: the DB clock supplies time, so submit passes a
-    // relative :delay_ms and the SQL computes run_at = now + delay.
-    const ins = {
-      id,
-      name: input.name,
-      queue: input.queue ?? "default",
-      payload: JSON.stringify(input.payload ?? {}),
-      metadata: JSON.stringify(input.metadata ?? {}),
-      max_attempts: input.maxAttempts ?? 3,
-      priority: input.priority ?? 0,
-      delay_ms: input.runAtDelayMs ?? 0,
-      parent_id: input.parentId ?? null,
-      root_id: input.rootId ?? id,
-      correlation_id: input.correlationId ?? null,
-    };
-    const key = input.key ?? null;
-    const conflict = input.conflict ?? "reuse";
-    if (key === null) return rowToTask((await this.run("insert_task", ins))[0]);
-
-    return this.tx(async (client) => {
-      const existing = (await this.runOn(client, "get_key", { key })) as { task_id: string }[];
-      if (existing.length) {
-        const exId = existing[0].task_id;
-        if (conflict === "reuse") return rowToTask((await this.runOn(client, "get", { id: exId }))[0]);
-        if (conflict === "reject") throw new AlreadyExists(key);
-        if (conflict === "replace") {
-          await this.runOn(client, "cancel", { id: exId });
-          const row = (await this.runOn(client, "insert_task", ins))[0];
-          await this.runOn(client, "upsert_key", { key, task_id: id });
-          return rowToTask(row);
-        }
-        throw new Error(`unknown conflict strategy: ${conflict}`);
-      }
-      const row = (await this.runOn(client, "insert_task", ins))[0];
-      await this.runOn(client, "upsert_key", { key, task_id: id });
-      return rowToTask(row);
+  /** A promise resolving on notification-or-timeout, deregistering either way.
+   * The timer is unref'd: while a listener exists its socket keeps the process
+   * alive, and dropListener wakes every waiter the moment it goes away. */
+  private wakeOn(key: string, timeoutMs: number): Promise<void> {
+    return new Promise((resolve) => {
+      let set = this.waiters.get(key);
+      if (!set) this.waiters.set(key, (set = new Set()));
+      const peers = set;
+      const waiter = () => {
+        clearTimeout(timer);
+        peers.delete(waiter);
+        if (peers.size === 0) this.waiters.delete(key);
+        resolve();
+      };
+      const timer = setTimeout(waiter, timeoutMs);
+      timer.unref?.();
+      peers.add(waiter);
     });
   }
 
-  async get(taskId: string): Promise<Task | null> {
-    await this.ensure();
-    const rows = await this.run("get", { id: taskId });
-    return rows.length ? rowToTask(rows[0]) : null;
-  }
-
-  async getByKey(key: string): Promise<Task | null> {
-    await this.ensure();
-    const rows = await this.run("get_by_key", { key });
-    return rows.length ? rowToTask(rows[0]) : null;
-  }
-
-  async list(input: ListInput = {}): Promise<Task[]> {
-    await this.ensure();
-    const rows = await this.run("list", {
-      status: input.status ?? null,
-      queue: input.queue ?? null,
-      name: input.name ?? null,
-      root_id: input.rootId ?? null,
-      correlation_id: input.correlationId ?? null,
-      limit: input.limit ?? 100,
-      offset: input.offset ?? 0,
-    });
-    return rows.map(rowToTask);
+  /** True once the LISTEN connection is up; starts connecting it otherwise
+   * (respecting the transient-failure backoff). Callers fall back to plain
+   * polling until it is ready (or forever, if it can't be established) —
+   * correctness never depends on it. */
+  private listenerReady(): boolean {
+    if (this.listener) return true;
+    if (!this.listenerUnavailable && !this.listenerConnecting && Date.now() >= this.listenerRetryAt) {
+      this.listenerConnecting = this.startListener();
+    }
+    return false;
   }
 
-  async cancel(taskId: string): Promise<Task | null> {
-    await this.ensure();
-    const rows = await this.run("cancel", { id: taskId });
-    return rows.length ? rowToTask(rows[0]) : null;
+  private async startListener(): Promise<void> {
+    try {
+      const pg = await loadPg();
+      // Built from the raw DSN: if `opts` ever grows connection-level settings
+      // (ssl, application_name), the listener must receive them too.
+      const client = new pg.Client({ connectionString: this.dsn });
+      try {
+        await client.connect();
+      } catch {
+        // Could not even connect — transient. Schedule a backed-off retry;
+        // polling covers the gap.
+        this.listenerRetryAt = Date.now() + this.listenerBackoffMs;
+        this.listenerBackoffMs = Math.min(LISTENER_RETRY_MAX_MS, this.listenerBackoffMs * 2);
+        return;
+      }
+      client.on("notification", (msg) => this.onNotification(msg.channel, msg.payload));
+      // A dropped listener degrades to polling; the next wake call reconnects.
+      client.on("error", () => this.dropListener());
+      try {
+        await client.query(`listen ${QUEUED_CHANNEL}; listen ${DONE_CHANNEL}`);
+      } catch {
+        // Connected, but LISTEN was refused (e.g. a transaction-mode pooler) —
+        // deterministic, so off for good. Polling covers it.
+        this.listenerUnavailable = true;
+        void client.end().catch(() => {});
+        return;
+      }
+      if (this.listenerUnavailable) {
+        void client.end().catch(() => {}); // closed while we were connecting
+        return;
+      }
+      this.listener = client;
+      this.listenerBackoffMs = LISTENER_RETRY_MS;
+    } catch {
+      // Anything unexpected (e.g. `pg` failed to load): off for good rather
+      // than a retry loop that cannot succeed.
+      this.listenerUnavailable = true;
+    } finally {
+      this.listenerConnecting = null;
+    }
   }
 
-  async cancelByKey(key: string): Promise<Task | null> {
-    await this.ensure();
-    return this.tx(async (client) => {
-      const existing = (await this.runOn(client, "get_key", { key })) as { task_id: string }[];
-      if (!existing.length) return null;
-      const rows = await this.runOn(client, "cancel", { id: existing[0].task_id });
-      return rows.length ? rowToTask(rows[0]) : null;
-    });
+  private onNotification(channel: string, payload: string | undefined): void {
+    let key: string;
+    if (channel === QUEUED_CHANNEL && payload) {
+      this.pendingQueues.add(payload);
+      key = "queued";
+    } else if (channel === DONE_CHANNEL && payload) {
+      key = `done:${payload}`;
+    } else {
+      return;
+    }
+    const set = this.waiters.get(key);
+    // A waiter only deletes itself, which Set iteration tolerates — no copy.
+    if (set) for (const w of set) w();
   }
 
-  async retry(taskId: string, opts: { resetAttempt?: boolean } = {}): Promise<Task | null> {
-    await this.ensure();
-    const rows = await this.run("retry", { id: taskId, reset_attempt: opts.resetAttempt ?? false });
-    return rows.length ? rowToTask(rows[0]) : null;
+  private dropListener(): void {
+    const client = this.listener;
+    this.listener = null;
+    if (client) void client.end().catch(() => {});
+    // Release everyone promptly; their fallback poll takes over.
+    for (const set of [...this.waiters.values()]) for (const w of [...set]) w();
   }
 
-  async retryByKey(key: string, opts: { resetAttempt?: boolean } = {}): Promise<Task | null> {
+  // ------------------------------------------------------------ dialect seam
+  protected async fetch(name: string, params: Params): Promise<any[]> {
     await this.ensure();
-    return this.tx(async (client) => {
-      const existing = (await this.runOn(client, "get_key", { key })) as { task_id: string }[];
-      if (!existing.length) return null;
-      const rows = await this.runOn(client, "retry", {
-        id: existing[0].task_id,
-        reset_attempt: opts.resetAttempt ?? false,
-      });
-      return rows.length ? rowToTask(rows[0]) : null;
-    });
+    const { text, values } = toPositional(this.statements[name], params);
+    return (await this.pool!.query(text, values)).rows;
   }
 
-  // ------------------------------------------------------------- worker side
-  async claim(input: {
-    queues: string[];
-    workerId: string;
-    leaseMs?: number;
-    limit?: number;
-  }): Promise<Task[]> {
+  protected async tx<T>(fn: (fetch: Fetch) => Promise<T>): Promise<T> {
     await this.ensure();
-    // No claimable_probe: PG readers don't block writers, so a plain transaction
-    // (recover then claim) is cheap even when idle. FOR UPDATE SKIP LOCKED in
-    // claim.sql gives true concurrent, non-contending dispatch.
-    return this.tx(async (client) => {
-      await this.runOn(client, "recover_leases", { lease_expired_error: LEASE_EXPIRED_ERROR_JSON });
-      const rows = await this.runOn(client, "claim", {
-        queues: input.queues,
-        worker_id: input.workerId,
-        lease_ms: input.leaseMs ?? 30_000,
-        limit: input.limit ?? 1,
+    const client = await this.pool!.connect();
+    try {
+      await client.query("begin");
+      const out = await fn(async (name, params) => {
+        const { text, values } = toPositional(this.statements[name], params);
+        return (await client.query(text, values)).rows;
       });
-      return rows.map(rowToTask);
-    });
-  }
-
-  async heartbeat(input: { taskId: string; workerId: string; leaseMs?: number }): Promise<Task> {
-    await this.ensure();
-    return this.ownedWrite("heartbeat", input.taskId, {
-      id: input.taskId,
-      worker_id: input.workerId,
-      lease_ms: input.leaseMs ?? 30_000,
-    });
-  }
-
-  async progress(input: {
-    taskId: string;
-    workerId: string;
-    progress: number | null;
-    message: string | null;
-  }): Promise<Task> {
-    await this.ensure();
-    return this.ownedWrite("progress", input.taskId, {
-      id: input.taskId,
-      worker_id: input.workerId,
-      progress: input.progress,
-      message: input.message,
-    });
-  }
-
-  async succeed(input: { taskId: string; workerId: string; result: unknown }): Promise<Task> {
-    await this.ensure();
-    return this.ownedWrite("succeed", input.taskId, {
-      id: input.taskId,
-      worker_id: input.workerId,
-      result: input.result == null ? null : JSON.stringify(input.result),
-      message: null,
-    });
-  }
-
-  async complete(input: { taskId: string; workerId: string; result: unknown }): Promise<Task> {
-    await this.ensure();
-    return this.ownedWrite("complete", input.taskId, {
-      id: input.taskId,
-      worker_id: input.workerId,
-      result: input.result == null ? null : JSON.stringify(input.result),
-    });
-  }
-
-  async fail(input: {
-    taskId: string;
-    workerId: string;
-    error: unknown;
-    retryable?: boolean;
-    delayMs?: number;
-  }): Promise<Task> {
-    await this.ensure();
-    return this.ownedWrite("fail", input.taskId, {
-      id: input.taskId,
-      worker_id: input.workerId,
-      error: JSON.stringify(input.error ?? {}),
-      retryable: input.retryable !== false,
-      delay_ms: input.delayMs ?? 0,
-    });
+      await client.query("commit");
+      return out;
+    } catch (e) {
+      await rollbackQuietly(client);
+      throw e;
+    } finally {
+      client.release();
+    }
   }
 }