diff --git a/internal-packages/run-engine/src/engine/waitpointCoordinator/scripts.ts b/internal-packages/run-engine/src/engine/waitpointCoordinator/scripts.ts index 1e60f6516..964924945 100644 --- a/internal-packages/run-engine/src/engine/waitpointCoordinator/scripts.ts +++ b/internal-packages/run-engine/src/engine/waitpointCoordinator/scripts.ts @@ -72,7 +72,9 @@ export function registerWaitpointCommands(redis: Redis): void { -- 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. - redis.call('HSET', watchers, ARGV[1], ARGV[2]) + -- + -- HSETNX: the first registration wins, mirroring the edge's ON CONFLICT DO NOTHING. + redis.call('HSETNX', watchers, ARGV[1], ARGV[2]) return { '${REGISTERED}' } `, }); @@ -115,6 +117,12 @@ export function registerWaitpointCommands(redis: Redis): void { 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 @@ -137,12 +145,16 @@ export function registerWaitpointCommands(redis: Redis): void { local pend, done, edge = KEYS[1], KEYS[2], KEYS[3] local n = tonumber(ARGV[1]) - -- countedPending and seenDelivered make both outputs DISTINCT BY ID. The count this - -- replaces was a COUNT(*) over waitpoint rows, so two edges for one waitpoint - -- contributed one. Counting per edge would inflate it. - local countedPending = {} + -- 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 * 4 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 pendingOfRequested = 0 local out = { '0', '0' } for i = 0, n - 1 do @@ -154,6 +166,7 @@ export function registerWaitpointCommands(redis: Redis): void { -- 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 reported ~= '' then -- Already COMPLETED when the watcher registered. It never becomes pending. @@ -176,14 +189,21 @@ export function registerWaitpointCommands(redis: Redis): void { end else redis.call('SADD', pend, id) - if not countedPending[id] then - countedPending[id] = true - pendingOfRequested = pendingOfRequested + 1 - end 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 @@ -234,6 +254,11 @@ export function registerWaitpointCommands(redis: Redis): void { 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}' } diff --git a/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.test.ts b/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.test.ts index 9ec5e73c7..0ad4f6e0f 100644 --- a/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.test.ts +++ b/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.test.ts @@ -3,12 +3,20 @@ import { createRedisClient, type RedisOptions } from "@internal/redis"; import { redisTest } from "@internal/testcontainers"; import { describe, expect } from "vitest"; -import { WaitpointKeyTagError } from "./keys.js"; +import { + edgeField, + idempotencyKey, + runBlockKeys, + watcherField, + WaitpointKeyTagError, +} from "./keys.js"; +import { registerWaitpointCommands } from "./scripts.js"; import { WaitpointNotFoundError, WaitpointStoreCoordinator, type WaitpointCompletion, type WaitpointRecordInput, + type WatcherEntry, } from "./storeCoordinator.js"; const ENV_ID = "env_1"; @@ -218,7 +226,7 @@ describe("registerOrReport", () => { const completed = await store.complete({ waitpointId: "w_a", completion: completion() }); expect(completed.watchers).toHaveLength(2); - expect(completed.watchers.map((w) => w.batchIndex).sort()).toEqual([0, 2]); + expect(completed.watchers.map((w) => w.batchIndex).sort((a, b) => a! - b!)).toEqual([0, 2]); } finally { await store.quit(); } @@ -306,6 +314,33 @@ describe("complete", () => { } }); + 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); @@ -324,6 +359,141 @@ describe("complete", () => { }); }); +// No coordinator method calls runAbsorbBlockers/runClear/wpIdemReserve yet — a later task +// wires those in. Registered directly on a raw client so the Lua itself is exercised now. +describe("runAbsorbBlockers (direct Lua)", () => { + 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, + "{}", + "", + "w_solo", + fieldB, + "{}", + 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, + "{}", + envelope, + "w_solo", + fieldA, + "{}", + "" + ); + + expect(reply).toEqual(["0", "0", "w_solo", 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 but only one group (4 ARGV entries) is supplied. + await expect( + client.runAbsorbBlockers(keys.pend, keys.done, keys.edge, "2", "w_solo", field, "{}", "") + ).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("the single-slot guard", () => { redisTest("rejects an invocation whose keys span two tags", async ({ redisOptions }) => { const store = coordinator(redisOptions); diff --git a/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.ts b/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.ts index f98436e6d..b8d4f8178 100644 --- a/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.ts +++ b/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.ts @@ -147,9 +147,13 @@ export class WaitpointStoreCoordinator { return command.call(this.redis, ...keys, ...argv); } - /** Exposed for the guard's own test. Asserts and returns; never invokes a script. */ - assertKeysForTest(operation: string, keys: string[]): void { - assertSingleSlot(operation, keys); + /** + * 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. + */ + assertKeysForTest(operation: string, keys: string[]) { + return this.#call(operation as ScriptName, keys); } async createIfAbsent(args: { @@ -171,10 +175,18 @@ export class WaitpointStoreCoordinator { 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: JSON.parse(reply[1] ?? "{}") as WaitpointRecordInput, - status: (reply[2] ?? "PENDING") as WaitpointStatus, + record, + status: reply[2] === "COMPLETED" ? "COMPLETED" : "PENDING", completion: parseJson(reply[3]), }; }