chore: import upstream snapshot with attribution

2026-07-13 13:32:57 +08:00
commit cd420f9332
4811 changed files with 884702 additions and 0 deletions
@@ -0,0 +1,263 @@
+import { describe, expect, it, vi } from "vitest";
+
+vi.mock("~/db.server", () => ({ prisma: {}, $replica: {} }));
+
+import { claimOrAwait, publishClaim, releaseClaim } from "~/v3/mollifier/idempotencyClaim.server";
+import type { IdempotencyClaimResult, MollifierBuffer } from "@trigger.dev/redis-worker";
+
+type ClaimState = {
+  value: string | null;
+  // Scripted return sequence for claimIdempotency calls. When set,
+  // overrides the default behaviour of returning based on `value`.
+  scriptedClaims?: IdempotencyClaimResult[];
+};
+
+function makeBuffer(initial: ClaimState = { value: null }): {
+  buffer: MollifierBuffer;
+  state: ClaimState;
+} {
+  const state = { ...initial };
+  const buffer = {
+    claimIdempotency: vi.fn(async (): Promise<IdempotencyClaimResult> => {
+      if (state.scriptedClaims && state.scriptedClaims.length > 0) {
+        return state.scriptedClaims.shift()!;
+      }
+      if (state.value === null) {
+        state.value = "pending";
+        return { kind: "claimed" };
+      }
+      if (state.value === "pending") return { kind: "pending" };
+      return { kind: "resolved", runId: state.value };
+    }),
+    readClaim: vi.fn(async (): Promise<IdempotencyClaimResult | null> => {
+      if (state.value === null) return null;
+      if (state.value === "pending") return { kind: "pending" };
+      return { kind: "resolved", runId: state.value };
+    }),
+    publishClaim: vi.fn(async ({ runId }: { runId: string }) => {
+      state.value = runId;
+    }),
+    releaseClaim: vi.fn(async () => {
+      state.value = null;
+    }),
+  } as unknown as MollifierBuffer;
+  return { buffer, state };
+}
+
+const baseInput = {
+  envId: "env_a",
+  taskIdentifier: "my-task",
+  idempotencyKey: "k-1",
+};
+
+describe("claimOrAwait", () => {
+  it("returns 'claimed' for the first caller — empty key wins SETNX", async () => {
+    const { buffer } = makeBuffer({ value: null });
+    const outcome = await claimOrAwait({
+      ...baseInput,
+      buffer,
+      generateToken: () => "token-1",
+    });
+    expect(outcome).toEqual({ kind: "claimed", token: "token-1" });
+  });
+
+  it("returns 'resolved' immediately when the key already holds a runId", async () => {
+    const { buffer } = makeBuffer({ value: "run_X" });
+    const outcome = await claimOrAwait({ ...baseInput, buffer });
+    expect(outcome).toEqual({ kind: "resolved", runId: "run_X" });
+  });
+
+  it("polls a pending key, then resolves when the runId is published", async () => {
+    const { buffer, state } = makeBuffer({ value: "pending" });
+    let nowValue = 0;
+    let pollCount = 0;
+    const outcome = await claimOrAwait({
+      ...baseInput,
+      buffer,
+      now: () => nowValue,
+      sleep: async (ms) => {
+        nowValue += ms;
+        pollCount += 1;
+        if (pollCount === 3) state.value = "run_X";
+      },
+      safetyNetMs: 1000,
+      pollStepMs: 25,
+    });
+    expect(outcome).toEqual({ kind: "resolved", runId: "run_X" });
+  });
+
+  it("returns 'timed_out' when the key stays pending past safetyNetMs", async () => {
+    const { buffer } = makeBuffer({ value: "pending" });
+    let nowValue = 0;
+    const outcome = await claimOrAwait({
+      ...baseInput,
+      buffer,
+      now: () => nowValue,
+      sleep: async (ms) => {
+        nowValue += ms;
+      },
+      safetyNetMs: 50,
+      pollStepMs: 25,
+    });
+    expect(outcome).toEqual({ kind: "timed_out" });
+  });
+
+  it("retries the claim when a polled key vanishes (claimant released)", async () => {
+    const { buffer, state } = makeBuffer({ value: "pending" });
+    let nowValue = 0;
+    let pollCount = 0;
+    // Scripted retry: on the second `claimIdempotency` call we win.
+    state.scriptedClaims = [
+      { kind: "pending" }, // first call (initial)
+      { kind: "claimed" }, // second call (retry after release)
+    ];
+    const outcome = await claimOrAwait({
+      ...baseInput,
+      buffer,
+      generateToken: () => "token-retry",
+      now: () => nowValue,
+      sleep: async (ms) => {
+        nowValue += ms;
+        pollCount += 1;
+        // First poll cycle: key vanishes (release).
+        if (pollCount === 1) state.value = null;
+      },
+      safetyNetMs: 1000,
+      pollStepMs: 25,
+    });
+    expect(outcome).toEqual({ kind: "claimed", token: "token-retry" });
+  });
+
+  it("fails open with 'claimed' when buffer is null (mollifier disabled)", async () => {
+    const outcome = await claimOrAwait({
+      ...baseInput,
+      buffer: null,
+      generateToken: () => "token-fallopen-null",
+    });
+    expect(outcome).toEqual({ kind: "claimed", token: "token-fallopen-null" });
+  });
+
+  it("fails open with 'claimed' if buffer.claimIdempotency throws (Redis down)", async () => {
+    const buffer = {
+      claimIdempotency: vi.fn(async () => {
+        throw new Error("ECONNREFUSED");
+      }),
+    } as unknown as MollifierBuffer;
+    const outcome = await claimOrAwait({
+      ...baseInput,
+      buffer,
+      generateToken: () => "token-fallopen-throw",
+    });
+    expect(outcome).toEqual({ kind: "claimed", token: "token-fallopen-throw" });
+  });
+
+  it("respects an aborted signal during the wait loop", async () => {
+    const { buffer } = makeBuffer({ value: "pending" });
+    const controller = new AbortController();
+    let nowValue = 0;
+    let pollCount = 0;
+    const outcome = await claimOrAwait({
+      ...baseInput,
+      buffer,
+      now: () => nowValue,
+      sleep: async (ms) => {
+        nowValue += ms;
+        pollCount += 1;
+        if (pollCount === 1) controller.abort();
+      },
+      abortSignal: controller.signal,
+      safetyNetMs: 5000,
+      pollStepMs: 25,
+    });
+    expect(outcome).toEqual({ kind: "timed_out" });
+  });
+});
+
+describe("publishClaim", () => {
+  it("writes the runId to the claim key", async () => {
+    const { buffer, state } = makeBuffer({ value: "pending" });
+    await publishClaim({ ...baseInput, token: "owner-token", runId: "run_X", buffer });
+    expect(state.value).toBe("run_X");
+    expect(buffer.publishClaim).toHaveBeenCalledOnce();
+  });
+
+  it("no-op when buffer is null", async () => {
+    await expect(
+      publishClaim({ ...baseInput, token: "owner-token", runId: "run_X", buffer: null })
+    ).resolves.toBeUndefined();
+  });
+
+  it("swallows errors so trigger pipeline isn't broken by Redis hiccups", async () => {
+    const buffer = {
+      publishClaim: vi.fn(async () => {
+        throw new Error("ECONNREFUSED");
+      }),
+    } as unknown as MollifierBuffer;
+    await expect(
+      publishClaim({ ...baseInput, token: "owner-token", runId: "run_X", buffer })
+    ).resolves.toBeUndefined();
+  });
+});
+
+describe("releaseClaim", () => {
+  it("DELs the claim so waiters can re-acquire", async () => {
+    const { buffer, state } = makeBuffer({ value: "pending" });
+    await releaseClaim({ ...baseInput, token: "owner-token", buffer });
+    expect(state.value).toBeNull();
+  });
+
+  it("no-op when buffer is null", async () => {
+    await expect(
+      releaseClaim({ ...baseInput, token: "owner-token", buffer: null })
+    ).resolves.toBeUndefined();
+  });
+});
+
+// End-to-end: the token from `claimOrAwait`'s `claimed` outcome must
+// reach `buffer.claimIdempotency` and round-trip through publishClaim /
+// releaseClaim. Without this the compare-and-act ownership protection
+// in the buffer is bypassed and the stale-claimant hazard returns.
+describe("claim ownership token wiring", () => {
+  it("threads the token from claimOrAwait into buffer.claimIdempotency", async () => {
+    const { buffer } = makeBuffer({ value: null });
+    const outcome = await claimOrAwait({
+      ...baseInput,
+      buffer,
+      generateToken: () => "owner-token-xyz",
+    });
+    expect(outcome).toEqual({ kind: "claimed", token: "owner-token-xyz" });
+    expect(buffer.claimIdempotency).toHaveBeenCalledWith({
+      ...baseInput,
+      token: "owner-token-xyz",
+      ttlSeconds: 30,
+    });
+  });
+
+  it("threads the token from publishClaim into buffer.publishClaim", async () => {
+    const { buffer } = makeBuffer({ value: "pending" });
+    await publishClaim({
+      ...baseInput,
+      token: "owner-token-xyz",
+      runId: "run_X",
+      buffer,
+    });
+    expect(buffer.publishClaim).toHaveBeenCalledWith(
+      expect.objectContaining({
+        token: "owner-token-xyz",
+        runId: "run_X",
+      })
+    );
+  });
+
+  it("threads the token from releaseClaim into buffer.releaseClaim", async () => {
+    const { buffer } = makeBuffer({ value: "pending" });
+    await releaseClaim({
+      ...baseInput,
+      token: "owner-token-xyz",
+      buffer,
+    });
+    expect(buffer.releaseClaim).toHaveBeenCalledWith(
+      expect.objectContaining({ token: "owner-token-xyz" })
+    );
+  });
+});