feat(run-engine): idempotency-keyed waitpoint creation, record before reservation

2026-08-21 20:03:43 +01:00
parent 79d14ca14a
commit 6d624d89dd
3 changed files with 171 additions and 1 deletions
@@ -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<string[]>
    ): Result<string[], Context>;
+    wpDiscard(
+      recordKey: string,
+      watchersKey: string,
+      callback?: Callback<string[]>
+    ): Result<string[], Context>;
    runAbsorbBlockers(
      pendKey: string,
      doneKey: string,
@@ -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);
@@ -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 };
+  }
 }