perf(run-engine): batch the resolver's deferred run reads, index its positions

Two costs in the completed-waitpoint resolver, both scaling with fan-in width,
which is the input the record set exists to make cheap.

The output hydration read one run per record, awaited in series inside the emit
loop. A batch parent resuming on 500 children therefore performed 500
sequential Postgres reads, where the hydration it replaces does one chunked
findMany. The reader is now plural: the resolver collects the distinct set of
runs its records defer to, reads them in one chunked batch before the loop, and
hydrates from the resulting map. A cycle that defers nothing reads nothing.

Positions came from a linear scan of the order per record, making the emit loop
quadratic -- about a million string comparisons for a thousand-wide wait. The
order is now indexed once into a map of id to positions. Absence still carries
the index-less case, so a wait with no batch index emits one entry with an
undefined index exactly as before.

Neither path has a production caller yet, so nothing changes for any
organisation today.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Author: d-cs

internal-packages/run-engine/src/engine/waitpointCoordinator/completedWaitpointEquivalence.test.ts

@@ -10,7 +10,7 @@ import { enhanceExecutionSnapshotWithWaitpoints } from "../systems/executionSnap
 import { buildCompletedWaitpointRecords } from "./completedWaitpointRecords.js";
 import {
   createCompletedWaitpointResolver,
-  createRunOutputReader,
+  createRunOutputsReader,
 } from "./completedWaitpointResolver.js";
 import { envelopeSourceFromWaitpointRow } from "./completionEnvelopeSource.js";
 import type { CompletionEnvelopeSource } from "./types.js";
@@ -89,7 +89,7 @@ async function bothPaths(
   const runStore = new PostgresRunStore({ prisma, readOnlyPrisma: prisma });
 
   const actual = await createCompletedWaitpointResolver({
-    readRunOutput: createRunOutputReader(runStore),
+    readRunOutputs: createRunOutputsReader(runStore),
   })({
     runId: RUN_ID,
     ...(batchId ? { batchId } : {}),

internal-packages/run-engine/src/engine/waitpointCoordinator/completedWaitpointResolver.runOutput.test.ts

@@ -10,10 +10,10 @@ import type { PrismaClient } from "@trigger.dev/database";
 import { describe, expect } from "vitest";
 import {
   createCompletedWaitpointResolver,
-  createRunOutputReader,
+  createRunOutputsReader,
   UnresolvableWaitpointId,
 } from "./completedWaitpointResolver.js";
-import { seedChildRunWithOutput } from "./testFixtures/childRun.js";
+import { seedChildRunsWithOutputs, seedChildRunWithOutput } from "./testFixtures/childRun.js";
 
 function deriveRecord(completedByTaskRunId: string): CompletedWaitpointRecord {
   return {
@@ -30,7 +30,30 @@ function deriveRecord(completedByTaskRunId: string): CompletedWaitpointRecord {
 
 function resolverFor(prisma: PrismaClient) {
   const runStore = new PostgresRunStore({ prisma, readOnlyPrisma: prisma });
-  return createCompletedWaitpointResolver({ readRunOutput: createRunOutputReader(runStore) });
+  return createCompletedWaitpointResolver({ readRunOutputs: createRunOutputsReader(runStore) });
+}
+
+/**
+ * A resolver that records the id set of every batched read and DELEGATES to the real reader, so
+ * the Postgres read still happens.
+ *
+ * Wrapping the collaborator rather than replacing it is deliberate: the assertion is about how
+ * many reads occur and what they ask for, and neither is observable from the resolved output.
+ */
+function countingResolver(prisma: PrismaClient) {
+  const runStore = new PostgresRunStore({ prisma, readOnlyPrisma: prisma });
+  const read = createRunOutputsReader(runStore);
+  const batches: string[][] = [];
+
+  return {
+    batches,
+    resolve: createCompletedWaitpointResolver({
+      readRunOutputs: async (ids) => {
+        batches.push(ids);
+        return read(ids);
+      },
+    }),
+  };
 }
 
 describe("the deriveFromRun branch", () => {
@@ -102,19 +125,9 @@ describe("the deriveFromRun branch", () => {
   // query per index.
   postgresTest("reads the run once for a record at several indexes", async ({ prisma }) => {
     const runId = await seedChildRunWithOutput(prisma, '{"value":42}');
-    const runStore = new PostgresRunStore({ prisma, readOnlyPrisma: prisma });
-    const reads: string[] = [];
-    const reader = createRunOutputReader(runStore);
-
-    // Counts calls and DELEGATES to the real reader, so the Postgres read still happens. This
-    // wraps the collaborator rather than replacing it: the assertion is about how many reads
-    // occur, which is not observable from the resolved output alone.
-    const result = await createCompletedWaitpointResolver({
-      readRunOutput: async (id) => {
-        reads.push(id);
-        return reader(id);
-      },
-    })({
+    const { resolve, batches } = countingResolver(prisma);
+
+    const result = await resolve({
       runId: "run_parent",
       pointer: { cycleSeq: 1, count: 2 },
       order: ["wp_run", "wp_run"],
@@ -123,7 +136,67 @@ describe("the deriveFromRun branch", () => {
     });
 
     expect(result).toHaveLength(2);
-    expect(reads).toEqual([runId]);
+    expect(batches).toEqual([[runId]]);
+  });
+
+  // The shape a batch fan-in produces. Every deferring record resolves in ONE read, not one
+  // read each: a per-record read put a serial round trip per child on the resume path, which is
+  // the cost the record set exists to remove.
+  postgresTest("reads every deferred run in one batch", async ({ prisma }) => {
+    const runIds = await seedChildRunsWithOutputs(
+      prisma,
+      Array.from({ length: 12 }, (_, i) => `{"value":${i}}`)
+    );
+    const { resolve, batches } = countingResolver(prisma);
+
+    const records = runIds.map((runId, i) => ({
+      ...deriveRecord(runId),
+      id: `wp_run_${i}`,
+      friendlyId: `waitpoint_wp_run_${i}`,
+    }));
+
+    const result = await resolve({
+      runId: "run_parent",
+      pointer: { cycleSeq: 1, count: records.length },
+      order: records.map((r) => r.id),
+      distinctIds: records.map((r) => r.id),
+      records,
+    });
+
+    expect(result).toHaveLength(12);
+    // One batch, holding every distinct run id.
+    expect(batches).toHaveLength(1);
+    expect(batches[0]?.slice().sort()).toEqual(runIds.slice().sort());
+    // And each output landed on its own waitpoint.
+    for (const [i, runId] of runIds.entries()) {
+      const entry = result.find((w) => w.id === `wp_run_${i}`);
+      expect(entry?.output).toBe(`{"value":${i}}`);
+      expect(runId).toBeTruthy();
+    }
+  });
+
+  // A cycle that defers nothing reads nothing, so an all-inline resume pays no Postgres round
+  // trip at all.
+  postgresTest("reads nothing when no record defers", async ({ prisma }) => {
+    const { resolve, batches } = countingResolver(prisma);
+
+    const result = await resolve({
+      runId: "run_parent",
+      pointer: { cycleSeq: 1, count: 0 },
+      order: [],
+      distinctIds: ["wp_inline"],
+      records: [
+        {
+          ...deriveRecord("run_unused"),
+          id: "wp_inline",
+          friendlyId: "waitpoint_wp_inline",
+          output: { inline: '{"ok":true}' },
+        },
+      ],
+    });
+
+    expect(result[0]?.output).toBe('{"ok":true}');
+    expect(batches).toEqual([]);
   });
 
   postgresTest("throws when a derive record arrives with no reader wired", async ({ prisma }) => {

internal-packages/run-engine/src/engine/waitpointCoordinator/completedWaitpointResolver.ts

@@ -39,26 +39,54 @@ export class UnresolvableWaitpointId extends Error {
 
 export type CompletedWaitpointResolverDeps = {
   /**
-   * Reads TaskRun.output. Returns undefined when the row is gone.
+   * Reads TaskRun.output for a SET of completing runs, keyed by run id.
+   *
+   * Plural on purpose. A batch parent resumes on every child at once, so a per-id reader made
+   * the resolver do one round trip per child -- 500 of them, in series, for a 500-wide fan-in,
+   * where the path this replaces did one chunked read. An id absent from the returned map is an
+   * absent output, which the caller refuses rather than resolving empty.
    *
    * Optional, because most cycles carry no `deriveFromRun` record and therefore never need it.
    * A cycle that DOES carry one without a reader is a wiring error, not a data condition, so it
    * throws rather than resolving empty.
    */
-  readRunOutput?(taskRunId: string): Promise<string | undefined>;
+  readRunOutputs?(taskRunIds: string[]): Promise<Map<string, string>>;
 };
 
+// Bounds one read, for the reason the envelope and waitpoint reads share: a run output can be
+// 100KB+, so a wide fan-in read whole can exceed Node's string conversion limits.
+const RUN_OUTPUT_CHUNK_SIZE = 100;
+
 /**
- * The production reader: TaskRun.output for the completing run, through the store so the read
+ * The production reader: TaskRun.output for the completing runs, through the store so each read
  * routes to the run's owning database.
+ *
+ * `findRunsByIds` is the store's own grouped replacement for `Promise.all(ids.map(findRun))`,
+ * and it forces `id` into the projection so the map keys correctly even though this select
+ * names only `output`.
  */
-export function createRunOutputReader(
-  runStore: Pick<RunStore, "findRun">,
+export function createRunOutputsReader(
+  runStore: Pick<RunStore, "findRunsByIds">,
   client?: ReadClient
-): (taskRunId: string) => Promise<string | undefined> {
-  return async (taskRunId) => {
-    const run = await runStore.findRun({ id: taskRunId }, { select: { output: true } }, client);
-    return run?.output ?? undefined;
+): (taskRunIds: string[]) => Promise<Map<string, string>> {
+  return async (taskRunIds) => {
+    const outputs = new Map<string, string>();
+
+    for (let i = 0; i < taskRunIds.length; i += RUN_OUTPUT_CHUNK_SIZE) {
+      const chunk = taskRunIds.slice(i, i + RUN_OUTPUT_CHUNK_SIZE);
+      const rows = await runStore.findRunsByIds(chunk, { select: { output: true } }, client);
+
+      for (const [id, row] of rows) {
+        // A row present with a null output is the same absence as a missing row: either way the
+        // value the waitpoint deferred is gone. Omitting it here keeps one absence rule, so the
+        // caller's refusal covers both.
+        if (row.output !== null) {
+          outputs.set(id, row.output);
+        }
+      }
+    }
+
+    return outputs;
   };
 }
 
@@ -100,13 +128,23 @@ export function createCompletedWaitpointResolver(deps: CompletedWaitpointResolve
       }
     }
 
+    // Every deferred output in ONE read, before the emit loop. Reading inside the loop meant a
+    // round trip per record, in series, which is the shape a batch fan-in punishes hardest: the
+    // wide wait this feature exists to make cheap is exactly the wide wait that paid most.
+    const runOutputs = await readDeferredOutputs(args.records, deps);
+
+    // Positions once, not once per record. `positionsOf` scanned the whole order for every
+    // record, so the emit loop was O(records x order) -- a million comparisons for a 1000-wide
+    // wait, growing with the same input as above.
+    const positions = positionsById(args.order);
+
     const out: CompletedWaitpoint[] = [];
 
     for (const record of args.records) {
-      const indexes = positionsOf(record.id, args.order);
+      const indexes = positions.get(record.id) ?? [undefined];
       // Hydrated once per record, not once per position, so a run at several batch indexes
-      // costs one read rather than one per index.
-      const output = await hydrateOutput(record, deps);
+      // resolves from one map entry rather than one per index.
+      const output = hydrateOutput(record, runOutputs);
 
       for (const index of indexes) {
         out.push({
@@ -144,24 +182,85 @@ export function createCompletedWaitpointResolver(deps: CompletedWaitpointResolve
   };
 }
 
-// An id with no position yields one entry with an undefined index, matching what the
-// existing hydration does for a wait that carried no batch index.
-function positionsOf(waitpointId: string, order: string[]): (number | undefined)[] {
-  const indexes: (number | undefined)[] = [];
+// Every id's positions in the order, built in one pass.
+//
+// An id ABSENT from this map has no position, and the caller emits it once with an undefined
+// index -- matching what the existing hydration does for a wait that carried no batch index.
+// Absence is how that case is carried, so this never stores an [undefined] entry itself.
+function positionsById(order: string[]): Map<string, number[]> {
+  const positions = new Map<string, number[]>();
 
   for (let i = 0; i < order.length; i++) {
-    if (order[i] === waitpointId) {
-      indexes.push(i);
+    const id = order[i];
+    if (id === undefined) {
+      continue;
+    }
+
+    const existing = positions.get(id);
+    if (existing) {
+      existing.push(i);
+    } else {
+      positions.set(id, [i]);
     }
   }
 
-  return indexes.length === 0 ? [undefined] : indexes;
+  return positions;
 }
 
-async function hydrateOutput(
-  record: CompletedWaitpointRecord,
+/**
+ * The output of every run a record defers to, in one batched read.
+ *
+ * Returns an empty map when no record defers, which is the common case: a cycle carrying only
+ * inline values, refs and BATCH records reads nothing at all.
+ */
+async function readDeferredOutputs(
+  records: CompletedWaitpointRecord[],
   deps: CompletedWaitpointResolverDeps
-): Promise<string | undefined> {
+): Promise<Map<string, string>> {
+  const runIds = new Set<string>();
+  let deferring: CompletedWaitpointRecord | undefined;
+
+  for (const record of records) {
+    const runId = deferredRunIdOf(record);
+    if (runId !== undefined) {
+      runIds.add(runId);
+      deferring ??= record;
+    }
+  }
+
+  if (runIds.size === 0) {
+    return new Map();
+  }
+
+  if (!deps.readRunOutputs) {
+    throw new Error(
+      `Waitpoint ${deferring?.id} defers its output to run ${deferring?.completedByTaskRunId}, but the resolver was built with no run-output reader.`
+    );
+  }
+
+  return deps.readRunOutputs([...runIds]);
+}
+
+// The run a record defers its output to, or undefined when it carries its own output or has
+// nothing to defer to. This is the single definition of "needs a run read", so the pre-pass and
+// the hydration cannot disagree about which records those are.
+function deferredRunIdOf(record: CompletedWaitpointRecord): string | undefined {
+  if (record.output === null) {
+    return undefined;
+  }
+
+  if ("inline" in record.output || "ref" in record.output) {
+    return undefined;
+  }
+
+  return record.completedByTaskRunId ?? undefined;
+}
+
+// Synchronous: every read this needs already happened in readDeferredOutputs.
+function hydrateOutput(
+  record: CompletedWaitpointRecord,
+  runOutputs: Map<string, string>
+): string | undefined {
   if (record.output === null) {
     return undefined;
   }
@@ -176,20 +275,15 @@ async function hydrateOutput(
     return record.output.ref;
   }
 
-  if (!record.completedByTaskRunId) {
+  const runId = deferredRunIdOf(record);
+  if (runId === undefined) {
     return undefined;
   }
 
-  if (!deps.readRunOutput) {
-    throw new Error(
-      `Waitpoint ${record.id} defers its output to run ${record.completedByTaskRunId}, but the resolver was built with no run-output reader.`
-    );
-  }
-
   // Postgres does not lose this: the back-reference nulls on delete but Waitpoint.output stays,
   // so the legacy path still emits it. Returning undefined here instead would resolve the
   // parent's triggerAndWait successfully with no output, which is silent wrong data.
-  const output = await deps.readRunOutput(record.completedByTaskRunId);
+  const output = runOutputs.get(runId);
   if (output === undefined) {
     throw new UnresolvableWaitpointId(record.id, "lost-run-output");
   }