github-featured/triggerdotdev--trigger.dev
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

refactor(run-store): hold RoutingRunStore's stores in a keyed shard map (#4752)

Browse Source 
## What

`RoutingRunStore` held two named store fields, `#new` and `#legacy`, and
took its routing policy from the order the statements happened to run
in. It now holds a `Map<ShardKey, RunStore>`, and the three policies
that were implicit are readable data:

- **`#probeOrder`** (`new` → `legacy`) — the sequential probe for a
lookup with no routable id. The first non-null result wins, and the
*last* entry owns the canonical not-found throw.
- **`#precedence`** (`legacy` → `new`) — ascending authority for a
merge, so the highest-authority shard wins a duplicate id.
- **`#idlessRouteShard`** (`new`) and **`#idlessWaitpointShard`**
(`legacy`) — the two id-less defaults, which differ by role and were
previously two unrelated literals in unrelated methods.

The two orders are the **reverse of each other**, which is why they are
separate fields rather than one ordering. Nine sites observe the
result-array order and must iterate `#probeOrder`; five decide a value
by which shard wins a duplicate and must iterate `#precedence`. Five
more sum counts and are order-independent, because addition commutes.

Four helpers absorb the twenty-six hand-written fan-outs —
`#probeFirst`, `#fanOut(order, fn)`, `#fanOutPartitioned`,
`#shardsExcept` — and `#shardKeyOf` replaces the inline
residency-to-store ternaries. `#fanOut` takes its order as an argument
so every call site states which policy it uses.

The constructor keeps its exact options type. No union arm, no `shards`
member: that would loosen the excess-property check and silently retire
the `@ts-expect-error onLegacyRead` lock in the test corpus. N-way
construction is a later change.

## One behaviour change

`findManyTaskRunWaitpoints` merged its edge rows NEW-first into a
last-wins dedupe, so a duplicate edge id resolved to the **legacy** row
— the opposite of the rule the other four merges follow, and the
opposite of what `dedupeEdgesById`'s own comment claimed. No test pinned
it in either direction.

It now resolves NEW-wins, consistent with every sibling merge, and a new
test pins the winner so it cannot drift back silently.

Reaching this case needs one edge id present on both stores at the same
time, with no routable `taskRunId`. That only arises from drain
mirroring. The drain seam is removed (`runOpsStore.test.ts`, "fan-out
spans NEW+LEGACY with no drain seam"), so **no new duplicates can be
created** — but removing the code does not delete rows it previously
wrote, and this class still carries comments treating mirrored rows as a
live data condition. Whether any historical duplicate edge rows persist
is an empirical question about production data, not something this diff
settles.

If such a row is hit, the two copies either agree — in which case the
winner is immaterial — or they have diverged, in which case NEW is the
authoritative copy by the router's own precedence rule. So the corrected
behaviour is at least as correct as the old one in every reachable case.

Everything else is behaviour-preserving.

## How it was verified

- **`internal-packages/run-store`: 69 files, 379 tests pass.** The
corpus is the regression gate for this refactor. 67 of the 68
pre-existing test files are byte-identical; the one that differs
(`runOpsStore.mixedResidency.test.ts`) changes only `//` comments.
- **`internal-packages/run-engine`: 12 files, 69 tests pass** — every
file that constructs the router, exercised at runtime.
- **The `@ts-expect-error onLegacyRead` lock still fires.**
`tsconfig.build.json` excludes `*.test.ts`, so a green typecheck does
not cover it. A scratch probe confirmed `tsc` still reports `TS2353` for
`onLegacyRead` and no error for the three real options.
- **All 48 construction sites outside the package compile unchanged.**
`tsconfig.check.json` also excludes `*.test.ts`, so the 25 webapp test
files were checked with the test exclusion dropped and compared against
the same check on the base commit: 614 errors before, 614 after, zero
present in one and not the other. Those 614 are pre-existing in
never-typechecked test files.
- `typecheck` passes for `run-store`, `run-engine` and `webapp`. `knip`
reports nothing in `run-store`.

## Also

Refreshes the sixteen stale `runOpsStore.ts` line references in
`runOpsStore.mixedResidency.test.ts`, each verified against the symbol
it names.

## Notes for the reviewer

- The riskiest possible mistake in this diff is a fan-out passing the
wrong order — the compiler cannot catch it, because both orders are
`readonly ShardKey[]`. The five `#precedence` sites are `#findRunsOpen`,
`findRunsByIdempotencyKeys`, `#collectManyWaitpoints`,
`findManyTaskRunWaitpoints` and `findManyWaitpointTags`. Those are the
lines worth the closest read.
- Four sites previously derived "the other store" by object identity
(`home === this.#new ? ...`). They now compare keys. The two are
equivalent: in single-database mode both keys map to the same store
object, and when the stores are distinct, identity and key comparison
agree.
- No changeset and no `.server-changes` note: the package is internal
and the one behaviour change is unreachable in production, so a release
note would tell a user nothing.
- Two CI checks fail for reasons that predate this branch and reproduce
on the base commit: `lint` (~16 unknown `react/*` rules make
`.oxlintrc.json` fail to parse, which disables oxlint entirely —
including the two `trigger-runops` fences) and `knip` (`unrun`, an
unused devDependency on the default branch). Both want their own fix.
This commit is contained in:
Daniel Sutton
2026-08-21 17:05:51 +01:00
committed by GitHub
parent c5c2ea92ca
commit aa9b888988
3 changed files with 505 additions and 303 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal-packages/run-store/src/runOpsStore.mixedResidency.test.ts
+14 -14
View File
@@ -195,10 +195,10 @@ async function seedSharedEnv(prisma14: PrismaClient, suffix: string) {
}
describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id #new coexisting)", () => {
// ── Case 1: findRuns by a MIXED bounded id-set (#findRunsByIdSet, runOpsStore.ts:294) ──
// ── Case 1: findRuns by a MIXED bounded id-set (#findRunsByIdSet) ──
// A list-hydrate id set spans cuid (legacy) + run-ops id (new) ids plus a run-ops id absent from legacy.
// Both resident runs returned; take/skip applied GLOBALLY post-merge; orderBy honored; the absent
// run-ops id short-circuits (never probed on LEGACY, :309).
// run-ops id short-circuits (never probed on LEGACY, #fanOutPartitioned).
heteroRunOpsPostgresTest(
"case 1: findRuns by a mixed id-set returns both DBs' runs, ordered, take/skip global",
async ({ prisma14, prisma17 }) => {
@@ -306,7 +306,7 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
}
);
// ── Case 2: findRuns by an OPEN predicate (#findRunsOpen, runOpsStore.ts:319) ──
// ── Case 2: findRuns by an OPEN predicate (#findRunsOpen) ──
// No id set → query BOTH stores, union, dedup by id (NEW wins). Filter by a shared scalar
// (runtimeEnvironmentId + status) that matches rows on both DBs.
heteroRunOpsPostgresTest(
@@ -355,7 +355,7 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
}
);
// ── Case 3: expireRunsBatch with a MIXED id list (runOpsStore.ts:474) ──
// ── Case 3: expireRunsBatch with a MIXED id list ──
// Partitions run-ops id→NEW / cuid→LEGACY; each leg called only when non-empty; counts summed; each row
// updated on its OWN DB only.
heteroRunOpsPostgresTest(
@@ -388,7 +388,7 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
}
);
// ── Case 4: clearIdempotencyKey fan-out arm (byFriendlyIds, runOpsStore.ts:358) ──
// ── Case 4: clearIdempotencyKey fan-out arm (byFriendlyIds) ──
// byFriendlyIds spans mixed residency → fan out to both, sum the count, each row cleared on its home.
heteroRunOpsPostgresTest(
"case 4: clearIdempotencyKey byFriendlyIds clears across both DBs and sums the count",
@@ -429,7 +429,7 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
}
);
// ── Case 5: countPendingWaitpoints scattered across both DBs (runOpsStore.ts:731) ──
// ── Case 5: countPendingWaitpoints scattered across both DBs ──
// A run's pending waitpoints can be split across both stores mid-drain → count on each and sum.
heteroRunOpsPostgresTest(
"case 5: countPendingWaitpoints sums PENDING waitpoints scattered across both DBs",
@@ -465,7 +465,7 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
}
);
// ── Case 6: findManyWaitpoints { id: { in: [...mixed...] } } (runOpsStore.ts:793) ──
// ── Case 6: findManyWaitpoints { id: { in: [...mixed...] } } ──
// Merge waitpoints from both DBs for a mixed id set.
heteroRunOpsPostgresTest(
"case 6: findManyWaitpoints merges a mixed id set from both DBs",
@@ -495,7 +495,7 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
// ── Case 8: findExecutionSnapshot / findManyExecutionSnapshots OPEN (no runId) where ──
// A by-snapshot-id-only lookup (snapshot ids are non-classifiable cuids) must fan out NEW→LEGACY
// (findExecutionSnapshot, :675) / merge both (findManyExecutionSnapshots, :688). Seed a snapshot on
// (findExecutionSnapshot) / merge both (findManyExecutionSnapshots). Seed a snapshot on
// EACH DB (one run-ops run on #new, one cuid run on #legacy) and read with a no-runId where.
heteroRunOpsPostgresTest(
"case 8: findExecutionSnapshot/findManyExecutionSnapshots with an open where reach both DBs",
@@ -544,7 +544,7 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
}
);
// ── Case 9a: findRun with an UNCLASSIFIABLE where (spanId) on a #legacy run (#findRunUnrouted, :213) ──
// ── Case 9a: findRun with an UNCLASSIFIABLE where (spanId) on a #legacy run (#findRunUnrouted) ──
// A run-ops run on #new and a cuid run on #legacy each carry a distinct spanId. A spanId where can't
// be id-classified → fan out NEW-first then LEGACY. The legacy-resident run must be found.
heteroRunOpsPostgresTest(
@@ -586,7 +586,7 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
}
);
// ── Case 9b: findRunOrThrow with an UNCLASSIFIABLE where (spanId) on a #legacy run (:593) ──
// ── Case 9b: findRunOrThrow with an UNCLASSIFIABLE where (spanId) on a #legacy run (#findRunOrThrowUnrouted) ──
// The throwing twin must match findRun's fan-out: an unclassifiable where whose only matching run
// lives on #legacy must NOT throw. A NEW-only fallback would miss the legacy run and throw.
heteroRunOpsPostgresTest(
@@ -628,7 +628,7 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
}
);
// ── Case 7: findManyTaskRunWaitpoints with edges whose relations STRADDLE DBs (runOpsStore.ts:876) ──
// ── Case 7: findManyTaskRunWaitpoints with edges whose relations STRADDLE DBs ──
// An edge co-locates with its RUN, but its `waitpoint`/`taskRun` relations can live on the OTHER DB
// (a cuid token blocking a run-ops run, and vice versa). The per-leg scalar query is stripped of the
// relation keys; the router re-hydrates `waitpoint`/`taskRun` across BOTH DBs. Exercises BOTH
@@ -721,7 +721,7 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
}
);
// ── Case 7b: the "blocking waitpoint not found on either DB" HARD ERROR (runOpsStore.ts:917) ──
// ── Case 7b: the "blocking waitpoint not found on either DB" HARD ERROR (#hydrateEdgeWaitpointsCrossDb) ──
// An edge whose `waitpointId` resolves on NEITHER DB must throw rather than leave a null status that
// would strand (hang) or wrongly unblock the run.
heteroRunOpsPostgresTest(
@@ -748,7 +748,7 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
}
);
// ── Case 10: findBatchTaskRunById / findBatchTaskRunByFriendlyId NEW-then-LEGACY probe (:1124,:1137) ──
// ── Case 10: findBatchTaskRunById / findBatchTaskRunByFriendlyId NEW-then-LEGACY probe ──
// A batch resident on #legacy AND a run-ops-id batch landed on #new (the control-plane window mints
// cuid ids, but a run-ops batch resides on #new) are BOTH found via the probe, regardless of id-shape.
heteroRunOpsPostgresTest(
@@ -792,7 +792,7 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
}
);
// ── Case 11a: updateManyWaitpoints with a NO-ID (batch) where fans out to both and sums (:822) ──
// ── Case 11a: updateManyWaitpoints with a NO-ID (batch) where fans out to both and sums ──
// A batch where (no single routable id, e.g. completedByTaskRunId IS NULL + status PENDING) must
// apply on BOTH DBs and sum the count.
heteroRunOpsPostgresTest(
internal-packages/run-store/src/runOpsStore.shardMap.test.ts
+180
View File
@@ -0,0 +1,180 @@
import { describe, expect, it } from "vitest";
import { RoutingRunStore } from "./runOpsStore.js";
import type { ReadClient, RunStore } from "./types.js";
// Pins the routing ALGEBRA: probe order, merge precedence, and the two id-less fallbacks that
// differ by role. No DB — each slot is a fake RunStore recording into ONE shared ordered log, so a
// sequential probe's leg order and a merge's winner are both observable. Real two-DB topology stays
// with the heteroRunOpsPostgresTest suites.
//
// MUST NOT assert invocation order for a PARALLEL fan-out: both legs are issued before either
// resolves, so the order they are created in is not a behaviour.
type Slot = "new" | "legacy";
type Call = { slot: Slot; method: string };
type FakeConfig = {
// Rows this store returns from findRun / findRuns / findRunOrThrow, regardless of filter.
runs?: Array<Record<string, unknown>>;
// Edge rows this store returns from findManyTaskRunWaitpoints, regardless of filter.
edges?: Array<Record<string, unknown>>;
// Waitpoint rows this store returns from findWaitpoint, regardless of filter.
waitpoint?: Record<string, unknown> | null;
};
type FakeStore = RunStore & {
slot: Slot;
primaryReadClient: { __primary: Slot };
};
function fakeStore(slot: Slot, log: Call[], config: FakeConfig = {}): FakeStore {
const record = (method: string) => log.push({ slot, method });
const runs = config.runs ?? [];
const store: Partial<FakeStore> = {
slot,
primaryReadClient: { __primary: slot },
findRun: ((_where: unknown, _args?: unknown) => {
record("findRun");
return Promise.resolve((runs[0] ?? null) as never);
}) as FakeStore["findRun"],
findRunOnPrimary: ((_where: unknown, _args?: unknown) => {
record("findRunOnPrimary");
return Promise.resolve((runs[0] ?? null) as never);
}) as FakeStore["findRunOnPrimary"],
findRunOrThrow: ((_where: unknown, _args?: unknown) => {
record("findRunOrThrow");
if (runs[0] === undefined) {
return Promise.reject(new Error(`no run on ${slot}`)) as never;
}
return Promise.resolve(runs[0] as never);
}) as FakeStore["findRunOrThrow"],
findRuns: ((_args: unknown, _client?: ReadClient) => {
record("findRuns");
return Promise.resolve(runs as never);
}) as FakeStore["findRuns"],
createRun: ((_params: unknown) => {
record("createRun");
return Promise.resolve({ slot } as never);
}) as FakeStore["createRun"],
createTaskRunCheckpoint: ((_args: unknown) => {
record("createTaskRunCheckpoint");
return Promise.resolve({ slot } as never);
}) as FakeStore["createTaskRunCheckpoint"],
findWaitpoint: ((_args: unknown, _client?: ReadClient) => {
record("findWaitpoint");
return Promise.resolve((config.waitpoint ?? null) as never);
}) as FakeStore["findWaitpoint"],
updateWaitpoint: ((_args: unknown) => {
record("updateWaitpoint");
return Promise.resolve({ slot } as never);
}) as FakeStore["updateWaitpoint"],
findManyTaskRunWaitpoints: ((_args: unknown, _client?: ReadClient) => {
record("findManyTaskRunWaitpoints");
return Promise.resolve((config.edges ?? []) as never);
}) as FakeStore["findManyTaskRunWaitpoints"],
};
return store as unknown as FakeStore;
}
// Deterministic residency by id prefix via the classify seam (no dependence on id-shape rules).
function buildRouter(newConfig: FakeConfig = {}, legacyConfig: FakeConfig = {}) {
const log: Call[] = [];
const newStore = fakeStore("new", log, newConfig);
const legacyStore = fakeStore("legacy", log, legacyConfig);
const router = new RoutingRunStore({
new: newStore,
legacy: legacyStore,
classify: (id: string) => (id.startsWith("new") ? "NEW" : "LEGACY"),
});
return { router, newStore, legacyStore, log };
}
const trace = (log: Call[]) => log.map((c) => `${c.slot}:${c.method}`);
describe("RoutingRunStore #probeOrder — new then legacy, sequential", () => {
it("probes new BEFORE legacy for an unrouted findRun", async () => {
const { router, log } = buildRouter();
await router.findRun({ spanId: "span_x" });
expect(trace(log)).toEqual(["new:findRun", "legacy:findRun"]);
});
it("stops at the first non-null leg and never consults legacy", async () => {
const { router, log } = buildRouter({ runs: [{ id: "r1" }] });
await router.findRun({ spanId: "span_x" });
expect(trace(log)).toEqual(["new:findRun"]);
});
it("gives the LAST probe leg the canonical not-found throw", async () => {
const { router, log } = buildRouter();
await expect(router.findRunOrThrow({ spanId: "span_x" })).rejects.toThrow("no run on legacy");
// new is probed with the nullable read; only legacy is asked to throw.
expect(trace(log)).toEqual(["new:findRun", "legacy:findRunOrThrow"]);
});
it("probes each store's own primary for a read-your-writes unrouted findRun", async () => {
const { router, log } = buildRouter();
await router.findRunOnPrimary({ spanId: "span_x" });
expect(trace(log)).toEqual(["new:findRunOnPrimary", "legacy:findRunOnPrimary"]);
});
});
describe("RoutingRunStore #precedence — NEW wins a merge", () => {
it("keeps the NEW row for a duplicate run id on an open predicate", async () => {
const { router } = buildRouter(
{ runs: [{ id: "dup", from: "new" }] },
{ runs: [{ id: "dup", from: "legacy" }] }
);
const rows = (await router.findRuns({
where: { runtimeEnvironmentId: "env_1" },
select: { id: true, from: true },
})) as Array<{ id: string; from: string }>;
expect(rows).toHaveLength(1);
expect(rows[0]?.from).toBe("new");
});
// Every merge in the router MUST resolve a duplicate id NEW-wins, edges included.
it("keeps the NEW row for a duplicate edge id on a waitpoint-keyed edge read", async () => {
const { router } = buildRouter(
{ edges: [{ id: "edge_dup", taskRunId: "new_run" }] },
{ edges: [{ id: "edge_dup", taskRunId: "legacy_run" }] }
);
const edges = (await router.findManyTaskRunWaitpoints({
where: { waitpointId: "waitpoint_x" },
select: { id: true, taskRunId: true },
})) as Array<{ id: string; taskRunId: string }>;
expect(edges).toHaveLength(1);
expect(edges[0]?.taskRunId).toBe("new_run");
});
});
describe("RoutingRunStore id-less fallbacks — the two defaults differ by role", () => {
it("routes an id-less create to new (#idlessRouteShard)", async () => {
const { router, log } = buildRouter();
await router.createRun({ data: {} } as never);
expect(trace(log)).toEqual(["new:createRun"]);
});
it("routes an id-less checkpoint create to new (#idlessRouteShard)", async () => {
const { router, log } = buildRouter();
await router.createTaskRunCheckpoint({ data: {} } as never);
expect(trace(log)).toEqual(["new:createTaskRunCheckpoint"]);
});
it("routes an id-less waitpoint update to legacy (#idlessWaitpointShard)", async () => {
const { router, log } = buildRouter();
await router.updateWaitpoint({ where: { idempotencyKey: "k" }, data: {} } as never);
expect(trace(log)).toEqual(["legacy:updateWaitpoint"]);
});
});
internal-packages/run-store/src/runOpsStore.ts
+311 -289
View File
@@ -7,7 +7,7 @@ import type {
TaskRunStatus,
WaitpointTag,
} from "@trigger.dev/database";
import { ownerEngine, type Residency } from "@trigger.dev/core/v3/isomorphic";
import { ownerEngine, type Residency, type ShardKey } from "@trigger.dev/core/v3/isomorphic";
import type { TaskRunError } from "@trigger.dev/core/v3/schemas";
import type {
ClearIdempotencyKeyInput,
@@ -34,23 +34,55 @@ import { isReadReplicaClient } from "./readReplicaClient.js";
import { CONNECTED_RUNS_LIMIT } from "./PostgresRunStore.js";
import { boundedIn } from "@trigger.dev/database";
// The two reserved gen-1 shard keys. They are multi-char, so a gen-2 single-char shard key can
// never collide with them.
const NEW_SHARD: ShardKey = "new";
const LEGACY_SHARD: ShardKey = "legacy";
/**
* Run-ops routing substrate for the TaskRun-core method group. Implements {@link RunStore}
* by selecting between a NEW store (the dedicated run-ops DB, where new runs are born) and
* a LEGACY store (the control-plane DB) via the residency classifier (`ownerEngine`:
* run-ops id→NEW, cuid→LEGACY). In single-DB both stores are the same, so routing is a no-op
* passthrough. Inert until the injecting seam wires it in under `isSplitEnabled()`; reads no
* flag here. The TaskRun-core methods (create/find/findRuns + updateMetadata/clearIdempotencyKey)
* route by residency; all other methods are mechanical residency-routing delegates.
* Run-ops routing substrate for the TaskRun-core method group. Implements {@link RunStore} over a
* map from shard key to store, selecting one by the residency classifier (`ownerEngine`: run-ops
* id→NEW, cuid→LEGACY). The compat constructor holds the two gen-1 shards — a NEW store (the
* dedicated run-ops DB, where new runs are born) and a LEGACY store (the control-plane DB).
* Inert until the injecting seam wires it in under `isSplitEnabled()`; reads no flag here.
*
* Every shard MUST be a distinct database. Single-DB does not construct this class at all — the
* injecting seam returns a bare PostgresRunStore — and split mode requires two configured run-ops
* URLs whose distinctness the boot sentinel enforces fail-closed. Two shard keys that resolve to
* ONE store would make the sum sites (#sumCounts, the counting fan-outs) count that store twice.
*
* Three policies are held as data rather than implied by statement order: {@link #probeOrder} for a
* lookup with no routable id, {@link #precedence} for a merge, and the two id-less fallbacks. A
* merge MUST iterate #precedence and a probe MUST iterate #probeOrder — the two are the reverse of
* each other, so swapping them changes behaviour.
*/
export class RoutingRunStore implements RunStore {
readonly #new: RunStore;
readonly #legacy: RunStore;
readonly #shards: ReadonlyMap<ShardKey, RunStore>;
// Sequential probe for a lookup with no routable id. The first non-null result wins, and the LAST
// entry owns the canonical not-found throw.
readonly #probeOrder: readonly ShardKey[];
// Ascending authority for a merge. The last write wins, so the highest-authority shard wins a
// duplicate id. Every merge in this class MUST use this order.
readonly #precedence: readonly ShardKey[];
// The two id-less defaults. They differ by role on purpose: a route with no id lands on the
// steady-state home, a waitpoint read with no id lands on the legacy store.
readonly #idlessRouteShard: ShardKey;
readonly #idlessWaitpointShard: ShardKey;
readonly #classify: (id: string) => Residency;
// Compat constructor: the two gen-1 stores, keyed by their reserved shard keys. The options type
// MUST stay closed — a union arm loosens the excess-property check and retires the
// `@ts-expect-error onLegacyRead` lock in the test corpus.
constructor(options: { new: RunStore; legacy: RunStore; classify?: (id: string) => Residency }) {
this.#new = options.new;
this.#legacy = options.legacy;
this.#shards = new Map<ShardKey, RunStore>([
[NEW_SHARD, options.new],
[LEGACY_SHARD, options.legacy],
]);
this.#probeOrder = [NEW_SHARD, LEGACY_SHARD];
this.#precedence = [LEGACY_SHARD, NEW_SHARD];
this.#idlessRouteShard = NEW_SHARD;
this.#idlessWaitpointShard = LEGACY_SHARD;
this.#classify = options.classify ?? ownerEngine;
}
@@ -71,16 +103,100 @@ export class RoutingRunStore implements RunStore {
return client != null && !isReadReplicaClient(client) ? store.primaryReadClient : undefined;
}
// The store for a shard key. Unreachable with the compat constructor — #shardKeyOfSafe yields only
// the two reserved keys — so this throw fires only if a caller wires a partial map.
#shardStore(key: ShardKey): RunStore {
const store = this.#shards.get(key);
if (store === undefined) {
throw new Error(`RoutingRunStore: no store is configured for shard key "${key}"`);
}
return store;
}
// The shard that owns an existing id. Throws only when an injected classifier throws.
#shardKeyOf(id: string): ShardKey {
return this.#classify(id) === "NEW" ? NEW_SHARD : LEGACY_SHARD;
}
// An unclassifiable id is treated as LEGACY (probe the control-plane DB rather than drop a
// real run), matching the read-through layer's policy.
#classifySafe(id: string): Residency {
#shardKeyOfSafe(id: string): ShardKey {
try {
return this.#classify(id);
return this.#shardKeyOf(id);
} catch {
return "LEGACY";
return LEGACY_SHARD;
}
}
// Sequential probe over #probeOrder, returning the first non-null result. `isLast` marks the leg
// that owns the canonical not-found throw, so a caller can swap in its throwing variant there.
async #probeFirst<R>(
fn: (store: RunStore, key: ShardKey, isLast: boolean) => Promise<R>
): Promise<R> {
const last = this.#probeOrder.length - 1;
for (let i = 0; i < last; i++) {
const key = this.#probeOrder[i]!;
const found = await fn(this.#shardStore(key), key, false);
if (found != null) {
return found;
}
}
const key = this.#probeOrder[last]!;
return fn(this.#shardStore(key), key, true);
}
// Run `fn` on every shard in parallel, returning the results in `order`. Pass #probeOrder where
// the result-array order is observable; pass #precedence where a duplicate id's winner decides
// the value. The two orders are the reverse of each other, so passing the wrong one is a
// behaviour change.
#fanOut<R>(
order: readonly ShardKey[],
fn: (store: RunStore, key: ShardKey) => Promise<R>
): Promise<R[]> {
return Promise.all(order.map((key) => fn(this.#shardStore(key), key)));
}
// Apply `fn` to every shard and sum the counts. A sum is order-independent, so this takes no order.
async #sumCounts(
fn: (store: RunStore, key: ShardKey) => Promise<{ count: number }>
): Promise<Prisma.BatchPayload> {
const legs = await this.#fanOut(this.#probeOrder, fn);
return { count: legs.reduce((sum, leg) => sum + leg.count, 0) };
}
// Group ids by owning shard, then run `fn` once per NON-EMPTY bucket, in `order`, in parallel. A
// shard with no id is never queried, so no leg issues an empty `IN ()`.
async #fanOutPartitioned<R>(
order: readonly ShardKey[],
ids: string[],
fn: (store: RunStore, shardIds: string[], key: ShardKey) => Promise<R>
): Promise<R[]> {
const byShard = new Map<ShardKey, string[]>();
for (const id of ids) {
const key = this.#shardKeyOfSafe(id);
const bucket = byShard.get(key);
if (bucket) bucket.push(id);
else byShard.set(key, [id]);
}
const legs: Array<Promise<R>> = [];
for (const key of order) {
const shardIds = byShard.get(key);
if (shardIds !== undefined) {
legs.push(fn(this.#shardStore(key), shardIds, key));
}
}
return Promise.all(legs);
}
// Every shard other than `key`, in probe order. With the compat constructor this yields exactly
// one entry, which is why each caller may take the first. At more than two shards a caller MUST
// fan out over all of them instead.
#shardsExcept(key: ShardKey): Array<{ key: ShardKey; store: RunStore }> {
return this.#probeOrder
.filter((k) => k !== key)
.map((k) => ({ key: k, store: this.#shardStore(k) }));
}
// A `findRuns` caller bound to the given store (preserves `this`; the overload set isn't
// assignable to a single call signature, so it's cast through the implementation shape). A
// caller-passed client resolves to the store's own primary (#ownPrimary) on every call.
@@ -98,23 +214,27 @@ export class RoutingRunStore implements RunStore {
// Route an existing run-ops id by residency. Throws on an unclassifiable id.
#route(id: string): RunStore {
return this.#classify(id) === "NEW" ? this.#new : this.#legacy;
return this.#shardStore(this.#shardKeyOf(id));
}
// Best-effort route; falls back to NEW (the steady-state home) when the id is absent.
// Classification is total (any id without the v1 version marker is LEGACY), so the
// catch below only guards injected classifiers that still throw.
#routeOrNew(id: string | undefined): RunStore {
// Best-effort shard key; falls back to #idlessRouteShard when the id is absent. Classification is
// total (any id without the v1 version marker is LEGACY), so the catch below only guards injected
// classifiers that still throw.
#routeKeyOrDefault(id: string | undefined): ShardKey {
if (typeof id !== "string") {
return this.#new;
return this.#idlessRouteShard;
}
try {
return this.#route(id);
return this.#shardKeyOf(id);
} catch {
return this.#new;
return this.#idlessRouteShard;
}
}
#routeOrNew(id: string | undefined): RunStore {
return this.#shardStore(this.#routeKeyOrDefault(id));
}
// WRITE routing is pure id-shape (cuid → LEGACY, run-ops id → NEW). A LEGACY-classified id is
// always LEGACY-resident; no marker check exists. Kept async so the many
// `await this.#routeForWrite(...)` call sites need no edits (awaiting a resolved store is
@@ -145,8 +265,8 @@ export class RoutingRunStore implements RunStore {
// `onPrimary` probes each store's own primary (read-your-writes callers; a fresh row may not
// be on the replica yet, which would mis-resolve the store).
async #resolveWaitpointStore(id: string | undefined, onPrimary = false): Promise<RunStore> {
const home =
typeof id === "string" && this.#classifySafe(id) === "NEW" ? this.#new : this.#legacy;
const homeKey = typeof id === "string" ? this.#shardKeyOfSafe(id) : this.#idlessWaitpointShard;
const home = this.#shardStore(homeKey);
if (typeof id !== "string") {
return home;
}
@@ -155,12 +275,15 @@ export class RoutingRunStore implements RunStore {
) {
return home;
}
const other = home === this.#new ? this.#legacy : this.#new;
return (await other.findWaitpoint(
const [other] = this.#shardsExcept(homeKey);
if (other === undefined) {
return home;
}
return (await other.store.findWaitpoint(
{ where: { id } },
onPrimary ? other.primaryReadClient : undefined
onPrimary ? other.store.primaryReadClient : undefined
))
? other
? other.store
: home;
}
@@ -249,14 +372,9 @@ export class RoutingRunStore implements RunStore {
onPrimary: boolean
): Promise<unknown> {
const method = onPrimary ? "findRunOnPrimary" : "findRun";
const fromNew = await (this.#new[method] as (...rest: unknown[]) => Promise<unknown>)(
where,
args
return this.#probeFirst((store) =>
(store[method] as (...rest: unknown[]) => Promise<unknown>)(where, args)
);
if (fromNew != null) {
return fromNew;
}
return (this.#legacy[method] as (...rest: unknown[]) => Promise<unknown>)(where, args);
}
findRuns<S extends Prisma.TaskRunSelect>(
@@ -339,15 +457,10 @@ export class RoutingRunStore implements RunStore {
): Promise<unknown[]> {
const { args: selArgs, addedFields } = ensureProjected(args);
const fan = { ...selArgs, take: undefined, skip: undefined };
const newIds = ids.filter((id) => this.#classifySafe(id) === "NEW");
const legacyIds = ids.filter((id) => this.#classifySafe(id) !== "NEW");
const findNew = this.#findManyOn(this.#new, client);
const findLegacy = this.#findManyOn(this.#legacy, client);
const [newRows, legacyRows] = await Promise.all([
newIds.length > 0 ? findNew(narrowToIds(fan, newIds)) : [],
legacyIds.length > 0 ? findLegacy(narrowToIds(fan, legacyIds)) : [],
]);
return finalizeRows([...newRows, ...legacyRows], args, addedFields);
const legs = await this.#fanOutPartitioned(this.#probeOrder, ids, (store, shardIds) =>
this.#findManyOn(store, client)(narrowToIds(fan, shardIds))
);
return finalizeRows(legs.flat(), args, addedFields);
}
// Open predicate (e.g. `{ batchId }`, `{ status, runtimeEnvironmentId }`): no id set to
@@ -355,12 +468,11 @@ export class RoutingRunStore implements RunStore {
async #findRunsOpen(args: FindRunsArgs, client?: ReadClient): Promise<unknown[]> {
const { args: selArgs, addedFields } = ensureProjected(args);
const fan = widenForMerge(selArgs);
const findNew = this.#findManyOn(this.#new, client);
const findLegacy = this.#findManyOn(this.#legacy, client);
const [newRows, legacyRows] = await Promise.all([findNew(fan), findLegacy(fan)]);
const legs = await this.#fanOut(this.#precedence, (store) =>
this.#findManyOn(store, client)(fan)
);
const byId = new Map<string, Record<string, unknown>>();
for (const r of legacyRows) byId.set(r.id as string, r);
for (const r of newRows) byId.set(r.id as string, r);
for (const r of legs.flat()) byId.set(r.id as string, r);
return finalizeRows([...byId.values()], args, addedFields);
}
@@ -427,18 +539,11 @@ export class RoutingRunStore implements RunStore {
if (args.idempotencyKeys.length === 0) {
return [];
}
const [newRows, legacyRows] = await Promise.all([
this.#new.findRunsByIdempotencyKeys(args, RoutingRunStore.#ownPrimary(this.#new, client)),
this.#legacy.findRunsByIdempotencyKeys(
args,
RoutingRunStore.#ownPrimary(this.#legacy, client)
),
]);
const legs = await this.#fanOut(this.#precedence, (store) =>
store.findRunsByIdempotencyKeys(args, RoutingRunStore.#ownPrimary(store, client))
);
const byKey = new Map<string, IdempotencyKeyRunMatch>();
for (const row of legacyRows) {
if (row.idempotencyKey != null) byKey.set(row.idempotencyKey, row);
}
for (const row of newRows) {
for (const row of legs.flat()) {
if (row.idempotencyKey != null) byKey.set(row.idempotencyKey, row);
}
return [...byKey.values()];
@@ -477,17 +582,16 @@ export class RoutingRunStore implements RunStore {
// so fall back to LEGACY when NEW matched nothing — otherwise the reset 404s and the stale legacy
// key keeps deduping. In the steady (fully-drained) state NEW matches and legacy is never touched.
if ("byPredicate" in params && params.byPredicate?.residency === "NEW") {
const fromNew = await this.#new.clearIdempotencyKey(params, undefined);
const fromNew = await this.#shardStore(NEW_SHARD).clearIdempotencyKey(params, undefined);
if (fromNew.count > 0) {
return fromNew;
}
const fromLegacy = await this.#legacy.clearIdempotencyKey(params);
return { count: fromNew.count + fromLegacy.count };
const rest = await Promise.all(
this.#shardsExcept(NEW_SHARD).map(({ store }) => store.clearIdempotencyKey(params))
);
return { count: rest.reduce((sum, leg) => sum + leg.count, fromNew.count) };
}
return Promise.all([
this.#new.clearIdempotencyKey(params),
this.#legacy.clearIdempotencyKey(params),
]).then(([fromNew, fromLegacy]) => ({ count: fromNew.count + fromLegacy.count }));
return this.#sumCounts((store) => store.clearIdempotencyKey(params));
}
// ---------------------------------------------------------------------------
@@ -625,15 +729,11 @@ export class RoutingRunStore implements RunStore {
data: { error: TaskRunError; now: Date },
tx?: PrismaClientOrTransaction
): Promise<number> {
// Partition by id-shape: run-ops id → NEW, everything else → LEGACY. Call each store
// only when its partition is non-empty (avoids an empty IN () clause). Sum counts.
const newIds = runIds.filter((id) => this.#classifySafe(id) === "NEW");
const legacyIds = runIds.filter((id) => this.#classifySafe(id) !== "NEW");
const [fromNew, fromLegacy] = await Promise.all([
newIds.length > 0 ? this.#new.expireRunsBatch(newIds, data) : 0,
legacyIds.length > 0 ? this.#legacy.expireRunsBatch(legacyIds, data) : 0,
]);
return fromNew + fromLegacy;
// Partition by id-shape, call each store only for its own ids, and sum the counts.
const legs = await this.#fanOutPartitioned(this.#probeOrder, runIds, (store, shardIds) =>
store.expireRunsBatch(shardIds, data)
);
return legs.reduce((sum, leg) => sum + leg, 0);
}
async lockRunToWorker(
@@ -771,16 +871,11 @@ export class RoutingRunStore implements RunStore {
onPrimary: boolean
): Promise<unknown> {
const probe = onPrimary ? "findRunOnPrimary" : "findRun";
const fromNew = await (this.#new[probe] as (...rest: unknown[]) => Promise<unknown>)(
where,
args
);
if (fromNew != null) {
return fromNew;
}
// LEGACY is the last leg probed, so it owns the canonical not-found throw when both DBs miss.
// The last probe leg owns the canonical not-found throw when every shard misses.
const throwMethod = onPrimary ? "findRunOrThrowOnPrimary" : "findRunOrThrow";
return (this.#legacy[throwMethod] as (...rest: unknown[]) => Promise<unknown>)(where, args);
return this.#probeFirst((store, _key, isLast) =>
(store[isLast ? throwMethod : probe] as (...rest: unknown[]) => Promise<unknown>)(where, args)
);
}
// Explicit read-your-writes entry points: route by residency to the owning store's PRIMARY
@@ -923,13 +1018,8 @@ export class RoutingRunStore implements RunStore {
const store = this.#routeOrNew(runId);
return store.findExecutionSnapshot(args, RoutingRunStore.#ownPrimary(store, client));
}
const fromNew = await this.#new.findExecutionSnapshot(
args,
RoutingRunStore.#ownPrimary(this.#new, client)
);
return (
fromNew ??
this.#legacy.findExecutionSnapshot(args, RoutingRunStore.#ownPrimary(this.#legacy, client))
return this.#probeFirst((store) =>
store.findExecutionSnapshot(args, RoutingRunStore.#ownPrimary(store, client))
);
}
@@ -943,14 +1033,10 @@ export class RoutingRunStore implements RunStore {
const store = this.#routeOrNew(runId);
return store.findManyExecutionSnapshots(args, RoutingRunStore.#ownPrimary(store, client));
}
const [fromNew, fromLegacy] = await Promise.all([
this.#new.findManyExecutionSnapshots(args, RoutingRunStore.#ownPrimary(this.#new, client)),
this.#legacy.findManyExecutionSnapshots(
args,
RoutingRunStore.#ownPrimary(this.#legacy, client)
),
]);
return [...fromNew, ...fromLegacy];
const legs = await this.#fanOut(this.#probeOrder, (store) =>
store.findManyExecutionSnapshots(args, RoutingRunStore.#ownPrimary(store, client))
);
return legs.flat();
}
async createExecutionSnapshot(
@@ -980,17 +1066,13 @@ export class RoutingRunStore implements RunStore {
RoutingRunStore.#ownPrimary(store, client)
);
}
const [fromNew, fromLegacy] = await Promise.all([
this.#new.findSnapshotCompletedWaitpointIds(
const legs = await this.#fanOut(this.#probeOrder, (store) =>
store.findSnapshotCompletedWaitpointIds(
snapshotId,
RoutingRunStore.#ownPrimary(this.#new, client)
),
this.#legacy.findSnapshotCompletedWaitpointIds(
snapshotId,
RoutingRunStore.#ownPrimary(this.#legacy, client)
),
]);
return uniqueStrings([...fromNew, ...fromLegacy]);
RoutingRunStore.#ownPrimary(store, client)
)
);
return uniqueStrings(legs.flat());
}
// As above: route to the run's store when the run id is threaded through, else fan out (the snapshot
@@ -1007,19 +1089,15 @@ export class RoutingRunStore implements RunStore {
RoutingRunStore.#ownPrimary(store, client)
);
}
const [fromNew, fromLegacy] = await Promise.all([
this.#new.findSnapshotCompletedWaitpointIdsWithPresence(
const legs = await this.#fanOut(this.#probeOrder, (store) =>
store.findSnapshotCompletedWaitpointIdsWithPresence(
snapshotId,
RoutingRunStore.#ownPrimary(this.#new, client)
),
this.#legacy.findSnapshotCompletedWaitpointIdsWithPresence(
snapshotId,
RoutingRunStore.#ownPrimary(this.#legacy, client)
),
]);
RoutingRunStore.#ownPrimary(store, client)
)
);
return {
present: fromNew.present || fromLegacy.present,
ids: uniqueStrings([...fromNew.ids, ...fromLegacy.ids]),
present: legs.some((leg) => leg.present),
ids: uniqueStrings(legs.flatMap((leg) => leg.ids)),
};
}
@@ -1029,34 +1107,23 @@ export class RoutingRunStore implements RunStore {
// row on each leg. Each sub-store already caps at CONNECTED_RUNS_LIMIT, but a disjoint run set on
// each side can still make the union exceed it, so slice again after the merge.
async findWaitpointConnectedRunIds(waitpointId: string, client?: ReadClient): Promise<string[]> {
const [fromNew, fromLegacy] = await Promise.all([
this.#new.findWaitpointConnectedRunIds(
waitpointId,
RoutingRunStore.#ownPrimary(this.#new, client)
),
this.#legacy.findWaitpointConnectedRunIds(
waitpointId,
RoutingRunStore.#ownPrimary(this.#legacy, client)
),
]);
return uniqueStrings([...fromNew, ...fromLegacy]).slice(0, CONNECTED_RUNS_LIMIT);
const legs = await this.#fanOut(this.#probeOrder, (store) =>
store.findWaitpointConnectedRunIds(waitpointId, RoutingRunStore.#ownPrimary(store, client))
);
return uniqueStrings(legs.flat()).slice(0, CONNECTED_RUNS_LIMIT);
}
async findWaitpointCompletedSnapshotIds(
waitpointId: string,
client?: ReadClient
): Promise<string[]> {
const [fromNew, fromLegacy] = await Promise.all([
this.#new.findWaitpointCompletedSnapshotIds(
const legs = await this.#fanOut(this.#probeOrder, (store) =>
store.findWaitpointCompletedSnapshotIds(
waitpointId,
RoutingRunStore.#ownPrimary(this.#new, client)
),
this.#legacy.findWaitpointCompletedSnapshotIds(
waitpointId,
RoutingRunStore.#ownPrimary(this.#legacy, client)
),
]);
return uniqueStrings([...fromNew, ...fromLegacy]);
RoutingRunStore.#ownPrimary(store, client)
)
);
return uniqueStrings(legs.flat());
}
async blockRunWithWaitpointEdges(params: {
@@ -1086,24 +1153,17 @@ export class RoutingRunStore implements RunStore {
runId?: string
): Promise<number> {
if (runId === undefined) {
const [fromNew, fromLegacy] = await Promise.all([
this.#new.countPendingWaitpoints(
waitpointIds,
RoutingRunStore.#ownPrimary(this.#new, client)
),
this.#legacy.countPendingWaitpoints(
waitpointIds,
RoutingRunStore.#ownPrimary(this.#legacy, client)
),
]);
return fromNew + fromLegacy;
const legs = await this.#fanOut(this.#probeOrder, (store) =>
store.countPendingWaitpoints(waitpointIds, RoutingRunStore.#ownPrimary(store, client))
);
return legs.reduce((sum, leg) => sum + leg, 0);
}
if (waitpointIds.length === 0) {
return 0;
}
const runStore = this.#routeOrNew(runId);
const otherStore = runStore === this.#new ? this.#legacy : this.#new;
const runKey = this.#routeKeyOrDefault(runId);
const runStore = this.#shardStore(runKey);
const { pendingIds, presentIds } = await runStore.countPendingWaitpointsWithPresence(
waitpointIds,
RoutingRunStore.#ownPrimary(runStore, client)
@@ -1113,9 +1173,13 @@ export class RoutingRunStore implements RunStore {
if (missing.length === 0) {
return pendingIds.length;
}
const otherPending = await otherStore.countPendingWaitpoints(
const [other] = this.#shardsExcept(runKey);
if (other === undefined) {
return pendingIds.length;
}
const otherPending = await other.store.countPendingWaitpoints(
missing,
RoutingRunStore.#ownPrimary(otherStore, client)
RoutingRunStore.#ownPrimary(other.store, client)
);
return pendingIds.length + otherPending;
}
@@ -1128,19 +1192,15 @@ export class RoutingRunStore implements RunStore {
waitpointIds: string[],
client?: ReadClient
): Promise<{ pendingIds: string[]; presentIds: string[] }> {
const [fromNew, fromLegacy] = await Promise.all([
this.#new.countPendingWaitpointsWithPresence(
const legs = await this.#fanOut(this.#probeOrder, (store) =>
store.countPendingWaitpointsWithPresence(
waitpointIds,
RoutingRunStore.#ownPrimary(this.#new, client)
),
this.#legacy.countPendingWaitpointsWithPresence(
waitpointIds,
RoutingRunStore.#ownPrimary(this.#legacy, client)
),
]);
RoutingRunStore.#ownPrimary(store, client)
)
);
return {
pendingIds: uniqueStrings([...fromNew.pendingIds, ...fromLegacy.pendingIds]),
presentIds: uniqueStrings([...fromNew.presentIds, ...fromLegacy.presentIds]),
pendingIds: uniqueStrings(legs.flatMap((leg) => leg.pendingIds)),
presentIds: uniqueStrings(legs.flatMap((leg) => leg.presentIds)),
};
}
@@ -1178,14 +1238,16 @@ export class RoutingRunStore implements RunStore {
waitpointId: string | undefined
): RunStore {
if (ownerId !== undefined) {
return this.#classifySafe(ownerId) === "NEW" ? this.#new : this.#legacy;
return this.#shardStore(this.#shardKeyOfSafe(ownerId));
}
if (residency !== undefined) {
return residency === "NEW" ? this.#new : this.#legacy;
return this.#shardStore(residency === "NEW" ? NEW_SHARD : LEGACY_SHARD);
}
return typeof waitpointId === "string" && this.#classifySafe(waitpointId) === "NEW"
? this.#new
: this.#legacy;
return this.#shardStore(
typeof waitpointId === "string"
? this.#shardKeyOfSafe(waitpointId)
: this.#idlessWaitpointShard
);
}
upsertWaitpoint<T extends Prisma.WaitpointUpsertArgs>(
@@ -1237,14 +1299,12 @@ export class RoutingRunStore implements RunStore {
scalarArgs as typeof args,
coLocatedDedup ? store.primaryReadClient : RoutingRunStore.#ownPrimary(store, client)
)) as Record<string, unknown> | null)
: (((await this.#new.findWaitpoint(
scalarArgs as typeof args,
RoutingRunStore.#ownPrimary(this.#new, client)
)) ??
(await this.#legacy.findWaitpoint(
: ((await this.#probeFirst((probed) =>
probed.findWaitpoint(
scalarArgs as typeof args,
RoutingRunStore.#ownPrimary(this.#legacy, client)
))) as Record<string, unknown> | null);
RoutingRunStore.#ownPrimary(probed, client)
)
)) as Record<string, unknown> | null);
if (row) {
await this.#reresolveWaitpointRelationsCrossDb(row, relations, client);
}
@@ -1258,7 +1318,12 @@ export class RoutingRunStore implements RunStore {
args: Prisma.SelectSubset<T, Prisma.WaitpointFindFirstArgs>
): Promise<Prisma.WaitpointGetPayload<T> | null> {
const id = RoutingRunStore.#waitpointId((args as { where?: unknown }).where);
const store = id !== undefined ? await this.#resolveWaitpointStore(id, true) : this.#new;
// The id-less arm takes the ROUTE default, not the waitpoint default. Only the unblock re-read
// calls this, and it always carries an id, so the arm is unreachable today.
const store =
id !== undefined
? await this.#resolveWaitpointStore(id, true)
: this.#shardStore(this.#idlessRouteShard);
return store.findWaitpointOnPrimary(args);
}
@@ -1297,7 +1362,8 @@ export class RoutingRunStore implements RunStore {
if (runId !== undefined) {
const requestedIds = idListFromWhere((scalarArgs.where ?? {}) as Prisma.TaskRunWhereInput);
if (requestedIds !== undefined) {
const runStore = this.#routeOrNew(runId);
const runKey = this.#routeKeyOrDefault(runId);
const runStore = this.#shardStore(runKey);
const fromRun = (await runStore.findManyWaitpoints(
scalarArgs as Prisma.WaitpointFindManyArgs,
RoutingRunStore.#ownPrimary(runStore, client)
@@ -1309,33 +1375,33 @@ export class RoutingRunStore implements RunStore {
if (missing.length === 0) {
return fromRun;
}
const otherStore = runStore === this.#new ? this.#legacy : this.#new;
const fromOther = (await otherStore.findManyWaitpoints(
const [other] = this.#shardsExcept(runKey);
if (other === undefined) {
return fromRun;
}
const fromOther = (await other.store.findManyWaitpoints(
narrowArgsToIds(scalarArgs, missing) as Prisma.WaitpointFindManyArgs,
RoutingRunStore.#ownPrimary(otherStore, client)
RoutingRunStore.#ownPrimary(other.store, client)
)) as Record<string, unknown>[];
return [...fromRun, ...fromOther];
}
// No bounded id set to partition on → fall through to the fan-out path.
}
const [fromNew, fromLegacy] = await Promise.all([
this.#new.findManyWaitpoints(
scalarArgs as Prisma.WaitpointFindManyArgs,
RoutingRunStore.#ownPrimary(this.#new, client)
) as Promise<Record<string, unknown>[]>,
this.#legacy.findManyWaitpoints(
scalarArgs as Prisma.WaitpointFindManyArgs,
RoutingRunStore.#ownPrimary(this.#legacy, client)
) as Promise<Record<string, unknown>[]>,
]);
// A token mirrored onto both DBs during drain appears in BOTH legs; dedup by id with NEW-wins
// (the NEW copy is authoritative once a run migrates), matching the router's NEW-wins invariant
// (#findRunsOpen). Without this, edge-waitpoint hydration could read a stale LEGACY status and
// strand the run. Rows whose projection omits `id` can't be deduped and pass through.
const legs = await this.#fanOut(
this.#precedence,
(store) =>
store.findManyWaitpoints(
scalarArgs as Prisma.WaitpointFindManyArgs,
RoutingRunStore.#ownPrimary(store, client)
) as Promise<Record<string, unknown>[]>
);
// A token mirrored onto both DBs during drain appears in BOTH legs; dedup by id in #precedence
// order, so the highest-authority copy wins. Without this, edge-waitpoint hydration could read a
// stale LEGACY status and strand the run. Rows whose projection omits `id` pass through.
const byId = new Map<string, Record<string, unknown>>();
const passthrough: Record<string, unknown>[] = [];
for (const w of [...fromLegacy, ...fromNew]) {
for (const w of legs.flat()) {
const id = w.id;
if (typeof id === "string") byId.set(id, w);
else passthrough.push(w);
@@ -1462,12 +1528,8 @@ export class RoutingRunStore implements RunStore {
const store = await this.#resolveWaitpointStore(id);
return store.updateManyWaitpoints(args, undefined);
}
// No single routable id (batch where): apply to both stores and sum.
const [fromNew, fromLegacy] = await Promise.all([
this.#new.updateManyWaitpoints(args),
this.#legacy.updateManyWaitpoints(args),
]);
return { count: fromNew.count + fromLegacy.count };
// No single routable id (batch where): apply to every store and sum.
return this.#sumCounts((store) => store.updateManyWaitpoints(args));
}
// Residency guard: selects the owning store by waitpointId.
@@ -1476,14 +1538,13 @@ export class RoutingRunStore implements RunStore {
context: ForWaitpointCompletionContext
): Promise<RunStore> {
// Preferred store: explicit legacy-authority pins first, else the waitpoint's id-shape.
const preferred =
const preferredKey =
context.treeOwnerResidency === "LEGACY" ||
context.isCrossTreeIdempotency === true ||
context.hasLegacyParent === true
? this.#legacy
: this.#classifySafe(waitpointId) === "NEW"
? this.#new
: this.#legacy;
? LEGACY_SHARD
: this.#shardKeyOfSafe(waitpointId);
const preferred = this.#shardStore(preferredKey);
// Resolve to where the waitpoint ACTUALLY lives: a migrated run's waitpoint can be on NEW
// with a LEGACY-classified id (or vice versa), so verify and fall back rather than route
// by id-shape alone and miss it (which leaves the blocked run stuck forever). This guard
@@ -1495,9 +1556,10 @@ export class RoutingRunStore implements RunStore {
) {
return preferred;
}
const other = preferred === this.#new ? this.#legacy : this.#new;
if (await other.findWaitpoint({ where: { id: waitpointId } }, other.primaryReadClient)) {
return other;
for (const { store } of this.#shardsExcept(preferredKey)) {
if (await store.findWaitpoint({ where: { id: waitpointId } }, store.primaryReadClient)) {
return store;
}
}
return preferred;
}
@@ -1534,17 +1596,13 @@ export class RoutingRunStore implements RunStore {
RoutingRunStore.#ownPrimary(store, client)
)) as Record<string, unknown>[];
} else {
const [fromNew, fromLegacy] = await Promise.all([
this.#new.findManyTaskRunWaitpoints(
const legs = await this.#fanOut(this.#precedence, (store) =>
store.findManyTaskRunWaitpoints(
scalarArgs as typeof args,
RoutingRunStore.#ownPrimary(this.#new, client)
),
this.#legacy.findManyTaskRunWaitpoints(
scalarArgs as typeof args,
RoutingRunStore.#ownPrimary(this.#legacy, client)
),
]);
edges = dedupeEdgesById([...fromNew, ...fromLegacy]) as Record<string, unknown>[];
RoutingRunStore.#ownPrimary(store, client)
)
);
edges = dedupeEdgesById(legs.flat()) as Record<string, unknown>[];
}
if (waitpoint) {
@@ -1623,11 +1681,7 @@ export class RoutingRunStore implements RunStore {
if (taskRunId !== undefined) {
return (await this.#routeOrNewForWrite(taskRunId)).deleteManyTaskRunWaitpoints(args);
}
const [fromNew, fromLegacy] = await Promise.all([
this.#new.deleteManyTaskRunWaitpoints(args),
this.#legacy.deleteManyTaskRunWaitpoints(args),
]);
return { count: fromNew.count + fromLegacy.count };
return this.#sumCounts((store) => store.deleteManyTaskRunWaitpoints(args));
}
findTaskRunAttempt<T extends Prisma.TaskRunAttemptFindFirstArgs>(
@@ -1650,14 +1704,9 @@ export class RoutingRunStore implements RunStore {
args: Prisma.SelectSubset<T, Prisma.TaskRunAttemptFindFirstArgs>,
client?: ReadClient
): Promise<Prisma.TaskRunAttemptGetPayload<T> | null> {
const fromNew = await this.#new.findTaskRunAttempt(
args,
RoutingRunStore.#ownPrimary(this.#new, client)
return this.#probeFirst((store) =>
store.findTaskRunAttempt(args, RoutingRunStore.#ownPrimary(store, client))
);
if (fromNew != null) {
return fromNew;
}
return this.#legacy.findTaskRunAttempt(args, RoutingRunStore.#ownPrimary(this.#legacy, client));
}
// Co-locate the checkpoint with its OWNING run so the run-routed snapshot's `checkpointId` FK
@@ -1714,16 +1763,8 @@ export class RoutingRunStore implements RunStore {
): Promise<Prisma.BatchTaskRunGetPayload<{ include: T }> | null> {
// Never forward the caller's client verbatim (a cross-DB probe with one shared client can
// only reach one DB); its presence resolves each leg to that store's OWN primary.
const fromNew = await this.#new.findBatchTaskRunById(
id,
args,
RoutingRunStore.#ownPrimary(this.#new, client)
);
if (fromNew != null) return fromNew;
return this.#legacy.findBatchTaskRunById(
id,
args,
RoutingRunStore.#ownPrimary(this.#legacy, client)
return this.#probeFirst((store) =>
store.findBatchTaskRunById(id, args, RoutingRunStore.#ownPrimary(store, client))
);
}
@@ -1736,18 +1777,13 @@ export class RoutingRunStore implements RunStore {
): Promise<Prisma.BatchTaskRunGetPayload<{ include: T }> | null> {
// Never forward the caller's client verbatim; its presence resolves each leg to that
// store's OWN primary.
const fromNew = await this.#new.findBatchTaskRunByFriendlyId(
friendlyId,
environmentId,
args,
RoutingRunStore.#ownPrimary(this.#new, client)
);
if (fromNew != null) return fromNew;
return this.#legacy.findBatchTaskRunByFriendlyId(
friendlyId,
environmentId,
args,
RoutingRunStore.#ownPrimary(this.#legacy, client)
return this.#probeFirst((store) =>
store.findBatchTaskRunByFriendlyId(
friendlyId,
environmentId,
args,
RoutingRunStore.#ownPrimary(store, client)
)
);
}
@@ -1766,18 +1802,13 @@ export class RoutingRunStore implements RunStore {
): Promise<Prisma.BatchTaskRunGetPayload<{ include: T }> | null> {
// Never forward the caller's client verbatim; its presence resolves each leg to that
// store's OWN primary.
const fromNew = await this.#new.findBatchTaskRunByIdempotencyKey(
environmentId,
idempotencyKey,
args,
RoutingRunStore.#ownPrimary(this.#new, client)
);
if (fromNew != null) return fromNew;
return this.#legacy.findBatchTaskRunByIdempotencyKey(
environmentId,
idempotencyKey,
args,
RoutingRunStore.#ownPrimary(this.#legacy, client)
return this.#probeFirst((store) =>
store.findBatchTaskRunByIdempotencyKey(
environmentId,
idempotencyKey,
args,
RoutingRunStore.#ownPrimary(store, client)
)
);
}
@@ -1791,11 +1822,7 @@ export class RoutingRunStore implements RunStore {
const store = this.#routeOrNew(id);
return store.updateManyBatchTaskRun(args, undefined);
}
const [fromNew, fromLegacy] = await Promise.all([
this.#new.updateManyBatchTaskRun(args),
this.#legacy.updateManyBatchTaskRun(args),
]);
return { count: fromNew.count + fromLegacy.count };
return this.#sumCounts((store) => store.updateManyBatchTaskRun(args));
}
// Items co-reside with their batch — route by `batchTaskRunId`, no fan-out.
@@ -1824,11 +1851,7 @@ export class RoutingRunStore implements RunStore {
const store = this.#routeOrNew(id);
return store.updateManyBatchTaskRunItems(args, undefined);
}
const [fromNew, fromLegacy] = await Promise.all([
this.#new.updateManyBatchTaskRunItems(args),
this.#legacy.updateManyBatchTaskRunItems(args),
]);
return { count: fromNew.count + fromLegacy.count };
return this.#sumCounts((store) => store.updateManyBatchTaskRunItems(args));
}
// An item co-resides with its batch AND its child run on one DB (both FKs local), so route by
@@ -1898,13 +1921,11 @@ export class RoutingRunStore implements RunStore {
skip: 0,
...(args.take != null ? { take: skip + args.take } : {}),
};
const [fromNew, fromLegacy] = await Promise.all([
this.#new.findManyWaitpointTags(perLeg, RoutingRunStore.#ownPrimary(this.#new, client)),
this.#legacy.findManyWaitpointTags(perLeg, RoutingRunStore.#ownPrimary(this.#legacy, client)),
]);
const legs = await this.#fanOut(this.#precedence, (store) =>
store.findManyWaitpointTags(perLeg, RoutingRunStore.#ownPrimary(store, client))
);
const byId = new Map<string, WaitpointTag>();
for (const tag of fromLegacy) byId.set(tag.id, tag);
for (const tag of fromNew) byId.set(tag.id, tag);
for (const tag of legs.flat()) byId.set(tag.id, tag);
const merged = args.orderBy
? (sortByOrderBy(
[...byId.values()] as unknown as Array<Record<string, unknown>>,
@@ -2034,8 +2055,9 @@ function narrowArgsToIds(args: Record<string, unknown>, ids: string[]): Record<s
};
}
// Merge edge rows from both stores, keeping one per edge `id` (NEW seen last wins). Rows whose
// projection omits `id` can't be deduped, so they pass through unchanged.
// Merge edge rows from every store, keeping one per edge `id`. The caller MUST supply the rows in
// #precedence order, so the highest-authority shard is seen last and wins. Rows whose projection
// omits `id` can't be deduped, so they pass through unchanged.
function dedupeEdgesById<R>(rows: R[]): R[] {
const byId = new Map<string, R>();
const passthrough: R[] = [];