diff --git a/internal-packages/run-engine/src/engine/bench/waitpointCoordinator.bench.test.ts b/internal-packages/run-engine/src/engine/bench/waitpointCoordinator.bench.test.ts new file mode 100644 index 00000000000..01393ba4f96 --- /dev/null +++ b/internal-packages/run-engine/src/engine/bench/waitpointCoordinator.bench.test.ts @@ -0,0 +1,266 @@ +/** + * Waitpoint coordination benchmark. Reports numbers; asserts nothing — on a shared runner + * the timings swing far more than any threshold worth gating on. + * + * Four groups, and only the first two are pairs: + * + * 1. Pending count — the store's SCARD gate against the previous path's + * `COUNT(*) ... WHERE status='PENDING'`, over the same population. Like for like. + * 2. Read amplification — the store's `readBlockState` against a full-payload `SELECT` + * of the same waitpoints. Like for like. + * 3. Store-only write paths — block+complete+deliver and K-watcher fan-out. Absolute + * numbers with NO Postgres counterpart: no single statement on the previous path + * corresponds to a Redis round trip that both blocks a run and delivers to watchers. + * 4. Register cost versus edge count — `registerBlocks` registers each edge with its own + * round trip before the single absorb. This measures whether that serial loop is a + * real cost at a wide fan-in, or a non-issue, at several fan-in widths. + * + * Every Postgres measurement here runs against rows this file inserts. A baseline over an + * empty table measures nothing. + * + * Knobs: BENCH_WP_ITERATIONS, BENCH_WP_FANIN, BENCH_WP_WATCHERS, BENCH_WP_REGISTER_WIDTHS, + * BENCH_WP_REGISTER_SAMPLES. + */ +import { containerTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import { + WaitpointStoreCoordinator, + type BlockEdge, + type WaitpointRecordInput, +} from "../waitpointCoordinator/storeCoordinator.js"; +import { setupAuthenticatedEnvironment } from "../tests/setup.js"; + +vi.setConfig({ testTimeout: 900_000 }); + +const ITERATIONS = Number(process.env.BENCH_WP_ITERATIONS ?? 100); +const FANIN = Number(process.env.BENCH_WP_FANIN ?? 1001); +const WATCHERS = Number(process.env.BENCH_WP_WATCHERS ?? 100); +const REGISTER_WIDTHS = (process.env.BENCH_WP_REGISTER_WIDTHS ?? "1,10,100,1001") + .split(",") + .map((raw) => Number(raw.trim())) + .filter((width) => Number.isFinite(width) && width > 0); +const REGISTER_SAMPLES = Number(process.env.BENCH_WP_REGISTER_SAMPLES ?? 20); +const NOW = new Date().toISOString(); + +type Sample = { label: string; count: number; p50: number; p99: number; totalMs: number }; + +function percentile(sorted: number[], p: number): number { + if (sorted.length === 0) return 0; + return sorted[Math.min(sorted.length - 1, Math.floor((p / 100) * sorted.length))]!; +} + +async function measure(label: string, count: number, run: (i: number) => Promise) { + const durations: number[] = []; + const started = Date.now(); + for (let i = 0; i < count; i++) { + const t0 = performance.now(); + await run(i); + durations.push(performance.now() - t0); + } + durations.sort((a, b) => a - b); + const sample: Sample = { + label, + count, + p50: percentile(durations, 50), + p99: percentile(durations, 99), + totalMs: Date.now() - started, + }; + console.log( + `[bench] ${sample.label} n=${sample.count} p50=${sample.p50.toFixed(2)}ms ` + + `p99=${sample.p99.toFixed(2)}ms total=${sample.totalMs}ms` + ); + return sample; +} + +function record(id: string, environmentId: string, projectId: string): WaitpointRecordInput { + return { + id, + friendlyId: `waitpoint_${id}`, + type: "MANUAL", + environmentId, + projectId, + createdAt: NOW, + updatedAt: NOW, + userProvidedIdempotencyKey: false, + tags: [], + }; +} + +const completion = { + completedAt: NOW, + outputType: "application/json", + outputIsError: false, + output: { inline: '{"ok":true}' }, +}; + +function edge(waitpointId: string, batchIndex?: number): BlockEdge { + return { waitpointId, batchIndex, createdAt: NOW, type: "MANUAL" }; +} + +async function insertWaitpoints( + prisma: PrismaClient, + ids: string[], + environmentId: string, + projectId: string +) { + await prisma.waitpoint.createMany({ + data: ids.map((id) => ({ + id, + friendlyId: `waitpoint_${id}`, + type: "MANUAL" as const, + idempotencyKey: id, + userProvidedIdempotencyKey: false, + projectId, + environmentId, + })), + }); +} + +containerTest( + "waitpoint coordination: pending count, read amplification, store write paths, register cost", + async ({ prisma, redisOptions }) => { + const env = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + const store = new WaitpointStoreCoordinator({ redisOptions }); + const samples: Sample[] = []; + const registerCost: Array<{ + width: number; + p50Ms: number; + p99Ms: number; + perEdgeMsP50: number; + }> = []; + + try { + const ids = Array.from({ length: FANIN }, (_, i) => `bench_w_${i}`); + + // Both stores get the SAME population. A Postgres baseline over an empty table + // measures an index probe against nothing. + await insertWaitpoints(prisma, ids, env.id, env.project.id); + for (const id of ids) { + await store.createIfAbsent({ + record: record(id, env.id, env.project.id), + status: "PENDING", + }); + } + await store.registerBlocks({ + runId: "bench_run_fanin", + edges: ids.map((id, index) => edge(id, index)), + }); + + // --- group 1: the pending-count gate, like for like --- + samples.push( + await measure("store.pendingCount", ITERATIONS, async () => { + await store.absorbBlockers({ runId: "bench_run_fanin", edges: [] }); + }) + ); + samples.push( + await measure("postgres.pendingCount", ITERATIONS, async () => { + await prisma.$queryRaw`SELECT COUNT(*) FROM "Waitpoint" WHERE id = ANY(${ids}::text[]) AND status = 'PENDING'`; + }) + ); + + // --- group 2: read amplification, like for like --- + samples.push( + await measure("store.readBlockState", ITERATIONS, async () => { + await store.readBlockState("bench_run_fanin"); + }) + ); + samples.push( + await measure("postgres.hydrateFullPayload", ITERATIONS, async () => { + // Every column of every waitpoint — the amplification the store removes. + await prisma.waitpoint.findMany({ where: { id: { in: ids } } }); + }) + ); + + // --- group 3: store-only write paths, no Postgres counterpart --- + samples.push( + await measure("store.block+complete+deliver", ITERATIONS, async (i) => { + const id = `bench_cycle_${i}`; + await store.createIfAbsent({ + record: record(id, env.id, env.project.id), + status: "PENDING", + }); + await store.registerBlocks({ runId: `bench_run_${i}`, edges: [edge(id)] }); + const done = await store.complete({ waitpointId: id, completion }); + for (const watcher of done.watchers) { + await store.deliverCompletion({ + runId: watcher.runId, + waitpointId: id, + completion: done.completion!, + }); + } + }) + ); + + const fanOutId = "bench_fanout_w"; + await store.createIfAbsent({ + record: record(fanOutId, env.id, env.project.id), + status: "PENDING", + }); + for (let i = 0; i < WATCHERS; i++) { + await store.registerBlocks({ runId: `bench_watcher_${i}`, edges: [edge(fanOutId)] }); + } + samples.push( + await measure(`store.complete+deliver(watchers=${WATCHERS})`, 1, async () => { + const done = await store.complete({ waitpointId: fanOutId, completion }); + // Serial on purpose: this is the worst case, and it is the number that says + // whether delivery needs to pipeline. + for (const watcher of done.watchers) { + await store.deliverCompletion({ + runId: watcher.runId, + waitpointId: fanOutId, + completion: done.completion!, + }); + } + }) + ); + + // --- group 4: register cost versus edge count --- + // registerBlocks registers each edge with its own round trip, serially, before the + // single absorb. A review flagged that a wide fan-in therefore serializes one round + // trip per edge. This measures the real cost at several widths rather than predicting + // it, so the decision about bounded concurrency is made against a number. + const registerPoolWidth = Math.max(0, ...REGISTER_WIDTHS); + const registerIds = Array.from({ length: registerPoolWidth }, (_, i) => `bench_reg_w_${i}`); + await insertWaitpoints(prisma, registerIds, env.id, env.project.id); + for (const id of registerIds) { + await store.createIfAbsent({ + record: record(id, env.id, env.project.id), + status: "PENDING", + }); + } + + for (const width of REGISTER_WIDTHS) { + const edges = registerIds.slice(0, width).map((id, index) => edge(id, index)); + let call = 0; + const sample = await measure( + `store.registerBlocks(edges=${width})`, + REGISTER_SAMPLES, + async () => { + await store.registerBlocks({ runId: `bench_register_${width}_${call++}`, edges }); + } + ); + samples.push(sample); + registerCost.push({ + width, + p50Ms: sample.p50, + p99Ms: sample.p99, + perEdgeMsP50: sample.p50 / width, + }); + console.log( + `[bench] store.registerBlocks(edges=${width}) implied per-edge cost ` + + `p50=${(sample.p50 / width).toFixed(3)}ms p99=${(sample.p99 / width).toFixed(3)}ms` + ); + } + + console.log( + `[bench] groups 1 and 2 are like-for-like pairs. Group 3 and the register-cost ` + + `group (4) have no Postgres counterpart: no single statement on the previous ` + + `path corresponds to a Redis round trip that blocks, completes and delivers, ` + + `or to a serial per-edge register loop.` + ); + console.log(`[bench] summary\n${JSON.stringify({ samples, registerCost }, null, 2)}`); + } finally { + await store.quit(); + } + } +); diff --git a/internal-packages/run-engine/src/engine/waitpointCoordinator/keys.test.ts b/internal-packages/run-engine/src/engine/waitpointCoordinator/keys.test.ts new file mode 100644 index 00000000000..463e1ecd388 --- /dev/null +++ b/internal-packages/run-engine/src/engine/waitpointCoordinator/keys.test.ts @@ -0,0 +1,132 @@ +import { describe, expect, it } from "vitest"; +import { + WaitpointKeyTagError, + assertSingleSlot, + edgeField, + idempotencyKey, + runBlockKeys, + waitpointIdFromEdgeField, + waitpointKeys, + watcherField, +} from "./keys.js"; + +describe("waitpointKeys", () => { + it("puts the record and its watchers under one hash tag", () => { + const k = waitpointKeys("abc123w"); + expect(k.record).toBe("wp:{abc123w}"); + expect(k.watchers).toBe("wp:{abc123w}:w"); + }); +}); + +describe("runBlockKeys", () => { + it("puts all three run keys under one hash tag", () => { + const k = runBlockKeys("run_abc"); + expect(k.pend).toBe("wp:run:{run_abc}:pend"); + expect(k.done).toBe("wp:run:{run_abc}:done"); + expect(k.edge).toBe("wp:run:{run_abc}:edge"); + }); +}); + +describe("idempotencyKey", () => { + it("tags by environment, so one environment's reservations share a slot", () => { + expect(idempotencyKey("env_1", "my-key")).toBe("wp:idem:{env_1}:my-key"); + }); +}); + +describe("edgeField", () => { + it("keys by waitpoint id and batch index, matching the Postgres unique key", () => { + expect(edgeField("w_a", 3)).toBe("w_a#3"); + }); + + it("collapses a null or absent batch index onto one field", () => { + expect(edgeField("w_a")).toBe("w_a#"); + expect(edgeField("w_a", null)).toBe("w_a#"); + }); + + it("distinguishes index 0 from an absent index", () => { + expect(edgeField("w_a", 0)).not.toBe(edgeField("w_a")); + }); +}); + +describe("waitpointIdFromEdgeField", () => { + it("round-trips back to the waitpoint id", () => { + for (const index of [undefined, null, 0, 7]) { + expect(waitpointIdFromEdgeField(edgeField("w_a", index))).toBe("w_a"); + } + }); + + it("returns undefined for a field with no separator", () => { + expect(waitpointIdFromEdgeField("nope")).toBeUndefined(); + }); + + it("splits on the last separator, tolerating a '#' inside the waitpoint id", () => { + expect(waitpointIdFromEdgeField("a#b#3")).toBe("a#b"); + }); +}); + +describe("watcherField", () => { + it("keys by run id and batch index, so one run can watch at several indexes", () => { + expect(watcherField("run_a", 2)).toBe("run_a#2"); + expect(watcherField("run_a")).toBe("run_a#"); + expect(watcherField("run_a", 0)).not.toBe(watcherField("run_a")); + }); +}); + +describe("assertSingleSlot", () => { + it("accepts keys that share one tag", () => { + const k = runBlockKeys("run_abc"); + expect(() => assertSingleSlot("runReadBlockState", [k.pend, k.done, k.edge])).not.toThrow(); + }); + + it("accepts a single tagged key", () => { + expect(() => assertSingleSlot("wpIdemReserve", [idempotencyKey("env_1", "k")])).not.toThrow(); + }); + + it("accepts an empty key list", () => { + expect(() => assertSingleSlot("noKeys", [])).not.toThrow(); + }); + + it("rejects keys from two different tags", () => { + const wp = waitpointKeys("w_a"); + const run = runBlockKeys("run_abc"); + expect(() => assertSingleSlot("bad", [wp.record, run.pend])).toThrow(WaitpointKeyTagError); + }); + + it("rejects an untagged key", () => { + expect(() => assertSingleSlot("bad", ["wp:no-tag"])).toThrow(WaitpointKeyTagError); + }); + + it("rejects an empty tag", () => { + expect(() => assertSingleSlot("bad", ["wp:{}"])).toThrow(WaitpointKeyTagError); + }); + + it("rejects an empty first pair, matching Redis rather than skipping to a later one", () => { + // Redis stops at the first `{`/`}` pair. An empty one means no tag at all, so it hashes + // the whole key. A regex would have found `a` here and wrongly claimed a shared slot. + expect(() => assertSingleSlot("bad", ["wp:{}{a}", "wp:{}{a}"])).toThrow(WaitpointKeyTagError); + }); + + it("takes the first pair when several are present", () => { + expect(() => assertSingleSlot("ok", ["wp:{a}{b}", "wp:{a}:w"])).not.toThrow(); + expect(() => assertSingleSlot("bad", ["wp:{a}{b}", "wp:{b}:w"])).toThrow(WaitpointKeyTagError); + }); + + it("does not degrade on a key made of many opening braces", () => { + const started = performance.now(); + expect(() => assertSingleSlot("bad", ["{".repeat(50_000)])).toThrow(WaitpointKeyTagError); + expect(performance.now() - started).toBeLessThan(1_000); + }); + + it("names the operation and the offending key in the error", () => { + const wp = waitpointKeys("w_a"); + const run = runBlockKeys("run_abc"); + try { + assertSingleSlot("myOperation", [wp.record, run.pend]); + throw new Error("should have thrown"); + } catch (error) { + expect(error).toBeInstanceOf(WaitpointKeyTagError); + expect((error as Error).message).toContain("myOperation"); + expect((error as Error).message).toContain(run.pend); + } + }); +}); diff --git a/internal-packages/run-engine/src/engine/waitpointCoordinator/keys.ts b/internal-packages/run-engine/src/engine/waitpointCoordinator/keys.ts new file mode 100644 index 00000000000..28eac087b4b --- /dev/null +++ b/internal-packages/run-engine/src/engine/waitpointCoordinator/keys.ts @@ -0,0 +1,98 @@ +/** + * Waitpoint coordination keyspace. Two hash tags, deliberately: + * + * - `wp:{waitpointId}` — the record, its status and completion envelope, plus the + * watcher hash. A waitpoint has N watcher runs, so it cannot live under any single + * run's tag. + * - `wp:run:{runId}:*` — one run's pending set, delivered set and edge set. The pending + * set's cardinality is the blocked-versus-unblocked signal, so it has to be readable + * atomically, which means one slot. + * + * Every script therefore touches exactly one tag, and assertSingleSlot enforces it on + * every invocation. A cluster would reject a cross-slot script; a single-node test server + * would not, so this assertion is the only thing standing between a cross-slot bug and + * production. + */ + +export type WaitpointKeys = { record: string; watchers: string }; +export type RunBlockKeys = { pend: string; done: string; edge: string }; + +export function waitpointKeys(waitpointId: string): WaitpointKeys { + const base = `wp:{${waitpointId}}`; + return { record: base, watchers: `${base}:w` }; +} + +export function runBlockKeys(runId: string): RunBlockKeys { + const base = `wp:run:{${runId}}`; + return { pend: `${base}:pend`, done: `${base}:done`, edge: `${base}:edge` }; +} + +export function idempotencyKey(environmentId: string, key: string): string { + return `wp:idem:{${environmentId}}:${key}`; +} + +// "#" separates the id from the index. An absent index collapses onto the empty suffix, +// which is how the partial unique index on a null batchIndex behaves; index 0 keeps its +// own field, because "0" and "" are different strings. The split back to an id below is +// taken from the LAST "#", not the first, so this stays unambiguous even if a waitpoint id +// or a run id ever contains "#" itself. +const SEPARATOR = "#"; + +export function edgeField(waitpointId: string, batchIndex?: number | null): string { + return `${waitpointId}${SEPARATOR}${batchIndex ?? ""}`; +} + +export function watcherField(runId: string, batchIndex?: number | null): string { + return `${runId}${SEPARATOR}${batchIndex ?? ""}`; +} + +// The last-"#" rule here is re-implemented as a Lua pattern in runClear (scripts.ts). This +// function has no caller besides its own test, so that test is what pins the rule as a +// specification the Lua mirrors, not just documentation of this helper. +export function waitpointIdFromEdgeField(field: string): string | undefined { + const separator = field.lastIndexOf(SEPARATOR); + return separator === -1 ? undefined : field.slice(0, separator); +} + +export class WaitpointKeyTagError extends Error { + constructor(operation: string, keys: string[], offending: string) { + super( + `Waitpoint operation ${operation} would span more than one cluster slot: ` + + `key ${JSON.stringify(offending)} does not share the tag of ${JSON.stringify(keys)}` + ); + this.name = "WaitpointKeyTagError"; + } +} + +// Redis's own keyHashSlot rule: the FIRST `{`, then the FIRST `}` after it. A missing brace +// or an empty pair means no tag, and Redis hashes the whole key. A regex would instead find +// the first NON-empty pair, disagreeing with Redis on `wp:{}{a}`. +function hashTag(key: string): string | undefined { + const open = key.indexOf("{"); + if (open === -1) return undefined; + + const close = key.indexOf("}", open + 1); + if (close === -1 || close === open + 1) return undefined; + + return key.slice(open + 1, close); +} + +/** + * Throw unless every key carries the same non-empty hash tag. Called on every script + * invocation, because the keys embed ids and are only known at call time. + */ +export function assertSingleSlot(operation: string, keys: string[]): void { + let tag: string | undefined; + + for (const key of keys) { + const found = hashTag(key); + if (!found) { + throw new WaitpointKeyTagError(operation, keys, key); + } + if (tag === undefined) { + tag = found; + } else if (found !== tag) { + throw new WaitpointKeyTagError(operation, keys, key); + } + } +} diff --git a/internal-packages/run-engine/src/engine/waitpointCoordinator/scripts.ts b/internal-packages/run-engine/src/engine/waitpointCoordinator/scripts.ts new file mode 100644 index 00000000000..820b4145f86 --- /dev/null +++ b/internal-packages/run-engine/src/engine/waitpointCoordinator/scripts.ts @@ -0,0 +1,388 @@ +import type { Callback, Redis, Result } from "@internal/redis"; + +/** + * Lua for the waitpoint coordination protocol. Three rules hold throughout: + * + * 1. Every key a script touches is declared in KEYS. No script builds a key name inside + * Lua. ioredis prefixes only the KEYS array, so a key minted in Lua would be + * unprefixed while the client wrote a prefixed one — and a script with a single + * declared key gives the caller's single-slot assertion nothing to compare. + * 2. Lua never parses JSON. Each script branches only on a short status string and moves + * opaque blobs, so every encoding decision stays in TypeScript. + * 3. A missing HGET returns Lua `false`, not `nil` — measured directly against a live + * Redis: `EVAL "return {'a', false, 'c'}"` and a table holding a missing-field HGET + * result both come back as 3 elements; only `EVAL "return {'a', nil, 'c'}"` comes back + * as 1. A `false` element converts to a reply-array null and does NOT shorten anything + * after it — only a genuine Lua nil truncates. Every returned slot is still coerced + * with `or ''` regardless, not to prevent truncation, but so an absent value arrives + * as `''` rather than `null`, giving the TypeScript one shape to decode instead of + * two. + * + * STORED_COMPLETED is the value written into the record's `status` field and is + * UPPERCASE. The outcome tokens below are lowercase and are a separate vocabulary: they + * name what a script DID, not what a record IS. Sharing one constant between the two + * makes an already-completed record invisible to every script. + */ + +const STORED_COMPLETED = "COMPLETED"; + +const MISSING = "missing"; +const CREATED = "created"; +const EXISTS = "exists"; +const REGISTERED = "registered"; +const DID_COMPLETE = "completed"; +const ALREADY = "already"; +const RESERVED = "reserved"; +const CLEARED = "cleared"; +const DRAINED = "drained"; +const DISCARDED = "discarded"; + +export function registerWaitpointCommands(redis: Redis): void { + // KEYS: record. ARGV: recordJson, status ('PENDING'|'COMPLETED'), completionJson (''). + redis.defineCommand("wpCreateIfAbsent", { + numberOfKeys: 1, + lua: ` + local record = KEYS[1] + + -- EXISTS-then-HSET inside one script, rather than a field-by-field HSETNX: the + -- record and its status must appear together or not at all. + if redis.call('EXISTS', record) == 1 then + local vals = redis.call('HMGET', record, 'r', 'status', 'c') + return { '${EXISTS}', vals[1] or '', vals[2] or '', vals[3] or '' } + end + + redis.call('HSET', record, 'r', ARGV[1], 'status', ARGV[2]) + if ARGV[3] ~= '' then + redis.call('HSET', record, 'c', ARGV[3]) + end + + return { '${CREATED}' } + `, + }); + + // KEYS: record, watchers. ARGV: watcherField, watcherJson. + redis.defineCommand("wpRegisterOrReport", { + numberOfKeys: 2, + lua: ` + local record, watchers = KEYS[1], KEYS[2] + + -- A missing waitpoint is never a silent no-op: the caller throws. Defaulting to + -- "not blocked" here would resume a run whose waitpoint never completed. + if redis.call('EXISTS', record) == 0 then + return { '${MISSING}' } + end + + if redis.call('HGET', record, 'status') == '${STORED_COMPLETED}' then + return { '${DID_COMPLETE}', redis.call('HGET', record, 'c') or '' } + end + + -- The watcher lands before any flip can read the watcher hash, because this script + -- and wpComplete are both atomic on this same shard. So a register either appears + -- in the flip's watcher list, or it observes COMPLETED above. + -- + -- HSETNX: the first registration wins, mirroring the edge's ON CONFLICT DO NOTHING. + redis.call('HSETNX', watchers, ARGV[1], ARGV[2]) + return { '${REGISTERED}' } + `, + }); + + // KEYS: record, watchers. ARGV: completionJson. + redis.defineCommand("wpComplete", { + numberOfKeys: 2, + lua: ` + local record, watchers = KEYS[1], KEYS[2] + + if redis.call('EXISTS', record) == 0 then + return { '${MISSING}' } + end + + local outcome = '${DID_COMPLETE}' + if redis.call('HGET', record, 'status') == '${STORED_COMPLETED}' then + -- Double completion is not an error, and the FIRST completion wins. This is the + -- guard a conditional UPDATE ... WHERE status = 'PENDING' used to provide. + outcome = '${ALREADY}' + else + redis.call('HSET', record, 'status', '${STORED_COMPLETED}', 'c', ARGV[1]) + end + + -- Returning the watchers here is what removes the reverse fan-out query. The + -- envelope comes back too, because delivery runs on each watcher's own shard and + -- cannot read this key. + local out = { outcome, redis.call('HGET', record, 'c') or '' } + local entries = redis.call('HVALS', watchers) + for i = 1, #entries do + out[#out + 1] = entries[i] + end + + return out + `, + }); + + // KEYS: idempotency key. ARGV: waitpointId, expiresAtMs ('' for no expiry). + redis.defineCommand("wpIdemReserve", { + numberOfKeys: 1, + lua: ` + local key = KEYS[1] + + -- Guard before the SET: a non-numeric expiry must not land a reservation that can + -- never expire because PEXPIREAT then errors out after the write already happened. + if ARGV[2] ~= '' and tonumber(ARGV[2]) == nil then + return redis.error_reply('wpIdemReserve: ARGV[2] must be numeric or empty') + end + + -- SET NX returns a status reply on success and false on conflict. + if redis.call('SET', key, ARGV[1], 'NX') then + -- Expiry only when the caller has one. A reservation with no expiry is the common + -- case and must never grow one here. + if ARGV[2] ~= '' then + redis.call('PEXPIREAT', key, tonumber(ARGV[2])) + end + return { '${RESERVED}', ARGV[1] } + end + + return { '${EXISTS}', redis.call('GET', key) or '' } + `, + }); + + // KEYS: record, watchers. No ARGV. Discards a losing reservation's orphan record. + redis.defineCommand("wpDiscard", { + numberOfKeys: 2, + lua: ` + redis.call('DEL', KEYS[1], KEYS[2]) + return { '${DISCARDED}' } + `, + }); + + // KEYS: pend, done, edge. + // ARGV: n, then n groups of 5 — waitpointId, edgeField, edgeJson, reportedFlag + // ('1'|'0'), reportedJson (''). reportedFlag, not the emptiness of reportedJson, is what + // decides the branch: a waitpoint can be reported COMPLETED with no completion envelope + // (see the FINISHED-healing path), and that case must still take the reported branch — + // flag '1', reportedJson '' — or the run would block forever on something already done. + redis.defineCommand("runAbsorbBlockers", { + numberOfKeys: 3, + lua: ` + local pend, done, edge = KEYS[1], KEYS[2], KEYS[3] + local n = tonumber(ARGV[1]) + + -- Guard before any write: a wrong n must not half-apply the script. HDEL/HSETNX below + -- are irreversible mid-script, and Redis does not roll back a script that errors. + if #ARGV ~= 1 + n * 5 then + return redis.error_reply('runAbsorbBlockers: arity mismatch') + end + + -- seenDelivered makes the delivered-pair output DISTINCT BY ID: two edges for one + -- waitpoint must contribute one pair, not two. + local requestedIds = {} + local seenDelivered = {} + local out = { '0', '0' } + + for i = 0, n - 1 do + local id = ARGV[2 + i * 5] + local field = ARGV[3 + i * 5] + local edgeJson = ARGV[4 + i * 5] + local reportedFlag = ARGV[5 + i * 5] + local reported = ARGV[6 + i * 5] + + -- HSETNX is the ON CONFLICT DO NOTHING of the edge write: a retry must not + -- overwrite the first attempt's metadata. + redis.call('HSETNX', edge, field, edgeJson) + requestedIds[id] = true + + if reportedFlag == '1' then + -- Already COMPLETED when the watcher registered. It never becomes pending, even + -- when reported ('' here) carries no envelope. + redis.call('HSET', done, id, reported) + redis.call('SREM', pend, id) + if not seenDelivered[id] then + seenDelivered[id] = true + out[#out + 1] = id + out[#out + 1] = reported + end + else + -- Check the delivered set FIRST. A completion that landed between register and + -- absorb has already delivered here, and that delivery wins. + local delivered = redis.call('HGET', done, id) + if delivered then + if not seenDelivered[id] then + seenDelivered[id] = true + out[#out + 1] = id + out[#out + 1] = delivered + end + else + redis.call('SADD', pend, id) + end + end + end + + -- Computed AFTER every write in this batch, as the count of distinct requested ids + -- with no entry in done. Counting incrementally during the loop is order-dependent: + -- a later group's completion for an id already counted as pending would leave the + -- count stale, reporting a waitpoint as both pending and delivered. + local pendingOfRequested = 0 + for id in pairs(requestedIds) do + if redis.call('HEXISTS', done, id) == 0 then + pendingOfRequested = pendingOfRequested + 1 + end + end + + out[1] = tostring(pendingOfRequested) + out[2] = tostring(redis.call('SCARD', pend)) + return out + `, + }); + + // KEYS: pend, done. ARGV: waitpointId, completionJson. + redis.defineCommand("runDeliverCompletion", { + numberOfKeys: 2, + lua: ` + local pend, done = KEYS[1], KEYS[2] + + redis.call('HSET', done, ARGV[1], ARGV[2]) + redis.call('SREM', pend, ARGV[1]) + + -- The caller treats this as a wakeup trigger, not as the resume decision: the + -- resume is decided under the run lock, and this count covers store-resident + -- blockers only. + return { tostring(redis.call('SCARD', pend)) } + `, + }); + + // KEYS: pend, done, edge. + redis.defineCommand("runReadBlockState", { + numberOfKeys: 3, + lua: ` + local pend, done, edge = KEYS[1], KEYS[2], KEYS[3] + + local pendIds = redis.call('SMEMBERS', pend) + -- HKEYS, never HGETALL: the delivered set's values are completion envelopes with + -- inline outputs, and materializing those inside a single-threaded script would + -- block the shard. + local doneIds = redis.call('HKEYS', done) + local edges = redis.call('HGETALL', edge) + + local out = { tostring(#pendIds), tostring(#doneIds), tostring(#edges) } + for i = 1, #pendIds do out[#out + 1] = pendIds[i] end + for i = 1, #doneIds do out[#out + 1] = doneIds[i] end + for i = 1, #edges do out[#out + 1] = edges[i] end + return out + `, + }); + + // KEYS: pend, done, edge. ARGV: n, then n edge fields. n = 0 clears everything. + redis.defineCommand("runClear", { + numberOfKeys: 3, + lua: ` + local pend, done, edge = KEYS[1], KEYS[2], KEYS[3] + local n = tonumber(ARGV[1]) + + -- Guard before any write, same reasoning as runAbsorbBlockers. + if #ARGV ~= 1 + n then + return redis.error_reply('runClear: arity mismatch') + end + + if n == 0 then + redis.call('DEL', pend, done, edge) + return { '${CLEARED}' } + end + + for i = 1, n do + redis.call('HDEL', edge, ARGV[1 + i]) + end + + -- Reconcile rather than delete by name. The edge set is the authority: after the + -- drain, pend and done may only hold ids that some surviving edge still references. + -- + -- Two reasons this is a superset of "remove the drained ids". First, one waitpoint + -- can hold several edges at different batch indexes, so a drained field must not + -- evict a delivery another edge still needs. Second, runDeliverCompletion writes + -- done[id] unconditionally, so a crash between register and absorb can leave a + -- delivered entry with no edge at all, which no name-derived drain could reach. + local remaining = {} + local fields = redis.call('HKEYS', edge) + for i = 1, #fields do + local sep = string.find(fields[i], '#[^#]*$') + if sep then + remaining[string.sub(fields[i], 1, sep - 1)] = true + end + end + + local doneIds = redis.call('HKEYS', done) + for i = 1, #doneIds do + if not remaining[doneIds[i]] then + redis.call('HDEL', done, doneIds[i]) + end + end + + local pendIds = redis.call('SMEMBERS', pend) + for i = 1, #pendIds do + if not remaining[pendIds[i]] then + redis.call('SREM', pend, pendIds[i]) + end + end + + return { '${DRAINED}' } + `, + }); +} + +declare module "@internal/redis" { + interface RedisCommander { + wpCreateIfAbsent( + recordKey: string, + recordJson: string, + status: string, + completionJson: string, + callback?: Callback + ): Result; + wpRegisterOrReport( + recordKey: string, + watchersKey: string, + watcherField: string, + watcherJson: string, + callback?: Callback + ): Result; + wpComplete( + recordKey: string, + watchersKey: string, + completionJson: string, + callback?: Callback + ): Result; + wpIdemReserve( + key: string, + waitpointId: string, + expiresAtMs: string, + callback?: Callback + ): Result; + wpDiscard( + recordKey: string, + watchersKey: string, + callback?: Callback + ): Result; + runAbsorbBlockers( + pendKey: string, + doneKey: string, + edgeKey: string, + ...args: Array> + ): Result; + runDeliverCompletion( + pendKey: string, + doneKey: string, + waitpointId: string, + completionJson: string, + callback?: Callback + ): Result; + runReadBlockState( + pendKey: string, + doneKey: string, + edgeKey: string, + callback?: Callback + ): Result; + runClear( + pendKey: string, + doneKey: string, + edgeKey: string, + ...args: Array> + ): Result; + } +} diff --git a/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.test.ts b/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.test.ts new file mode 100644 index 00000000000..f0e9c0c297d --- /dev/null +++ b/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.test.ts @@ -0,0 +1,1837 @@ +// Redis-only suite: the coordinator holds no Prisma reference, so no Postgres container +// is needed. redisTest FLUSHALLs before every test, so ids may be reused across describes. +import { createRedisClient, type RedisOptions } from "@internal/redis"; +import { redisTest } from "@internal/testcontainers"; +import { + deriveWaitpointIdFromAnchor, + generateRunOpsId, + generateWaitpointId, +} from "@trigger.dev/core/v3/isomorphic"; +import { describe, expect } from "vitest"; +import { + edgeField, + idempotencyKey, + runBlockKeys, + watcherField, + WaitpointKeyTagError, +} from "./keys.js"; +import { registerWaitpointCommands } from "./scripts.js"; +import { + WaitpointNotFoundError, + WaitpointStoreCoordinator, + type BlockEdge, + type WaitpointCompletion, + type WaitpointRecordInput, + type WatcherEntry, +} from "./storeCoordinator.js"; + +const ENV_ID = "env_1"; +const PROJECT_ID = "proj_1"; +const NOW = "2026-08-21T12:00:00.000Z"; + +function coordinator(redisOptions: RedisOptions) { + return new WaitpointStoreCoordinator({ redisOptions }); +} + +function record(id: string, overrides: Partial = {}): WaitpointRecordInput { + return { + id, + friendlyId: `waitpoint_${id}`, + type: "MANUAL", + environmentId: ENV_ID, + projectId: PROJECT_ID, + createdAt: NOW, + updatedAt: NOW, + userProvidedIdempotencyKey: false, + tags: [], + ...overrides, + }; +} + +function completion(overrides: Partial = {}): WaitpointCompletion { + return { + completedAt: NOW, + outputType: "application/json", + outputIsError: false, + output: { inline: '{"ok":true}' }, + ...overrides, + }; +} + +describe("createIfAbsent", () => { + redisTest("creates a PENDING record and reports created", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + const result = await store.createIfAbsent({ record: record("w_a"), status: "PENDING" }); + expect(result.outcome).toBe("created"); + } finally { + await store.quit(); + } + }); + + redisTest("returns the existing record on a second call", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.createIfAbsent({ record: record("w_a"), status: "PENDING" }); + const second = await store.createIfAbsent({ + record: record("w_a", { friendlyId: "waitpoint_DIFFERENT" }), + status: "PENDING", + }); + + expect(second.outcome).toBe("exists"); + if (second.outcome !== "exists") throw new Error("unreachable"); + // The first write wins: a retry must not overwrite the stored record. + expect(second.record.friendlyId).toBe("waitpoint_w_a"); + expect(second.status).toBe("PENDING"); + expect(second.completion).toBeUndefined(); + } finally { + await store.quit(); + } + }); + + redisTest("preserves every record field through a round trip", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + const full = record("w_a", { + type: "RUN", + idempotencyKey: "key-1", + userProvidedIdempotencyKey: true, + idempotencyKeyExpiresAt: NOW, + completedAfter: NOW, + completedByTaskRunId: "run_child", + completedByBatchId: "batch_1", + tags: ["one", "two"], + }); + + await store.createIfAbsent({ record: full, status: "PENDING" }); + const read = await store.createIfAbsent({ record: record("w_a"), status: "PENDING" }); + + expect(read.outcome).toBe("exists"); + if (read.outcome !== "exists") throw new Error("unreachable"); + // Every field the frozen return shapes need must survive the blob round trip. + expect(read.record).toEqual(full); + } finally { + await store.quit(); + } + }); + + redisTest( + "can create an already-COMPLETED record with no completion envelope", + async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + // This is the shape that catches a status-casing mismatch: the record is stored + // COMPLETED, and a register must see it as completed rather than pending. + await store.createIfAbsent({ record: record("w_a"), status: "COMPLETED" }); + + const reported = await store.registerOrReport({ + waitpointId: "w_a", + runId: "run_1", + createdAt: NOW, + }); + + expect(reported.outcome).toBe("completed"); + if (reported.outcome !== "completed") throw new Error("unreachable"); + expect(reported.completion).toBeUndefined(); + } finally { + await store.quit(); + } + } + ); + + redisTest( + "can create an already-COMPLETED record with a completion", + async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.createIfAbsent({ + record: record("w_a", { type: "RUN" }), + status: "COMPLETED", + completion: completion(), + }); + + const reported = await store.registerOrReport({ + waitpointId: "w_a", + runId: "run_1", + createdAt: NOW, + }); + + expect(reported.outcome).toBe("completed"); + if (reported.outcome !== "completed") throw new Error("unreachable"); + expect(reported.completion?.output).toEqual({ inline: '{"ok":true}' }); + } finally { + await store.quit(); + } + } + ); + + redisTest( + "reads a COMPLETED record back through createIfAbsent, with an envelope", + async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.createIfAbsent({ + record: record("w_a"), + status: "COMPLETED", + completion: completion(), + }); + + const second = await store.createIfAbsent({ record: record("w_a"), status: "PENDING" }); + + expect(second.outcome).toBe("exists"); + if (second.outcome !== "exists") throw new Error("unreachable"); + expect(second.status).toBe("COMPLETED"); + expect(second.completion?.output).toEqual({ inline: '{"ok":true}' }); + } finally { + await store.quit(); + } + } + ); + + redisTest( + "reads a COMPLETED record back through createIfAbsent, with no envelope", + async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.createIfAbsent({ record: record("w_a"), status: "COMPLETED" }); + + const second = await store.createIfAbsent({ record: record("w_a"), status: "PENDING" }); + + expect(second.outcome).toBe("exists"); + if (second.outcome !== "exists") throw new Error("unreachable"); + expect(second.status).toBe("COMPLETED"); + expect(second.completion).toBeUndefined(); + } finally { + await store.quit(); + } + } + ); +}); + +describe("registerOrReport", () => { + redisTest("registers a watcher against a PENDING waitpoint", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.createIfAbsent({ record: record("w_a"), status: "PENDING" }); + const result = await store.registerOrReport({ + waitpointId: "w_a", + runId: "run_1", + createdAt: NOW, + }); + expect(result.outcome).toBe("registered"); + } finally { + await store.quit(); + } + }); + + redisTest("reports the completion inline for a COMPLETED waitpoint", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.createIfAbsent({ record: record("w_a"), status: "PENDING" }); + await store.complete({ waitpointId: "w_a", completion: completion() }); + + const result = await store.registerOrReport({ + waitpointId: "w_a", + runId: "run_1", + createdAt: NOW, + }); + + expect(result.outcome).toBe("completed"); + if (result.outcome !== "completed") throw new Error("unreachable"); + expect(result.completion?.output).toEqual({ inline: '{"ok":true}' }); + } finally { + await store.quit(); + } + }); + + redisTest("throws for a waitpoint that does not exist", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await expect( + store.registerOrReport({ waitpointId: "w_missing", runId: "run_1", createdAt: NOW }) + ).rejects.toThrow(WaitpointNotFoundError); + } finally { + await store.quit(); + } + }); + + redisTest("keeps one watcher entry per batch index", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.createIfAbsent({ record: record("w_a"), status: "PENDING" }); + await store.registerOrReport({ + waitpointId: "w_a", + runId: "run_1", + batchIndex: 0, + createdAt: NOW, + }); + await store.registerOrReport({ + waitpointId: "w_a", + runId: "run_1", + batchIndex: 2, + createdAt: NOW, + }); + + const completed = await store.complete({ waitpointId: "w_a", completion: completion() }); + expect(completed.watchers).toHaveLength(2); + expect(completed.watchers.map((w) => w.batchIndex).sort((a, b) => a! - b!)).toEqual([0, 2]); + } finally { + await store.quit(); + } + }); + + redisTest("carries spanIdToComplete through to the watcher entry", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.createIfAbsent({ record: record("w_a"), status: "PENDING" }); + await store.registerOrReport({ + waitpointId: "w_a", + runId: "run_1", + spanIdToComplete: "span_abc", + createdAt: NOW, + }); + + const completed = await store.complete({ waitpointId: "w_a", completion: completion() }); + expect(completed.watchers[0]!.spanIdToComplete).toBe("span_abc"); + expect(completed.watchers[0]!.runId).toBe("run_1"); + expect(completed.watchers[0]!.createdAt).toBe(NOW); + } finally { + await store.quit(); + } + }); + + redisTest("keeps the first registration's watcher on a re-register", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.createIfAbsent({ record: record("w_a"), status: "PENDING" }); + await store.registerOrReport({ + waitpointId: "w_a", + runId: "run_1", + spanIdToComplete: "span_first", + createdAt: NOW, + }); + // Same run, same (absent) batch index, so the watcher field collides. HSETNX must + // not let this second registration overwrite the first one's span. + await store.registerOrReport({ + waitpointId: "w_a", + runId: "run_1", + spanIdToComplete: "span_second", + createdAt: NOW, + }); + + const completed = await store.complete({ waitpointId: "w_a", completion: completion() }); + expect(completed.watchers).toHaveLength(1); + expect(completed.watchers[0]!.spanIdToComplete).toBe("span_first"); + } finally { + await store.quit(); + } + }); +}); + +describe("complete", () => { + redisTest("flips PENDING to COMPLETED and returns the watchers", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.createIfAbsent({ record: record("w_a"), status: "PENDING" }); + await store.registerOrReport({ waitpointId: "w_a", runId: "run_1", createdAt: NOW }); + await store.registerOrReport({ waitpointId: "w_a", runId: "run_2", createdAt: NOW }); + + const result = await store.complete({ waitpointId: "w_a", completion: completion() }); + + expect(result.outcome).toBe("completed"); + expect(result.watchers.map((w) => w.runId).sort()).toEqual(["run_1", "run_2"]); + } finally { + await store.quit(); + } + }); + + redisTest("is idempotent and returns the watchers again", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.createIfAbsent({ record: record("w_a"), status: "PENDING" }); + await store.registerOrReport({ waitpointId: "w_a", runId: "run_1", createdAt: NOW }); + + const first = await store.complete({ waitpointId: "w_a", completion: completion() }); + const second = await store.complete({ + waitpointId: "w_a", + completion: completion({ output: { inline: '{"second":true}' } }), + }); + + expect(first.outcome).toBe("completed"); + expect(second.outcome).toBe("already"); + // The FIRST completion wins, matching the guard on status = PENDING. + expect(second.completion?.output).toEqual({ inline: '{"ok":true}' }); + expect(second.watchers.map((w) => w.runId)).toEqual(["run_1"]); + } finally { + await store.quit(); + } + }); + + redisTest("throws for a waitpoint that does not exist", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await expect( + store.complete({ waitpointId: "w_missing", completion: completion() }) + ).rejects.toThrow(WaitpointNotFoundError); + } finally { + await store.quit(); + } + }); + + redisTest("returns an empty watcher list when nobody is blocked", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.createIfAbsent({ record: record("w_a"), status: "PENDING" }); + const result = await store.complete({ waitpointId: "w_a", completion: completion() }); + expect(result.watchers).toEqual([]); + } finally { + await store.quit(); + } + }); + + redisTest( + "keeps the watcher list intact when the completion field is absent on an already-completed record", + async ({ redisOptions }) => { + const store = coordinator(redisOptions); + const probe = createRedisClient(redisOptions); + try { + // registerOrReport never lets a watcher land once status is COMPLETED, so this + // shape is forced by hand: it pins that an absent 'c' field decodes to an + // undefined completion without disturbing the watchers that follow it in the + // reply array. + await store.createIfAbsent({ record: record("w_a"), status: "COMPLETED" }); + const watcher: WatcherEntry = { runId: "run_1", createdAt: NOW }; + await probe.hset("wp:{w_a}:w", watcherField("run_1"), JSON.stringify(watcher)); + + const result = await store.complete({ waitpointId: "w_a", completion: completion() }); + + expect(result.outcome).toBe("already"); + expect(result.completion).toBeUndefined(); + expect(result.watchers).toHaveLength(1); + expect(result.watchers[0]!.runId).toBe("run_1"); + } finally { + probe.disconnect(); + await store.quit(); + } + } + ); + + redisTest("sets no TTL on the record or the watcher key", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + const probe = createRedisClient(redisOptions); + try { + await store.createIfAbsent({ record: record("w_a"), status: "PENDING" }); + await store.registerOrReport({ waitpointId: "w_a", runId: "run_1", createdAt: NOW }); + await store.complete({ waitpointId: "w_a", completion: completion() }); + + // -1 means the key exists with no expiry. Anything >= 0 breaks the retention rule. + expect(await probe.pttl("wp:{w_a}")).toBe(-1); + expect(await probe.pttl("wp:{w_a}:w")).toBe(-1); + } finally { + probe.disconnect(); + await store.quit(); + } + }); +}); + +// A coordinator method now drives each of these scripts, but this block stays: it is the +// only place asserting the RAW reply shape, so a Lua/TypeScript framing change made on +// both sides at once would still fail here even though every class-level test passed. +describe("reply framing (direct Lua — pins the wire shape the coordinator decodes)", () => { + const envelope = JSON.stringify(completion()); + + redisTest( + "does not double-count a waitpoint reported pending then delivered in the same batch", + async ({ redisOptions }) => { + const client = createRedisClient(redisOptions); + registerWaitpointCommands(client); + try { + const keys = runBlockKeys("run_1"); + const fieldA = edgeField("w_solo", 0); + const fieldB = edgeField("w_solo", 1); + + // Group 0 arrives unreported (still pending); group 1 for the SAME waitpoint + // arrives already reported. This is the straddle that broke pendingOfRequested. + const reply = await client.runAbsorbBlockers( + keys.pend, + keys.done, + keys.edge, + "2", + "w_solo", + fieldA, + "{}", + "0", + "", + "w_solo", + fieldB, + "{}", + "1", + envelope + ); + + expect(reply).toEqual(["0", "0", "w_solo", envelope]); + expect(await client.scard(keys.pend)).toBe(0); + } finally { + client.disconnect(); + } + } + ); + + redisTest( + "produces the identical result when the same two groups arrive in reverse order", + async ({ redisOptions }) => { + const client = createRedisClient(redisOptions); + registerWaitpointCommands(client); + try { + const keys = runBlockKeys("run_1"); + const fieldA = edgeField("w_solo", 0); + const fieldB = edgeField("w_solo", 1); + + const reply = await client.runAbsorbBlockers( + keys.pend, + keys.done, + keys.edge, + "2", + "w_solo", + fieldB, + "{}", + "1", + envelope, + "w_solo", + fieldA, + "{}", + "0", + "" + ); + + expect(reply).toEqual(["0", "0", "w_solo", envelope]); + expect(await client.scard(keys.pend)).toBe(0); + } finally { + client.disconnect(); + } + } + ); + + redisTest("counts two distinct unreported ids as fully pending", async ({ redisOptions }) => { + const client = createRedisClient(redisOptions); + registerWaitpointCommands(client); + try { + const keys = runBlockKeys("run_1"); + + const reply = await client.runAbsorbBlockers( + keys.pend, + keys.done, + keys.edge, + "2", + "w_a", + edgeField("w_a", 0), + "{}", + "0", + "", + "w_b", + edgeField("w_b", 0), + "{}", + "0", + "" + ); + + expect(reply).toEqual(["2", "2"]); + expect(await client.scard(keys.pend)).toBe(2); + } finally { + client.disconnect(); + } + }); + + redisTest( + "counts one reported and one unreported id as one pending, one delivered", + async ({ redisOptions }) => { + const client = createRedisClient(redisOptions); + registerWaitpointCommands(client); + try { + const keys = runBlockKeys("run_1"); + + const reply = await client.runAbsorbBlockers( + keys.pend, + keys.done, + keys.edge, + "2", + "w_a", + edgeField("w_a", 0), + "{}", + "0", + "", + "w_b", + edgeField("w_b", 0), + "{}", + "1", + envelope + ); + + expect(reply).toEqual(["1", "1", "w_b", envelope]); + expect(await client.scard(keys.pend)).toBe(1); + } finally { + client.disconnect(); + } + } + ); + + redisTest( + "reported flag '1' with an empty envelope still delivers, not pends", + async ({ redisOptions }) => { + const client = createRedisClient(redisOptions); + registerWaitpointCommands(client); + try { + const keys = runBlockKeys("run_1"); + + // The bug this task fixed: COMPLETED-with-no-envelope must take the reported + // branch on the flag alone, not on the envelope being non-empty. + const reply = await client.runAbsorbBlockers( + keys.pend, + keys.done, + keys.edge, + "1", + "w_a", + edgeField("w_a", 0), + "{}", + "1", + "" + ); + + expect(reply).toEqual(["0", "0", "w_a", ""]); + expect(await client.scard(keys.pend)).toBe(0); + } finally { + client.disconnect(); + } + } + ); + + redisTest( + "counts the same unreported id passed twice as one pending, not two", + async ({ redisOptions }) => { + const client = createRedisClient(redisOptions); + registerWaitpointCommands(client); + try { + const keys = runBlockKeys("run_1"); + + const reply = await client.runAbsorbBlockers( + keys.pend, + keys.done, + keys.edge, + "2", + "w_a", + edgeField("w_a", 0), + "{}", + "0", + "", + "w_a", + edgeField("w_a", 1), + "{}", + "0", + "" + ); + + expect(reply).toEqual(["1", "1"]); + expect(await client.scard(keys.pend)).toBe(1); + } finally { + client.disconnect(); + } + } + ); + + redisTest( + "counts an id already in done, passed unreported, as delivered rather than pending", + async ({ redisOptions }) => { + const client = createRedisClient(redisOptions); + registerWaitpointCommands(client); + try { + const keys = runBlockKeys("run_1"); + // A completion that landed between register and absorb — the delivered set + // already has this id before the absorb call ever sees it. + await client.hset(keys.done, "w_a", envelope); + + const reply = await client.runAbsorbBlockers( + keys.pend, + keys.done, + keys.edge, + "1", + "w_a", + edgeField("w_a", 0), + "{}", + "0", + "" + ); + + expect(reply).toEqual(["0", "0", "w_a", envelope]); + expect(await client.scard(keys.pend)).toBe(0); + } finally { + client.disconnect(); + } + } + ); + + redisTest("rejects an arity mismatch before writing anything", async ({ redisOptions }) => { + const client = createRedisClient(redisOptions); + registerWaitpointCommands(client); + try { + const keys = runBlockKeys("run_1"); + const field = edgeField("w_solo", 0); + + // n says 2 groups (1 + 2 * 5 = 11 ARGV entries expected) but only one group (5 + // ARGV entries) is supplied. + await expect( + client.runAbsorbBlockers( + keys.pend, + keys.done, + keys.edge, + "2", + "w_solo", + field, + "{}", + "0", + "" + ) + ).rejects.toThrow(); + + expect(await client.exists(keys.pend)).toBe(0); + expect(await client.exists(keys.done)).toBe(0); + expect(await client.exists(keys.edge)).toBe(0); + } finally { + client.disconnect(); + } + }); + + redisTest( + "runClear rejects an arity mismatch before writing anything", + async ({ redisOptions }) => { + const client = createRedisClient(redisOptions); + registerWaitpointCommands(client); + try { + const keys = runBlockKeys("run_1"); + const field = edgeField("w_solo", 0); + await client.hset(keys.edge, field, "{}"); + await client.sadd(keys.pend, "w_solo"); + + // n says 2 fields but only one field is supplied. + await expect( + client.runClear(keys.pend, keys.done, keys.edge, "2", field) + ).rejects.toThrow(); + + expect(await client.hexists(keys.edge, field)).toBe(1); + expect(await client.sismember(keys.pend, "w_solo")).toBe(1); + } finally { + client.disconnect(); + } + } + ); + + redisTest( + "wpIdemReserve rejects a non-numeric expiry and does not create the reservation", + async ({ redisOptions }) => { + const client = createRedisClient(redisOptions); + registerWaitpointCommands(client); + try { + const key = idempotencyKey(ENV_ID, "key-1"); + + await expect(client.wpIdemReserve(key, "w_a", "not-a-number")).rejects.toThrow(); + + expect(await client.exists(key)).toBe(0); + } finally { + client.disconnect(); + } + } + ); +}); + +describe("createWithIdempotencyKey", () => { + // Real minted ids. The method rejects anything but a standalone DATETIME/MANUAL id, because + // its loser-discard is only safe for an id that was never handed out. + const idA = generateWaitpointId("MANUAL"); + const idB = generateWaitpointId("MANUAL"); + redisTest("creates the waitpoint and wins the reservation", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + const result = await store.createWithIdempotencyKey({ + record: record(idA, { idempotencyKey: "key-1", userProvidedIdempotencyKey: true }), + environmentId: ENV_ID, + idempotencyKey: "key-1", + }); + + expect(result).toEqual({ waitpointId: idA, created: true }); + } finally { + await store.quit(); + } + }); + + redisTest("returns the winner's id and deletes the loser", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + const probe = createRedisClient(redisOptions); + try { + await store.createWithIdempotencyKey({ + record: record(idA, { idempotencyKey: "key-1", userProvidedIdempotencyKey: true }), + environmentId: ENV_ID, + idempotencyKey: "key-1", + }); + + const second = await store.createWithIdempotencyKey({ + record: record(idB, { idempotencyKey: "key-1", userProvidedIdempotencyKey: true }), + environmentId: ENV_ID, + idempotencyKey: "key-1", + }); + + expect(second).toEqual({ waitpointId: idA, created: false }); + // The loser cleans up after itself: nothing ever referenced its id. + expect(await probe.exists(`wp:{${idB}}`)).toBe(0); + expect(await probe.exists(`wp:{${idA}}`)).toBe(1); + } finally { + probe.disconnect(); + await store.quit(); + } + }); + + redisTest( + "the original creator's own retry does not discard its own record", + async ({ redisOptions }) => { + const store = coordinator(redisOptions); + const probe = createRedisClient(redisOptions); + try { + const withKey = record(idA, { + idempotencyKey: "key-1", + userProvidedIdempotencyKey: true, + }); + + const first = await store.createWithIdempotencyKey({ + record: withKey, + environmentId: ENV_ID, + idempotencyKey: "key-1", + }); + expect(first).toEqual({ waitpointId: idA, created: true }); + + // The SAME caller, retrying with the SAME record id and the SAME key — not a + // different id racing for the same reservation. + const retry = await store.createWithIdempotencyKey({ + record: withKey, + environmentId: ENV_ID, + idempotencyKey: "key-1", + }); + + expect(retry).toEqual({ waitpointId: idA, created: false }); + // The record must survive: a wrongly-discarded record would delete this too. + expect(await probe.exists(`wp:{${idA}}`)).toBe(1); + + // The real proof: something usable is still there for every later caller that + // blocks on this id. + const registered = await store.registerOrReport({ + waitpointId: idA, + runId: "run_1", + createdAt: NOW, + }); + expect(registered.outcome).toBe("registered"); + } finally { + probe.disconnect(); + await store.quit(); + } + } + ); + + redisTest("sets no expiry when the record carries none", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + const probe = createRedisClient(redisOptions); + try { + await store.createWithIdempotencyKey({ + record: record(idA, { idempotencyKey: "key-1", userProvidedIdempotencyKey: true }), + environmentId: ENV_ID, + idempotencyKey: "key-1", + }); + + // The common case. An expiry appearing here would be a retention rule violation. + expect(await probe.pttl(`wp:idem:{${ENV_ID}}:key-1`)).toBe(-1); + } finally { + probe.disconnect(); + await store.quit(); + } + }); + + redisTest("sets the expiry the record carries", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + const probe = createRedisClient(redisOptions); + try { + await store.createWithIdempotencyKey({ + record: record(idA, { + idempotencyKey: "key-1", + userProvidedIdempotencyKey: true, + idempotencyKeyExpiresAt: new Date(Date.now() + 60_000).toISOString(), + }), + environmentId: ENV_ID, + idempotencyKey: "key-1", + }); + + const ttl = await probe.pttl(`wp:idem:{${ENV_ID}}:key-1`); + // Wide band, deliberately: the deadline is computed from the test process's clock + // and applied as an absolute PEXPIREAT, while PTTL is computed against the Redis + // server's own clock. A few ms of disagreement between those two clocks is normal + // and shows up as overshoot on this read, not as a bug in the reservation. The + // band still catches every failure worth catching — wrong units, no expiry + // applied, a negative TTL — without re-asserting that two independent clocks + // agree to the millisecond. + expect(ttl).toBeGreaterThan(55_000); + expect(ttl).toBeLessThanOrEqual(65_000); + } finally { + probe.disconnect(); + await store.quit(); + } + }); + + redisTest("scopes reservations by environment", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.createWithIdempotencyKey({ + record: record(idA, { idempotencyKey: "key-1" }), + environmentId: "env_1", + idempotencyKey: "key-1", + }); + + const other = await store.createWithIdempotencyKey({ + record: record(idB, { idempotencyKey: "key-1", environmentId: "env_2" }), + environmentId: "env_2", + idempotencyKey: "key-1", + }); + + expect(other).toEqual({ waitpointId: idB, created: true }); + } finally { + await store.quit(); + } + }); +}); + +redisTest( + "rejects a derived RUN id, whose loser-discard would be unsafe", + async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + // A derived id is recomputable from its anchor, so another caller can register a + // watcher on it. Discarding one could delete a record already in use. + const derived = deriveWaitpointIdFromAnchor(`run_${generateRunOpsId()}`, "RUN")!; + await expect( + store.createWithIdempotencyKey({ + record: record(derived, { type: "RUN", idempotencyKey: "key-1" }), + environmentId: ENV_ID, + idempotencyKey: "key-1", + }) + ).rejects.toThrow(/freshly minted DATETIME or MANUAL/); + } finally { + await store.quit(); + } + } +); + +describe("the single-slot guard", () => { + redisTest("rejects an invocation whose keys span two tags", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + // Reaches the same wrapper every operation goes through, so this proves the guard + // is live at the call path and not only in the pure unit test. + expect(() => + store.assertKeysForTest("wpComplete", ["wp:{w_a}", "wp:run:{run_1}:pend"]) + ).toThrow(WaitpointKeyTagError); + } finally { + await store.quit(); + } + }); +}); + +const RUN_ID = "run_1"; + +function edge(waitpointId: string, overrides: Partial = {}): BlockEdge { + return { waitpointId, createdAt: NOW, type: "MANUAL", ...overrides }; +} + +describe("absorbBlockers", () => { + redisTest("counts pending blockers and reports the store total", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + const result = await store.absorbBlockers({ + runId: RUN_ID, + edges: [edge("w_a"), edge("w_b")], + }); + + expect(result.pendingOfRequested).toBe(2); + expect(result.storePendingTotal).toBe(2); + expect(result.alreadyDelivered).toEqual([]); + } finally { + await store.quit(); + } + }); + + redisTest( + "counts a repeated waitpoint id once, matching a count over distinct rows", + async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + const result = await store.absorbBlockers({ + runId: RUN_ID, + edges: [edge("w_a", { batchIndex: 0 }), edge("w_a", { batchIndex: 2 })], + }); + + // The count this replaces was a COUNT(*) over waitpoint rows, so two edges for + // one waitpoint contributed one. Both numbers must say 1, not 2. + expect(result.pendingOfRequested).toBe(1); + expect(result.storePendingTotal).toBe(1); + + const state = await store.readBlockState(RUN_ID); + expect(state.edges).toHaveLength(2); + expect(state.edges.map((e) => e.batchIndex).sort()).toEqual([0, 2]); + } finally { + await store.quit(); + } + } + ); + + redisTest( + "does not add a reported-complete blocker to the pending set", + async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + const result = await store.absorbBlockers({ + runId: RUN_ID, + edges: [edge("w_a", { reported: { completion: completion() } }), edge("w_b")], + }); + + expect(result.pendingOfRequested).toBe(1); + expect(result.storePendingTotal).toBe(1); + expect(result.alreadyDelivered.map((d) => d.waitpointId)).toEqual(["w_a"]); + } finally { + await store.quit(); + } + } + ); + + redisTest( + "a later absorb reads back the stored envelope, not a bare flag", + async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + const envelope = completion({ output: { inline: '{"first":true}' } }); + + // Reported once, with an envelope — this write is what's under test. + await store.absorbBlockers({ + runId: RUN_ID, + edges: [edge("w_a", { reported: { completion: envelope } })], + }); + + // Same waitpoint id, arriving unreported this time: takes the "read `done` back" + // path, exposing whatever the first call actually stored under that id. + const second = await store.absorbBlockers({ runId: RUN_ID, edges: [edge("w_a")] }); + + expect(second.alreadyDelivered).toHaveLength(1); + expect(second.alreadyDelivered[0]!.completion).toEqual(envelope); + } finally { + await store.quit(); + } + } + ); + + redisTest( + "a later absorb for a no-envelope delivery reads back no completion", + async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.absorbBlockers({ + runId: RUN_ID, + edges: [edge("w_a", { reported: {} })], + }); + + const second = await store.absorbBlockers({ runId: RUN_ID, edges: [edge("w_a")] }); + + expect(second.alreadyDelivered).toHaveLength(1); + expect(second.alreadyDelivered[0]!.completion).toBeUndefined(); + } finally { + await store.quit(); + } + } + ); + + redisTest("reports a repeated already-delivered id once", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + const result = await store.absorbBlockers({ + runId: RUN_ID, + edges: [ + edge("w_a", { batchIndex: 0, reported: { completion: completion() } }), + edge("w_a", { batchIndex: 1, reported: { completion: completion() } }), + ], + }); + + expect(result.alreadyDelivered).toHaveLength(1); + } finally { + await store.quit(); + } + }); + + redisTest("lets a delivery that raced ahead of the absorb win", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.deliverCompletion({ + runId: RUN_ID, + waitpointId: "w_a", + completion: completion(), + }); + + const result = await store.absorbBlockers({ runId: RUN_ID, edges: [edge("w_a")] }); + + expect(result.pendingOfRequested).toBe(0); + expect(result.storePendingTotal).toBe(0); + expect(result.alreadyDelivered.map((d) => d.waitpointId)).toEqual(["w_a"]); + } finally { + await store.quit(); + } + }); + + redisTest("is idempotent when run twice", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + const first = await store.absorbBlockers({ runId: RUN_ID, edges: [edge("w_a")] }); + const second = await store.absorbBlockers({ runId: RUN_ID, edges: [edge("w_a")] }); + + expect(first.storePendingTotal).toBe(1); + expect(second.storePendingTotal).toBe(1); + expect((await store.readBlockState(RUN_ID)).edges).toHaveLength(1); + } finally { + await store.quit(); + } + }); + + redisTest("keeps the first edge's metadata on a retry", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.absorbBlockers({ + runId: RUN_ID, + edges: [edge("w_a", { spanIdToComplete: "span_first" })], + }); + await store.absorbBlockers({ + runId: RUN_ID, + edges: [edge("w_a", { spanIdToComplete: "span_second" })], + }); + + expect((await store.readBlockState(RUN_ID)).edges[0]!.spanIdToComplete).toBe("span_first"); + } finally { + await store.quit(); + } + }); + + redisTest("reports the run's real total for an empty edge list", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.absorbBlockers({ runId: RUN_ID, edges: [edge("w_a")] }); + + const result = await store.absorbBlockers({ runId: RUN_ID, edges: [] }); + + // pendingOfRequested is 0 because nothing was requested. storePendingTotal is the + // run's whole store-resident set, which is NOT empty. + expect(result.pendingOfRequested).toBe(0); + expect(result.storePendingTotal).toBe(1); + expect(result.alreadyDelivered).toEqual([]); + } finally { + await store.quit(); + } + }); + + redisTest( + "reports a smaller pendingOfRequested than storePendingTotal when an unrelated blocker is already pending", + async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + // w_x is a live blocker from an earlier absorb, unrelated to this call's request. + await store.absorbBlockers({ runId: RUN_ID, edges: [edge("w_x")] }); + + const result = await store.absorbBlockers({ + runId: RUN_ID, + edges: [edge("w_a", { reported: { completion: completion() } })], + }); + + // Nothing THIS call requested is pending (w_a arrived already delivered), but the + // run's whole store-resident set still holds w_x — a divergence for a different + // reason than an empty request list, so a reply[0]/reply[1] swap or a + // re-derived-in-TypeScript pendingOfRequested would both be caught here too. + expect(result.pendingOfRequested).toBe(0); + expect(result.storePendingTotal).toBe(1); + } finally { + await store.quit(); + } + } + ); + + redisTest("sets no TTL on any run key", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + const probe = createRedisClient(redisOptions); + try { + await store.absorbBlockers({ runId: RUN_ID, edges: [edge("w_a")] }); + await store.deliverCompletion({ + runId: RUN_ID, + waitpointId: "w_a", + completion: completion(), + }); + + // -1 is "exists, no expiry"; -2 is "no key". Neither is a TTL. `pend` is emptied by + // the delivery, and Redis deletes an empty set, so -2 is expected there. + for (const key of [ + `wp:run:{${RUN_ID}}:pend`, + `wp:run:{${RUN_ID}}:done`, + `wp:run:{${RUN_ID}}:edge`, + ]) { + expect(await probe.pttl(key)).toBeLessThan(0); + } + } finally { + probe.disconnect(); + await store.quit(); + } + }); +}); + +describe("deliverCompletion", () => { + redisTest("removes the blocker and returns the new store total", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.absorbBlockers({ runId: RUN_ID, edges: [edge("w_a"), edge("w_b")] }); + + expect( + ( + await store.deliverCompletion({ + runId: RUN_ID, + waitpointId: "w_a", + completion: completion(), + }) + ).storePendingTotal + ).toBe(1); + + expect( + ( + await store.deliverCompletion({ + runId: RUN_ID, + waitpointId: "w_b", + completion: completion(), + }) + ).storePendingTotal + ).toBe(0); + } finally { + await store.quit(); + } + }); + + redisTest("is idempotent", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.absorbBlockers({ runId: RUN_ID, edges: [edge("w_a")] }); + await store.deliverCompletion({ + runId: RUN_ID, + waitpointId: "w_a", + completion: completion(), + }); + const again = await store.deliverCompletion({ + runId: RUN_ID, + waitpointId: "w_a", + completion: completion(), + }); + + expect(again.storePendingTotal).toBe(0); + } finally { + await store.quit(); + } + }); +}); + +describe("readBlockState", () => { + redisTest( + "returns the pending ids, the delivered ids and the edges", + async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.absorbBlockers({ + runId: RUN_ID, + edges: [ + edge("w_a", { batchIndex: 0, completedAfter: NOW, type: "DATETIME" }), + edge("w_b"), + ], + }); + await store.deliverCompletion({ + runId: RUN_ID, + waitpointId: "w_a", + completion: completion(), + }); + + const state = await store.readBlockState(RUN_ID); + + expect(state.pendingIds).toEqual(["w_b"]); + expect(state.deliveredIds).toEqual(["w_a"]); + expect(state.edges).toHaveLength(2); + + const datetime = state.edges.find((e) => e.waitpointId === "w_a"); + // type and completedAfter must ride the edge: a frozen return type needs them, and + // they live on the waitpoint's own shard, which this read cannot touch. + expect(datetime?.type).toBe("DATETIME"); + expect(datetime?.completedAfter).toBe(NOW); + expect(datetime?.edgeId).toBe("w_a#0"); + } finally { + await store.quit(); + } + } + ); + + redisTest("returns empty collections for a run with no blockers", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + expect(await store.readBlockState("run_unknown")).toEqual({ + pendingIds: [], + deliveredIds: [], + edges: [], + }); + } finally { + await store.quit(); + } + }); +}); + +describe("clearBlockState", () => { + redisTest("drains the named edges and reconciles", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.absorbBlockers({ runId: RUN_ID, edges: [edge("w_a"), edge("w_b")] }); + await store.deliverCompletion({ + runId: RUN_ID, + waitpointId: "w_a", + completion: completion(), + }); + + expect((await store.clearBlockState({ runId: RUN_ID, edgeIds: ["w_a#"] })).outcome).toBe( + "drained" + ); + + const state = await store.readBlockState(RUN_ID); + expect(state.edges.map((e) => e.waitpointId)).toEqual(["w_b"]); + expect(state.deliveredIds).toEqual([]); + expect(state.pendingIds).toEqual(["w_b"]); + } finally { + await store.quit(); + } + }); + + redisTest( + "keeps a waitpoint's delivery while another edge for it survives", + async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.absorbBlockers({ + runId: RUN_ID, + edges: [edge("w_a", { batchIndex: 0 }), edge("w_a", { batchIndex: 1 })], + }); + await store.deliverCompletion({ + runId: RUN_ID, + waitpointId: "w_a", + completion: completion(), + }); + + await store.clearBlockState({ runId: RUN_ID, edgeIds: ["w_a#0"] }); + + // One edge remains, so the delivery must remain too — dropping it would make the + // surviving edge look undelivered. + const state = await store.readBlockState(RUN_ID); + expect(state.edges.map((e) => e.edgeId)).toEqual(["w_a#1"]); + expect(state.deliveredIds).toEqual(["w_a"]); + } finally { + await store.quit(); + } + } + ); + + redisTest("reaps a delivered entry that no edge references", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + // The register-before-absorb window: a delivery can land for a waitpoint whose edge + // was never written. A name-derived drain could never reach it. + await store.absorbBlockers({ runId: RUN_ID, edges: [edge("w_kept")] }); + await store.deliverCompletion({ + runId: RUN_ID, + waitpointId: "w_orphan", + completion: completion(), + }); + + expect((await store.readBlockState(RUN_ID)).deliveredIds).toEqual(["w_orphan"]); + + await store.clearBlockState({ runId: RUN_ID, edgeIds: ["w_nothing#"] }); + + const state = await store.readBlockState(RUN_ID); + expect(state.deliveredIds).toEqual([]); + expect(state.edges.map((e) => e.waitpointId)).toEqual(["w_kept"]); + } finally { + await store.quit(); + } + }); + + redisTest("clears everything when no edge ids are given", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.absorbBlockers({ runId: RUN_ID, edges: [edge("w_a"), edge("w_b")] }); + + expect((await store.clearBlockState({ runId: RUN_ID })).outcome).toBe("cleared"); + expect(await store.readBlockState(RUN_ID)).toEqual({ + pendingIds: [], + deliveredIds: [], + edges: [], + }); + } finally { + await store.quit(); + } + }); + + redisTest( + "is a no-op for an explicitly empty edge id list, unlike an omitted one", + async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.absorbBlockers({ runId: RUN_ID, edges: [edge("w_a"), edge("w_b")] }); + + // Omitting edgeIds reaches the Lua's n === 0 branch and clears everything (proven + // above). A caller-computed EMPTY array must not collapse onto that: it means + // "nothing to drain", not "clear the run". + expect((await store.clearBlockState({ runId: RUN_ID, edgeIds: [] })).outcome).toBe("noop"); + + const state = await store.readBlockState(RUN_ID); + expect(state.edges.map((e) => e.waitpointId).sort()).toEqual(["w_a", "w_b"]); + expect(state.pendingIds.sort()).toEqual(["w_a", "w_b"]); + } finally { + await store.quit(); + } + } + ); +}); + +describe("registerBlocks: a COMPLETED waitpoint with no envelope never blocks (regression)", () => { + redisTest( + "created COMPLETED with no envelope: registerBlocks does not block the run", + async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + // No `completion` at all — the FINISHED-healing shape from Task 4's "can create an + // already-COMPLETED record with no completion envelope" test. + await store.createIfAbsent({ record: record("w_a"), status: "COMPLETED" }); + + const result = await store.registerBlocks({ runId: RUN_ID, edges: [edge("w_a")] }); + + expect(result.pendingOfRequested).toBe(0); + expect(result.storePendingTotal).toBe(0); + expect(result.alreadyDelivered.map((d) => d.waitpointId)).toEqual(["w_a"]); + // The whole point: no fabricated envelope, and the delivery is real on the run + // shard, not just absent from pending. + expect(result.alreadyDelivered[0]!.completion).toBeUndefined(); + expect((await store.readBlockState(RUN_ID)).deliveredIds).toEqual(["w_a"]); + } finally { + await store.quit(); + } + } + ); + + redisTest( + "created COMPLETED with an envelope: behaves identically with respect to blocking", + async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.createIfAbsent({ + record: record("w_a"), + status: "COMPLETED", + completion: completion(), + }); + + const result = await store.registerBlocks({ runId: RUN_ID, edges: [edge("w_a")] }); + + expect(result.pendingOfRequested).toBe(0); + expect(result.storePendingTotal).toBe(0); + expect(result.alreadyDelivered.map((d) => d.waitpointId)).toEqual(["w_a"]); + } finally { + await store.quit(); + } + } + ); +}); + +describe("registerBlocks: the two orderings", () => { + redisTest("block first, then complete: the run blocks, then wakes", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.createIfAbsent({ record: record("w_a"), status: "PENDING" }); + + const blocked = await store.registerBlocks({ runId: RUN_ID, edges: [edge("w_a")] }); + expect(blocked.pendingOfRequested).toBe(1); + expect(blocked.storePendingTotal).toBe(1); + + const completed = await store.complete({ waitpointId: "w_a", completion: completion() }); + expect(completed.watchers.map((w) => w.runId)).toEqual([RUN_ID]); + + const delivered = await store.deliverCompletion({ + runId: RUN_ID, + waitpointId: "w_a", + completion: completed.completion!, + }); + expect(delivered.storePendingTotal).toBe(0); + } finally { + await store.quit(); + } + }); + + redisTest("complete first, then block: the run never goes pending", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.createIfAbsent({ record: record("w_a"), status: "PENDING" }); + await store.complete({ waitpointId: "w_a", completion: completion() }); + + const result = await store.registerBlocks({ runId: RUN_ID, edges: [edge("w_a")] }); + + expect(result.pendingOfRequested).toBe(0); + expect(result.storePendingTotal).toBe(0); + expect(result.alreadyDelivered.map((d) => d.waitpointId)).toEqual(["w_a"]); + + const state = await store.readBlockState(RUN_ID); + expect(state.pendingIds).toEqual([]); + expect(state.deliveredIds).toEqual(["w_a"]); + } finally { + await store.quit(); + } + }); + + redisTest("throws when a blocking waitpoint does not exist", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await expect( + store.registerBlocks({ runId: RUN_ID, edges: [edge("w_missing")] }) + ).rejects.toThrow(WaitpointNotFoundError); + } finally { + await store.quit(); + } + }); + + redisTest( + "a throw mid-loop leaves the earlier watcher registered, and that residue is safe", + async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.createIfAbsent({ record: record("w_ok"), status: "PENDING" }); + + await expect( + store.registerBlocks({ runId: RUN_ID, edges: [edge("w_ok"), edge("w_missing")] }) + ).rejects.toThrow(WaitpointNotFoundError); + + // registerBlocks throws before absorbBlockers ever runs, so the run's own shard + // is untouched. + const state = await store.readBlockState(RUN_ID); + expect(state.pendingIds).toEqual([]); + expect(state.edges).toEqual([]); + + // But w_ok's watcher WAS registered on w_ok's own shard before the throw. + const completed = await store.complete({ waitpointId: "w_ok", completion: completion() }); + expect(completed.watchers.map((w) => w.runId)).toEqual([RUN_ID]); + + // Delivering it writes a `done` entry for a run that was never blocked on it — + // inert residue, not a false resume: no edge ever named it, and clearBlockState's + // reconcile would drop it the moment this run's block state is next drained. + const delivered = await store.deliverCompletion({ + runId: RUN_ID, + waitpointId: "w_ok", + completion: completed.completion!, + }); + expect(delivered.storePendingTotal).toBe(0); + expect((await store.readBlockState(RUN_ID)).deliveredIds).toEqual(["w_ok"]); + } finally { + await store.quit(); + } + } + ); + + redisTest("is idempotent when run twice", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.createIfAbsent({ record: record("w_a"), status: "PENDING" }); + + const first = await store.registerBlocks({ runId: RUN_ID, edges: [edge("w_a")] }); + const second = await store.registerBlocks({ runId: RUN_ID, edges: [edge("w_a")] }); + + expect(first.storePendingTotal).toBe(1); + expect(second.storePendingTotal).toBe(1); + expect((await store.readBlockState(RUN_ID)).edges).toHaveLength(1); + } finally { + await store.quit(); + } + }); + + redisTest( + "mixed set: one pending and one already complete blocks the run once", + async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.createIfAbsent({ record: record("w_pending"), status: "PENDING" }); + await store.createIfAbsent({ record: record("w_done"), status: "PENDING" }); + await store.complete({ waitpointId: "w_done", completion: completion() }); + + const result = await store.registerBlocks({ + runId: RUN_ID, + edges: [edge("w_pending"), edge("w_done")], + }); + + expect(result.pendingOfRequested).toBe(1); + expect(result.storePendingTotal).toBe(1); + expect(result.alreadyDelivered.map((d) => d.waitpointId)).toEqual(["w_done"]); + } finally { + await store.quit(); + } + } + ); +}); + +describe("multi-index merge, end to end into the executor shape", () => { + redisTest( + "a run blocked on one waitpoint at two indexes resolves to two entries", + async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.createIfAbsent({ + record: record("w_child", { type: "RUN", completedByTaskRunId: "run_child" }), + status: "PENDING", + }); + + await store.registerBlocks({ + runId: RUN_ID, + edges: [ + edge("w_child", { batchIndex: 0, batchId: "batch_1", type: "RUN" }), + edge("w_child", { batchIndex: 2, batchId: "batch_1", type: "RUN" }), + ], + }); + + const completed = await store.complete({ + waitpointId: "w_child", + completion: completion({ output: null }), + }); + // The cross-shard fact this test claims to prove: two registers for the same + // waitpoint at different indexes fanned out into two distinct watcher entries. + expect( + completed.watchers.map((w) => w.batchIndex).sort((a, b) => (a ?? 0) - (b ?? 0)) + ).toEqual([0, 2]); + await store.deliverCompletion({ + runId: RUN_ID, + waitpointId: "w_child", + completion: completed.completion!, + }); + + const state = await store.readBlockState(RUN_ID); + expect(state.pendingIds).toEqual([]); + expect(state.deliveredIds).toEqual(["w_child"]); + + // Derive the cycle's ordered id list the way the read path does: keep only edges + // that carry a batch index, sort ascending, map to id. Derived inline on purpose — + // another lane owns the order rule and its resolver, and this test's job is to + // prove the COORDINATOR preserved the edge multiplicity across two shards, not to + // own that rule. + const order = state.edges + .filter((e) => e.batchIndex !== undefined && e.batchIndex !== null) + .sort((a, b) => a.batchIndex! - b.batchIndex!) + .map((e) => e.waitpointId); + + // One waitpoint, two edges, so the id repeats — that repeat is what expands into + // two entries for the executor, and losing it would silently drop a batch item. + expect(order).toEqual(["w_child", "w_child"]); + expect(state.edges.map((e) => e.edgeId).sort()).toEqual(["w_child#0", "w_child#2"]); + expect(state.edges.every((e) => e.batchId === "batch_1")).toBe(true); + } finally { + await store.quit(); + } + } + ); +}); + +describe("the resume cycle drains and can start again", () => { + redisTest("a second wait on the same waitpoint blocks nothing", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.createIfAbsent({ record: record("w_a"), status: "PENDING" }); + await store.registerBlocks({ runId: RUN_ID, edges: [edge("w_a")] }); + + const completed = await store.complete({ waitpointId: "w_a", completion: completion() }); + await store.deliverCompletion({ + runId: RUN_ID, + waitpointId: "w_a", + completion: completed.completion!, + }); + + const first = await store.readBlockState(RUN_ID); + await store.clearBlockState({ runId: RUN_ID, edgeIds: first.edges.map((e) => e.edgeId) }); + + // Cycle two. The waitpoint is COMPLETED for good, so the register reports it and the + // run is never blocked. + const second = await store.registerBlocks({ + runId: RUN_ID, + edges: [edge("w_a", { batchIndex: 5 })], + }); + + expect(second.storePendingTotal).toBe(0); + expect(second.alreadyDelivered.map((d) => d.waitpointId)).toEqual(["w_a"]); + expect((await store.readBlockState(RUN_ID)).edges.map((e) => e.edgeId)).toEqual(["w_a#5"]); + } finally { + await store.quit(); + } + }); +}); + +// Every test above is a sequence of awaits. Redis guarantees atomicity WITHIN a script, so +// those tests can only ever prove single-script invariants. These races drive real +// concurrent calls (Promise.all over N copies) against the multi-script TypeScript +// sequences, and assert an invariant that holds regardless of who wins — never a timing. +describe("genuine concurrency", () => { + const CONCURRENCY = 8; + + redisTest( + "exactly one of N concurrent completers wins, and every caller sees its completion", + async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.createIfAbsent({ record: record("w_a"), status: "PENDING" }); + await store.registerOrReport({ waitpointId: "w_a", runId: "run_1", createdAt: NOW }); + await store.registerOrReport({ waitpointId: "w_a", runId: "run_2", createdAt: NOW }); + + const results = await Promise.all( + Array.from({ length: CONCURRENCY }, (_, i) => + store.complete({ + waitpointId: "w_a", + completion: completion({ output: { inline: `{"racer":${i}}` } }), + }) + ) + ); + + const winners = results.filter((r) => r.outcome === "completed"); + const losers = results.filter((r) => r.outcome === "already"); + expect(winners).toHaveLength(1); + expect(losers).toHaveLength(CONCURRENCY - 1); + + // Every caller, winner and losers alike, reads back the SAME stored completion. + const stored = winners[0]!.completion; + for (const r of results) { + expect(r.completion).toEqual(stored); + } + + // And every caller returns the full watcher list — a race must never truncate it. + for (const r of results) { + expect(r.watchers.map((w) => w.runId).sort()).toEqual(["run_1", "run_2"]); + } + } finally { + await store.quit(); + } + } + ); + + redisTest( + "a pre-existing registration survives N concurrent attempts to re-register its field", + async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.createIfAbsent({ record: record("w_a"), status: "PENDING" }); + await store.registerOrReport({ + waitpointId: "w_a", + runId: "run_1", + spanIdToComplete: "span_first", + createdAt: NOW, + }); + + // Same run, same (absent) batch index as the registration above, so every one of + // these collides on the exact same watcher field. + await Promise.all( + Array.from({ length: CONCURRENCY }, (_, i) => + store.registerOrReport({ + waitpointId: "w_a", + runId: "run_1", + spanIdToComplete: `span_racer_${i}`, + createdAt: NOW, + }) + ) + ); + + const completed = await store.complete({ waitpointId: "w_a", completion: completion() }); + const forRun1 = completed.watchers.filter((w) => w.runId === "run_1"); + expect(forRun1).toHaveLength(1); + expect(forRun1[0]!.spanIdToComplete).toBe("span_first"); + } finally { + await store.quit(); + } + } + ); + + redisTest( + "exactly one of N concurrent idempotency-keyed creators wins, and every loser cleans up", + async ({ redisOptions }) => { + const store = coordinator(redisOptions); + const probe = createRedisClient(redisOptions); + try { + const ids = Array.from({ length: CONCURRENCY }, () => generateWaitpointId("MANUAL")); + + const results = await Promise.all( + ids.map((id) => + store.createWithIdempotencyKey({ + record: record(id, { idempotencyKey: "key-1", userProvidedIdempotencyKey: true }), + environmentId: ENV_ID, + idempotencyKey: "key-1", + }) + ) + ); + + const winners = results.filter((r) => r.created); + expect(winners).toHaveLength(1); + + const winnerId = winners[0]!.waitpointId; + for (const r of results) { + expect(r.waitpointId).toBe(winnerId); + } + expect(await probe.exists(`wp:{${winnerId}}`)).toBe(1); + + for (const id of ids) { + if (id === winnerId) continue; + expect(await probe.exists(`wp:{${id}}`)).toBe(0); + expect(await probe.exists(`wp:{${id}}:w`)).toBe(0); + } + } finally { + probe.disconnect(); + await store.quit(); + } + } + ); + + redisTest( + "registerBlocks racing complete never leaves a waitpoint double-booked or the pending count negative", + async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + for (let i = 0; i < 30; i++) { + const waitpointId = `w_race_${i}`; + const runId = `run_race_${i}`; + await store.createIfAbsent({ record: record(waitpointId), status: "PENDING" }); + + // Two edges for the SAME waitpoint: registerBlocks registers them one at a + // time, so a concurrent complete() has a real window to land between the two + // registrations — the exact straddle that makes absorbBlockers' per-group + // reported/unreported split matter, rather than racing a single all-or-nothing + // group. + const [blocked] = await Promise.all([ + store.registerBlocks({ + runId, + edges: [edge(waitpointId, { batchIndex: 0 }), edge(waitpointId, { batchIndex: 1 })], + }), + store.complete({ waitpointId, completion: completion() }), + ]); + + const state = await store.readBlockState(runId); + const delivered = state.deliveredIds.includes(waitpointId); + const pending = state.pendingIds.includes(waitpointId); + + expect(delivered && pending).toBe(false); + expect(blocked.storePendingTotal).toBeGreaterThanOrEqual(0); + expect(blocked.storePendingTotal).toBeLessThanOrEqual(1); + } + } finally { + await store.quit(); + } + } + ); +}); diff --git a/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.ts b/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.ts new file mode 100644 index 00000000000..723552c57ab --- /dev/null +++ b/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.ts @@ -0,0 +1,538 @@ +import { createRedisClient, type Redis, type RedisOptions } from "@internal/redis"; +import { Logger } from "@trigger.dev/core/logger"; +import { parseWaitpointId } from "@trigger.dev/core/v3/isomorphic"; +import { + assertSingleSlot, + edgeField, + idempotencyKey, + runBlockKeys, + waitpointKeys, + watcherField, +} from "./keys.js"; +import { registerWaitpointCommands } from "./scripts.js"; + +/** The values written into a record's `status` field. Uppercase, and never a token. */ +export type WaitpointStatus = "PENDING" | "COMPLETED"; + +/** Every script this coordinator may invoke. The wrapper below is the only entry point. */ +type ScriptName = + | "wpCreateIfAbsent" + | "wpRegisterOrReport" + | "wpComplete" + | "wpIdemReserve" + | "wpDiscard" + | "runAbsorbBlockers" + | "runDeliverCompletion" + | "runReadBlockState" + | "runClear"; + +/** + * The immutable half of a waitpoint, written once at creation. Carries every field the + * legacy-shaped return types need, including the two that gate the executor-visible + * idempotency key and the token surface. + */ +export type WaitpointRecordInput = { + id: string; + friendlyId: string; + type: "RUN" | "BATCH" | "DATETIME" | "MANUAL"; + environmentId: string; + projectId: string; + createdAt: string; + updatedAt: string; + userProvidedIdempotencyKey: boolean; + tags: string[]; + idempotencyKey?: string; + idempotencyKeyExpiresAt?: string; + completedAfter?: string; + completedByTaskRunId?: string; + completedByBatchId?: string; +}; + +/** + * A stored output: a small inline value, an already-offloaded reference, or null when the + * value is re-derivable from a business fact and is therefore never copied forward. + */ +export type WaitpointCompletionOutput = { inline: string } | { ref: string } | null; + +/** + * The completion half of a waitpoint, written at the flip. + * + * This is the coordinator's OWN type, deliberately not a projection of any frozen record + * type. The store treats a completion as an opaque blob: it writes it, returns it, and + * never inspects a field. Whoever owns the read-time resolver maps between this and the + * frozen record shape, so the two can evolve without a type dependency in either + * direction. + */ +export type WaitpointCompletion = { + /** ISO 8601. */ + completedAt: string; + outputType: string; + outputIsError: boolean; + output: WaitpointCompletionOutput; +}; + +export type WatcherEntry = { + runId: string; + batchIndex?: number; + spanIdToComplete?: string; + createdAt: string; +}; + +export type CreateIfAbsentResult = + | { outcome: "created" } + | { + outcome: "exists"; + record: WaitpointRecordInput; + status: WaitpointStatus; + completion?: WaitpointCompletion; + }; + +export type RegisterOrReportResult = + | { outcome: "registered" } + | { outcome: "completed"; completion?: WaitpointCompletion }; + +export type CompleteResult = { + outcome: "completed" | "already"; + completion?: WaitpointCompletion; + watchers: WatcherEntry[]; +}; + +/** + * One run-to-waitpoint edge. The metadata a frozen return type — an existing API response + * shape this store must keep reproducing — needs travels here. + */ +export type BlockEdge = { + waitpointId: string; + batchIndex?: number | null; + batchId?: string; + spanIdToComplete?: string; + createdAt: string; + type: WaitpointRecordInput["type"]; + completedAfter?: string; + // Set when the register step already reported this waitpoint COMPLETED. The box, not + // `completion`, carries the "reported" fact: box present + no completion means + // COMPLETED-with-no-envelope, box absent means never reported. + reported?: { completion?: WaitpointCompletion }; +}; + +export type AbsorbResult = { + /** + * How many DISTINCT requested ids were still pending. Equivalent to the count the + * previous path took over this call's ids, which was a COUNT over waitpoint rows — so + * two edges for one waitpoint contribute one. This is the number a caller should use to + * keep today's block-time gate unchanged. + */ + pendingOfRequested: number; + /** + * The run's whole pending set, counting STORE-RESIDENT blockers only. A run can also be + * blocked by a legacy waitpoint, which this number cannot see, so it is never on its own + * a decision to resume. + */ + storePendingTotal: number; + alreadyDelivered: Array<{ waitpointId: string; completion?: WaitpointCompletion }>; +}; + +// absorbBlockers strips `reported` before writing the edge blob, so a value read back +// here can never carry it — Omit says so instead of inheriting a field that is always +// undefined. +export type BlockStateEdge = Omit & { edgeId: string }; + +export type BlockState = { + pendingIds: string[]; + deliveredIds: string[]; + edges: BlockStateEdge[]; +}; + +export class WaitpointNotFoundError extends Error { + constructor(waitpointId: string) { + super(`Waitpoint ${waitpointId} is not present in the store`); + this.name = "WaitpointNotFoundError"; + } +} + +export type WaitpointStoreCoordinatorOptions = { + redisOptions: RedisOptions; + logger?: Logger; +}; + +// Lua returns '' for an absent value, never nil, because every reply slot is coerced to +// keep the array from truncating. So a nullish check would not fire and JSON.parse('') +// throws. One helper, used at every decode site. +function parseJson(raw: string | undefined): T | undefined { + return raw ? (JSON.parse(raw) as T) : undefined; +} + +export class WaitpointStoreCoordinator { + private readonly redis: Redis; + private readonly logger: Logger; + #quit?: Promise; + + constructor(options: WaitpointStoreCoordinatorOptions) { + this.logger = options.logger ?? new Logger("WaitpointStoreCoordinator", "debug"); + this.redis = createRedisClient(options.redisOptions, { + onError: (error) => + this.logger.error("WaitpointStoreCoordinator redis client error", { error }), + }); + registerWaitpointCommands(this.redis); + } + + // Idempotent and error-swallowing: every test calls this in a finally, and a double quit + // must never mask the real assertion failure. + async quit(): Promise { + if (!this.#quit) { + this.#quit = this.redis.quit().then( + () => undefined, + () => undefined + ); + } + await this.#quit; + } + + /** + * The ONLY way this class invokes a script. Routing every call through one place is what + * makes the single-slot guard un-forgettable: a method added later cannot reach a script + * without passing its keys through this assertion. + * + * Every script's signature is (...keys, ...argv) => string[], so one cast covers them + * all. The typed RedisCommander augmentation in scripts.ts documents each shape. + */ + #call(script: ScriptName, keys: string[], ...argv: string[]): Promise { + assertSingleSlot(script, keys); + const command = this.redis[script] as (...args: string[]) => Promise; + return command.call(this.redis, ...keys, ...argv); + } + + /** + * Exposed for the guard's own test. Delegates through #call rather than calling + * assertSingleSlot directly, so a mutation to the guard inside #call fails this test too + * — not only the tests that happen to exercise a real script. + * + * With cross-tag (invalid) keys, assertSingleSlot throws synchronously inside #call, + * before any promise exists, and that throw propagates straight out of this method. With + * same-tag (valid) keys, #call would go on to dispatch a real script call; this method + * never returns or awaits that promise, and swallows whatever it eventually settles to, + * so a valid-key call here can never surface as an unhandled rejection in the caller. + */ + assertKeysForTest(operation: string, keys: string[]): void { + this.#call(operation as ScriptName, keys).catch(() => undefined); + } + + async createIfAbsent(args: { + record: WaitpointRecordInput; + status: WaitpointStatus; + completion?: WaitpointCompletion; + }): Promise { + const keys = waitpointKeys(args.record.id); + + const reply = await this.#call( + "wpCreateIfAbsent", + [keys.record], + JSON.stringify(args.record), + args.status, + args.completion ? JSON.stringify(args.completion) : "" + ); + + if (reply[0] === "created") { + return { outcome: "created" }; + } + + // reply[1] is '' only if the record hash exists with no 'r' field, which should never + // happen — but ?? never fires on '', so a bare JSON.parse('') would throw an + // undiagnosable SyntaxError instead of naming the waitpoint. + const record = parseJson(reply[1]); + if (!record) { + throw new Error(`Waitpoint ${args.record.id} exists in the store with no record blob`); + } + + return { + outcome: "exists", + record, + status: reply[2] === "COMPLETED" ? "COMPLETED" : "PENDING", + completion: parseJson(reply[3]), + }; + } + + async registerOrReport(args: { + waitpointId: string; + runId: string; + batchIndex?: number | null; + spanIdToComplete?: string; + createdAt: string; + }): Promise { + const keys = waitpointKeys(args.waitpointId); + + // batchIndex is nullable at the boundary (matching the column) and undefined inside, + // because JSON.stringify drops an undefined field but keeps a null one. + const watcher: WatcherEntry = { + runId: args.runId, + batchIndex: args.batchIndex ?? undefined, + spanIdToComplete: args.spanIdToComplete, + createdAt: args.createdAt, + }; + + const reply = await this.#call( + "wpRegisterOrReport", + [keys.record, keys.watchers], + watcherField(args.runId, args.batchIndex), + JSON.stringify(watcher) + ); + + if (reply[0] === "missing") { + throw new WaitpointNotFoundError(args.waitpointId); + } + if (reply[0] === "completed") { + return { outcome: "completed", completion: parseJson(reply[1]) }; + } + + return { outcome: "registered" }; + } + + async complete(args: { + waitpointId: string; + completion: WaitpointCompletion; + }): Promise { + const keys = waitpointKeys(args.waitpointId); + + const reply = await this.#call( + "wpComplete", + [keys.record, keys.watchers], + JSON.stringify(args.completion) + ); + + if (reply[0] === "missing") { + throw new WaitpointNotFoundError(args.waitpointId); + } + + return { + outcome: reply[0] as "completed" | "already", + completion: parseJson(reply[1]), + watchers: reply.slice(2).map((entry) => JSON.parse(entry) as WatcherEntry), + }; + } + + /** + * Create a waitpoint under an idempotency key. + * + * The reservation and the record sit under different hash tags, so no script spans + * them. That makes the ORDER load-bearing: create first, then reserve. + * + * Reserve-first would mean a crash between the two steps leaves a reservation naming a + * waitpoint that does not exist. Every later request with that key loses the + * reservation, blocks on the winner's id, and throws when it registers — correctly, but + * forever, because an idempotency key commonly carries no expiry to clear it. + * + * Create-first inverts the failure: a crash leaves an orphan record that nothing ever + * referenced, because its id is random and unpublished. No caller hangs, but nothing + * currently reclaims that record either: the backstop collector the wider plan + * describes is keyed off a run's status, and this orphan has no owning run, so that + * collector never sees it. The record is harmless — inert, unreferenced, never + * returned to anyone — but it is a real leak until a later ticket adds a reaper for + * standalone idempotency-keyed orphans specifically. + */ + async createWithIdempotencyKey(args: { + record: WaitpointRecordInput; + environmentId: string; + idempotencyKey: string; + // `created` means THIS CALL won the reservation, not that the id is new. A retry by the + // original creator reports false, because the reservation it is losing to is its own. A + // caller must not gate one-time side effects on it without handling that. + }): Promise<{ waitpointId: string; created: boolean }> { + // Standalone ids only. The discard below deletes this call's own record, and that is + // only safe because a freshly minted id was never handed out, so nothing can reference + // it. A RUN or BATCH id is DERIVED from its anchor, so any caller can recompute it and + // register a watcher on it — discarding one could delete a record already in use. + const parsed = parseWaitpointId(args.record.id); + if (parsed.format !== "b32hexW" || (parsed.type !== "DATETIME" && parsed.type !== "MANUAL")) { + throw new Error( + `createWithIdempotencyKey requires a freshly minted DATETIME or MANUAL id, got ${args.record.id}` + ); + } + + await this.createIfAbsent({ record: args.record, status: "PENDING" }); + + const expiresAtMs = args.record.idempotencyKeyExpiresAt + ? String(new Date(args.record.idempotencyKeyExpiresAt).getTime()) + : ""; + + const reply = await this.#call( + "wpIdemReserve", + [idempotencyKey(args.environmentId, args.idempotencyKey)], + args.record.id, + expiresAtMs + ); + + if (reply[0] === "reserved") { + return { waitpointId: args.record.id, created: true }; + } + + const winner = reply[1]; + if (winner !== args.record.id) { + // Safe to discard: this id is random and was never handed to any caller, so no + // watcher can reference it. Both keys share the record's tag. + const keys = waitpointKeys(args.record.id); + await this.#call("wpDiscard", [keys.record, keys.watchers]); + } + + return { waitpointId: winner, created: false }; + } + + async absorbBlockers(args: { runId: string; edges: BlockEdge[] }): Promise { + const keys = runBlockKeys(args.runId); + + // No fast path for an empty list: storePendingTotal is defined as the run's WHOLE + // store-resident pending set, so it has to be read even when nothing is requested. + const argv: string[] = [String(args.edges.length)]; + for (const item of args.edges) { + const { reported, ...stored } = item; + const reportedFlag = reported !== undefined ? "1" : "0"; + const reportedJson = reported?.completion ? JSON.stringify(reported.completion) : ""; + argv.push( + item.waitpointId, + edgeField(item.waitpointId, item.batchIndex), + JSON.stringify(stored), + reportedFlag, + reportedJson + ); + } + + const reply = await this.#call("runAbsorbBlockers", [keys.pend, keys.done, keys.edge], ...argv); + + const alreadyDelivered: AbsorbResult["alreadyDelivered"] = []; + for (let i = 2; i < reply.length; i += 2) { + alreadyDelivered.push({ + waitpointId: reply[i]!, + completion: parseJson(reply[i + 1]), + }); + } + + return { + pendingOfRequested: Number(reply[0]), + storePendingTotal: Number(reply[1]), + alreadyDelivered, + }; + } + + /** + * Block a run on a set of waitpoints. + * + * Register on every waitpoint's own shard FIRST, then absorb on the run's shard. The + * order is the protocol: a completion that lands in between finds the watcher already + * registered, so it delivers onto the run's shard, and the absorb sees that delivery and + * never marks the waitpoint pending. + * + * The register keys the decision to skip the pending set on OUTCOME, never on whether a + * completion envelope came back — a waitpoint can be reported COMPLETED with none. + * + * A throw partway through (a missing waitpoint) intentionally leaves any + * already-registered watchers in place rather than unwinding them. That's safe: a later + * `complete` on one of those waitpoints still delivers correctly, and if it lands before + * this run ever retries `registerBlocks`, the stray `done` entry it writes is inert until + * a future absorb or `clearBlockState`'s reconcile reads it — never a false resume. + */ + async registerBlocks(args: { runId: string; edges: BlockEdge[] }): Promise { + const registered: BlockEdge[] = []; + + for (const item of args.edges) { + const result = await this.registerOrReport({ + waitpointId: item.waitpointId, + runId: args.runId, + batchIndex: item.batchIndex, + spanIdToComplete: item.spanIdToComplete, + createdAt: item.createdAt, + }); + + registered.push( + result.outcome === "completed" + ? { ...item, reported: { completion: result.completion } } + : item + ); + } + + return this.absorbBlockers({ runId: args.runId, edges: registered }); + } + + async deliverCompletion(args: { + runId: string; + waitpointId: string; + completion: WaitpointCompletion; + }): Promise<{ storePendingTotal: number }> { + const keys = runBlockKeys(args.runId); + + const reply = await this.#call( + "runDeliverCompletion", + [keys.pend, keys.done], + args.waitpointId, + JSON.stringify(args.completion) + ); + + return { storePendingTotal: Number(reply[0]) }; + } + + async readBlockState(runId: string): Promise { + const keys = runBlockKeys(runId); + const reply = await this.#call("runReadBlockState", [keys.pend, keys.done, keys.edge]); + + // Slots 0 and 1 are true element counts, but slot 2 is the FLAT length of the edge + // HGETALL — two entries per edge, field then value. The cursor arithmetic below relies + // on that asymmetry, so do not "normalise" it without changing the Lua too. + const pendCount = Number(reply[0]); + const doneCount = Number(reply[1]); + const edgeCount = Number(reply[2]); + + let cursor = 3; + const pendingIds = reply.slice(cursor, cursor + pendCount); + cursor += pendCount; + const deliveredIds = reply.slice(cursor, cursor + doneCount); + cursor += doneCount; + + const edges: BlockStateEdge[] = []; + for (let i = 0; i < edgeCount; i += 2) { + const edgeId = reply[cursor + i]!; + // An edge value is always a non-empty JSON.stringify, so a missing slot here means + // the cursor walked off the end of the reply. That must fail loudly, not decode a + // BlockEdge with no waitpointId — the exact off-by-one this task's arithmetic guards + // against. + const edgeJson = reply[cursor + i + 1]; + if (!edgeJson) { + throw new Error( + `readBlockState(${runId}): missing edge payload at reply index ${cursor + i + 1}` + ); + } + const stored = JSON.parse(edgeJson) as BlockEdge; + edges.push({ ...stored, edgeId }); + } + + return { pendingIds, deliveredIds, edges }; + } + + /** + * Drain one cycle's edges, or clear the run entirely when no edge ids are given. + * + * The selective form RECONCILES: any pending or delivered entry that no surviving edge + * references goes too, not only the named ones. See runClear in scripts.ts for why. + */ + async clearBlockState(args: { + runId: string; + edgeIds?: string[]; + }): Promise<{ outcome: "cleared" | "drained" | "noop" }> { + // `omitted` and `explicitly empty` must not collapse onto each other: the Lua's + // n === 0 means "clear the whole run", so an omitted edgeIds stays the terminal clear, + // but a caller that computed zero edges to drain gets a genuine no-op that never + // reaches Redis. + if (args.edgeIds && args.edgeIds.length === 0) { + return { outcome: "noop" }; + } + + const keys = runBlockKeys(args.runId); + const edgeIds = args.edgeIds ?? []; + + const reply = await this.#call( + "runClear", + [keys.pend, keys.done, keys.edge], + String(edgeIds.length), + ...edgeIds + ); + + return { outcome: reply[0] as "cleared" | "drained" }; + } +} diff --git a/internal-packages/run-engine/src/index.ts b/internal-packages/run-engine/src/index.ts index 2c54e4c20c0..2c98edf6866 100644 --- a/internal-packages/run-engine/src/index.ts +++ b/internal-packages/run-engine/src/index.ts @@ -38,3 +38,26 @@ export type { ProcessBatchItemCallback, BatchCompletionCallback, } from "./batch-queue/types.js"; + +// Waitpoint store coordinator. Exported but not yet wired: a later ticket routes +// WaitpointSystem onto it behind a per-organisation flag. +export { + WaitpointStoreCoordinator, + WaitpointNotFoundError, +} from "./engine/waitpointCoordinator/storeCoordinator.js"; +export type { + AbsorbResult, + BlockEdge, + BlockState, + BlockStateEdge, + CompleteResult, + CreateIfAbsentResult, + RegisterOrReportResult, + WaitpointCompletion, + WaitpointCompletionOutput, + WaitpointRecordInput, + WaitpointStatus, + WaitpointStoreCoordinatorOptions, + WatcherEntry, +} from "./engine/waitpointCoordinator/storeCoordinator.js"; +export { WaitpointKeyTagError } from "./engine/waitpointCoordinator/keys.js"; diff --git a/packages/core/src/v3/isomorphic/friendlyId.test.ts b/packages/core/src/v3/isomorphic/friendlyId.test.ts index 2e3ba4d83a5..b5ea7a51971 100644 --- a/packages/core/src/v3/isomorphic/friendlyId.test.ts +++ b/packages/core/src/v3/isomorphic/friendlyId.test.ts @@ -11,13 +11,20 @@ import { RUN_OPS_ID_VERSION, RUN_OPS_ID_VERSION_2, RUN_OPS_ID_VERSION_INDEX, + WAITPOINT_ID_TYPE_INDEX, + WAITPOINT_ID_VERSION, base32hexDecode, base32hexEncode, + deriveWaitpointIdFromAnchor, + generateFriendlyId, generateRunOpsId, generateRunOpsIdV2, + generateWaitpointId, parseRunId, parseRunOpsIdBody, parseRunOpsIdV2Body, + parseWaitpointId, + type WaitpointIdType, } from "./friendlyId.js"; /** Every legal gen-2 shard char: the full DNS-safe lowercase range. */ @@ -410,3 +417,158 @@ describe("parseRunId — v2 arm", () => { expect(parseRunId(`waitpoint_${generateRunOpsIdV2("a")}`).format).toBe("legacy"); }); }); + +describe("waitpoint ids: run-ops format with version char w", () => { + it("mints a 26-char body per type, with the type char at index 24 and version w at 25", () => { + const cases: Array<[WaitpointIdType, string]> = [ + ["RUN", "r"], + ["BATCH", "b"], + ["DATETIME", "d"], + ["MANUAL", "m"], + ]; + + for (const [type, typeChar] of cases) { + const body = generateWaitpointId(type); + expect(body.length).toBe(RUN_OPS_ID_LENGTH); + expect(body[WAITPOINT_ID_TYPE_INDEX]).toBe(typeChar); + expect(body[RUN_OPS_ID_VERSION_INDEX]).toBe(WAITPOINT_ID_VERSION); + } + }); + + it("round-trips every type char through parseWaitpointId", () => { + for (const type of ["RUN", "BATCH", "DATETIME", "MANUAL"] as WaitpointIdType[]) { + const parsed = parseWaitpointId(generateWaitpointId(type)); + expect(parsed.format).toBe("b32hexW"); + if (parsed.format !== "b32hexW") throw new Error("unreachable"); + expect(parsed.type).toBe(type); + } + }); + + it("classifies both the prefixed and the bare form identically", () => { + const body = generateWaitpointId("MANUAL"); + const bare = parseWaitpointId(body); + const prefixed = parseWaitpointId(`waitpoint_${body}`); + expect(bare).toEqual(prefixed); + expect(bare).toEqual({ format: "b32hexW", type: "MANUAL", timestamp: expect.any(Date) }); + }); + + it("recovers the mint timestamp from the core", () => { + vi.useFakeTimers(); + try { + vi.setSystemTime(new Date("2026-08-21T12:00:00.000Z")); + const parsed = parseWaitpointId(generateWaitpointId("DATETIME")); + if (parsed.format !== "b32hexW") throw new Error("unreachable"); + expect(parsed.timestamp.toISOString()).toBe("2026-08-21T12:00:00.000Z"); + } finally { + vi.useRealTimers(); + } + }); + + it("classifies every legacy shape as legacy", () => { + const legacy = [ + WaitpointId.generate().id, + WaitpointId.generate().friendlyId, + generateFriendlyId("waitpoint"), + "", + "waitpoint_", + "a".repeat(27), + "a".repeat(26), + ]; + + for (const id of legacy) { + expect(parseWaitpointId(id).format).toBe("legacy"); + } + }); + + it("rejects a 26-char body whose version is w but whose type char is not r/b/d/m", () => { + const body = generateWaitpointId("RUN"); + const bad = `${body.slice(0, WAITPOINT_ID_TYPE_INDEX)}x${WAITPOINT_ID_VERSION}`; + expect(parseWaitpointId(bad).format).toBe("legacy"); + }); + + it("rejects a body whose core is outside the base32hex alphabet", () => { + const body = generateWaitpointId("RUN"); + // "w" is outside [0-9a-v], so the core no longer decodes. + expect(parseWaitpointId(`w${body.slice(1)}`).format).toBe("legacy"); + }); + + it("never parses a run id as a waitpoint id, or the reverse", () => { + expect(parseWaitpointId(generateRunOpsId()).format).toBe("legacy"); + expect(parseWaitpointId(generateRunOpsIdV2("7")).format).toBe("legacy"); + expect(parseRunId(`run_${generateWaitpointId("RUN")}`).format).toBe("legacy"); + }); + + it("rejects a well-formed waitpoint body wearing a foreign prefix", () => { + const body = `${"0".repeat(24)}rw`; // valid core + RUN type char + version w + expect(parseWaitpointId(`run_${body}`).format).toBe("legacy"); + expect(parseWaitpointId(`batch_${body}`).format).toBe("legacy"); + expect(parseWaitpointId(`waitpoint_${body}`)).toEqual({ + format: "b32hexW", + type: "RUN", + timestamp: expect.any(Date), + }); + expect(parseWaitpointId(body).format).toBe("b32hexW"); + }); + + it("handles a bare body that happens to contain an underscore sanely (never throws, never misclassifies)", () => { + const body = generateWaitpointId("BATCH"); + const withUnderscore = `_${body.slice(1)}`; + expect(() => parseWaitpointId(withUnderscore)).not.toThrow(); + // "_" is outside the base32hex alphabet, so this can never be a real waitpoint id. + expect(parseWaitpointId(withUnderscore).format).toBe("legacy"); + }); +}); + +describe("deriveWaitpointIdFromAnchor", () => { + it("is deterministic: the same anchor and type always give the same id", () => { + const anchor = `run_${generateRunOpsId("us-east-1")}`; + const first = deriveWaitpointIdFromAnchor(anchor, "RUN"); + expect(first).toBeDefined(); + expect(first).toBe(deriveWaitpointIdFromAnchor(anchor, "RUN")); + }); + + it("shares the anchor's 24-char core and replaces the region and version chars", () => { + const anchorBody = generateRunOpsId("us-east-1"); + const derived = deriveWaitpointIdFromAnchor(`run_${anchorBody}`, "RUN"); + expect(derived).toBeDefined(); + expect(derived!.slice(0, WAITPOINT_ID_TYPE_INDEX)).toBe( + anchorBody.slice(0, WAITPOINT_ID_TYPE_INDEX) + ); + expect(derived![WAITPOINT_ID_TYPE_INDEX]).toBe("r"); + expect(derived![RUN_OPS_ID_VERSION_INDEX]).toBe(WAITPOINT_ID_VERSION); + }); + + it("accepts a bare anchor body as well as a prefixed one", () => { + const anchorBody = generateRunOpsId(); + expect(deriveWaitpointIdFromAnchor(anchorBody, "RUN")).toBe( + deriveWaitpointIdFromAnchor(`run_${anchorBody}`, "RUN") + ); + }); + + it("accepts a gen-2 anchor", () => { + const derived = deriveWaitpointIdFromAnchor(`run_${generateRunOpsIdV2("7")}`, "RUN"); + expect(derived).toBeDefined(); + expect(parseWaitpointId(derived!).format).toBe("b32hexW"); + }); + + it("derives a BATCH id from a run-ops format batch anchor", () => { + const derived = deriveWaitpointIdFromAnchor(`batch_${generateRunOpsId()}`, "BATCH"); + expect(derived).toBeDefined(); + const parsed = parseWaitpointId(derived!); + if (parsed.format !== "b32hexW") throw new Error("unreachable"); + expect(parsed.type).toBe("BATCH"); + }); + + it("returns undefined for a legacy anchor, so the caller falls back to a legacy mint", () => { + expect(deriveWaitpointIdFromAnchor(RunId.generate().friendlyId, "RUN")).toBeUndefined(); + expect(deriveWaitpointIdFromAnchor("run_", "RUN")).toBeUndefined(); + expect(deriveWaitpointIdFromAnchor("", "RUN")).toBeUndefined(); + }); + + it("gives a different id per type from one anchor", () => { + const anchor = `run_${generateRunOpsId()}`; + expect(deriveWaitpointIdFromAnchor(anchor, "RUN")).not.toBe( + deriveWaitpointIdFromAnchor(anchor, "BATCH") + ); + }); +}); diff --git a/packages/core/src/v3/isomorphic/friendlyId.ts b/packages/core/src/v3/isomorphic/friendlyId.ts index c468de65319..2f436b93a3e 100644 --- a/packages/core/src/v3/isomorphic/friendlyId.ts +++ b/packages/core/src/v3/isomorphic/friendlyId.ts @@ -238,6 +238,105 @@ export function parseRunId(id: string): ParsedRunId { return LEGACY_RUN_ID; } +// Waitpoint ids reuse the run-ops body layout — 24-char base32hex core, then a +// positional char, then a version char — so the body parses positionally instead of +// splitting on "_". Index 24 carries the TYPE (the slot a run uses for its region or +// shard char), which leaves room to move to a shard char under a later version. +export const WAITPOINT_ID_VERSION = "w"; +export const WAITPOINT_ID_TYPE_INDEX = RUN_OPS_ID_REGION_INDEX; + +export type WaitpointIdType = "RUN" | "BATCH" | "DATETIME" | "MANUAL"; + +// "w" sits OUTSIDE the base32hex alphabet [0-9a-v], so the version char can never be +// mistaken for a core char, and it can never collide with a numeric run generation. +const WAITPOINT_TYPE_CHARS: Readonly> = { + RUN: "r", + BATCH: "b", + DATETIME: "d", + MANUAL: "m", +}; + +const WAITPOINT_TYPES_BY_CHAR: Readonly> = { + r: "RUN", + b: "BATCH", + d: "DATETIME", + m: "MANUAL", +}; + +export type ParsedWaitpointId = + | { format: "b32hexW"; type: WaitpointIdType; timestamp: Date } + | { format: "legacy" }; + +const LEGACY_WAITPOINT_ID: ParsedWaitpointId = { format: "legacy" }; + +/** + * Mint a standalone waitpoint id body (26 chars, no prefix) for DATETIME and MANUAL: a + * fresh core, the type char, then the waitpoint version char. + */ +export function generateWaitpointId(type: WaitpointIdType): string { + return `${mintRunOpsIdCore()}${WAITPOINT_TYPE_CHARS[type]}${WAITPOINT_ID_VERSION}`; +} + +/** + * Derive the 1:1 waitpoint id body for a RUN or BATCH anchor by reusing the anchor's + * 24-char core. Pure, so create-if-absent is idempotent without a lock. Returns + * undefined when the anchor is not a run-ops id, which is the caller's signal to mint a + * legacy waitpoint instead. + * + * Only the core survives: the anchor's region or shard char and its version char are + * both replaced. So the anchor id is NOT recoverable from the waitpoint id — the reverse + * direction uses the completedBy* back-pointer. + */ +export function deriveWaitpointIdFromAnchor( + anchorId: string, + type: WaitpointIdType +): string | undefined { + const body = stripAnchorPrefix(anchorId); + if (!parseRunOpsIdBody(body) && !parseRunOpsIdV2Body(body)) { + return undefined; + } + + return `${body.slice(0, RUN_OPS_ID_CORE_LENGTH)}${WAITPOINT_TYPE_CHARS[type]}${WAITPOINT_ID_VERSION}`; +} + +/** + * Classify a waitpoint id. Accepts the prefixed (`waitpoint_`) and bare forms, but + * NOT another entity's prefix (`run_`, `batch_`, ...) — this is the discriminator a + * later ticket uses to route a possibly customer-supplied id, so a foreign prefix must + * classify legacy rather than have its body reinterpreted as a waitpoint id. Total: + * never throws. + */ +export function parseWaitpointId(id: string): ParsedWaitpointId { + const body = stripWaitpointIdPrefix(id); + if (body.length !== RUN_OPS_ID_LENGTH) return LEGACY_WAITPOINT_ID; + if (body[RUN_OPS_ID_VERSION_INDEX] !== WAITPOINT_ID_VERSION) return LEGACY_WAITPOINT_ID; + + const type = WAITPOINT_TYPES_BY_CHAR[body[WAITPOINT_ID_TYPE_INDEX] ?? ""]; + if (!type) return LEGACY_WAITPOINT_ID; + + const timestamp = parseRunOpsIdCoreTimestamp(body); + if (timestamp === undefined) return LEGACY_WAITPOINT_ID; + + return { format: "b32hexW", type, timestamp }; +} + +// Strip any `_` if present. Prefix-agnostic is correct ONLY here: the caller +// already knows anchorId names a run or batch anchor, so there is no foreign prefix to +// guard against. Do not reuse for parseWaitpointId — see stripWaitpointIdPrefix. +function stripAnchorPrefix(id: string): string { + const underscore = id.indexOf("_"); + return underscore === -1 ? id : id.slice(underscore + 1); +} + +const WAITPOINT_ID_PREFIX = "waitpoint_"; + +// Strip the `waitpoint_` prefix if present; any other prefix, or a bare body, is left +// as-is. Unlike stripAnchorPrefix, this must never strip a foreign prefix down to a body +// that then happens to pass the run-ops shape check. +function stripWaitpointIdPrefix(id: string): string { + return id.startsWith(WAITPOINT_ID_PREFIX) ? id.slice(WAITPOINT_ID_PREFIX.length) : id; +} + export function generateInternalId(): string { return cuid(); }