feat(run-engine): run shard operations for absorb, deliver, read and clear

2026-08-21 20:23:10 +01:00
parent 6aea717ade
commit 77f0e4dd2b
2 changed files with 526 additions and 1 deletions
@@ -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> = {}): 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();
+    }
+  });
+});
@@ -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<AbsorbResult> {
+    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<WaitpointCompletion>(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<BlockState> {
+    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" };
+  }
 }