diff --git a/apps/webapp/app/v3/runStore.server.ts b/apps/webapp/app/v3/runStore.server.ts index 9ccf84b5117..15e860e65c3 100644 --- a/apps/webapp/app/v3/runStore.server.ts +++ b/apps/webapp/app/v3/runStore.server.ts @@ -1,5 +1,12 @@ -import { PostgresRunStore, RoutingRunStore, type RunStore } from "@internal/run-store"; -import { ownerEngine, type Residency } from "@trigger.dev/core/v3/isomorphic"; +import { + PostgresRunStore, + RoutingRunStore, + type RoutingStoreMetrics, + type RunStore, +} from "@internal/run-store"; +import { resolveShard, type ShardKey } from "@trigger.dev/core/v3/isomorphic"; +import { Counter } from "prom-client"; +import { metricsRegister } from "~/metrics.server"; import type { PrismaClient, PrismaReplicaClient } from "@trigger.dev/database"; import type { RunOpsPrismaClient } from "@internal/run-ops-database"; import { @@ -29,8 +36,8 @@ type BuildRunStoreDeps = { /** Single-DB store handles (control-plane pair). Used verbatim when split is OFF. */ singleWriter: PrismaClient; singleReplica: PrismaReplicaClient; - /** Residency classifier; defaults to ownerEngine inside RoutingRunStore. */ - classify?: (id: string) => Residency; + /** Id-to-shard-key resolver; defaults to the core resolveShard inside RoutingRunStore. */ + resolveShard?: (id: string) => ShardKey; /** Per-pool transaction-resilience configs threaded into the store(s) this builds (IoC). */ singleResilience?: TransactionResilienceConfig; newResilience?: TransactionResilienceConfig; @@ -82,10 +89,31 @@ export function buildRunStore(deps: BuildRunStoreDeps): RunStore { return new RoutingRunStore({ new: newStore, legacy: legacyStore, - classify: deps.classify ?? ownerEngine, + resolveShard: deps.resolveShard ?? resolveShard, + metrics: routingStoreMetrics, }); } +// singleton: module-scope Counter registration double-registers under dev HMR. +const routingStoreMetrics: RoutingStoreMetrics = singleton("routingStoreMetrics", () => { + const duplicateId = new Counter({ + name: "runops_shard_duplicate_id_total", + help: "One id was returned by two run-ops shards that must be disjoint (a routing-invariant violation).", + labelNames: ["shard_keys"], + registers: [metricsRegister], + }); + const probeFallback = new Counter({ + name: "runops_waitpoint_probe_fallback_total", + help: "A waitpoint was not on the run-ops store its id named and was found by a fallback probe.", + labelNames: ["from", "to"], + registers: [metricsRegister], + }); + return { + recordDuplicateId: (shardKeys) => duplicateId.inc({ shard_keys: shardKeys.join(",") }), + recordWaitpointProbeFallback: (from, to) => probeFallback.inc({ from, to }), + }; +}); + // Build the routing store whenever BOTH run-ops DBs are configured, independent of // RUN_OPS_SPLIT_ENABLED. Reads must fan out across both DBs so a run that lives on the new // DB stays visible even with the flag off (matches the db.server topology factory). The flag diff --git a/internal-packages/run-store/src/PostgresRunStore.ts b/internal-packages/run-store/src/PostgresRunStore.ts index b7d5086431b..33604d01148 100644 --- a/internal-packages/run-store/src/PostgresRunStore.ts +++ b/internal-packages/run-store/src/PostgresRunStore.ts @@ -1826,7 +1826,7 @@ export class PostgresRunStore implements RunStore { const branches = args.idempotencyKeys.map((key) => { const base = params.length; params.push(args.runtimeEnvironmentId, args.taskIdentifier, key); - return `SELECT "friendlyId", "idempotencyKey", "idempotencyKeyExpiresAt" FROM "TaskRun" WHERE "runtimeEnvironmentId" = $${base + 1} AND "taskIdentifier" = $${base + 2} AND "idempotencyKey" = $${base + 3}`; + return `SELECT "id", "createdAt", "friendlyId", "idempotencyKey", "idempotencyKeyExpiresAt" FROM "TaskRun" WHERE "runtimeEnvironmentId" = $${base + 1} AND "taskIdentifier" = $${base + 2} AND "idempotencyKey" = $${base + 3}`; }); return prisma.$queryRawUnsafe( branches.join(" UNION ALL "), diff --git a/internal-packages/run-store/src/index.ts b/internal-packages/run-store/src/index.ts index 8893975cf12..a4109ab0104 100644 --- a/internal-packages/run-store/src/index.ts +++ b/internal-packages/run-store/src/index.ts @@ -3,4 +3,5 @@ export * from "./PostgresRunStore.js"; export * from "./runOpsStore.js"; export * from "./readReplicaClient.js"; export * from "./redisSnapshotStore.js"; +export * from "./routingStoreMetrics.js"; export * from "./snapshotComparator.js"; diff --git a/internal-packages/run-store/src/routingStoreMetrics.ts b/internal-packages/run-store/src/routingStoreMetrics.ts new file mode 100644 index 00000000000..d4aa454d903 --- /dev/null +++ b/internal-packages/run-store/src/routingStoreMetrics.ts @@ -0,0 +1,16 @@ +/** + * Counters the routing store emits. Injected the same way RedisSnapshotStore takes its metrics, + * so the package stays free of a metrics dependency and a test can assert on a fake. + * + * runops_shard_duplicate_id_total — one id returned by two shards that should be disjoint + * runops_waitpoint_probe_fallback_total — a waitpoint was not on the store its id named + */ +export type RoutingStoreMetrics = { + recordDuplicateId(shardKeys: string[]): void; + recordWaitpointProbeFallback(from: string, to: string): void; +}; + +export const noopRoutingStoreMetrics: RoutingStoreMetrics = { + recordDuplicateId() {}, + recordWaitpointProbeFallback() {}, +}; diff --git a/internal-packages/run-store/src/runOpsStore.nShardMatrix.test.ts b/internal-packages/run-store/src/runOpsStore.nShardMatrix.test.ts new file mode 100644 index 00000000000..3f42e36171a --- /dev/null +++ b/internal-packages/run-store/src/runOpsStore.nShardMatrix.test.ts @@ -0,0 +1,305 @@ +// FOUR-STORE MATRIX — proves RoutingRunStore is correct across legacy + new + two gen-2 shards +// (a, b) against REAL databases (makeNShardRunOpsPostgresTest). NEVER mocked. This is where the +// §3.4 disjoint-sum fix is proven end-to-end: a double count here strands a blocked run forever. +// +// runOpsStore.mixedResidency.test.ts is the TWO-store invariant lock and stays byte-identical; this +// file is the N-store extension and lives separately. + +import { makeNShardRunOpsPostgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import type { RunOpsPrismaClient } from "@internal/run-ops-database"; +import { resolveShard } from "@trigger.dev/core/v3/isomorphic"; +import { describe, expect } from "vitest"; +import { PostgresRunStore } from "./PostgresRunStore.js"; +import { RoutingRunStore } from "./runOpsStore.js"; +import type { CreateRunInput } from "./types.js"; + +type AnyClient = PrismaClient | RunOpsPrismaClient; + +const matrixTest = makeNShardRunOpsPostgresTest(2); + +// A gen-2 id: 24-char base32hex core, the shard char at index 24, version "2" at index 25. +// resolveShard(gen2("a", ...)) === "a". A bare cuid-length id classifies "legacy". +function gen2(shardChar: string, seed: string): string { + const core = (seed.replace(/[^0-9a-v]/g, "0") + "k".repeat(24)).slice(0, 24); + return `${core}${shardChar}2`; +} +function cuid(seed: string): string { + return (seed + "c".repeat(25)).slice(0, 25); // 25 chars → LEGACY +} + +function makeStore(prisma: AnyClient, variant: "legacy" | "dedicated") { + return new PostgresRunStore({ + prisma: prisma as never, + readOnlyPrisma: prisma as never, + schemaVariant: variant, + }); +} + +// The real four-store split: legacy (full schema) + new + gen-2 a + gen-2 b (dedicated subset), +// routed by the REAL core resolveShard. +function makeMatrixRouter( + legacyPrisma: PrismaClient, + newPrisma: RunOpsPrismaClient, + shardPrismas: RunOpsPrismaClient[] +) { + return new RoutingRunStore({ + new: makeStore(newPrisma, "dedicated"), + legacy: makeStore(legacyPrisma, "legacy"), + shards: [ + { key: "a", store: makeStore(shardPrismas[0]!, "dedicated") }, + { key: "b", store: makeStore(shardPrismas[1]!, "dedicated") }, + ], + resolveShard, + }); +} + +async function seedLegacyEnv(prisma: PrismaClient, suffix: string) { + const organization = await prisma.organization.create({ + data: { title: `Org ${suffix}`, slug: `org-${suffix}` }, + }); + const project = await prisma.project.create({ + data: { + name: `Project ${suffix}`, + slug: `project-${suffix}`, + externalRef: `proj_${suffix}`, + organizationId: organization.id, + }, + }); + const environment = await prisma.runtimeEnvironment.create({ + data: { + type: "DEVELOPMENT", + slug: "dev", + projectId: project.id, + organizationId: organization.id, + apiKey: `tr_dev_${suffix}`, + pkApiKey: `pk_dev_${suffix}`, + shortcode: `short_${suffix}`, + }, + }); + return { + organizationId: organization.id, + projectId: project.id, + runtimeEnvironmentId: environment.id, + environmentId: environment.id, + }; +} + +function buildRun(params: { + runId: string; + runtimeEnvironmentId: string; + organizationId: string; + projectId: string; + createdAt?: Date; +}): CreateRunInput { + return { + data: { + id: params.runId, + engine: "V2", + status: "PENDING", + friendlyId: `run_${params.runId}`, + runtimeEnvironmentId: params.runtimeEnvironmentId, + environmentType: "DEVELOPMENT", + organizationId: params.organizationId, + projectId: params.projectId, + taskIdentifier: "my-task", + payload: "{}", + payloadType: "application/json", + context: {}, + traceContext: {}, + traceId: `trace_${params.runId}`, + spanId: `span_${params.runId}`, + runTags: [], + queue: "task/my-task", + isTest: false, + taskEventStore: "taskEvent", + depth: 0, + createdAt: params.createdAt ?? new Date("2024-01-01T00:00:00.000Z"), + }, + snapshot: { + engine: "V2", + executionStatus: "RUN_CREATED", + description: "Run was created", + runStatus: "PENDING", + environmentId: params.runtimeEnvironmentId, + environmentType: "DEVELOPMENT", + projectId: params.projectId, + organizationId: params.organizationId, + }, + }; +} + +async function seedPendingWaitpoint( + prisma: AnyClient, + params: { id: string; projectId: string; environmentId: string } +) { + await (prisma as PrismaClient).waitpoint.create({ + data: { + id: params.id, + friendlyId: `wp_${params.id}`, + type: "MANUAL", + status: "PENDING", + idempotencyKey: `idem_${params.id}`, + userProvidedIdempotencyKey: false, + projectId: params.projectId, + environmentId: params.environmentId, + }, + }); +} + +describe("RoutingRunStore four-store matrix — disjoint sum on real databases", () => { + matrixTest( + "countPendingWaitpoints unions across a gen-2 shard and the gen-1 pair with no double count", + async ({ legacyPrisma, newPrisma, shardPrismas }) => { + const env = await seedLegacyEnv(legacyPrisma, "disjoint"); + const router = makeMatrixRouter(legacyPrisma, newPrisma, shardPrismas); + + // The blocked run lives on shard a. + const runId = gen2("a", "run"); + // Its blocking waitpoints: a gen-2 waitpoint on shard b, a cuid on legacy, and a cuid MIRRORED + // onto both gen-1 stores (the drain-mirror case that must count once). + const wpB = gen2("b", "wpb"); + const wpCuid = cuid("wpcuid"); + const wpMirror = cuid("wpmirror"); + + const dedicatedEnv = { projectId: env.projectId, environmentId: env.environmentId }; + await seedPendingWaitpoint(shardPrismas[1]!, { id: wpB, ...dedicatedEnv }); + await seedPendingWaitpoint(legacyPrisma, { id: wpCuid, ...dedicatedEnv }); + await seedPendingWaitpoint(legacyPrisma, { id: wpMirror, ...dedicatedEnv }); + await seedPendingWaitpoint(newPrisma, { id: wpMirror, ...dedicatedEnv }); + + // b:wpB (1) + legacy:wpCuid (1) + wpMirror (once, though on both gen-1 stores) = 3. + const count = await router.countPendingWaitpoints([wpB, wpCuid, wpMirror], undefined, runId); + expect(count).toBe(3); + } + ); + + matrixTest( + "a gen-2 waitpoint on the run's own shard contributes exactly once, not twice", + async ({ legacyPrisma, newPrisma, shardPrismas }) => { + const env = await seedLegacyEnv(legacyPrisma, "ownshard"); + const router = makeMatrixRouter(legacyPrisma, newPrisma, shardPrismas); + const runId = gen2("a", "run2"); + const wpA = gen2("a", "wpa"); + await seedPendingWaitpoint(shardPrismas[0]!, { + id: wpA, + projectId: env.projectId, + environmentId: env.environmentId, + }); + // wpA lives on the run's own shard a → found by the presence query, never re-queried elsewhere. + expect(await router.countPendingWaitpoints([wpA], undefined, runId)).toBe(1); + } + ); +}); + +describe("RoutingRunStore four-store matrix — alias topology", () => { + matrixTest( + "an aliased gen-2 shard counts its database ONCE in a sum (declaration, not identity)", + async ({ legacyPrisma, newPrisma, shardPrismas }) => { + const env = await seedLegacyEnv(legacyPrisma, "alias"); + // Shard "a" aliases "new" over the SAME database, but via a SEPARATE store object built over + // the same client — exactly how the wiring layer will construct it. Identity dedupe would see + // two objects and double-count; declaration dedupe counts the database once. + const newStore = makeStore(newPrisma, "dedicated"); + const aStoreSameDb = makeStore(newPrisma, "dedicated"); // distinct object, same DB + const router = new RoutingRunStore({ + new: newStore, + legacy: makeStore(legacyPrisma, "legacy"), + shards: [{ key: "a", store: aStoreSameDb, aliasOf: "new" }], + resolveShard, + }); + + const wp = cuid("aliaswp"); + await seedPendingWaitpoint(newPrisma, { + id: wp, + projectId: env.projectId, + environmentId: env.environmentId, + }); + // No runId → the id-less sum fans over DISTINCT databases. The aliased "a" must not add a + // second leg over the "new" database, or the one pending waitpoint counts twice. + expect(await router.countPendingWaitpoints([wp])).toBe(1); + } + ); +}); + +describe("RoutingRunStore four-store matrix — mixed gen-1 and gen-2 reads", () => { + matrixTest( + "findRunsByIds hydrates a mixed id set across legacy, new and both gen-2 shards", + async ({ legacyPrisma, newPrisma, shardPrismas }) => { + const env = await seedLegacyEnv(legacyPrisma, "mixed"); + const router = makeMatrixRouter(legacyPrisma, newPrisma, shardPrismas); + + const legacyId = cuid("mixleg"); + // A v1 run-ops id (version "1") routes to "new"; gen-2 ids route to their shard char. + const newId = ("mixnew".replace(/[^0-9a-v]/g, "0") + "k".repeat(24)).slice(0, 24) + "01"; + const aId = gen2("a", "mixa"); + const bId = gen2("b", "mixb"); + const all = [legacyId, newId, aId, bId]; + + for (const runId of all) { + await router.createRun(buildRun({ runId, ...env })); + } + + const found = await router.findRunsByIds(all, { select: { id: true } }); + expect(new Set([...found.keys()])).toEqual(new Set(all)); + } + ); +}); + +describe("RoutingRunStore four-store matrix — cross-tree completion across gen-2 shards", () => { + matrixTest( + "a gen-2 waitpoint completes on its own shard even under the cross-tree legacy pin", + async ({ legacyPrisma, newPrisma, shardPrismas }) => { + const env = await seedLegacyEnv(legacyPrisma, "crosstree"); + const router = makeMatrixRouter(legacyPrisma, newPrisma, shardPrismas); + // The waitpoint is owned by a run on shard b; the blocked run is on shard a (cross-tree). + const wpB = gen2("b", "ctwp"); + await seedPendingWaitpoint(shardPrismas[1]!, { + id: wpB, + projectId: env.projectId, + environmentId: env.environmentId, + }); + // isCrossTreeIdempotency pins gen-1 flows to legacy; a gen-2 id must OVERRIDE that pin, or the + // completion write lands on legacy, matches zero rows, and strands the run. + const store = await router.forWaitpointCompletion(wpB, { + isCrossTreeIdempotency: true, + } as never); + // The returned store finds wpB on its primary — only shard b holds it, so the override worked. + const found = await store.findWaitpoint({ where: { id: wpB } }, store.primaryReadClient); + expect(found?.id).toBe(wpB); + } + ); +}); + +describe("RoutingRunStore four-store matrix — pagination merge", () => { + matrixTest( + "findRuns merges an open-predicate page across all four stores in orderBy order", + async ({ legacyPrisma, newPrisma, shardPrismas }) => { + const env = await seedLegacyEnv(legacyPrisma, "paginate"); + const router = makeMatrixRouter(legacyPrisma, newPrisma, shardPrismas); + // One run per store, distinct createdAt so the global sort order is unambiguous. + const rows = [ + { id: cuid("pgleg"), at: new Date("2024-01-01T00:00:00Z") }, + { + id: ("pgnew".replace(/[^0-9a-v]/g, "0") + "k".repeat(24)).slice(0, 24) + "01", + at: new Date("2024-01-02T00:00:00Z"), + }, + { id: gen2("a", "pga"), at: new Date("2024-01-03T00:00:00Z") }, + { id: gen2("b", "pgb"), at: new Date("2024-01-04T00:00:00Z") }, + ]; + for (const r of rows) { + await router.createRun(buildRun({ runId: r.id, ...env, createdAt: r.at })); + } + // Open predicate (no id set) → fan out + merge; take 2 skip 1 over createdAt desc. + const page = (await router.findRuns({ + where: { runtimeEnvironmentId: env.runtimeEnvironmentId }, + select: { id: true, createdAt: true }, + orderBy: { createdAt: "desc" }, + take: 2, + skip: 1, + })) as Array<{ id: string }>; + // Global desc order is b, a, new, legacy; skip 1 take 2 → [a, new]. + expect(page.map((r) => r.id)).toEqual([rows[2]!.id, rows[1]!.id]); + } + ); +}); diff --git a/internal-packages/run-store/src/runOpsStore.runKeyedRouting.test.ts b/internal-packages/run-store/src/runOpsStore.runKeyedRouting.test.ts index e29ecb11cb8..e5cb229ef06 100644 --- a/internal-packages/run-store/src/runOpsStore.runKeyedRouting.test.ts +++ b/internal-packages/run-store/src/runOpsStore.runKeyedRouting.test.ts @@ -319,8 +319,9 @@ describe("RoutingRunStore.countPendingWaitpoints — route by runId then partiti "legacy_run" ); expect(count).toBe(1); - // Fallback queried the other store with ONLY the id missing on the run's store. - expect(newStore.calls.map((c) => c.method)).toEqual(["countPendingWaitpoints"]); + // Fallback queried the other store with ONLY the id missing on the run's store. It uses the + // presence variant so the results can be unioned by id (a drain mirror counts once at N). + expect(newStore.calls.map((c) => c.method)).toEqual(["countPendingWaitpointsWithPresence"]); expect(newStore.calls[0]?.args[0]).toEqual(["waitpoint_crosstree"]); }); diff --git a/internal-packages/run-store/src/runOpsStore.shardMap.test.ts b/internal-packages/run-store/src/runOpsStore.shardMap.test.ts index c25574e7561..b6ea71e9f60 100644 --- a/internal-packages/run-store/src/runOpsStore.shardMap.test.ts +++ b/internal-packages/run-store/src/runOpsStore.shardMap.test.ts @@ -10,7 +10,7 @@ import type { ReadClient, RunStore } from "./types.js"; // 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 Slot = string; type Call = { slot: Slot; method: string }; @@ -19,8 +19,14 @@ type FakeConfig = { runs?: Array>; // Edge rows this store returns from findManyTaskRunWaitpoints, regardless of filter. edges?: Array>; + // Batch row this store returns from findBatchTaskRunById, regardless of filter. + batch?: Record | null; // Waitpoint rows this store returns from findWaitpoint, regardless of filter. waitpoint?: Record | null; + // Rows this store returns from findRunsByIdempotencyKeys, regardless of filter. + idempotencyMatches?: Array>; + // Waitpoint ids this store reports as pending (present = pending here) for count/collect probes. + pendingWaitpointIds?: string[]; }; type FakeStore = RunStore & { @@ -79,10 +85,42 @@ function fakeStore(slot: Slot, log: Call[], config: FakeConfig = {}): FakeStore return Promise.resolve({ slot } as never); }) as FakeStore["updateWaitpoint"], + createWaitpoint: ((_args: unknown) => { + record("createWaitpoint"); + return Promise.resolve({ slot } as never); + }) as FakeStore["createWaitpoint"], + + runInTransaction: ((_runId: unknown, fn: (store: unknown, tx: unknown) => unknown) => { + record("runInTransaction"); + return Promise.resolve(fn(store, {})); + }) as FakeStore["runInTransaction"], + findManyTaskRunWaitpoints: ((_args: unknown, _client?: ReadClient) => { record("findManyTaskRunWaitpoints"); return Promise.resolve((config.edges ?? []) as never); }) as FakeStore["findManyTaskRunWaitpoints"], + + updateManyWaitpoints: ((_args: unknown) => { + record("updateManyWaitpoints"); + return Promise.resolve({ count: 1 } as never); + }) as FakeStore["updateManyWaitpoints"], + + findRunsByIdempotencyKeys: ((_args: unknown, _client?: ReadClient) => { + record("findRunsByIdempotencyKeys"); + return Promise.resolve((config.idempotencyMatches ?? []) as never); + }) as FakeStore["findRunsByIdempotencyKeys"], + + findBatchTaskRunById: ((_id: unknown, _args?: unknown, _client?: ReadClient) => { + record("findBatchTaskRunById"); + return Promise.resolve((config.batch ?? null) as never); + }) as FakeStore["findBatchTaskRunById"], + + countPendingWaitpointsWithPresence: ((waitpointIds: string[], _client?: ReadClient) => { + record("countPendingWaitpointsWithPresence"); + const pending = new Set(config.pendingWaitpointIds ?? []); + const found = waitpointIds.filter((id) => pending.has(id)); + return Promise.resolve({ pendingIds: found, presentIds: found } as never); + }) as FakeStore["countPendingWaitpointsWithPresence"], }; return store as unknown as FakeStore; @@ -103,6 +141,12 @@ function buildRouter(newConfig: FakeConfig = {}, legacyConfig: FakeConfig = {}) const trace = (log: Call[]) => log.map((c) => `${c.slot}:${c.method}`); +// A parallel fan-out issues every leg before any resolves, so the ORDER legs appear in the log is +// NOT a behaviour and MUST NOT be asserted. Compare the multiset of slots instead. Order assertions +// via `trace` are valid only for the sequential (two-or-fewer-store) probe. +const slots = (log: Call[], method?: string) => + (method ? log.filter((c) => c.method === method) : log).map((c) => c.slot).sort(); + describe("RoutingRunStore #probeOrder — new then legacy, sequential", () => { it("probes new BEFORE legacy for an unrouted findRun", async () => { const { router, log } = buildRouter(); @@ -166,10 +210,19 @@ describe("RoutingRunStore id-less fallbacks — the two defaults differ by role" expect(trace(log)).toEqual(["new:createRun"]); }); - it("routes an id-less checkpoint create to new (#idlessRouteShard)", async () => { + it("throws for an id-less checkpoint create rather than defaulting to new", async () => { const { router, log } = buildRouter(); - await router.createTaskRunCheckpoint({ data: {} } as never); - expect(trace(log)).toEqual(["new:createTaskRunCheckpoint"]); + await expect(router.createTaskRunCheckpoint({ data: {} } as never)).rejects.toThrow( + "createTaskRunCheckpoint requires ownerRunId to route" + ); + expect(trace(log)).toEqual([]); + }); + + it("throws for a batch create with no id rather than defaulting to new", async () => { + const { router } = buildRouter(); + await expect(router.createBatchTaskRun({} as never)).rejects.toThrow( + "createBatchTaskRun requires data.id to route" + ); }); it("routes an id-less waitpoint update to legacy (#idlessWaitpointShard)", async () => { @@ -178,3 +231,618 @@ describe("RoutingRunStore id-less fallbacks — the two defaults differ by role" expect(trace(log)).toEqual(["legacy:updateWaitpoint"]); }); }); + +describe("RoutingRunStore id-to-shard-key seam", () => { + it("defaults to the core resolveShard when neither seam is injected", async () => { + const log: Call[] = []; + const router = new RoutingRunStore({ + new: fakeStore("new", log), + legacy: fakeStore("legacy", log), + }); + // A gen-1 v1 body (version "1" at index 25) routes to new. + await router.findRun({ id: "a".repeat(24) + "01" }); + expect(trace(log)).toEqual(["new:findRun"]); + }); + + it("keeps the legacy classify seam working, so the corpus stays green", async () => { + const { router, log } = buildRouter(); + await router.findRun({ id: "new_run_1" }); + expect(trace(log)).toEqual(["new:findRun"]); + }); + + it("prefers an injected resolveShard over classify", async () => { + const log: Call[] = []; + const router = new RoutingRunStore({ + new: fakeStore("new", log), + legacy: fakeStore("legacy", log), + classify: () => "NEW", + resolveShard: () => "legacy", + }); + await router.findRun({ id: "anything" }); + expect(trace(log)).toEqual(["legacy:findRun"]); + }); + + // An id naming a shard nobody configured must fail loud rather than fall back to a default + // store, which would be a silent read against the wrong database. The throw is SYNCHRONOUS: + // routing happens before any query is issued, and `await store.findRun(...)` propagates it + // identically. Only a `.catch()`-style caller would see the difference. + it("throws for an id resolving to an unconfigured shard key", () => { + const log: Call[] = []; + const router = new RoutingRunStore({ + new: fakeStore("new", log), + legacy: fakeStore("legacy", log), + resolveShard: () => "a", + }); + expect(() => router.findRun({ id: "anything" })).toThrow( + 'no store is configured for shard key "a"' + ); + expect(trace(log)).toEqual([]); + }); +}); + +function buildNShardRouter(shardKeys: string[], opts: { aliasOf?: Record } = {}) { + const log: Call[] = []; + const newStore = fakeStore("new", log); + const legacyStore = fakeStore("legacy", log); + const byKey: Record = { new: newStore, legacy: legacyStore }; + const shards = shardKeys.map((key) => { + const aliasOf = opts.aliasOf?.[key]; + const store = aliasOf ? byKey[aliasOf]! : fakeStore(key as Slot, log); + byKey[key] = store; + return aliasOf ? { key, store, aliasOf } : { key, store }; + }); + const router = new RoutingRunStore({ + new: newStore, + legacy: legacyStore, + shards, + resolveShard: (id: string) => id.split(":")[0]!, + }); + return { router, log, byKey }; +} + +describe("RoutingRunStore #distinctStores — one entry per database", () => { + it("routes an id to its gen-2 shard", async () => { + const { router, log } = buildNShardRouter(["a", "b"]); + await router.findRun({ id: "a:run_1" }); + expect(trace(log)).toEqual(["a:findRun"]); + }); + + it("counts an aliased shard's database ONCE in a sum", async () => { + // "a" aliases "new": two keys, one database. + const { router, log } = buildNShardRouter(["a"], { aliasOf: { a: "new" } }); + const result = await router.updateManyWaitpoints({ + where: { status: "PENDING" }, + data: {}, + } as never); + expect(trace(log)).toEqual(["new:updateManyWaitpoints", "legacy:updateManyWaitpoints"]); + expect(result.count).toBe(2); + }); + + it("still routes an id whose key is an alias", async () => { + const { router, log, byKey } = buildNShardRouter(["a"], { aliasOf: { a: "new" } }); + expect(byKey.a).toBe(byKey.new); + await router.findRun({ id: "a:run_1" }); + expect(trace(log)).toEqual(["new:findRun"]); + }); + + it("keeps #fanOutPartitioned key-driven so an aliased bucket is not dropped", async () => { + const { router, log } = buildNShardRouter(["a"], { aliasOf: { a: "new" } }); + await router.findRunsByIds(["a:r1", "legacy:r2"]); + // Both buckets get a leg. The aliased bucket routes onto the shared store. + expect(log.filter((c) => c.method === "findRuns")).toHaveLength(2); + }); + + it("rejects an aliasOf naming an unconfigured key", () => { + const log: Call[] = []; + expect( + () => + new RoutingRunStore({ + new: fakeStore("new", log), + legacy: fakeStore("legacy", log), + shards: [{ key: "a", store: fakeStore("a" as Slot, log), aliasOf: "nope" }], + }) + ).toThrow('aliasOf "nope"'); + }); + + it("rejects a shard key that reuses a reserved key (new/legacy)", () => { + const log: Call[] = []; + expect( + () => + new RoutingRunStore({ + new: fakeStore("new", log), + legacy: fakeStore("legacy", log), + shards: [{ key: "new", store: fakeStore("shadow" as Slot, log) }], + }) + ).toThrow("must be unique"); + }); + + it("rejects a duplicate custom shard key", () => { + const log: Call[] = []; + expect( + () => + new RoutingRunStore({ + new: fakeStore("new", log), + legacy: fakeStore("legacy", log), + shards: [ + { key: "a", store: fakeStore("a1" as Slot, log) }, + { key: "a", store: fakeStore("a2" as Slot, log) }, + ], + }) + ).toThrow("must be unique"); + }); + + it("rejects a self-alias (a -> a), which would drop its database from every fan-out", () => { + const log: Call[] = []; + expect( + () => + new RoutingRunStore({ + new: fakeStore("new", log), + legacy: fakeStore("legacy", log), + shards: [{ key: "a", store: fakeStore("a" as Slot, log), aliasOf: "a" }], + }) + ).toThrow("must name a non-aliased store"); + }); + + it("rejects an alias chain (a -> b where b is itself aliased)", () => { + const log: Call[] = []; + expect( + () => + new RoutingRunStore({ + new: fakeStore("new", log), + legacy: fakeStore("legacy", log), + shards: [ + { key: "a", store: fakeStore("a" as Slot, log), aliasOf: "b" }, + { key: "b", store: fakeStore("b" as Slot, log), aliasOf: "new" }, + ], + }) + ).toThrow("must name a non-aliased store"); + }); + + it("rejects an alias cycle (a -> b, b -> a), which would drop both databases", () => { + const log: Call[] = []; + expect( + () => + new RoutingRunStore({ + new: fakeStore("new", log), + legacy: fakeStore("legacy", log), + shards: [ + { key: "a", store: fakeStore("a" as Slot, log), aliasOf: "b" }, + { key: "b", store: fakeStore("b" as Slot, log), aliasOf: "a" }, + ], + }) + ).toThrow("must name a non-aliased store"); + }); +}); + +describe("RoutingRunStore probe at N", () => { + it("keeps the sequential short circuit at two distinct stores", async () => { + const { router, log } = buildRouter({ runs: [{ id: "r1" }] }); + await router.findRun({ spanId: "span_x" }); + expect(trace(log)).toEqual(["new:findRun"]); + }); + + it("issues every leg in parallel above two distinct stores", async () => { + // "b" carries the only hit. A sequential probe would stop the moment it found a result; a + // true parallel fan-out queries every OTHER leg too, since all legs are issued before any + // resolves. Miss-path (every leg misses) throw semantics are covered separately below. + const log: Call[] = []; + const router = new RoutingRunStore({ + new: fakeStore("new", log), + legacy: fakeStore("legacy", log), + shards: [ + { key: "a", store: fakeStore("a", log) }, + { key: "b", store: fakeStore("b", log, { runs: [{ id: "r1" }] }) }, + ], + resolveShard: (id: string) => id.split(":")[0]!, + }); + await router.findRun({ spanId: "span_x" }); + expect(slots(log)).toEqual(["a", "b", "legacy", "new"]); + }); + + it("gives the legacy leg the canonical throw when every leg misses", async () => { + const { router } = buildNShardRouter(["a"]); + await expect(router.findRunOrThrow({ spanId: "span_x" })).rejects.toThrow("no run on legacy"); + }); + + it("tolerates a failing leg when another leg wins", async () => { + const log: Call[] = []; + const broken = fakeStore("a", log); + (broken as { findRun: unknown }).findRun = () => Promise.reject(new Error("shard a is down")); + const router = new RoutingRunStore({ + new: fakeStore("new", log, { runs: [{ id: "r1" }] }), + legacy: fakeStore("legacy", log), + shards: [{ key: "a", store: broken }], + resolveShard: (id: string) => id.split(":")[0]!, + }); + await expect(router.findRun({ spanId: "span_x" })).resolves.toMatchObject({ id: "r1" }); + }); + + it("surfaces a leg failure when no leg wins", async () => { + const log: Call[] = []; + const broken = fakeStore("a", log); + (broken as { findRun: unknown }).findRun = () => Promise.reject(new Error("shard a is down")); + const router = new RoutingRunStore({ + new: fakeStore("new", log), + legacy: fakeStore("legacy", log), + shards: [{ key: "a", store: broken }], + resolveShard: (id: string) => id.split(":")[0]!, + }); + await expect(router.findRun({ spanId: "span_x" })).rejects.toThrow("shard a is down"); + }); +}); + +describe("RoutingRunStore merge precedence and duplicate alarm", () => { + const spy = () => { + const seen: string[][] = []; + return { + metrics: { + recordDuplicateId: (k: string[]) => seen.push(k), + recordWaitpointProbeFallback() {}, + }, + seen, + }; + }; + + it("stays silent for a duplicate run id across the gen-1 pair", async () => { + const { metrics, seen } = spy(); + const log: Call[] = []; + const router = new RoutingRunStore({ + new: fakeStore("new", log, { runs: [{ id: "dup", from: "new" }] }), + legacy: fakeStore("legacy", log, { runs: [{ id: "dup", from: "legacy" }] }), + metrics, + }); + const rows = (await router.findRuns({ + where: { runtimeEnvironmentId: "env_1" }, + select: { id: true, from: true }, + })) as Array<{ from: string }>; + expect(rows).toHaveLength(1); + expect(rows[0]?.from).toBe("new"); // NEW wins the precedence merge + expect(seen).toEqual([]); + }); + + it("alarms for a duplicate involving a gen-2 shard and still picks deterministically", async () => { + const { metrics, seen } = spy(); + const log: Call[] = []; + const router = new RoutingRunStore({ + new: fakeStore("new", log), + legacy: fakeStore("legacy", log), + shards: [ + { key: "a", store: fakeStore("a", log, { runs: [{ id: "dup", from: "a" }] }) }, + { key: "b", store: fakeStore("b", log, { runs: [{ id: "dup", from: "b" }] }) }, + ], + resolveShard: (id: string) => id.split(":")[0]!, + metrics, + }); + const rows = (await router.findRuns({ + where: { runtimeEnvironmentId: "env_1" }, + select: { id: true, from: true }, + })) as Array<{ from: string }>; + expect(rows).toHaveLength(1); + expect(rows[0]?.from).toBe("b"); // last in #precedence [legacy, new, a, b] wins + expect(seen).toEqual([["a", "b"]]); + }); + + it("passes through an edge row whose projection omits id", async () => { + const log: Call[] = []; + const router = new RoutingRunStore({ + new: fakeStore("new", log, { edges: [{ taskRunId: "r1" }] }), + legacy: fakeStore("legacy", log, { edges: [{ taskRunId: "r2" }] }), + }); + const edges = (await router.findManyTaskRunWaitpoints({ + where: { waitpointId: "w" }, + select: { taskRunId: true }, + })) as Array<{ taskRunId: string }>; + expect(edges).toHaveLength(2); + }); +}); + +describe("RoutingRunStore findRunsByIdempotencyKeys tiebreak", () => { + const older = new Date("2026-01-01T00:00:00Z"); + const newer = new Date("2026-01-02T00:00:00Z"); + const match = (id: string, createdAt: Date) => ({ + id, + createdAt, + friendlyId: `run_${id}`, + idempotencyKey: "k", + idempotencyKeyExpiresAt: null, + }); + + it("keeps NEW-wins across the gen-1 pair even when legacy is older", async () => { + const log: Call[] = []; + const router = new RoutingRunStore({ + new: fakeStore("new", log, { idempotencyMatches: [match("n1", newer)] }), + legacy: fakeStore("legacy", log, { idempotencyMatches: [match("l1", older)] }), + }); + const rows = await router.findRunsByIdempotencyKeys({ + runtimeEnvironmentId: "env", + taskIdentifier: "t", + idempotencyKeys: ["k"], + }); + expect(rows.map((r) => r.id)).toEqual(["n1"]); + }); + + it("takes the earliest createdAt once a gen-2 shard is involved", async () => { + const log: Call[] = []; + const router = new RoutingRunStore({ + new: fakeStore("new", log, { idempotencyMatches: [match("n1", newer)] }), + legacy: fakeStore("legacy", log), + shards: [ + { key: "a", store: fakeStore("a", log, { idempotencyMatches: [match("a1", older)] }) }, + ], + resolveShard: (id: string) => id.split(":")[0]!, + }); + const rows = await router.findRunsByIdempotencyKeys({ + runtimeEnvironmentId: "env", + taskIdentifier: "t", + idempotencyKeys: ["k"], + }); + expect(rows.map((r) => r.id)).toEqual(["a1"]); + }); +}); + +describe("RoutingRunStore countPendingWaitpoints — disjoint-sum partition", () => { + // resolveShard: "x:..." -> "x" (a gen-2 shard); a bare cuid -> "legacy". + function partitionRouter(pending: Record, spy?: (k: string[]) => void) { + const log: Call[] = []; + const mk = (slot: string) => fakeStore(slot, log, { pendingWaitpointIds: pending[slot] ?? [] }); + const router = new RoutingRunStore({ + new: mk("new"), + legacy: mk("legacy"), + shards: [ + { key: "a", store: mk("a") }, + { key: "b", store: mk("b") }, + ], + resolveShard: (id: string) => (id.includes(":") ? id.split(":")[0]! : "legacy"), + ...(spy ? { metrics: { recordDuplicateId: spy, recordWaitpointProbeFallback() {} } } : {}), + }); + return { router, log }; + } + const countCalls = (log: Call[]) => slots(log, "countPendingWaitpointsWithPresence"); + + it("sends a gen-2 absent id to its own shard only", async () => { + const { router, log } = partitionRouter({ b: ["b:w1"] }); + expect(await router.countPendingWaitpoints(["b:w1"], undefined, "a:run")).toBe(1); + expect(countCalls(log)).toEqual(["a", "b"]); // run shard a (presence) + fallback b + }); + + it("contributes zero for a gen-2 id whose shard IS the run's shard", async () => { + const { router, log } = partitionRouter({}); + expect(await router.countPendingWaitpoints(["a:w1"], undefined, "a:run")).toBe(0); + expect(countCalls(log)).toEqual(["a"]); // absent on a; no fallback leg (b-bucket empty, a skipped) + }); + + it("probes BOTH gen-1 stores for a cuid absent id when the run is on a gen-2 shard", async () => { + const { router, log } = partitionRouter({ legacy: ["cuid_w1"] }); + expect(await router.countPendingWaitpoints(["cuid_w1"], undefined, "a:run")).toBe(1); + expect(countCalls(log)).toEqual(["a", "legacy", "new"]); + }); + + it("probes only legacy for a cuid when the run is on new", async () => { + const { router, log } = partitionRouter({ legacy: ["cuid_w1"] }); + expect(await router.countPendingWaitpoints(["cuid_w1"], undefined, "new:run")).toBe(1); + // run shard resolves "new:run" -> "new" (presence); cuid -> {legacy, new} minus new = legacy + expect(countCalls(log)).toEqual(["legacy", "new"]); + }); + + it("probes only new for a cuid when the run is on legacy", async () => { + const { router, log } = partitionRouter({ new: ["cuid_w1"] }); + expect(await router.countPendingWaitpoints(["cuid_w1"], undefined, "cuid_run")).toBe(1); + // run shard resolves cuid -> "legacy"; cuid -> {legacy, new} minus legacy = new only + expect(countCalls(log)).toEqual(["legacy", "new"]); + }); + + it("counts a drain-mirrored cuid ONCE and stays silent", async () => { + const seen: string[][] = []; + const { router } = partitionRouter({ new: ["cuid_w1"], legacy: ["cuid_w1"] }, (k) => + seen.push(k) + ); + expect(await router.countPendingWaitpoints(["cuid_w1"], undefined, "a:run")).toBe(1); + expect(seen).toEqual([]); // the gen-1 mirror is expected, never alarmed + }); + + it("fails loud when an absent id resolves to an unconfigured shard key", async () => { + // "c:w1" resolves to shard "c", which is not configured. Silently dropping it would under-count + // a pending waitpoint and prematurely unblock the run — so it must throw, not skip. + const { router } = partitionRouter({}); + await expect(router.countPendingWaitpoints(["c:w1"], undefined, "a:run")).rejects.toThrow( + 'unconfigured shard key "c"' + ); + }); + + it("returns zero for an id absent everywhere", async () => { + const { router } = partitionRouter({}); + expect(await router.countPendingWaitpoints(["b:w9"], undefined, "a:run")).toBe(0); + }); + + it("id-less count unions a drain-mirrored cuid to one, never sums it to two", async () => { + // No runId: fans over distinct stores. A cuid pending on BOTH gen-1 stores is one waitpoint. + const { router } = partitionRouter({ new: ["cuid_w1"], legacy: ["cuid_w1"] }); + expect(await router.countPendingWaitpoints(["cuid_w1"])).toBe(1); + }); + + it("returns the true total for a mixed gen-2 and cuid set with no double count", async () => { + const { router } = partitionRouter({ + b: ["b:w1"], + legacy: ["cuid_w1"], + new: ["cuid_w1", "cuid_w2"], + }); + const count = await router.countPendingWaitpoints( + ["b:w1", "cuid_w1", "cuid_w2", "b:w9"], + undefined, + "a:run" + ); + expect(count).toBe(3); // b:w1 + cuid_w1 (mirror, once) + cuid_w2; b:w9 absent + }); +}); + +describe("RoutingRunStore waitpoint probes at N", () => { + it("routes a gen-2 waitpoint directly, with no probe", async () => { + const { router, log } = buildNShardRouter(["a", "b"]); + await router.updateWaitpoint({ where: { id: "b:w1" }, data: {} } as never); + expect(trace(log)).toEqual(["b:updateWaitpoint"]); + }); + + it("keeps a cuid waitpoint on the gen-1 pair and never probes a gen-2 shard", async () => { + const log: Call[] = []; + const router = new RoutingRunStore({ + new: fakeStore("new", log, { waitpoint: { id: "cuid_w1" } }), + legacy: fakeStore("legacy", log), + shards: [{ key: "a", store: fakeStore("a", log) }], + resolveShard: (id: string) => (id.includes(":") ? id.split(":")[0]! : "legacy"), + }); + await router.updateWaitpoint({ where: { id: "cuid_w1" }, data: {} } as never); + expect(log.map((c) => c.slot)).not.toContain("a"); + }); + + it("records a probe fallback when the waitpoint is not on the store its id names", async () => { + const falls: Array<[string, string]> = []; + const log: Call[] = []; + const router = new RoutingRunStore({ + new: fakeStore("new", log, { waitpoint: { id: "cuid_w1" } }), + legacy: fakeStore("legacy", log, { waitpoint: null }), + metrics: { + recordDuplicateId() {}, + recordWaitpointProbeFallback: (from, to) => falls.push([from, to]), + }, + }); + await router.updateWaitpoint({ where: { id: "cuid_w1" }, data: {} } as never); + expect(falls).toEqual([["legacy", "new"]]); + }); + + it("lets a gen-2 waitpoint id beat the cross-tree legacy pin", async () => { + const { router, log } = buildNShardRouter(["a", "b"]); + const store = await router.forWaitpointCompletion("b:w1", { + isCrossTreeIdempotency: true, + } as never); + expect((store as FakeStore).slot).toBe("b"); + expect(log.map((c) => c.slot)).not.toContain("legacy"); + }); + + it("keeps the legacy pin for a cuid waitpoint in a cross-tree completion", async () => { + const log: Call[] = []; + const router = new RoutingRunStore({ + new: fakeStore("new", log), + legacy: fakeStore("legacy", log, { waitpoint: { id: "cuid_w1" } }), + shards: [{ key: "a", store: fakeStore("a", log) }], + resolveShard: (id: string) => (id.includes(":") ? id.split(":")[0]! : "legacy"), + }); + const store = await router.forWaitpointCompletion("cuid_w1", { + isCrossTreeIdempotency: true, + } as never); + expect((store as FakeStore).slot).toBe("legacy"); + expect(log.map((c) => c.slot)).not.toContain("a"); + }); +}); + +describe("RoutingRunStore gen-2 shard refuses a co-located cuid waitpoint", () => { + // resolveShard: "x:..." -> "x" (gen-2 shard); a bare id (no colon) -> "legacy". + function coLocateRouter() { + const log: Call[] = []; + const router = new RoutingRunStore({ + new: fakeStore("new", log), + legacy: fakeStore("legacy", log), + shards: [{ key: "a", store: fakeStore("a", log) }], + resolveShard: (id: string) => (id.includes(":") ? id.split(":")[0]! : "legacy"), + }); + return { router, log }; + } + + it("throws when a cuid waitpoint is co-located onto a gen-2 shard", () => { + const { router, log } = coLocateRouter(); + // Routing throws SYNCHRONOUSLY, before any write is issued. + expect(() => + router.createWaitpoint({ data: { id: "cuid_w1" } } as never, undefined, { + coLocateWithRunId: "a:run_1", + }) + ).toThrow('onto gen-2 shard "a"'); + expect(trace(log)).toEqual([]); + }); + + it("throws when an id-less waitpoint is co-located onto a gen-2 shard", () => { + // Prisma's @default(cuid()) would otherwise mint a cuid on the gen-2 shard AFTER the write, + // leaving it unroutable for its own completion (CodeRabbit finding). Reject it up front. + const { router, log } = coLocateRouter(); + expect(() => + router.createWaitpoint({ data: {} } as never, undefined, { coLocateWithRunId: "a:run_1" }) + ).toThrow('onto gen-2 shard "a"'); + expect(trace(log)).toEqual([]); + }); + + it("allows a gen-2 waitpoint co-located onto its own shard", async () => { + const { router, log } = coLocateRouter(); + await router.createWaitpoint({ data: { id: "a:w1" } } as never, undefined, { + coLocateWithRunId: "a:run_1", + }); + expect(trace(log)).toEqual(["a:createWaitpoint"]); + }); + + it("allows a cuid waitpoint co-located onto a gen-1 store", async () => { + const { router, log } = coLocateRouter(); + await router.createWaitpoint({ data: { id: "cuid_w1" } } as never, undefined, { + coLocateWithRunId: "legacy_run", + }); + expect(trace(log)).toEqual(["legacy:createWaitpoint"]); + }); +}); + +describe("RoutingRunStore id-less read/route defaults hold at N (never a gen-2 shard)", () => { + // With gen-2 shards a and b configured, each id-less default must still resolve to its named + // gen-1 store, never leak to a gen-2 shard. + it("#routeOrNew falls back to new for an id-less create", async () => { + const { router, log } = buildNShardRouter(["a", "b"]); + await router.createRun({ data: {} } as never); + expect(trace(log)).toEqual(["new:createRun"]); + }); + + it("#routeOrNew falls back to new for an id-less runInTransaction", async () => { + const { router, log } = buildNShardRouter(["a", "b"]); + await router.runInTransaction(undefined, async () => undefined); + expect(trace(log)).toEqual(["new:runInTransaction"]); + }); + + it("#resolveWaitpointStore(undefined) falls back to legacy for an id-less update", async () => { + const { router, log } = buildNShardRouter(["a", "b"]); + await router.updateWaitpoint({ where: { idempotencyKey: "k" }, data: {} } as never); + expect(trace(log)).toEqual(["legacy:updateWaitpoint"]); + }); + + it("#waitpointWriteStore with no owner and no residency falls back to legacy", async () => { + const { router, log } = buildNShardRouter(["a", "b"]); + await router.createWaitpoint({ data: {} } as never); + expect(trace(log)).toEqual(["legacy:createWaitpoint"]); + }); +}); + +describe("RoutingRunStore batch probe tolerates legitimate dual-residency", () => { + it("does NOT alarm when a batch is found on a gen-2 shard AND legacy", async () => { + const seen: string[][] = []; + const log: Call[] = []; + const router = new RoutingRunStore({ + new: fakeStore("new", log), + legacy: fakeStore("legacy", log, { batch: { id: "batch_dup", from: "legacy" } }), + shards: [{ key: "a", store: fakeStore("a", log, { batch: { id: "batch_dup", from: "a" } }) }], + resolveShard: (id: string) => id.split(":")[0]!, + metrics: { recordDuplicateId: (k) => seen.push(k), recordWaitpointProbeFallback() {} }, + }); + const batch = (await router.findBatchTaskRunById("batch_dup")) as { from: string } | null; + // batchTriggerV3 writes raw to the control plane while runEngine routes by id, so this is a + // legitimate dual-residency, not a routing-invariant violation — no alarm. + expect(seen).toEqual([]); + // Precedence still picks deterministically (gen-2 shard 'a' outranks legacy). + expect(batch?.from).toBe("a"); + }); + + it("still alarms when a RUN is found on a gen-2 shard AND legacy (real violation)", async () => { + const seen: string[][] = []; + const log: Call[] = []; + const router = new RoutingRunStore({ + new: fakeStore("new", log), + legacy: fakeStore("legacy", log, { runs: [{ id: "dup", from: "legacy" }] }), + shards: [{ key: "a", store: fakeStore("a", log, { runs: [{ id: "dup", from: "a" }] }) }], + resolveShard: (id: string) => id.split(":")[0]!, + metrics: { recordDuplicateId: (k) => seen.push(k), recordWaitpointProbeFallback() {} }, + }); + await router.findRun({ spanId: "span_x" }); + expect(seen).toEqual([["legacy", "a"]]); + }); +}); diff --git a/internal-packages/run-store/src/runOpsStore.ts b/internal-packages/run-store/src/runOpsStore.ts index 27bd78866d4..53089da21a5 100644 --- a/internal-packages/run-store/src/runOpsStore.ts +++ b/internal-packages/run-store/src/runOpsStore.ts @@ -7,7 +7,11 @@ import type { TaskRunStatus, WaitpointTag, } from "@trigger.dev/database"; -import { ownerEngine, type Residency, type ShardKey } from "@trigger.dev/core/v3/isomorphic"; +import { + resolveShard as coreResolveShard, + type Residency, + type ShardKey, +} from "@trigger.dev/core/v3/isomorphic"; import type { TaskRunError } from "@trigger.dev/core/v3/schemas"; import type { ClearIdempotencyKeyInput, @@ -30,8 +34,10 @@ import type { TaskRunWithWaitpoint, WaitpointColocationOptions, } from "./types.js"; +import { Logger } from "@trigger.dev/core/logger"; import { isReadReplicaClient } from "./readReplicaClient.js"; import { CONNECTED_RUNS_LIMIT } from "./PostgresRunStore.js"; +import { noopRoutingStoreMetrics, type RoutingStoreMetrics } from "./routingStoreMetrics.js"; import { boundedIn } from "@trigger.dev/database"; @@ -42,20 +48,19 @@ const LEGACY_SHARD: ShardKey = "legacy"; /** * 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 + * map from shard key to store, selecting one by `resolveShard` (gen-2 id→its shard char, gen-1 + * 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. + * Two shard keys MAY resolve to one store only through a declared `aliasOf`. #distinctStores then + * holds one entry per database, so a sum never counts a database twice. An undeclared duplicate + * store is still a configuration error. * - * 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. + * #probeOrder and #precedence each list one key per distinct store. At two shards they are exact + * reverses. At N they are not: both put gen-2 shards last, so a probe finds the gen-1 pair first + * and a merge lets a gen-2 shard win. A probe MUST iterate #probeOrder and a merge MUST iterate + * #precedence. */ export class RoutingRunStore implements RunStore { readonly #shards: ReadonlyMap; @@ -65,25 +70,90 @@ export class RoutingRunStore implements RunStore { // 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[]; + // One entry per distinct database, in precedence order. A fan-out sum iterates this, never #shards. + readonly #distinctStores: ReadonlyArray<{ key: ShardKey; store: RunStore }>; // 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; + // Id to shard key. `resolveShard` is the gen-2 seam. `classify` is the gen-1 seam, kept because + // five sites inject it and three of those are the regression corpus. A gen-1 classifier can + // only ever name the two reserved keys, so it can never reach a gen-2 shard. + readonly #resolveShardKey: (id: string) => ShardKey; + readonly #metrics: RoutingStoreMetrics; + readonly #logger: Logger; // 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 }) { + constructor(options: { + new: RunStore; + legacy: RunStore; + classify?: (id: string) => Residency; + resolveShard?: (id: string) => ShardKey; + shards?: ReadonlyArray<{ key: ShardKey; store: RunStore; aliasOf?: ShardKey }>; + metrics?: RoutingStoreMetrics; + logger?: Logger; + }) { + const shards = options.shards ?? []; + // Keys must be unique across the reserved pair and every configured shard. A key of "new" or + // "legacy" would overwrite the reserved #shards entry; a repeated custom key would appear twice + // in #precedence and #distinctStores, so a fan-out would query one database twice. + const shardKeys = [NEW_SHARD, LEGACY_SHARD, ...shards.map((s) => s.key)]; + if (new Set(shardKeys).size !== shardKeys.length) { + throw new Error( + "RoutingRunStore: shard keys must be unique and cannot reuse the reserved 'new' or 'legacy' keys" + ); + } + const configured = new Set(shardKeys); + const aliasedKeys = new Set(shards.filter((s) => s.aliasOf !== undefined).map((s) => s.key)); + for (const shard of shards) { + if (shard.aliasOf === undefined) { + continue; + } + // An alias must name a REAL root store, so #distinctStores keeps exactly one entry per + // database. A target that is itself aliased (a chain or a cycle) would drop every key in the + // cycle from #distinctStores, and that database would vanish from every read and write. + if (!configured.has(shard.aliasOf)) { + throw new Error( + `RoutingRunStore: shard "${shard.key}" declares aliasOf "${shard.aliasOf}", which is not configured` + ); + } + if (shard.aliasOf === shard.key || aliasedKeys.has(shard.aliasOf)) { + throw new Error( + `RoutingRunStore: shard "${shard.key}" aliasOf "${shard.aliasOf}" must name a non-aliased store; chains and cycles are not allowed` + ); + } + } + this.#shards = new Map([ [NEW_SHARD, options.new], [LEGACY_SHARD, options.legacy], + ...shards.map((s) => [s.key, s.store] as const), ]); - this.#probeOrder = [NEW_SHARD, LEGACY_SHARD]; - this.#precedence = [LEGACY_SHARD, NEW_SHARD]; + + const gen2Keys = shards.map((s) => s.key); + this.#probeOrder = [NEW_SHARD, LEGACY_SHARD, ...gen2Keys]; + this.#precedence = [LEGACY_SHARD, NEW_SHARD, ...gen2Keys]; this.#idlessRouteShard = NEW_SHARD; this.#idlessWaitpointShard = LEGACY_SHARD; - this.#classify = options.classify ?? ownerEngine; + const classify = options.classify; + this.#resolveShardKey = + options.resolveShard ?? + (classify !== undefined + ? (id: string) => (classify(id) === "NEW" ? NEW_SHARD : LEGACY_SHARD) + : coreResolveShard); + + // One entry per physical database, in precedence order. A declared alias contributes none: + // it shares its target's database, and a second leg over one database double-counts a sum. + // The discriminator is the DECLARATION, not object identity — the wiring layer may build a + // second store object over a shared client. + this.#distinctStores = this.#precedence + .filter((key) => !aliasedKeys.has(key)) + .map((key) => ({ key, store: this.#shardStore(key) })); + + this.#metrics = options.metrics ?? noopRoutingStoreMetrics; + this.#logger = options.logger ?? new Logger("RoutingRunStore", "warn"); } // A routing store spans two databases and has no single primary — routed reads resolve the @@ -113,9 +183,99 @@ export class RoutingRunStore implements RunStore { return store; } - // The shard that owns an existing id. Throws only when an injected classifier throws. + // A duplicate id is EXPECTED across the gen-1 pair (drain mirrors a token onto both). Any other + // combination breaks id-determinism: alarm, but keep the deterministic pick. + #reportDuplicateId(id: string, shardKeys: ShardKey[]): void { + if (shardKeys.every((key) => key === NEW_SHARD || key === LEGACY_SHARD)) { + return; + } + this.#metrics.recordDuplicateId(shardKeys); + this.#logger.error("RoutingRunStore: one id returned by two shards", { id, shardKeys }); + } + + // Merge key-tagged legs, keeping one row per id. Legs MUST arrive in #precedence order, so the + // highest-authority copy is written last and wins. A winner keeps the POSITION of its first + // sighting: callers observe row order whenever `orderBy` is absent. Rows whose projection omits + // `id` cannot be deduped and pass through unchanged. A duplicate whose reporting keys leave the + // gen-1 pair alarms via #reportDuplicateId. + #mergeById>(legs: Array<{ key: ShardKey; rows: R[] }>): R[] { + const byId = new Map(); + const keysById = new Map(); + const passthrough: R[] = []; + for (const { key, rows } of legs) { + for (const row of rows) { + const id = row.id; + if (typeof id !== "string") { + passthrough.push(row); + continue; + } + byId.set(id, row); + const keys = keysById.get(id); + if (keys) keys.push(key); + else keysById.set(id, [key]); + } + } + for (const [id, keys] of keysById) { + if (keys.length > 1) this.#reportDuplicateId(id, keys); + } + return [...byId.values(), ...passthrough]; + } + + // Where to look for waitpoint ids the run's own shard did not return. A gen-2 id names exactly + // one shard, so it goes there and nowhere else — that is what keeps the legs disjoint and the sum + // sound. A cuid names no shard: drain can mirror it onto NEW while it keeps its id, and a gen-2 + // run can block on a pre-gen-2 cuid token, so there is no single gen-1 partner. Both gen-1 stores + // are probed and the results are de-duped by id. A target equal to `runKey` is skipped (already + // probed); an id resolving to an unconfigured shard fails loud rather than being dropped. + #partitionAbsentIds(runKey: ShardKey, ids: string[]): Array<{ key: ShardKey; ids: string[] }> { + const byKey = new Map(); + const push = (key: ShardKey, id: string) => { + // The run's own shard was already probed, so skip it. But an id resolving to a shard nobody + // configured is UnknownShardKey: silently dropping it here would UNDER-count a pending + // waitpoint and prematurely unblock the run — the exact failure this method guards against. + // Fail loud instead (§7 append-only rule). + if (key === runKey) return; + if (!this.#shards.has(key)) { + throw new Error( + `RoutingRunStore: waitpoint "${id}" resolves to unconfigured shard key "${key}"` + ); + } + const bucket = byKey.get(key); + if (bucket) bucket.push(id); + else byKey.set(key, [id]); + }; + for (const id of ids) { + const home = this.#shardKeyOfSafe(id); + if (home === LEGACY_SHARD) { + push(LEGACY_SHARD, id); + push(NEW_SHARD, id); + } else { + push(home, id); + } + } + return this.#precedence + .filter((key) => byKey.has(key)) + .map((key) => ({ key, ids: byKey.get(key)! })); + } + + // A cuid is deliberately probed on BOTH gen-1 stores, so the same id from both is the expected + // drain mirror, not a violation. Any other id maps to one shard, so a two-leg return is a bug. + #isGen1MirrorProbe(id: string): boolean { + return this.#shardKeyOfSafe(id) === LEGACY_SHARD; + } + + // The gen-1 pair members other than `key`. A cuid waitpoint can only ever be drain-relocated + // BETWEEN the two gen-1 stores, so its "where does it really live" probe stays confined here and + // never touches a gen-2 shard. + #gen1PairExcept(key: ShardKey): Array<{ key: ShardKey; store: RunStore }> { + return [NEW_SHARD, LEGACY_SHARD] + .filter((k) => k !== key) + .map((k) => ({ key: k, store: this.#shardStore(k) })); + } + + // The shard that owns an existing id. Throws only when an injected resolver throws. #shardKeyOf(id: string): ShardKey { - return this.#classify(id) === "NEW" ? NEW_SHARD : LEGACY_SHARD; + return this.#resolveShardKey(id); } // An unclassifiable id is treated as LEGACY (probe the control-plane DB rather than drop a @@ -128,32 +288,88 @@ export class RoutingRunStore implements RunStore { } } - // 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. + // #distinctStores sorted by `order`. Shared by #probeFirst and #fanOut so every ordered walk + // over the distinct-store set goes through one sort. + #orderedLegs(order: readonly ShardKey[]): ReadonlyArray<{ key: ShardKey; store: RunStore }> { + const rank = new Map(order.map((key, i) => [key, i])); + return [...this.#distinctStores].sort( + (a, b) => (rank.get(a.key) ?? 0) - (rank.get(b.key) ?? 0) + ); + } + + // A lookup with no routable id. At two distinct stores this is the sequential short circuit, + // byte-identical to before: the second leg is never queried when the first answers. Above two, + // a sequential walk would cost N round trips, so every leg is issued in parallel and the winner + // comes from #precedence. `isLast` marks the leg that owns the canonical not-found throw. async #probeFirst( - fn: (store: RunStore, key: ShardKey, isLast: boolean) => Promise + fn: (store: RunStore, key: ShardKey, isLast: boolean) => Promise, + opts?: { alarmOnDuplicate?: boolean } ): Promise { - 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 legs = this.#orderedLegs(this.#probeOrder); + const lastIndex = legs.length - 1; + + if (legs.length <= 2) { + for (let i = 0; i < lastIndex; i++) { + const { store, key } = legs[i]!; + const found = await fn(store, key, false); + if (found != null) { + return found; + } } + const { store, key } = legs[lastIndex]!; + return fn(store, key, true); } - const key = this.#probeOrder[last]!; - return fn(this.#shardStore(key), key, true); + + // Parallel. Every leg takes the NON-throwing arm, so one leg cannot reject a lookup another + // leg answers. A rejection is held and only surfaces when nothing was found. + const settled = await Promise.allSettled(legs.map(({ store, key }) => fn(store, key, false))); + + const hits: Array<{ key: ShardKey; value: Awaited }> = []; + let firstRejection: unknown; + settled.forEach((outcome, i) => { + if (outcome.status === "rejected") { + firstRejection ??= outcome.reason; + return; + } + if (outcome.value != null) { + hits.push({ key: legs[i]!.key, value: outcome.value }); + } + }); + + // A row on two stores is a routing-invariant violation for entities that live on exactly one + // store (runs, waitpoints, attempts, snapshots). Batches are the exception: `batchTriggerV3` + // writes raw to the control plane while runEngine routes by id, so a batch is legitimately + // dual-resident. Those callers pass `alarmOnDuplicate: false` so a batch on legacy + a gen-2 + // shard is not mistaken for a violation. + if (hits.length > 1 && opts?.alarmOnDuplicate !== false) { + this.#reportDuplicateId( + String((hits[0]!.value as { id?: unknown })?.id ?? "unknown"), + hits.map((h) => h.key) + ); + } + if (hits.length > 0) { + // Highest authority wins: #precedence ascends, so take the last hit in that order. + const rank = new Map(this.#precedence.map((key, i) => [key, i])); + hits.sort((a, b) => (rank.get(a.key) ?? 0) - (rank.get(b.key) ?? 0)); + return hits[hits.length - 1]!.value; + } + if (firstRejection !== undefined) { + throw firstRejection; + } + // Nothing found anywhere. The LEGACY leg owns the canonical not-found throw, so give it the + // throwing arm. One extra query, on the miss path only. + const legacy = this.#shardStore(LEGACY_SHARD); + return fn(legacy, LEGACY_SHARD, 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. + // Run `fn` on every DISTINCT store in parallel. `order` selects the ordering; membership is + // always one entry per database, so an aliased key never contributes a second leg. #fanOut( order: readonly ShardKey[], fn: (store: RunStore, key: ShardKey) => Promise ): Promise { - return Promise.all(order.map((key) => fn(this.#shardStore(key), key))); + const legs = this.#orderedLegs(order); + return Promise.all(legs.map(({ store, key }) => fn(store, key))); } // Apply `fn` to every shard and sum the counts. A sum is order-independent, so this takes no order. @@ -174,6 +390,11 @@ export class RoutingRunStore implements RunStore { const byShard = new Map(); for (const id of ids) { const key = this.#shardKeyOfSafe(id); + // An id resolving to a shard nobody configured is UnknownShardKey. Dropping it would silently + // omit a row from the hydrated set, so fail loud (§7 append-only rule). + if (!this.#shards.has(key)) { + throw new Error(`RoutingRunStore: id "${id}" resolves to unconfigured shard key "${key}"`); + } const bucket = byShard.get(key); if (bucket) bucket.push(id); else byShard.set(key, [id]); @@ -188,13 +409,11 @@ export class RoutingRunStore implements RunStore { 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. + // Every distinct store other than `key`'s, in precedence 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) })); + return this.#distinctStores.filter((s) => s.key !== key); } // A `findRuns` caller bound to the given store (preserves `this`; the overload set isn't @@ -270,21 +489,27 @@ export class RoutingRunStore implements RunStore { if (typeof id !== "string") { return home; } + // A gen-2 waitpoint carries its shard in its id and its row lives there. No probe. + if (homeKey !== NEW_SHARD && homeKey !== LEGACY_SHARD) { + return home; + } if ( await home.findWaitpoint({ where: { id } }, onPrimary ? home.primaryReadClient : undefined) ) { return home; } - const [other] = this.#shardsExcept(homeKey); - if (other === undefined) { - return home; + for (const { key, store } of this.#gen1PairExcept(homeKey)) { + if ( + await store.findWaitpoint( + { where: { id } }, + onPrimary ? store.primaryReadClient : undefined + ) + ) { + this.#metrics.recordWaitpointProbeFallback(homeKey, key); + return store; + } } - return (await other.store.findWaitpoint( - { where: { id } }, - onPrimary ? other.store.primaryReadClient : undefined - )) - ? other.store - : home; + return home; } static #waitpointId(clause: unknown): string | undefined { @@ -468,12 +693,11 @@ export class RoutingRunStore implements RunStore { async #findRunsOpen(args: FindRunsArgs, client?: ReadClient): Promise { const { args: selArgs, addedFields } = ensureProjected(args); const fan = widenForMerge(selArgs); - const legs = await this.#fanOut(this.#precedence, (store) => - this.#findManyOn(store, client)(fan) - ); - const byId = new Map>(); - for (const r of legs.flat()) byId.set(r.id as string, r); - return finalizeRows([...byId.values()], args, addedFields); + const legs = await this.#fanOut(this.#precedence, async (store, key) => ({ + key, + rows: (await this.#findManyOn(store, client)(fan)) as Record[], + })); + return finalizeRows(this.#mergeById(legs), args, addedFields); } // Canonical grouped replacement for `Promise.all(ids.map(id => readThroughRun(id)))`: reuses @@ -539,14 +763,45 @@ export class RoutingRunStore implements RunStore { if (args.idempotencyKeys.length === 0) { return []; } - const legs = await this.#fanOut(this.#precedence, (store) => - store.findRunsByIdempotencyKeys(args, RoutingRunStore.#ownPrimary(store, client)) + const legs = await this.#fanOut(this.#precedence, (store, key) => + store + .findRunsByIdempotencyKeys(args, RoutingRunStore.#ownPrimary(store, client)) + .then((rows) => ({ key, rows })) ); - const byKey = new Map(); - for (const row of legs.flat()) { - if (row.idempotencyKey != null) byKey.set(row.idempotencyKey, row); + // Dedupe by KEY, not id: two runs on two shards can legitimately share a global key. Across the + // gen-1 pair the winner stays today's precedence result (NEW wins), which the duplicate-guard + // contract depends on. Once a gen-2 shard supplies a candidate, order by creation instead, so + // the winner does not depend on configured shard order. + const byKey = new Map>(); + for (const { key, rows } of legs) { + for (const row of rows) { + if (row.idempotencyKey == null) continue; + const bucket = byKey.get(row.idempotencyKey); + if (bucket) bucket.push({ key, row }); + else byKey.set(row.idempotencyKey, [{ key, row }]); + } } - return [...byKey.values()]; + const out: IdempotencyKeyRunMatch[] = []; + for (const candidates of byKey.values()) { + const gen1Only = candidates.every(({ key }) => key === NEW_SHARD || key === LEGACY_SHARD); + if (gen1Only) { + out.push(candidates[candidates.length - 1]!.row); + continue; + } + out.push( + [...candidates].sort((a, b) => { + const byCreated = a.row.createdAt.getTime() - b.row.createdAt.getTime(); + return byCreated !== 0 + ? byCreated + : a.row.id < b.row.id + ? -1 + : a.row.id > b.row.id + ? 1 + : 0; + })[0]!.row + ); + } + return out; } // --------------------------------------------------------------------------- @@ -1153,10 +1408,20 @@ export class RoutingRunStore implements RunStore { runId?: string ): Promise { if (runId === undefined) { + // No run id to partition on: query every distinct store and UNION by id, matching the routed + // path below. A drain-mirrored cuid pending on both gen-1 stores must count once, not twice — + // summing raw counts here would reintroduce the double count this method exists to remove. const legs = await this.#fanOut(this.#probeOrder, (store) => - store.countPendingWaitpoints(waitpointIds, RoutingRunStore.#ownPrimary(store, client)) + store.countPendingWaitpointsWithPresence( + waitpointIds, + RoutingRunStore.#ownPrimary(store, client) + ) ); - return legs.reduce((sum, leg) => sum + leg, 0); + const union = new Set(); + for (const leg of legs) { + for (const id of leg.pendingIds) union.add(id); + } + return union.size; } if (waitpointIds.length === 0) { @@ -1173,15 +1438,41 @@ export class RoutingRunStore implements RunStore { if (missing.length === 0) { return pendingIds.length; } - const [other] = this.#shardsExcept(runKey); - if (other === undefined) { + const plan = this.#partitionAbsentIds(runKey, missing); + if (plan.length === 0) { return pendingIds.length; } - const otherPending = await other.store.countPendingWaitpoints( - missing, - RoutingRunStore.#ownPrimary(other.store, client) + const legs = await Promise.all( + plan.map(async ({ key, ids }) => { + const store = this.#shardStore(key); + const { pendingIds: found } = await store.countPendingWaitpointsWithPresence( + ids, + RoutingRunStore.#ownPrimary(store, client) + ); + return { key, found }; + }) ); - return pendingIds.length + otherPending; + // UNION by id, never a sum of counts. A cuid mirrored onto both gen-1 stores appears twice; + // summing it is exactly the double count that leaves pendingCount above zero forever and never + // unblocks the run. The run store's pending set seeds the union; missing ids are disjoint from + // it by construction. #reportDuplicateId is the tripwire for a NON-mirror two-leg return, which + // a consistent resolveShard makes unreachable — it guards a future partition bug. + const union = new Set(pendingIds); + const seenFrom = new Map(); + for (const { key, found } of legs) { + for (const id of found) { + const keys = seenFrom.get(id); + if (keys) keys.push(key); + else seenFrom.set(id, [key]); + union.add(id); + } + } + for (const [id, keys] of seenFrom) { + if (keys.length > 1 && !this.#isGen1MirrorProbe(id)) { + this.#reportDuplicateId(id, keys); + } + } + return union.size; } // Fan out and union: an id lives on exactly one store in steady state (a drain-mirror can put it on @@ -1238,7 +1529,19 @@ export class RoutingRunStore implements RunStore { waitpointId: string | undefined ): RunStore { if (ownerId !== undefined) { - return this.#shardStore(this.#shardKeyOfSafe(ownerId)); + const key = this.#shardKeyOfSafe(ownerId); + // A gen-2 shard holds only ids stamped for that shard, because a waitpoint completes on the + // shard its own id names. Anything else stranded the blocked run: a cuid (routes to the gen-1 + // pair on completion), an id for a DIFFERENT gen-2 shard, or NO id at all — Prisma's + // @default(cuid()) then mints a cuid on the gen-2 shard after the write. The mint layer must + // stamp the owner's shard onto the waitpoint id, so fail loud rather than write an orphan. + const isGen2 = key !== NEW_SHARD && key !== LEGACY_SHARD; + if (isGen2 && (waitpointId === undefined || this.#shardKeyOfSafe(waitpointId) !== key)) { + throw new Error( + `RoutingRunStore: refusing to co-locate waitpoint "${waitpointId ?? ""}" onto gen-2 shard "${key}"; its id must be stamped for that shard` + ); + } + return this.#shardStore(key); } if (residency !== undefined) { return this.#shardStore(residency === "NEW" ? NEW_SHARD : LEGACY_SHARD); @@ -1350,10 +1653,11 @@ export class RoutingRunStore implements RunStore { return rows; } - // Collect the scalar waitpoint rows (relation re-resolution happens in the caller). With a run id in - // scope and a bounded id set to partition on, route to the run's store and fall back to the other DB - // for ONLY the ids missing there (a rare cross-tree token) — the two legs are disjoint by - // construction, so no dedup is needed. Otherwise fan out to BOTH and dedup by id NEW-wins. + // Collect the scalar waitpoint rows (relation re-resolution happens in the caller). With a run id + // in scope and a bounded id set, route to the run's store and fall back for ONLY the ids missing + // there. The fallback targets come from #partitionAbsentIds: a gen-2 id to its own shard, a cuid + // to BOTH gen-1 stores. The cuid legs are not disjoint, so the fallback rows are merged by id. + // Otherwise (no bounded id set) fan out to every store and dedup by id NEW-wins. async #collectManyWaitpoints( scalarArgs: Record, client: ReadClient | undefined, @@ -1375,38 +1679,41 @@ export class RoutingRunStore implements RunStore { if (missing.length === 0) { return fromRun; } - const [other] = this.#shardsExcept(runKey); - if (other === undefined) { + // Same partition as countPendingWaitpoints: a gen-2 missing id goes to its own shard, a + // cuid to both gen-1 stores. The cuid legs are NOT disjoint, so merge by id (a drain mirror + // appears once) rather than concatenate. + const plan = this.#partitionAbsentIds(runKey, missing); + if (plan.length === 0) { return fromRun; } - const fromOther = (await other.store.findManyWaitpoints( - narrowArgsToIds(scalarArgs, missing) as Prisma.WaitpointFindManyArgs, - RoutingRunStore.#ownPrimary(other.store, client) - )) as Record[]; - return [...fromRun, ...fromOther]; + const legs = await Promise.all( + plan.map(async ({ key, ids }) => { + const store = this.#shardStore(key); + return { + key, + rows: (await store.findManyWaitpoints( + narrowArgsToIds(scalarArgs, ids) as Prisma.WaitpointFindManyArgs, + RoutingRunStore.#ownPrimary(store, client) + )) as Record[], + }; + }) + ); + return [...fromRun, ...this.#mergeById(legs)]; } // No bounded id set to partition on → fall through to the fan-out path. } - const legs = await this.#fanOut( - this.#precedence, - (store) => - store.findManyWaitpoints( - scalarArgs as Prisma.WaitpointFindManyArgs, - RoutingRunStore.#ownPrimary(store, client) - ) as Promise[]> - ); - // 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>(); - const passthrough: Record[] = []; - for (const w of legs.flat()) { - const id = w.id; - if (typeof id === "string") byId.set(id, w); - else passthrough.push(w); - } - return [...byId.values(), ...passthrough]; + // A token mirrored onto both DBs during drain appears in BOTH legs; #mergeById dedups 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. + const legs = await this.#fanOut(this.#precedence, async (store, key) => ({ + key, + rows: (await store.findManyWaitpoints( + scalarArgs as Prisma.WaitpointFindManyArgs, + RoutingRunStore.#ownPrimary(store, client) + )) as Record[], + })); + return this.#mergeById(legs); } // Re-resolve a waitpoint's group-A relations across BOTH DBs and attach them to `row`. Each target @@ -1537,27 +1844,36 @@ export class RoutingRunStore implements RunStore { waitpointId: string, context: ForWaitpointCompletionContext ): Promise { - // Preferred store: explicit legacy-authority pins first, else the waitpoint's id-shape. + // A gen-2 waitpoint's row lives on the shard its id names. The three legacy pins encode "the one + // non-NEW store", a gen-1 idea, so a gen-2 id OVERRIDES them: honouring the pin would send the + // completion write to legacy, match zero rows, and strand the blocked run. A gen-2 id is also + // directly routable, so it takes no probe. + const idKey = this.#shardKeyOfSafe(waitpointId); + const isGen2 = idKey !== NEW_SHARD && idKey !== LEGACY_SHARD; + if (isGen2) { + return this.#shardStore(idKey); + } const preferredKey = context.treeOwnerResidency === "LEGACY" || context.isCrossTreeIdempotency === true || context.hasLegacyParent === true ? LEGACY_SHARD - : this.#shardKeyOfSafe(waitpointId); + : idKey; 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 - // selects the store a WRITE (updateManyWaitpoints) then lands on, so it must probe each - // store's PRIMARY (mirroring #resolveWaitpointStore's onPrimary): a just-created waitpoint the - // replica hasn't caught up on would otherwise mis-resolve the owner and strand the run. + // Resolve to where a CUID 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 across the gen-1 pair + // rather than route by id-shape alone and miss it (which leaves the blocked run stuck forever). + // This guard selects the store a WRITE (updateManyWaitpoints) then lands on, so it must probe + // each store's PRIMARY: a just-created waitpoint the replica has not caught up on would + // otherwise mis-resolve the owner and strand the run. if ( await preferred.findWaitpoint({ where: { id: waitpointId } }, preferred.primaryReadClient) ) { return preferred; } - for (const { store } of this.#shardsExcept(preferredKey)) { + for (const { key, store } of this.#gen1PairExcept(preferredKey)) { if (await store.findWaitpoint({ where: { id: waitpointId } }, store.primaryReadClient)) { + this.#metrics.recordWaitpointProbeFallback(preferredKey, key); return store; } } @@ -1596,13 +1912,14 @@ export class RoutingRunStore implements RunStore { RoutingRunStore.#ownPrimary(store, client) )) as Record[]; } else { - const legs = await this.#fanOut(this.#precedence, (store) => - store.findManyTaskRunWaitpoints( + const legs = await this.#fanOut(this.#precedence, async (store, key) => ({ + key, + rows: (await store.findManyTaskRunWaitpoints( scalarArgs as typeof args, RoutingRunStore.#ownPrimary(store, client) - ) - ); - edges = dedupeEdgesById(legs.flat()) as Record[]; + )) as Record[], + })); + edges = this.#mergeById(legs); } if (waitpoint) { @@ -1717,7 +2034,13 @@ export class RoutingRunStore implements RunStore { ownerRunId?: string, tx?: PrismaClientOrTransaction ): Promise> { - const store = this.#routeOrNew(ownerRunId); + // A create is a mint decision the mint layer owns. Defaulting to NEW was harmless with one + // dedicated store; at N it is a silent write to the wrong shard, and the run-routed snapshot's + // checkpointId FK then resolves on a different database. Fail loud instead. + if (ownerRunId === undefined) { + throw new Error("createTaskRunCheckpoint requires ownerRunId to route"); + } + const store = this.#route(ownerRunId); return store.createTaskRunCheckpoint(args, ownerRunId, undefined); } @@ -1731,9 +2054,12 @@ export class RoutingRunStore implements RunStore { ): Promise { // Route by the batch's classifiable internal id: run-ops id→NEW, cuid→LEGACY. The caller's // `tx` is never forwarded — the create runs on the owning store's own client so the batch and - // its co-resident child runs/items land on the same DB. Mirrors the by-id waitpoint-write routing / - // updateBatchTaskRun. - const store = await this.#routeOrNewForWrite(data.id); + // its co-resident child runs/items land on the same DB. A create with no id is a mint decision + // the mint layer must have made; at N a silent NEW default is a wrong-shard write, so fail loud. + if (data.id === undefined) { + throw new Error("createBatchTaskRun requires data.id to route"); + } + const store = await this.#routeForWrite(data.id); return store.createBatchTaskRun(data, undefined); } @@ -1763,8 +2089,9 @@ export class RoutingRunStore implements RunStore { ): Promise | 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. - return this.#probeFirst((store) => - store.findBatchTaskRunById(id, args, RoutingRunStore.#ownPrimary(store, client)) + return this.#probeFirst( + (store) => store.findBatchTaskRunById(id, args, RoutingRunStore.#ownPrimary(store, client)), + { alarmOnDuplicate: false } ); } @@ -1777,13 +2104,15 @@ export class RoutingRunStore implements RunStore { ): Promise | null> { // Never forward the caller's client verbatim; its presence resolves each leg to that // store's OWN primary. - return this.#probeFirst((store) => - store.findBatchTaskRunByFriendlyId( - friendlyId, - environmentId, - args, - RoutingRunStore.#ownPrimary(store, client) - ) + return this.#probeFirst( + (store) => + store.findBatchTaskRunByFriendlyId( + friendlyId, + environmentId, + args, + RoutingRunStore.#ownPrimary(store, client) + ), + { alarmOnDuplicate: false } ); } @@ -1802,13 +2131,15 @@ export class RoutingRunStore implements RunStore { ): Promise | null> { // Never forward the caller's client verbatim; its presence resolves each leg to that // store's OWN primary. - return this.#probeFirst((store) => - store.findBatchTaskRunByIdempotencyKey( - environmentId, - idempotencyKey, - args, - RoutingRunStore.#ownPrimary(store, client) - ) + return this.#probeFirst( + (store) => + store.findBatchTaskRunByIdempotencyKey( + environmentId, + idempotencyKey, + args, + RoutingRunStore.#ownPrimary(store, client) + ), + { alarmOnDuplicate: false } ); } @@ -1921,17 +2252,20 @@ export class RoutingRunStore implements RunStore { skip: 0, ...(args.take != null ? { take: skip + args.take } : {}), }; - const legs = await this.#fanOut(this.#precedence, (store) => - store.findManyWaitpointTags(perLeg, RoutingRunStore.#ownPrimary(store, client)) - ); - const byId = new Map(); - for (const tag of legs.flat()) byId.set(tag.id, tag); + const legs = await this.#fanOut(this.#precedence, async (store, key) => ({ + key, + rows: (await store.findManyWaitpointTags( + perLeg, + RoutingRunStore.#ownPrimary(store, client) + )) as unknown as Array>, + })); + const deduped = this.#mergeById(legs) as unknown as WaitpointTag[]; const merged = args.orderBy ? (sortByOrderBy( - [...byId.values()] as unknown as Array>, + deduped as unknown as Array>, args.orderBy as unknown as NonNullable ) as unknown as WaitpointTag[]) - : [...byId.values()]; + : deduped; return merged.slice(skip, args.take != null ? skip + args.take : undefined); } @@ -2055,20 +2389,6 @@ function narrowArgsToIds(args: Record, ids: string[]): Record(rows: R[]): R[] { - const byId = new Map(); - const passthrough: R[] = []; - for (const row of rows) { - const id = (row as { id?: unknown }).id; - if (typeof id === "string") byId.set(id, row); - else passthrough.push(row); - } - return [...byId.values(), ...passthrough]; -} - // A caller sub-select for an edge relation: `{ select?, include? }`, `true` for a bare `key: true`, // or undefined when not requested. type SubProjection = { select?: any; include?: any } | true | undefined; diff --git a/internal-packages/run-store/src/types.ts b/internal-packages/run-store/src/types.ts index 7c7f9566893..3ceb0a462dd 100644 --- a/internal-packages/run-store/src/types.ts +++ b/internal-packages/run-store/src/types.ts @@ -23,6 +23,8 @@ import type { Residency } from "@trigger.dev/core/v3/isomorphic"; export type ReadClient = PrismaClientOrTransaction | PrismaReplicaClient; export type IdempotencyKeyRunMatch = { + id: string; + createdAt: Date; friendlyId: string; idempotencyKey: string | null; idempotencyKeyExpiresAt: Date | null; diff --git a/internal-packages/testcontainers/src/index.ts b/internal-packages/testcontainers/src/index.ts index 8cdaee8571b..e72a63766e9 100644 --- a/internal-packages/testcontainers/src/index.ts +++ b/internal-packages/testcontainers/src/index.ts @@ -566,6 +566,93 @@ export const threeDbRunOpsPostgresTest = test.extend + test.extend({ + legacyUri: async ({}, use) => { + const container = await getWorkerPostgresContainer(); + const baseUri = container.getConnectionUri(); + const cloneDb = `nShardLegacy_${pgCloneCounter++}`; + await createDatabaseFromTemplate(baseUri, cloneDb); + try { + await use(postgresUriWithDatabase(baseUri, cloneDb)); + } finally { + await dropCloneDatabase(baseUri, cloneDb); + } + }, + newUri: async ({}, use) => { + const container = await getRunOpsWorkerPostgresContainer17(); + const baseUri = container.getConnectionUri(); + const cloneDb = `nShardNew_${pgCloneCounter++}`; + await createDatabaseFromTemplate(baseUri, cloneDb); + try { + await use(postgresUriWithDatabase(baseUri, cloneDb)); + } finally { + await dropCloneDatabase(baseUri, cloneDb); + } + }, + shardUris: async ({}, use) => { + const container = await getRunOpsWorkerPostgresContainer17(); + const baseUri = container.getConnectionUri(); + const clones: string[] = []; + try { + for (let i = 0; i < gen2ShardCount; i++) { + const cloneDb = `nShardGen2_${pgCloneCounter++}`; + await createDatabaseFromTemplate(baseUri, cloneDb); + clones.push(cloneDb); + } + await use(clones.map((db) => postgresUriWithDatabase(baseUri, db))); + } finally { + for (const db of clones) { + await dropCloneDatabase(baseUri, db); + } + } + }, + legacyPrisma: async ({ legacyUri }, use) => { + const prisma = new PrismaClient({ datasources: { db: { url: legacyUri } } }); + try { + await use(prisma); + } finally { + await prisma.$disconnect(); + } + }, + newPrisma: async ({ newUri }, use) => { + const prisma = new RunOpsPrismaClient({ datasources: { db: { url: newUri } } }); + try { + await use(prisma); + } finally { + await prisma.$disconnect(); + } + }, + shardPrismas: async ({ shardUris }, use) => { + const clients = shardUris.map( + (url) => new RunOpsPrismaClient({ datasources: { db: { url } } }) + ); + try { + await use(clients); + } finally { + for (const c of clients) { + await c.$disconnect(); + } + } + }, + }); + export const redisContainer = async ( { network, task }: { network: StartedNetwork } & TestContext, use: Use diff --git a/internal-packages/testcontainers/src/nShardFixture.test.ts b/internal-packages/testcontainers/src/nShardFixture.test.ts new file mode 100644 index 00000000000..1fd3449616f --- /dev/null +++ b/internal-packages/testcontainers/src/nShardFixture.test.ts @@ -0,0 +1,27 @@ +import { expect } from "vitest"; +import { makeNShardRunOpsPostgresTest } from "./index.js"; + +const nShardTest = makeNShardRunOpsPostgresTest(2); + +// Booting the PG14 + PG17 containers and cloning four databases on a cold runner far exceeds +// vitest's 5s default (this package sets no global testTimeout), so pass a generous per-test one. +nShardTest( + "builds 4 distinct databases (legacy + new + 2 gen-2 shards)", + async ({ legacyUri, newUri, shardUris }) => { + expect(shardUris).toHaveLength(2); + const all = [legacyUri, newUri, ...shardUris]; + expect(new Set(all).size).toBe(4); + }, + 120_000 +); + +nShardTest( + "each gen-2 clone carries the run-ops subset schema", + async ({ shardPrismas }) => { + expect(shardPrismas).toHaveLength(2); + for (const prisma of shardPrismas) { + await expect(prisma.taskRun.count()).resolves.toBe(0); + } + }, + 120_000 +);