diff --git a/internal-packages/run-engine/src/engine/waitpointCoordinator/scripts.ts b/internal-packages/run-engine/src/engine/waitpointCoordinator/scripts.ts index 6818fd029..5608b7226 100644 --- a/internal-packages/run-engine/src/engine/waitpointCoordinator/scripts.ts +++ b/internal-packages/run-engine/src/engine/waitpointCoordinator/scripts.ts @@ -35,6 +35,7 @@ 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 (''). @@ -143,6 +144,15 @@ export function registerWaitpointCommands(redis: Redis): void { `, }); + // 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 4 — waitpointId, edgeField, edgeJson, reportedJson (''). redis.defineCommand("runAbsorbBlockers", { @@ -338,6 +348,11 @@ declare module "@internal/redis" { expiresAtMs: string, callback?: Callback ): Result; + wpDiscard( + recordKey: string, + watchersKey: string, + callback?: Callback + ): Result; runAbsorbBlockers( pendKey: string, doneKey: string, 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 2a62b491a..b6aabb690 100644 --- a/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.test.ts +++ b/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.test.ts @@ -614,6 +614,111 @@ describe("runAbsorbBlockers, runClear and wpIdemReserve (direct Lua)", () => { ); }); +describe("createWithIdempotencyKey", () => { + redisTest("creates the waitpoint and wins the reservation", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + const result = await store.createWithIdempotencyKey({ + record: record("w_a", { idempotencyKey: "key-1", userProvidedIdempotencyKey: true }), + environmentId: ENV_ID, + idempotencyKey: "key-1", + }); + + expect(result).toEqual({ waitpointId: "w_a", 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("w_first", { idempotencyKey: "key-1", userProvidedIdempotencyKey: true }), + environmentId: ENV_ID, + idempotencyKey: "key-1", + }); + + const second = await store.createWithIdempotencyKey({ + record: record("w_second", { idempotencyKey: "key-1", userProvidedIdempotencyKey: true }), + environmentId: ENV_ID, + idempotencyKey: "key-1", + }); + + expect(second).toEqual({ waitpointId: "w_first", created: false }); + // The loser cleans up after itself: nothing ever referenced its id. + expect(await probe.exists("wp:{w_second}")).toBe(0); + expect(await probe.exists("wp:{w_first}")).toBe(1); + } 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("w_a", { 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("w_a", { + 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`); + expect(ttl).toBeGreaterThan(0); + expect(ttl).toBeLessThanOrEqual(60_000); + } finally { + probe.disconnect(); + await store.quit(); + } + }); + + redisTest("scopes reservations by environment", async ({ redisOptions }) => { + const store = coordinator(redisOptions); + try { + await store.createWithIdempotencyKey({ + record: record("w_a", { idempotencyKey: "key-1" }), + environmentId: "env_1", + idempotencyKey: "key-1", + }); + + const other = await store.createWithIdempotencyKey({ + record: record("w_b", { idempotencyKey: "key-1", environmentId: "env_2" }), + environmentId: "env_2", + idempotencyKey: "key-1", + }); + + expect(other).toEqual({ waitpointId: "w_b", created: true }); + } finally { + await store.quit(); + } + }); +}); + 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 506849fd8..d2c3f0b50 100644 --- a/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.ts +++ b/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.ts @@ -1,6 +1,6 @@ import { createRedisClient, type Redis, type RedisOptions } from "@internal/redis"; import { Logger } from "@trigger.dev/core/logger"; -import { assertSingleSlot, waitpointKeys, watcherField } from "./keys.js"; +import { assertSingleSlot, idempotencyKey, waitpointKeys, watcherField } from "./keys.js"; import { registerWaitpointCommands } from "./scripts.js"; /** The values written into a record's `status` field. Uppercase, and never a token. */ @@ -12,6 +12,7 @@ type ScriptName = | "wpRegisterOrReport" | "wpComplete" | "wpIdemReserve" + | "wpDiscard" | "runAbsorbBlockers" | "runDeliverCompletion" | "runReadBlockState" @@ -254,4 +255,53 @@ export class WaitpointStoreCoordinator { 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, and the orphan + * is reaped by the store's own garbage collection rather than by an expiry, which + * pending keys never carry. + */ + async createWithIdempotencyKey(args: { + record: WaitpointRecordInput; + environmentId: string; + idempotencyKey: string; + }): Promise<{ waitpointId: string; created: boolean }> { + 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] ?? args.record.id; + 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 }; + } }