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 23dca517a..0f951ed21 100644 --- a/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.test.ts +++ b/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.test.ts @@ -14,6 +14,7 @@ import { registerWaitpointCommands } from "./scripts.js"; import { WaitpointNotFoundError, WaitpointStoreCoordinator, + type BlockEdge, type WaitpointCompletion, type WaitpointRecordInput, type WatcherEntry, @@ -740,3 +741,380 @@ describe("the single-slot guard", () => { } }); }); + +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); + + // The edges themselves stay distinct — that multiplicity is what produces the + // repeats in the cycle's ordered id list. + 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() }), 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("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() }), + edge("w_a", { batchIndex: 1, reported: 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 { + // The completion landed between register and absorb, so it is already delivered. + 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("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(); + } + }); +}); diff --git a/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.ts b/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.ts index 0ef39b522..32a00e355 100644 --- a/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.ts +++ b/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.ts @@ -1,6 +1,13 @@ import { createRedisClient, type Redis, type RedisOptions } from "@internal/redis"; import { Logger } from "@trigger.dev/core/logger"; -import { assertSingleSlot, idempotencyKey, waitpointKeys, watcherField } from "./keys.js"; +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. */ @@ -89,6 +96,44 @@ export type CompleteResult = { watchers: WatcherEntry[]; }; +/** One run-to-waitpoint edge. The metadata a frozen return type 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. */ + reported?: 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 }>; +}; + +export type BlockStateEdge = BlockEdge & { 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`); @@ -307,4 +352,106 @@ export class WaitpointStoreCoordinator { 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; + argv.push( + item.waitpointId, + edgeField(item.waitpointId, item.batchIndex), + JSON.stringify(stored), + reported ? JSON.stringify(reported) : "" + ); + } + + 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, + }; + } + + 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]!; + const stored = JSON.parse(reply[cursor + i + 1] ?? "{}") 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: after the named edges go, any pending or delivered + * entry that no surviving edge references goes too. That is wider than deleting the + * named ids, and it has to be — a delivery is written unconditionally, so the window + * between register and absorb can leave a delivered entry with no edge at all. + */ + async clearBlockState(args: { + runId: string; + edgeIds?: string[]; + }): Promise<{ outcome: "cleared" | "drained" }> { + 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" }; + } }