Skip to content

Commit 90ea1e0

Browse files
d-csd-cs
committed
feat(webapp,run-store): resolve idempotency and seam reads through the shard map
Idempotency resolved its client with a binary `ownerEngine(id) === "NEW"` ternary, which cannot name a gen-2 shard. Both call sites now resolve through one shard-keyed map, so they cannot disagree about which store owns an id. An absent key takes an explicit logged branch to the fallback rather than a silent legacy default. `clientForShardKey` is the single place an id becomes a client. `ShardKey` collapses to `string`, so the compiler cannot catch a wrong key here; the `classify` seam is retyped to return a `ShardKey` so a `Residency` value ("NEW") can no longer be fed into a shard-key parameter, which the two differ from each other only by case. `resolveIdempotencyDedupClient` keeps its policy. The mint-kind branch has no id to decode and still resolves to the gen-1 pair. Delete the `isMigrated` branch. Nothing implements it, and the one production comment recorded that omitting it was deliberate. The two cross-seam batch hydration sites classified with the binary `ownerEngine` too. A gen-2 id joined the gen-1 `new` group, missed there, and — classifying dedicated-family — never reached the legacy probe either, so it was dropped from a bulk-action page and from batch results with no error. Both now partition ids by shard key and read each configured shard once. `PostgresRunStore._residency` widens to `ShardKey` for call-site consistency. It stays unused; the store stays unaware of its siblings. Inert while RUN_OPS_SHARDS is unset: the shard map holds only the two reserved gen-1 keys, so every partition falls through to today's paths.
1 parent 77cc7bb commit 90ea1e0

11 files changed

Lines changed: 402 additions & 80 deletions

‎apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts‎

Lines changed: 44 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import type { BatchTaskRunExecutionResult } from "@trigger.dev/core/v3";
2-
import { ownerEngine } from "@trigger.dev/core/v3/isomorphic";
2+
import { resolveShard, type ShardKey } from "@trigger.dev/core/v3/isomorphic";
33
import {
44
$replica,
55
type PrismaClientOrTransaction,
@@ -13,6 +13,7 @@ import { runStore as defaultRunStore } from "~/v3/runStore.server";
1313
import { BasePresenter } from "./basePresenter.server";
1414

1515
import { boundedIn } from "@trigger.dev/database";
16+
import { runOpsShardReplicas } from "~/v3/runOpsMigration/shardHandles.server";
1617
/**
1718
* Run-ops read-through wiring. All optional; absent (or `splitEnabled` falsy) collapses `call` to
1819
* passthrough. `legacyReplica` is a READ REPLICA handle only — there is NO legacy-primary field.
@@ -21,6 +22,8 @@ type ApiBatchResultsReadThroughDeps = {
2122
splitEnabled?: boolean;
2223
newClient?: PrismaReplicaClient;
2324
legacyReplica?: PrismaReplicaClient;
25+
/** Gen-2 shard replicas by shard char; empty (RUN_OPS_SHARDS unset) keeps today's behaviour. */
26+
shardReplicas?: ReadonlyMap<ShardKey, PrismaReplicaClient>;
2427
isPastRetention?: (runId: string) => boolean;
2528
};
2629

@@ -181,16 +184,48 @@ export class ApiBatchResultsPresenter extends BasePresenter {
181184

182185
const taskRunIds = batchRun.items.map((item) => item.taskRunId);
183186

184-
const newRows = (await newClient.taskRun.findMany({
185-
where: { id: { in: boundedIn(taskRunIds) } },
186-
select: memberRunSelect,
187-
})) as TaskRunWithAttempts[];
187+
// A gen-2 id is directly routable to its own shard, so it must not join the gen-1 read:
188+
// it would miss there, and (being dedicated-family) never reach the legacy probe either.
189+
const shardReplicas = this.readThrough?.shardReplicas ?? runOpsShardReplicas;
190+
const genOneIds: string[] = [];
191+
const idsByShard = new Map<ShardKey, string[]>();
192+
for (const id of taskRunIds) {
193+
const shardKey = resolveShard(id);
194+
if (shardKey !== "new" && shardKey !== "legacy" && shardReplicas.has(shardKey)) {
195+
const group = idsByShard.get(shardKey);
196+
group ? group.push(id) : idsByShard.set(shardKey, [id]);
197+
} else {
198+
genOneIds.push(id);
199+
}
200+
}
201+
202+
const newRows = (
203+
genOneIds.length > 0
204+
? ((await newClient.taskRun.findMany({
205+
where: { id: { in: boundedIn(genOneIds) } },
206+
select: memberRunSelect,
207+
})) as TaskRunWithAttempts[])
208+
: []
209+
).concat(
210+
(
211+
await Promise.all(
212+
[...idsByShard.entries()].map(
213+
async ([shardKey, ids]) =>
214+
(await shardReplicas.get(shardKey)!.taskRun.findMany({
215+
where: { id: { in: boundedIn(ids) } },
216+
select: memberRunSelect,
217+
})) as TaskRunWithAttempts[]
218+
)
219+
)
220+
).flat()
221+
);
188222
const runsById = new Map(newRows.map((run) => [run.id, run]));
189223

190-
// A run-ops id can only live on NEW, so only misses that AREN'T run-ops-shaped are candidates
191-
// for the legacy probe — mirrors readThroughRun's per-id "NEW residency skips legacy" rule.
192-
const legacyCandidateIds = taskRunIds.filter(
193-
(id) => !runsById.has(id) && ownerEngine(id) !== "NEW"
224+
// A dedicated-family id (gen-1 v1 or gen-2) can only live on its own store, so only
225+
// misses that AREN'T dedicated-shaped are candidates for the legacy probe — mirrors
226+
// readThroughRun's per-id "dedicated residency skips legacy" rule.
227+
const legacyCandidateIds = genOneIds.filter(
228+
(id) => !runsById.has(id) && resolveShard(id) === "legacy"
194229
);
195230
if (legacyCandidateIds.length > 0) {
196231
const legacyRows = (await legacyReplica.taskRun.findMany({

‎apps/webapp/app/runEngine/concerns/idempotencyKeys.server.ts‎

Lines changed: 27 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { ownerEngine, RunId } from "@trigger.dev/core/v3/isomorphic";
1+
import { resolveShard, RunId, type ShardKey } from "@trigger.dev/core/v3/isomorphic";
22
import type { PrismaClientOrTransaction, TaskRun, Waitpoint } from "@trigger.dev/database";
33
import { env } from "~/env.server";
44
import { logger } from "~/services/logger.server";
@@ -13,9 +13,10 @@ import { computeClaimTtlSeconds } from "~/v3/mollifier/claimTtl";
1313
import { makeResolveMollifierFlag } from "~/v3/mollifier/mollifierGate.server";
1414
import { runStore } from "~/v3/runStore.server";
1515
import { runOpsLegacyPrisma, runOpsNewPrisma } from "~/db.server";
16+
import { runOpsShardWriters } from "~/v3/runOpsMigration/shardHandles.server";
1617
import { isSplitEnabled } from "~/v3/runOpsMigration/splitMode.server";
1718
import { resolveRunIdMintKind } from "~/v3/engineVersion.server";
18-
import { resolveIdempotencyDedupClient } from "./idempotencyResidency.server";
19+
import { clientForShardKey, resolveIdempotencyDedupClient } from "./idempotencyResidency.server";
1920
import type { TraceEventConcern, TriggerTaskRequest } from "../types";
2021

2122
// In-memory per-org mollifier-enabled check, shared with `evaluateGate`
@@ -32,6 +33,19 @@ const resolveOrgMollifierFlag = makeResolveMollifierFlag();
3233
// PG's unique index as the backstop.
3334
const MAX_CLEARED_WINNER_REACQUIRES = 5;
3435

36+
// Every run-ops store keyed by shard key. Both idempotency call sites resolve through this
37+
// one map, so they cannot disagree about which store owns an id.
38+
const idempotencyShardClients: ReadonlyMap<ShardKey, PrismaClientOrTransaction> = new Map<
39+
ShardKey,
40+
PrismaClientOrTransaction
41+
>([
42+
["legacy", runOpsLegacyPrisma],
43+
["new", runOpsNewPrisma],
44+
...[...runOpsShardWriters.entries()].map(
45+
([key, writer]) => [key, writer as PrismaClientOrTransaction] as const
46+
),
47+
]);
48+
3549
// Claim ownership context returned to the caller when the
3650
// IdempotencyKeyConcern won a pre-gate claim. Caller MUST publish the
3751
// winning runId on pipeline success (`publishClaim`) or release the
@@ -172,12 +186,9 @@ export class IdempotencyKeyConcern {
172186
{
173187
isSplitEnabled,
174188
fallbackClient: this.prisma,
175-
newClient: runOpsNewPrisma,
176-
legacyClient: runOpsLegacyPrisma,
189+
clients: idempotencyShardClients,
177190
resolveMintKind: resolveRunIdMintKind,
178-
// `isMigrated` is intentionally omitted: until a child of a swept
179-
// legacy-id parent can be born on the new DB, the swept-marker override
180-
// would never change the answer, so a child routes by parent id-shape.
191+
logger,
181192
}
182193
);
183194

@@ -640,12 +651,15 @@ export class IdempotencyKeyConcern {
640651
} catch {
641652
return null;
642653
}
643-
let client: PrismaClientOrTransaction;
644-
try {
645-
client = ownerEngine(internalId) === "NEW" ? runOpsNewPrisma : runOpsLegacyPrisma;
646-
} catch {
647-
client = this.prisma;
648-
}
654+
// The routing store routes by id and never forwards this object, so its identity only
655+
// signals read-your-writes. Resolving it through the shard map keeps the two idempotency
656+
// call sites in agreement and stops this reading as gen-2-unaware.
657+
const client = clientForShardKey(
658+
resolveShard(internalId),
659+
idempotencyShardClients,
660+
this.prisma,
661+
logger
662+
);
649663
return runStore.findRun(
650664
{ id: internalId, runtimeEnvironmentId: environmentId },
651665
{ include: { associatedWaitpoint: true } },

‎apps/webapp/app/runEngine/concerns/idempotencyResidency.server.test.ts‎

Lines changed: 49 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import { describe, expect, it } from "vitest";
22
import { RunId } from "@trigger.dev/core/v3/isomorphic";
33
import {
4+
clientForShardKey,
45
resolveIdempotencyDedupClient,
56
type ResolveIdempotencyClientDeps,
67
} from "./idempotencyResidency.server";
@@ -9,20 +10,30 @@ import {
910
const FALLBACK = { __tag: "fallback" } as never;
1011
const NEW_CLIENT = { __tag: "new" } as never;
1112
const LEGACY_CLIENT = { __tag: "legacy" } as never;
13+
const SHARD_A_CLIENT = { __tag: "shard-a" } as never;
14+
15+
function clientMap() {
16+
return new Map([
17+
["new", NEW_CLIENT],
18+
["legacy", LEGACY_CLIENT],
19+
["a", SHARD_A_CLIENT],
20+
]);
21+
}
1222

1323
function makeDeps(over: Partial<ResolveIdempotencyClientDeps>): ResolveIdempotencyClientDeps {
1424
return {
1525
isSplitEnabled: async () => true,
1626
fallbackClient: FALLBACK,
17-
newClient: NEW_CLIENT,
18-
legacyClient: LEGACY_CLIENT,
27+
clients: clientMap(),
1928
resolveMintKind: async () => "runOpsId",
29+
// Kept as an injected seam: the real resolveShard is total, so only an injected
30+
// classifier can exercise the throw-to-fallback arm below.
2031
classify: (id) => {
21-
if (id.length === 26 && id[25] === "1") return "NEW";
22-
if (id.length === 25) return "LEGACY";
32+
if (id.length === 26 && id[25] === "2") return id[24]!;
33+
if (id.length === 26 && id[25] === "1") return "new";
34+
if (id.length === 25) return "legacy";
2335
throw new Error(`unclassifiable: ${id.length}`);
2436
},
25-
isMigrated: undefined,
2637
...over,
2738
};
2839
}
@@ -72,29 +83,49 @@ describe("resolveIdempotencyDedupClient", () => {
7283
expect(client).toBe(LEGACY_CLIENT);
7384
});
7485

75-
it("routes a swept (migrated) cuid-parent child to the NEW client", async () => {
76-
const cuidParent = RunId.toFriendlyId("c".repeat(25));
86+
it("falls back to the fallback client when a present parent id is unclassifiable", async () => {
7787
const client = await resolveIdempotencyDedupClient(
78-
{ environmentForMint: env, parentRunFriendlyId: cuidParent },
79-
makeDeps({ isMigrated: async () => true })
88+
{ environmentForMint: env, parentRunFriendlyId: "run_not-a-valid-length" },
89+
makeDeps({})
8090
);
81-
expect(client).toBe(NEW_CLIENT);
91+
expect(client).toBe(FALLBACK);
8292
});
8393

84-
it("routes a non-migrated cuid-parent child to the LEGACY client even when isMigrated is provided", async () => {
85-
const cuidParent = RunId.toFriendlyId("d".repeat(25));
94+
it("routes a child to its OWN SHARD client when the parent is a gen-2 id", async () => {
95+
const genTwoParent = RunId.toFriendlyId("e".repeat(24) + "a2");
8696
const client = await resolveIdempotencyDedupClient(
87-
{ environmentForMint: env, parentRunFriendlyId: cuidParent },
88-
makeDeps({ isMigrated: async () => false })
97+
{ environmentForMint: env, parentRunFriendlyId: genTwoParent },
98+
makeDeps({ resolveMintKind: async () => "cuid" }) // mint flag must NOT win for a child
8999
);
90-
expect(client).toBe(LEGACY_CLIENT);
100+
expect(client).toBe(SHARD_A_CLIENT);
91101
});
92102

93-
it("falls back to the fallback client when a present parent id is unclassifiable", async () => {
103+
it("falls back and logs when a gen-2 parent names an unconfigured shard key", async () => {
104+
const errors: unknown[] = [];
105+
const genTwoParent = RunId.toFriendlyId("f".repeat(24) + "z2");
94106
const client = await resolveIdempotencyDedupClient(
95-
{ environmentForMint: env, parentRunFriendlyId: "run_not-a-valid-length" },
96-
makeDeps({})
107+
{ environmentForMint: env, parentRunFriendlyId: genTwoParent },
108+
makeDeps({ logger: { error: (_m, meta) => errors.push(meta) } })
97109
);
98110
expect(client).toBe(FALLBACK);
111+
expect(errors).toHaveLength(1);
112+
});
113+
});
114+
115+
describe("clientForShardKey", () => {
116+
it("selects the same client the map holds for each reserved key and shard key", () => {
117+
const clients = clientMap();
118+
expect(clientForShardKey("new", clients, FALLBACK)).toBe(NEW_CLIENT);
119+
expect(clientForShardKey("legacy", clients, FALLBACK)).toBe(LEGACY_CLIENT);
120+
expect(clientForShardKey("a", clients, FALLBACK)).toBe(SHARD_A_CLIENT);
121+
});
122+
123+
it("returns the fallback and logs for a key the map does not hold", () => {
124+
const errors: unknown[] = [];
125+
const client = clientForShardKey("z", clientMap(), FALLBACK, {
126+
error: (_m, meta) => errors.push(meta),
127+
});
128+
expect(client).toBe(FALLBACK);
129+
expect(errors).toHaveLength(1);
99130
});
100131
});

‎apps/webapp/app/runEngine/concerns/idempotencyResidency.server.ts‎

Lines changed: 39 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,46 @@
1-
import { ownerEngine, RunId, type Residency } from "@trigger.dev/core/v3/isomorphic";
1+
import { resolveShard, RunId, type ShardKey } from "@trigger.dev/core/v3/isomorphic";
22
import type { PrismaClientOrTransaction } from "@trigger.dev/database";
33

44
type MintKind = "cuid" | "runOpsId";
55

6+
type Logger = { error: (message: string, meta?: Record<string, unknown>) => void };
7+
68
export type ResolveIdempotencyClientDeps = {
79
isSplitEnabled: () => Promise<boolean>;
810
fallbackClient: PrismaClientOrTransaction;
9-
newClient: PrismaClientOrTransaction;
10-
legacyClient: PrismaClientOrTransaction;
11+
/** Every store keyed by shard key: the reserved `legacy`/`new` plus one entry per gen-2 shard. */
12+
clients: ReadonlyMap<ShardKey, PrismaClientOrTransaction>;
1113
resolveMintKind: (environment: {
1214
organizationId: string;
1315
id: string;
1416
orgFeatureFlags?: unknown;
1517
}) => Promise<MintKind>;
16-
classify?: (id: string) => Residency;
17-
isMigrated?: (id: string) => Promise<boolean>;
18+
classify?: (id: string) => ShardKey;
19+
logger?: Logger;
1820
};
1921

22+
/**
23+
* The one place an id becomes a client. `ShardKey` collapses to `string`, so the compiler
24+
* cannot catch a wrong key here — an absent key takes an explicit logged branch to the
25+
* fallback rather than a silent `?? legacy`.
26+
*/
27+
export function clientForShardKey(
28+
shardKey: ShardKey,
29+
clients: ReadonlyMap<ShardKey, PrismaClientOrTransaction>,
30+
fallback: PrismaClientOrTransaction,
31+
logger?: Logger
32+
): PrismaClientOrTransaction {
33+
const client = clients.get(shardKey);
34+
if (client === undefined) {
35+
logger?.error("idempotency: no client configured for shard key", {
36+
shardKey,
37+
configured: [...clients.keys()],
38+
});
39+
return fallback;
40+
}
41+
return client;
42+
}
43+
2044
export async function resolveIdempotencyDedupClient(
2145
args: {
2246
environmentForMint: { organizationId: string; id: string; orgFeatureFlags?: unknown };
@@ -28,9 +52,9 @@ export async function resolveIdempotencyDedupClient(
2852
return deps.fallbackClient;
2953
}
3054

31-
const classify = deps.classify ?? ownerEngine;
32-
const clientFor = (residency: Residency): PrismaClientOrTransaction =>
33-
residency === "NEW" ? deps.newClient : deps.legacyClient;
55+
const classify = deps.classify ?? resolveShard;
56+
const clientFor = (shardKey: ShardKey): PrismaClientOrTransaction =>
57+
clientForShardKey(shardKey, deps.clients, deps.fallbackClient, deps.logger);
3458

3559
if (args.parentRunFriendlyId) {
3660
let parentInternalId: string;
@@ -39,18 +63,18 @@ export async function resolveIdempotencyDedupClient(
3963
} catch {
4064
return deps.fallbackClient;
4165
}
42-
let residency: Residency;
66+
let shardKey: ShardKey;
4367
try {
44-
residency = classify(parentInternalId);
68+
shardKey = classify(parentInternalId);
4569
} catch {
4670
return deps.fallbackClient;
4771
}
48-
if (residency === "LEGACY" && deps.isMigrated && (await deps.isMigrated(parentInternalId))) {
49-
return deps.newClient;
50-
}
51-
return clientFor(residency);
72+
return clientFor(shardKey);
5273
}
5374

75+
// Mint kind, not an id: there is no shard to decode, so this keeps resolving to the
76+
// gen-1 pair exactly as before. Which shard a gen-2 env mints into is the mint layer's
77+
// decision, and this client is a read-your-writes signal rather than a correctness gate.
5478
const kind = await deps.resolveMintKind(args.environmentForMint);
55-
return clientFor(kind === "runOpsId" ? "NEW" : "LEGACY");
79+
return clientFor(kind === "runOpsId" ? "new" : "legacy");
5680
}

‎apps/webapp/app/v3/runOpsMigration/readThrough.server.test.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,9 @@ describe("readThroughRun (legacy replica + new DB)", () => {
8888
},
8989
});
9090

91-
expect(pastRetentionResult.found === false && pastRetentionResult.reason).toBe("past-retention");
91+
expect(pastRetentionResult.found === false && pastRetentionResult.reason).toBe(
92+
"past-retention"
93+
);
9294

9395
// A run that is simply absent (not past retention) yields not-found.
9496
const notFoundResult = await readThroughRun({

‎apps/webapp/app/v3/runOpsMigration/readThrough.server.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ type ReadThroughDeps = {
5353
/** Resolved boot constant; never `await`ed per-request when supplied. */
5454
splitEnabled?: boolean;
5555
isPastRetention?: (id: string) => boolean;
56-
logger?: { error: (m: string, meta?: unknown) => void };
56+
logger?: { error: (m: string, meta?: Record<string, unknown>) => void };
5757
/** Saturation-signal emit hook: called on each legacy-replica hit. */
5858
onLegacyReplicaRead?: (id: string) => void;
5959
};

‎apps/webapp/app/v3/runOpsMigration/shardHandles.server.ts‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,17 @@ export function buildShardHandleMaps(handles: ShardHandle[]): {
3030

3131
// A gen-2 shard is the same dedicated subset schema as the gen-1 new store, so these casts
3232
// carry exactly the precedent (and the same residual risk) as `runOpsNewPrisma`'s.
33-
const maps = buildShardHandleMaps(runOpsShardHandles ?? []);
33+
// The try/catch mirrors `runStore.server.ts`'s handle resolution: a minimal `db.server` mock
34+
// does not define this export at all, and accessing an undefined mock export throws.
35+
function resolveShardHandles(): ShardHandle[] {
36+
try {
37+
return runOpsShardHandles ?? [];
38+
} catch {
39+
return [];
40+
}
41+
}
42+
43+
const maps = buildShardHandleMaps(resolveShardHandles());
3444

3545
export const runOpsShardReplicas = maps.replicas;
3646
export const runOpsShardWriters = maps.writers;

0 commit comments

Comments
 (0)