diff --git a/apps/webapp/app/env.server.ts b/apps/webapp/app/env.server.ts index 2b1fba86980..0c3a03148ea 100644 --- a/apps/webapp/app/env.server.ts +++ b/apps/webapp/app/env.server.ts @@ -2020,6 +2020,12 @@ const EnvironmentSchema = z // (stale or fresh) resolves to the same kind for the whole window. See mintFlipGrace.ts. RUN_OPS_MINT_FLIP_GRACE_MS: z.coerce.number().int().default(90_000), + // Per-organization waitpoint coordinator cutover. The org's waitpointSystem flag wins; + // this is the fallback when the org has no override. Read only at waitpoint mint time. + WAITPOINT_SYSTEM_DEFAULT: z.enum(["legacy", "redis"]).default("legacy"), + WAITPOINT_MINT_FLAG_CACHE_TTL_MS: z.coerce.number().int().default(30_000), + WAITPOINT_MINT_FLAG_CACHE_MAX_ENTRIES: z.coerce.number().int().default(10_000), + // Session replication (Postgres → ClickHouse sessions_v1). Shares Redis // with the runs replicator for leader locking but has its own slot and // publication so the two consume independently. diff --git a/apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts b/apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts index 67ef45ebd27..abe04948083 100644 --- a/apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts @@ -1,5 +1,5 @@ import type { BatchTaskRunExecutionResult } from "@trigger.dev/core/v3"; -import { ownerEngine } from "@trigger.dev/core/v3/isomorphic"; +import { resolveShard, type ShardKey } from "@trigger.dev/core/v3/isomorphic"; import { $replica, type PrismaClientOrTransaction, @@ -13,6 +13,8 @@ import { runStore as defaultRunStore } from "~/v3/runStore.server"; import { BasePresenter } from "./basePresenter.server"; import { boundedIn } from "@trigger.dev/database"; +import { runOpsShardReplicas } from "~/v3/runOpsMigration/shardHandles.server"; +import { logger } from "~/services/logger.server"; /** * Run-ops read-through wiring. All optional; absent (or `splitEnabled` falsy) collapses `call` to * passthrough. `legacyReplica` is a READ REPLICA handle only — there is NO legacy-primary field. @@ -21,6 +23,8 @@ type ApiBatchResultsReadThroughDeps = { splitEnabled?: boolean; newClient?: PrismaReplicaClient; legacyReplica?: PrismaReplicaClient; + /** Gen-2 shard replicas by shard char; empty (RUN_OPS_SHARDS unset) keeps today's behaviour. */ + shardReplicas?: ReadonlyMap; isPastRetention?: (runId: string) => boolean; }; @@ -181,16 +185,57 @@ export class ApiBatchResultsPresenter extends BasePresenter { const taskRunIds = batchRun.items.map((item) => item.taskRunId); - const newRows = (await newClient.taskRun.findMany({ - where: { id: { in: boundedIn(taskRunIds) } }, - select: memberRunSelect, - })) as TaskRunWithAttempts[]; + // A gen-2 id is directly routable to its own shard, so it must not join the gen-1 read: + // it would miss there, and (being dedicated-family) never reach the legacy probe either. + const shardReplicas = this.readThrough?.shardReplicas ?? runOpsShardReplicas; + const genOneIds: string[] = []; + const idsByShard = new Map(); + for (const id of taskRunIds) { + const shardKey = resolveShard(id); + if (shardKey === "new" || shardKey === "legacy") { + genOneIds.push(id); + } else if (shardReplicas.has(shardKey)) { + const group = idsByShard.get(shardKey); + group ? group.push(id) : idsByShard.set(shardKey, [id]); + } else { + // Not routable and not a gen-1 shape. A gen-1 store is the wrong database, and a + // dedicated-family id never reaches the legacy probe, so falling back there would + // drop the member silently. Drop it loudly instead. + logger.error("ApiBatchResultsPresenter: gen-2 member on an unconfigured shard key", { + runId: id, + shardKey, + configured: [...shardReplicas.keys()], + }); + } + } + + const newRows = ( + genOneIds.length > 0 + ? ((await newClient.taskRun.findMany({ + where: { id: { in: boundedIn(genOneIds) } }, + select: memberRunSelect, + })) as TaskRunWithAttempts[]) + : [] + ).concat( + ( + await Promise.all( + [...idsByShard.entries()].map( + async ([shardKey, ids]) => + (await shardReplicas.get(shardKey)!.taskRun.findMany({ + where: { id: { in: boundedIn(ids) } }, + select: memberRunSelect, + })) as TaskRunWithAttempts[] + ) + ) + ).flat() + ); const runsById = new Map(newRows.map((run) => [run.id, run])); - // A run-ops id can only live on NEW, so only misses that AREN'T run-ops-shaped are candidates - // for the legacy probe — mirrors readThroughRun's per-id "NEW residency skips legacy" rule. - const legacyCandidateIds = taskRunIds.filter( - (id) => !runsById.has(id) && ownerEngine(id) !== "NEW" + // A dedicated-family id (gen-1 v1 or gen-2) can only live on its own store, so only + // misses that AREN'T dedicated-shaped are candidates for the legacy probe — mirrors + // readThroughRun's per-id "dedicated residency skips legacy" rule. + const legacyCandidateIds = genOneIds.filter( + (id) => !runsById.has(id) && resolveShard(id) === "legacy" ); if (legacyCandidateIds.length > 0) { const legacyRows = (await legacyReplica.taskRun.findMany({ diff --git a/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.input-streams.wait.ts b/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.input-streams.wait.ts index 024779ac666..3d28861fbb9 100644 --- a/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.input-streams.wait.ts +++ b/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.input-streams.wait.ts @@ -1,4 +1,5 @@ import { json } from "@remix-run/server-runtime"; +import { resolveWaitpointMintKind } from "~/v3/waitpointMigration/waitpointMintKind.server"; import { z } from "zod"; import { CreateInputStreamWaitpointRequestBody, @@ -82,7 +83,14 @@ const { action, loader } = createActionApiRoute( // Create the waitpoint. Co-locate it with the owning run (run-ops split) so a run-ops id // run's input-stream waitpoint lands on the run's DB and its block edge resolves. + const waitpointMintKind = await resolveWaitpointMintKind({ + organizationId: authentication.environment.organizationId, + id: authentication.environment.id, + orgFeatureFlags: authentication.environment.organization.featureFlags, + }); + const result = await engine.createManualWaitpoint({ + waitpointMintKind, runId: run.id, environmentId: authentication.environment.id, projectId: authentication.environment.projectId, diff --git a/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.session-streams.wait.ts b/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.session-streams.wait.ts index c00ff51b3be..001a25e892a 100644 --- a/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.session-streams.wait.ts +++ b/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.session-streams.wait.ts @@ -1,4 +1,5 @@ import { json } from "@remix-run/server-runtime"; +import { resolveWaitpointMintKind } from "~/v3/waitpointMigration/waitpointMintKind.server"; import { CreateSessionStreamWaitpointRequestBody, type CreateSessionStreamWaitpointResponseBody, @@ -103,7 +104,14 @@ const { action, loader } = createActionApiRoute( // Create the waitpoint. Co-locate it with the owning run (run-ops split) so a run-ops id // run's session-stream waitpoint lands on the run's DB and its block edge resolves. + const waitpointMintKind = await resolveWaitpointMintKind({ + organizationId: authentication.environment.organizationId, + id: authentication.environment.id, + orgFeatureFlags: authentication.environment.organization.featureFlags, + }); + const result = await engine.createManualWaitpoint({ + waitpointMintKind, runId: run.id, environmentId: authentication.environment.id, projectId: authentication.environment.projectId, diff --git a/apps/webapp/app/routes/api.v1.waitpoints.tokens.ts b/apps/webapp/app/routes/api.v1.waitpoints.tokens.ts index 62322c527c7..66b43eeb1ed 100644 --- a/apps/webapp/app/routes/api.v1.waitpoints.tokens.ts +++ b/apps/webapp/app/routes/api.v1.waitpoints.tokens.ts @@ -1,4 +1,5 @@ import { json } from "@remix-run/server-runtime"; +import { resolveWaitpointMintKind } from "~/v3/waitpointMigration/waitpointMintKind.server"; import { CreateWaitpointTokenRequestBody, type CreateWaitpointTokenResponseBody, @@ -93,7 +94,14 @@ const { action } = createActionApiRoute( } } + const waitpointMintKind = await resolveWaitpointMintKind({ + organizationId: authentication.environment.organizationId, + id: authentication.environment.id, + orgFeatureFlags: authentication.environment.organization.featureFlags, + }); + const result = await engine.createManualWaitpoint({ + waitpointMintKind, environmentId: authentication.environment.id, projectId: authentication.environment.projectId, idempotencyKey: body.idempotencyKey, diff --git a/apps/webapp/app/routes/engine.v1.runs.$runFriendlyId.wait.duration.ts b/apps/webapp/app/routes/engine.v1.runs.$runFriendlyId.wait.duration.ts index c7a8c3c5619..957f8ba1a74 100644 --- a/apps/webapp/app/routes/engine.v1.runs.$runFriendlyId.wait.duration.ts +++ b/apps/webapp/app/routes/engine.v1.runs.$runFriendlyId.wait.duration.ts @@ -1,4 +1,5 @@ import type { TypedResponse } from "@remix-run/server-runtime"; +import { resolveWaitpointMintKind } from "~/v3/waitpointMigration/waitpointMintKind.server"; import { json } from "@remix-run/server-runtime"; import type { WaitForDurationResponseBody } from "@trigger.dev/core/v3"; import { WaitForDurationRequestBody } from "@trigger.dev/core/v3"; @@ -41,7 +42,14 @@ const { action } = createActionApiRoute( ? resolveIdempotencyKeyTTL(body.idempotencyKeyTTL) : undefined; + const waitpointMintKind = await resolveWaitpointMintKind({ + organizationId: authentication.environment.organizationId, + id: authentication.environment.id, + orgFeatureFlags: authentication.environment.organization.featureFlags, + }); + const { waitpoint } = await engine.createDateTimeWaitpoint({ + waitpointMintKind, // Co-locate the waitpoint with the run that blocks on it (run-ops split): a run-ops run lives // on the dedicated DB, but the minted waitpoint id is always a cuid, so without the run id // the waitpoint would route to the control-plane DB and the block edge would never resolve. diff --git a/apps/webapp/app/routes/engine.v1.runs.$runFriendlyId.waitpoints.tokens.$waitpointFriendlyId.wait.ts b/apps/webapp/app/routes/engine.v1.runs.$runFriendlyId.waitpoints.tokens.$waitpointFriendlyId.wait.ts index ea1ebab0679..566ffc05876 100644 --- a/apps/webapp/app/routes/engine.v1.runs.$runFriendlyId.waitpoints.tokens.$waitpointFriendlyId.wait.ts +++ b/apps/webapp/app/routes/engine.v1.runs.$runFriendlyId.waitpoints.tokens.$waitpointFriendlyId.wait.ts @@ -38,7 +38,14 @@ const { action } = createActionApiRoute( }); if (!waitpoint) { - throw json({ error: "Waitpoint not found" }, { status: 404 }); + // Retryable: a miss here can be replica lag. resolveWaitpointThroughReadThrough + // deliberately does not read the legacy primary, so it relies on the caller retrying. + // A plain 404 is not retried by the SDK, which would turn a transient miss into a + // permanent failure. + throw json( + { error: "Waitpoint not found" }, + { status: 404, headers: { "x-should-retry": "true" } } + ); } const _result = await engine.blockRunWithWaitpoint({ @@ -55,6 +62,11 @@ const { action } = createActionApiRoute( { status: 200 } ); } catch (error) { + // A Response thrown inside the try is a deliberate status (the 404 above), not a + // failure. Re-throw it untouched, or every intentional 4xx here becomes a 500. + if (error instanceof Response) { + throw error; + } logger.error("Failed to wait for waitpoint", { runId, waitpointId, error }); throw json({ error: "Failed to wait for waitpoint token" }, { status: 500 }); } diff --git a/apps/webapp/app/runEngine/concerns/idempotencyKeys.server.ts b/apps/webapp/app/runEngine/concerns/idempotencyKeys.server.ts index f6696865e94..dfa2d4f5845 100644 --- a/apps/webapp/app/runEngine/concerns/idempotencyKeys.server.ts +++ b/apps/webapp/app/runEngine/concerns/idempotencyKeys.server.ts @@ -1,4 +1,4 @@ -import { ownerEngine, RunId } from "@trigger.dev/core/v3/isomorphic"; +import { resolveShard, RunId, type ShardKey } from "@trigger.dev/core/v3/isomorphic"; import type { PrismaClientOrTransaction, TaskRun, Waitpoint } from "@trigger.dev/database"; import { env } from "~/env.server"; import { logger } from "~/services/logger.server"; @@ -13,9 +13,10 @@ import { computeClaimTtlSeconds } from "~/v3/mollifier/claimTtl"; import { makeResolveMollifierFlag } from "~/v3/mollifier/mollifierGate.server"; import { runStore } from "~/v3/runStore.server"; import { runOpsLegacyPrisma, runOpsNewPrisma } from "~/db.server"; +import { runOpsShardWriters } from "~/v3/runOpsMigration/shardHandles.server"; import { isSplitEnabled } from "~/v3/runOpsMigration/splitMode.server"; import { resolveRunIdMintKind } from "~/v3/engineVersion.server"; -import { resolveIdempotencyDedupClient } from "./idempotencyResidency.server"; +import { clientForShardKey, resolveIdempotencyDedupClient } from "./idempotencyResidency.server"; import type { TraceEventConcern, TriggerTaskRequest } from "../types"; // In-memory per-org mollifier-enabled check, shared with `evaluateGate` @@ -32,6 +33,16 @@ const resolveOrgMollifierFlag = makeResolveMollifierFlag(); // PG's unique index as the backstop. const MAX_CLEARED_WINNER_REACQUIRES = 5; +// The store that owns a shard key. A function, not a map: the handles are module constants and +// `runOpsShardWriters` is already keyed, so a second structure would add an allocation and, if +// memoised, mutable module state. Reading them lazily also keeps this module importable by +// triggerTask under a `~/db.server` mock that omits them. +function idempotencyClientFor(shardKey: ShardKey): PrismaClientOrTransaction | undefined { + if (shardKey === "legacy") return runOpsLegacyPrisma; + if (shardKey === "new") return runOpsNewPrisma; + return runOpsShardWriters.get(shardKey); +} + // Claim ownership context returned to the caller when the // IdempotencyKeyConcern won a pre-gate claim. Caller MUST publish the // winning runId on pipeline success (`publishClaim`) or release the @@ -172,12 +183,9 @@ export class IdempotencyKeyConcern { { isSplitEnabled, fallbackClient: this.prisma, - newClient: runOpsNewPrisma, - legacyClient: runOpsLegacyPrisma, + clientFor: idempotencyClientFor, resolveMintKind: resolveRunIdMintKind, - // `isMigrated` is intentionally omitted: until a child of a swept - // legacy-id parent can be born on the new DB, the swept-marker override - // would never change the answer, so a child routes by parent id-shape. + logger, } ); @@ -640,12 +648,15 @@ export class IdempotencyKeyConcern { } catch { return null; } - let client: PrismaClientOrTransaction; - try { - client = ownerEngine(internalId) === "NEW" ? runOpsNewPrisma : runOpsLegacyPrisma; - } catch { - client = this.prisma; - } + // The routing store routes by id and never forwards this object, so its identity only + // signals read-your-writes. Resolving it through the shard map keeps the two idempotency + // call sites in agreement and stops this reading as gen-2-unaware. + const client = clientForShardKey( + resolveShard(internalId), + idempotencyClientFor, + this.prisma, + logger + ); return runStore.findRun( { id: internalId, runtimeEnvironmentId: environmentId }, { include: { associatedWaitpoint: true } }, diff --git a/apps/webapp/app/runEngine/concerns/idempotencyResidency.server.test.ts b/apps/webapp/app/runEngine/concerns/idempotencyResidency.server.test.ts index 39b806a0f71..976a4ffd784 100644 --- a/apps/webapp/app/runEngine/concerns/idempotencyResidency.server.test.ts +++ b/apps/webapp/app/runEngine/concerns/idempotencyResidency.server.test.ts @@ -1,6 +1,7 @@ import { describe, expect, it } from "vitest"; import { RunId } from "@trigger.dev/core/v3/isomorphic"; import { + clientForShardKey, resolveIdempotencyDedupClient, type ResolveIdempotencyClientDeps, } from "./idempotencyResidency.server"; @@ -9,20 +10,30 @@ import { const FALLBACK = { __tag: "fallback" } as never; const NEW_CLIENT = { __tag: "new" } as never; const LEGACY_CLIENT = { __tag: "legacy" } as never; +const SHARD_A_CLIENT = { __tag: "shard-a" } as never; + +function clientMap() { + return new Map([ + ["new", NEW_CLIENT], + ["legacy", LEGACY_CLIENT], + ["a", SHARD_A_CLIENT], + ]); +} function makeDeps(over: Partial): ResolveIdempotencyClientDeps { return { isSplitEnabled: async () => true, fallbackClient: FALLBACK, - newClient: NEW_CLIENT, - legacyClient: LEGACY_CLIENT, + clientFor: (key) => clientMap().get(key), resolveMintKind: async () => "runOpsId", + // Kept as an injected seam: the real resolveShard is total, so only an injected + // classifier can exercise the throw-to-fallback arm below. classify: (id) => { - if (id.length === 26 && id[25] === "1") return "NEW"; - if (id.length === 25) return "LEGACY"; + if (id.length === 26 && id[25] === "2") return id[24]!; + if (id.length === 26 && id[25] === "1") return "new"; + if (id.length === 25) return "legacy"; throw new Error(`unclassifiable: ${id.length}`); }, - isMigrated: undefined, ...over, }; } @@ -72,29 +83,51 @@ describe("resolveIdempotencyDedupClient", () => { expect(client).toBe(LEGACY_CLIENT); }); - it("routes a swept (migrated) cuid-parent child to the NEW client", async () => { - const cuidParent = RunId.toFriendlyId("c".repeat(25)); + it("falls back to the fallback client when a present parent id is unclassifiable", async () => { const client = await resolveIdempotencyDedupClient( - { environmentForMint: env, parentRunFriendlyId: cuidParent }, - makeDeps({ isMigrated: async () => true }) + { environmentForMint: env, parentRunFriendlyId: "run_not-a-valid-length" }, + makeDeps({}) ); - expect(client).toBe(NEW_CLIENT); + expect(client).toBe(FALLBACK); }); - it("routes a non-migrated cuid-parent child to the LEGACY client even when isMigrated is provided", async () => { - const cuidParent = RunId.toFriendlyId("d".repeat(25)); + it("routes a child to its OWN SHARD client when the parent is a gen-2 id", async () => { + const genTwoParent = RunId.toFriendlyId("e".repeat(24) + "a2"); const client = await resolveIdempotencyDedupClient( - { environmentForMint: env, parentRunFriendlyId: cuidParent }, - makeDeps({ isMigrated: async () => false }) + { environmentForMint: env, parentRunFriendlyId: genTwoParent }, + makeDeps({ resolveMintKind: async () => "cuid" }) // mint flag must NOT win for a child ); - expect(client).toBe(LEGACY_CLIENT); + expect(client).toBe(SHARD_A_CLIENT); }); - it("falls back to the fallback client when a present parent id is unclassifiable", async () => { + it("falls back and logs when a gen-2 parent names an unconfigured shard key", async () => { + const errors: unknown[] = []; + const genTwoParent = RunId.toFriendlyId("f".repeat(24) + "z2"); const client = await resolveIdempotencyDedupClient( - { environmentForMint: env, parentRunFriendlyId: "run_not-a-valid-length" }, - makeDeps({}) + { environmentForMint: env, parentRunFriendlyId: genTwoParent }, + makeDeps({ logger: { error: (_m, meta) => errors.push(meta) } }) ); expect(client).toBe(FALLBACK); + expect(errors).toHaveLength(1); + }); +}); + +describe("clientForShardKey", () => { + it("selects the same client the map holds for each reserved key and shard key", () => { + const clients = clientMap(); + const clientFor = (key: string) => clients.get(key); + expect(clientForShardKey("new", clientFor, FALLBACK)).toBe(NEW_CLIENT); + expect(clientForShardKey("legacy", clientFor, FALLBACK)).toBe(LEGACY_CLIENT); + expect(clientForShardKey("a", clientFor, FALLBACK)).toBe(SHARD_A_CLIENT); + }); + + it("returns the fallback and logs for a key the map does not hold", () => { + const errors: unknown[] = []; + const map = clientMap(); + const client = clientForShardKey("z", (key) => map.get(key), FALLBACK, { + error: (_m, meta) => errors.push(meta), + }); + expect(client).toBe(FALLBACK); + expect(errors).toHaveLength(1); }); }); diff --git a/apps/webapp/app/runEngine/concerns/idempotencyResidency.server.ts b/apps/webapp/app/runEngine/concerns/idempotencyResidency.server.ts index 86f1435654b..f2a731e61ca 100644 --- a/apps/webapp/app/runEngine/concerns/idempotencyResidency.server.ts +++ b/apps/webapp/app/runEngine/concerns/idempotencyResidency.server.ts @@ -1,22 +1,44 @@ -import { ownerEngine, RunId, type Residency } from "@trigger.dev/core/v3/isomorphic"; +import { resolveShard, RunId, type ShardKey } from "@trigger.dev/core/v3/isomorphic"; import type { PrismaClientOrTransaction } from "@trigger.dev/database"; type MintKind = "cuid" | "runOpsId"; +type Logger = { error: (message: string, meta?: Record) => void }; + export type ResolveIdempotencyClientDeps = { isSplitEnabled: () => Promise; fallbackClient: PrismaClientOrTransaction; - newClient: PrismaClientOrTransaction; - legacyClient: PrismaClientOrTransaction; + /** The store that owns a shard key: the reserved `legacy`/`new`, or a gen-2 shard. */ + clientFor: (shardKey: ShardKey) => PrismaClientOrTransaction | undefined; resolveMintKind: (environment: { organizationId: string; id: string; orgFeatureFlags?: unknown; }) => Promise; - classify?: (id: string) => Residency; - isMigrated?: (id: string) => Promise; + classify?: (id: string) => ShardKey; + logger?: Logger; }; +/** + * The one place an id becomes a client. `ShardKey` collapses to `string`, so the compiler + * cannot catch a wrong key here — an absent key takes an explicit logged branch to the + * fallback rather than a silent `?? legacy`. The configured set is not repeated in the log: + * boot already prints the shard table. + */ +export function clientForShardKey( + shardKey: ShardKey, + clientFor: (shardKey: ShardKey) => PrismaClientOrTransaction | undefined, + fallback: PrismaClientOrTransaction, + logger?: Logger +): PrismaClientOrTransaction { + const client = clientFor(shardKey); + if (client === undefined) { + logger?.error("idempotency: no client configured for shard key", { shardKey }); + return fallback; + } + return client; +} + export async function resolveIdempotencyDedupClient( args: { environmentForMint: { organizationId: string; id: string; orgFeatureFlags?: unknown }; @@ -28,9 +50,9 @@ export async function resolveIdempotencyDedupClient( return deps.fallbackClient; } - const classify = deps.classify ?? ownerEngine; - const clientFor = (residency: Residency): PrismaClientOrTransaction => - residency === "NEW" ? deps.newClient : deps.legacyClient; + const classify = deps.classify ?? resolveShard; + const clientFor = (shardKey: ShardKey): PrismaClientOrTransaction => + clientForShardKey(shardKey, deps.clientFor, deps.fallbackClient, deps.logger); if (args.parentRunFriendlyId) { let parentInternalId: string; @@ -39,18 +61,18 @@ export async function resolveIdempotencyDedupClient( } catch { return deps.fallbackClient; } - let residency: Residency; + let shardKey: ShardKey; try { - residency = classify(parentInternalId); + shardKey = classify(parentInternalId); } catch { return deps.fallbackClient; } - if (residency === "LEGACY" && deps.isMigrated && (await deps.isMigrated(parentInternalId))) { - return deps.newClient; - } - return clientFor(residency); + return clientFor(shardKey); } + // Mint kind, not an id: there is no shard to decode, so this keeps resolving to the + // gen-1 pair exactly as before. Which shard a gen-2 env mints into is the mint layer's + // decision, and this client is a read-your-writes signal rather than a correctness gate. const kind = await deps.resolveMintKind(args.environmentForMint); - return clientFor(kind === "runOpsId" ? "NEW" : "LEGACY"); + return clientFor(kind === "runOpsId" ? "new" : "legacy"); } diff --git a/apps/webapp/app/runEngine/concerns/resolveWaitpointThroughReadThrough.server.ts b/apps/webapp/app/runEngine/concerns/resolveWaitpointThroughReadThrough.server.ts index ec5adc13a6c..b1c7a2bd05d 100644 --- a/apps/webapp/app/runEngine/concerns/resolveWaitpointThroughReadThrough.server.ts +++ b/apps/webapp/app/runEngine/concerns/resolveWaitpointThroughReadThrough.server.ts @@ -1,3 +1,4 @@ +import { resolveShard, type ShardKey } from "@trigger.dev/core/v3/isomorphic"; import type { PrismaReplicaClient } from "~/db.server"; import { runOpsLegacyReplica as defaultLegacyReplica, @@ -5,12 +6,18 @@ import { runOpsNewReplica as defaultNewClient, runOpsSplitReadEnabled as defaultSplitReadEnabled, } from "~/db.server"; +import { + runOpsShardReplicas as defaultShardReplicas, + runOpsShardWriters as defaultShardWriters, +} from "~/v3/runOpsMigration/shardHandles.server"; import { readThroughRun } from "~/v3/runOpsMigration/readThrough.server"; type ResolveWaitpointDeps = { newClient?: PrismaReplicaClient; legacyReplica?: PrismaReplicaClient; newPrimary?: PrismaReplicaClient; + shardReplicas?: ReadonlyMap; + shardWriters?: ReadonlyMap; splitEnabled?: boolean; isPastRetention?: (id: string) => boolean; }; @@ -21,6 +28,8 @@ export type ResolveWaitpointReadThroughDefaults = { newClient: PrismaReplicaClient; legacyReplica: PrismaReplicaClient; newPrimary: PrismaReplicaClient; + shardReplicas: ReadonlyMap; + shardWriters: ReadonlyMap; splitEnabled: boolean; }; @@ -28,6 +37,8 @@ const productionDefaults: ResolveWaitpointReadThroughDefaults = { newClient: defaultNewClient, legacyReplica: defaultLegacyReplica, newPrimary: defaultNewPrimary as unknown as PrismaReplicaClient, + shardReplicas: defaultShardReplicas, + shardWriters: defaultShardWriters as unknown as ReadonlyMap, splitEnabled: defaultSplitReadEnabled, }; @@ -43,7 +54,8 @@ export async function resolveWaitpointThroughReadThrough(opts: { const splitEnabled = opts.deps?.splitEnabled ?? defaults.splitEnabled; const result = await readThroughRun({ - runId: opts.waitpointId, + id: opts.waitpointId, + idKind: "waitpoint", environmentId: opts.environmentId, readNew: (client) => opts.read(client), readLegacy: (replica) => opts.read(replica), @@ -51,22 +63,31 @@ export async function resolveWaitpointThroughReadThrough(opts: { splitEnabled, newClient: opts.deps?.newClient ?? defaults.newClient, legacyReplica: opts.deps?.legacyReplica ?? defaults.legacyReplica, + shardReplicas: opts.deps?.shardReplicas ?? defaults.shardReplicas, isPastRetention: opts.deps?.isPastRetention, }, }); - if (result.source === "new" || result.source === "legacy-replica") { + if (result.found) { return result.value; } // past-retention is an intentional not-found: the token is gone. - if (result.source === "past-retention") { + if (result.reason === "past-retention") { return null; } // Read-your-writes fallback for a token completed immediately after mint, before it replicated: - // re-read from the run-ops PRIMARY only. We deliberately never read the control-plane/legacy + // re-read from the owning store's PRIMARY only. We deliberately never read the control-plane/legacy // primary here (that is the load the replica-only read-through exists to shed), so a legacy-resident // token that misses its replica stays a miss and the caller retries, rather than adding primary load. + const shardKey = resolveShard(opts.waitpointId); + if (shardKey !== "new" && shardKey !== "legacy") { + // A gen-2 token's primary is its OWN shard's writer. The gen-1 new writer is a different + // database, so reading it would miss and silently disable read-your-writes here. + const shardWriter = (opts.deps?.shardWriters ?? defaults.shardWriters).get(shardKey); + return shardWriter ? await opts.read(shardWriter) : null; + } + const fromNewPrimary = await opts.read(opts.deps?.newPrimary ?? defaults.newPrimary); if (fromNewPrimary != null) { return fromNewPrimary; diff --git a/apps/webapp/app/runEngine/services/batchTrigger.server.ts b/apps/webapp/app/runEngine/services/batchTrigger.server.ts index 5e29d158925..59d5f594462 100644 --- a/apps/webapp/app/runEngine/services/batchTrigger.server.ts +++ b/apps/webapp/app/runEngine/services/batchTrigger.server.ts @@ -1,3 +1,4 @@ +import { resolveWaitpointMintKind } from "~/v3/waitpointMigration/waitpointMintKind.server"; import { type BatchTriggerTaskV2RequestBody, type BatchTriggerTaskV3RequestBody, @@ -194,6 +195,11 @@ export class RunEngineBatchTriggerService extends WithRunEngine { environmentId: environment.id, projectId: environment.projectId, organizationId: environment.organizationId, + waitpointMintKind: await resolveWaitpointMintKind({ + organizationId: environment.organizationId, + id: environment.id, + orgFeatureFlags: environment.organization.featureFlags, + }), }); } @@ -285,6 +291,11 @@ export class RunEngineBatchTriggerService extends WithRunEngine { environmentId: environment.id, projectId: environment.projectId, organizationId: environment.organizationId, + waitpointMintKind: await resolveWaitpointMintKind({ + organizationId: environment.organizationId, + id: environment.id, + orgFeatureFlags: environment.organization.featureFlags, + }), }); } diff --git a/apps/webapp/app/runEngine/services/createBatch.server.ts b/apps/webapp/app/runEngine/services/createBatch.server.ts index 0289e68e2c7..b3f652b2f89 100644 --- a/apps/webapp/app/runEngine/services/createBatch.server.ts +++ b/apps/webapp/app/runEngine/services/createBatch.server.ts @@ -1,4 +1,5 @@ import type { InitializeBatchOptions } from "@internal/run-engine"; +import { resolveWaitpointMintKind } from "~/v3/waitpointMigration/waitpointMintKind.server"; import { type CreateBatchRequestBody, type CreateBatchResponse } from "@trigger.dev/core/v3"; import { RunId } from "@trigger.dev/core/v3/isomorphic"; import { type BatchTaskRun, Prisma } from "@trigger.dev/database"; @@ -137,6 +138,11 @@ export class CreateBatchService extends WithRunEngine { environmentId: environment.id, projectId: environment.projectId, organizationId: environment.organizationId, + waitpointMintKind: await resolveWaitpointMintKind({ + organizationId: environment.organizationId, + id: environment.id, + orgFeatureFlags: environment.organization.featureFlags, + }), }); } diff --git a/apps/webapp/app/runEngine/services/triggerTask.server.ts b/apps/webapp/app/runEngine/services/triggerTask.server.ts index 8e9e99d7f09..00d7c3990f0 100644 --- a/apps/webapp/app/runEngine/services/triggerTask.server.ts +++ b/apps/webapp/app/runEngine/services/triggerTask.server.ts @@ -29,6 +29,10 @@ import { removeNullBytesFromKey } from "~/utils/nullBytes"; import { handleMetadataPacket } from "~/utils/packets"; import { startSpan } from "~/v3/tracing.server"; import { resolveRunIdMintKind } from "~/v3/engineVersion.server"; +import { + resolveWaitpointMintKind, + type WaitpointMintKind, +} from "~/v3/waitpointMigration/waitpointMintKind.server"; import { resolveInheritedMintKind } from "~/v3/runOpsMigration/resolveInheritedMintKind.server"; import { mintFriendlyIdForKind } from "~/v3/runOpsMigration/mintAnchoredRunFriendlyId.server"; import type { @@ -652,10 +656,16 @@ export class RunEngineTriggerTaskService { event.setAttribute("taskRunId", runFriendlyId); const payloadPacket = await this.payloadProcessor.process(triggerRequest); + const waitpointMintKind = await resolveWaitpointMintKind({ + organizationId: environment.organizationId, + id: environment.id, + orgFeatureFlags: environment.organization.featureFlags, + }); const engineTriggerInput = this.#buildEngineTriggerInput({ runFriendlyId, environment, + waitpointMintKind, idempotencyKey, idempotencyKeyExpiresAt, body, @@ -732,10 +742,16 @@ export class RunEngineTriggerTaskService { } const payloadPacket = await this.payloadProcessor.process(triggerRequest); + const waitpointMintKind = await resolveWaitpointMintKind({ + organizationId: environment.organizationId, + id: environment.id, + orgFeatureFlags: environment.organization.featureFlags, + }); const baseEngineInput = this.#buildEngineTriggerInput({ runFriendlyId, environment, + waitpointMintKind, idempotencyKey, idempotencyKeyExpiresAt, body, @@ -898,6 +914,7 @@ export class RunEngineTriggerTaskService { #buildEngineTriggerInput(args: { runFriendlyId: string; environment: AuthenticatedEnvironment; + waitpointMintKind: WaitpointMintKind; idempotencyKey?: string; idempotencyKeyExpiresAt?: Date; body: TriggerTaskRequest["body"]; @@ -987,6 +1004,7 @@ export class RunEngineTriggerTaskService { ? { id: args.options.batchId, index: args.options.batchIndex ?? 0 } : undefined, resumeParentOnCompletion: args.body.options?.resumeParentOnCompletion, + waitpointMintKind: args.waitpointMintKind, depth: args.depth, metadata: args.metadataPacket?.data, metadataType: args.metadataPacket?.dataType, diff --git a/apps/webapp/app/services/routeBuilders/apiBuilder.server.ts b/apps/webapp/app/services/routeBuilders/apiBuilder.server.ts index a5a808e3195..1f0004dda1a 100644 --- a/apps/webapp/app/services/routeBuilders/apiBuilder.server.ts +++ b/apps/webapp/app/services/routeBuilders/apiBuilder.server.ts @@ -25,13 +25,14 @@ import { getApiVersion } from "~/api/versions"; import { WORKER_HEADERS } from "@trigger.dev/core/v3/runEngineWorker"; import { ServiceValidationError } from "~/v3/services/common.server"; import { EngineServiceValidationError } from "@internal/run-engine"; +import { unroutableIdResponse } from "./unroutableId.server"; import { tenantContext, tenantContextFromAuthEnvironment } from "~/services/tenantContext.server"; // Client aborts and service-level validation errors aren't bugs — they're // expected at API boundaries. Log them at `warn` so they stay in stdout // without flowing to Sentry via Logger.onError. function logBoundaryError( - message: "Error in loader" | "Error in action", + message: "Error in loader" | "Error in action" | "Unroutable id", error: unknown, url: string ) { @@ -451,6 +452,12 @@ export function createLoaderApiRoute< return await wrapResponse(request, error, corsStrategy !== "none"); } + const unroutable = unroutableIdResponse(error); + if (unroutable) { + logBoundaryError("Unroutable id", error, request.url); + return await wrapResponse(request, unroutable, corsStrategy !== "none"); + } + logBoundaryError("Error in loader", error, request.url); return await wrapResponse( @@ -722,6 +729,12 @@ export function createLoaderPATApiRoute< if (error instanceof Response) { return await wrapResponse(request, error, corsStrategy !== "none"); } + + const unroutable = unroutableIdResponse(error); + if (unroutable) { + logBoundaryError("Unroutable id", error, request.url); + return await wrapResponse(request, unroutable, corsStrategy !== "none"); + } return await wrapResponse( request, json({ error: "Internal Server Error" }, { status: 500 }), @@ -996,6 +1009,12 @@ export function createActionPATApiRoute< return await wrapResponse(request, error, corsStrategy !== "none"); } + const unroutable = unroutableIdResponse(error); + if (unroutable) { + logBoundaryError("Unroutable id", error, request.url); + return await wrapResponse(request, unroutable, corsStrategy !== "none"); + } + logBoundaryError("Error in action", error, request.url); // Typed validation errors map to their own status (default 400); @@ -1346,6 +1365,12 @@ export function createActionApiRoute< return await wrapResponse(request, error, corsStrategy !== "none"); } + const unroutable = unroutableIdResponse(error); + if (unroutable) { + logBoundaryError("Unroutable id", error, request.url); + return await wrapResponse(request, unroutable, corsStrategy !== "none"); + } + logBoundaryError("Error in action", error, request.url); return await wrapResponse( @@ -1612,6 +1637,12 @@ export function createMultiMethodApiRoute< return await wrapResponse(request, error, corsStrategy !== "none"); } + const unroutable = unroutableIdResponse(error); + if (unroutable) { + logBoundaryError("Unroutable id", error, request.url); + return await wrapResponse(request, unroutable, corsStrategy !== "none"); + } + logBoundaryError("Error in action", error, request.url); return await wrapResponse( diff --git a/apps/webapp/app/services/routeBuilders/unroutableId.server.ts b/apps/webapp/app/services/routeBuilders/unroutableId.server.ts new file mode 100644 index 00000000000..dcc2d46d7b6 --- /dev/null +++ b/apps/webapp/app/services/routeBuilders/unroutableId.server.ts @@ -0,0 +1,19 @@ +import { json } from "@remix-run/server-runtime"; +import { UnknownShardKey } from "@internal/run-store"; + +/** + * An id naming a shard the topology has no store for cannot be routed, so a read cannot locate + * the row: that is a 404, and matches what an absent gen-1 or cuid id already returns. It must + * not be a 500 — `resolveShard` is pure id-shape, so any base32hex core plus `[a-z0-9]` plus "2" + * parses as gen-2, which lets any caller induce a 5xx, and a 5xx on a read trips canary rollbacks. + * + * The router still throws. Callers log it before returning this, so a genuine misconfiguration — + * a shard key dropped from a config that is meant to be append-only — still alarms. + */ +export function unroutableIdResponse(error: unknown): Response | undefined { + // Explicitly NOT retryable: an id naming an unconfigured shard is not a transient miss, and + // no number of retries makes a topology grow a store. + return error instanceof UnknownShardKey + ? json({ error: "Not Found" }, { status: 404, headers: { "x-should-retry": "false" } }) + : undefined; +} diff --git a/apps/webapp/app/v3/featureFlags.ts b/apps/webapp/app/v3/featureFlags.ts index 3a88beb54bc..2b527a367d1 100644 --- a/apps/webapp/app/v3/featureFlags.ts +++ b/apps/webapp/app/v3/featureFlags.ts @@ -36,6 +36,9 @@ export const FEATURE_FLAG = { runOpsMintShardSetFlippedAt: "runOpsMintShardSetFlippedAt", // Fleet-wide pin for the complete cutover. Beats every per-org and per-env pin. runOpsMintShardOverride: "runOpsMintShardOverride", + // Per-organization waitpoint coordinator selection. Read ONLY at waitpoint mint time; + // every later operation on a waitpoint routes by its id shape, never by this flag. + waitpointSystem: "waitpointSystem", queueMetricsUiEnabled: "queueMetricsUiEnabled", // Per-organization rollout for creating additional environment API keys. additionalApiKeysEnabled: "additionalApiKeysEnabled", @@ -95,6 +98,7 @@ export const FeatureFlagCatalog = { // Per-org run-ops-id mint cutover. Defaults to "cuid"; only honored when // RUN_OPS_MINT_ENABLED is on AND isSplitEnabled() is true. [FEATURE_FLAG.runOpsMintKind]: z.enum(["cuid", "runOpsId"]), + [FEATURE_FLAG.waitpointSystem]: z.enum(["legacy", "redis"]), // Grace-linger stamp: the previously-effective kind and the flip timestamp, written // by stampMintKindFlip on a genuine flip. Display-only (see ORG_LOCKED_FLAGS). [FEATURE_FLAG.runOpsMintKindPrev]: z.enum(["cuid", "runOpsId"]), diff --git a/apps/webapp/app/v3/runEngineHandlersShared.server.ts b/apps/webapp/app/v3/runEngineHandlersShared.server.ts index 4ce8cc2de8a..d8999e2332a 100644 --- a/apps/webapp/app/v3/runEngineHandlersShared.server.ts +++ b/apps/webapp/app/v3/runEngineHandlersShared.server.ts @@ -35,7 +35,8 @@ export async function readRunForEvent( deps: EventReadDeps ): Promise | null> { const result = await readThroughRun>({ - runId, + id: runId, + idKind: "run", environmentId, readNew: (client) => deps.store.findRun({ id: runId }, { select }, client), readLegacy: (replica) => deps.store.findRun({ id: runId }, { select }, replica), @@ -47,7 +48,7 @@ export async function readRunForEvent( }, }); - return result.source === "not-found" || result.source === "past-retention" ? null : result.value; + return result.found ? result.value : null; } /** diff --git a/apps/webapp/app/v3/runOpsMigration/readThrough.server.test.ts b/apps/webapp/app/v3/runOpsMigration/readThrough.server.test.ts index f7f7c43a530..8a657060ef9 100644 --- a/apps/webapp/app/v3/runOpsMigration/readThrough.server.test.ts +++ b/apps/webapp/app/v3/runOpsMigration/readThrough.server.test.ts @@ -13,6 +13,21 @@ vi.setConfig({ testTimeout: 60_000 }); // 25-char cuid body → LEGACY residency. 26-char v1 body (version "1" at index 25) → NEW residency. const LEGACY_RUN_ID = "run_" + "a".repeat(25); const NEW_RUN_ID = "run_" + "b".repeat(24) + "01"; +// 26-char gen-2 body: shard char at index 24, version "2" at index 25. +const SHARD_A_RUN_ID = "run_" + "c".repeat(24) + "a2"; +const SHARD_Z_RUN_ID = "run_" + "c".repeat(24) + "z2"; +const LEGACY_WAITPOINT_ID = "waitpoint_" + "d".repeat(25); + +function throwingClient(label: string) { + return vi.fn(async (): Promise<{ marker: number } | null> => { + throw new Error(`${label} must never be read`); + }); +} + +function collectingLogger() { + const errors: { message: string; meta?: unknown }[] = []; + return { errors, error: (message: string, meta?: unknown) => errors.push({ message, meta }) }; +} // Lightweight real read: a trivial `$queryRaw` that genuinely hits the given container. // `hit` controls whether the read "finds" the run, so we exercise routing without @@ -28,14 +43,7 @@ async function realRead( // A presenter-shaped mapping: both "not-found" and "past-retention" collapse to the // same 404-ish surface, so an old run after termination yields the normal response. function toHttpish(result: ReadThroughResult): { status: number; value?: T } { - switch (result.source) { - case "new": - case "legacy-replica": - return { status: 200, value: result.value }; - case "not-found": - case "past-retention": - return { status: 404 }; - } + return result.found ? { status: 200, value: result.value } : { status: 404 }; } describe("readThroughRun (legacy replica + new DB)", () => { @@ -46,7 +54,8 @@ describe("readThroughRun (legacy replica + new DB)", () => { // read resolving through `legacyReplica` (prisma14) IS the structural guarantee // that the primary is never touched. const result = await readThroughRun({ - runId: LEGACY_RUN_ID, + id: LEGACY_RUN_ID, + idKind: "run", environmentId: "env_1", readNew: (c) => realRead(c, false), readLegacy: (c) => realRead(c, true), @@ -57,7 +66,7 @@ describe("readThroughRun (legacy replica + new DB)", () => { }, }); - expect(result.source).toBe("legacy-replica"); + expect(result.found && result.source).toBe("legacy-replica"); expect(toHttpish(result).status).toBe(200); } ); @@ -66,7 +75,8 @@ describe("readThroughRun (legacy replica + new DB)", () => { "post-termination past-retention returns the normal not-found surface", async ({ prisma14, prisma17 }) => { const pastRetentionResult = await readThroughRun({ - runId: LEGACY_RUN_ID, + id: LEGACY_RUN_ID, + idKind: "run", environmentId: "env_1", readNew: (c) => realRead(c, false), readLegacy: (c) => realRead(c, false), // legacy gone / retention elapsed @@ -78,11 +88,14 @@ describe("readThroughRun (legacy replica + new DB)", () => { }, }); - expect(pastRetentionResult.source).toBe("past-retention"); + expect(pastRetentionResult.found === false && pastRetentionResult.reason).toBe( + "past-retention" + ); // A run that is simply absent (not past retention) yields not-found. const notFoundResult = await readThroughRun({ - runId: LEGACY_RUN_ID, + id: LEGACY_RUN_ID, + idKind: "run", environmentId: "env_1", readNew: (c) => realRead(c, false), readLegacy: (c) => realRead(c, false), @@ -94,7 +107,7 @@ describe("readThroughRun (legacy replica + new DB)", () => { }, }); - expect(notFoundResult.source).toBe("not-found"); + expect(notFoundResult.found === false && notFoundResult.reason).toBe("not-found"); // Both collapse to the same 404-ish surface. expect(toHttpish(pastRetentionResult).status).toBe(toHttpish(notFoundResult).status); expect(toHttpish(pastRetentionResult).status).toBe(404); @@ -110,7 +123,8 @@ describe("readThroughRun (legacy replica + new DB)", () => { const newRead = vi.fn((c: PrismaReplicaClient) => realRead(c, true)); const result = await readThroughRun({ - runId: LEGACY_RUN_ID, + id: LEGACY_RUN_ID, + idKind: "run", environmentId: "env_1", readNew: newRead, readLegacy: throwingLegacy, @@ -121,7 +135,7 @@ describe("readThroughRun (legacy replica + new DB)", () => { }, }); - expect(result.source).toBe("new"); + expect(result.found && result.source).toBe("new"); expect(newRead).toHaveBeenCalledTimes(1); expect(throwingLegacy).not.toHaveBeenCalled(); } @@ -135,7 +149,152 @@ describe("readThroughRun (legacy replica + new DB)", () => { }); const result = await readThroughRun({ - runId: NEW_RUN_ID, + id: NEW_RUN_ID, + idKind: "run", + environmentId: "env_1", + readNew: (c) => realRead(c, true), + readLegacy: throwingLegacy, + deps: { + splitEnabled: true, + newClient: prisma17 as unknown as PrismaReplicaClient, + legacyReplica: prisma14 as unknown as PrismaReplicaClient, + }, + }); + + expect(result.found && result.source).toBe("new"); + expect(throwingLegacy).not.toHaveBeenCalled(); + } + ); + + heteroPostgresTest( + "gen-2 id reads its OWN shard replica once and probes no other store", + async ({ prisma14, prisma17 }) => { + const throwingNew = throwingClient("the gen-1 new store"); + const throwingLegacy = throwingClient("the legacy replica"); + const shardRead = vi.fn((c: PrismaReplicaClient) => realRead(c, true)); + + const result = await readThroughRun({ + id: SHARD_A_RUN_ID, + idKind: "run", + environmentId: "env_1", + // One closure serves both the gen-1 new store and a shard: a shard is the same + // dedicated schema. The throwing clients prove WHICH client it was handed. + readNew: (c) => shardRead(c), + readLegacy: throwingLegacy, + deps: { + splitEnabled: true, + newClient: throwingNew as unknown as PrismaReplicaClient, + legacyReplica: prisma14 as unknown as PrismaReplicaClient, + shardReplicas: new Map([["a", prisma17 as unknown as PrismaReplicaClient]]), + }, + }); + + expect(result.found && result.source).toBe("shard:a"); + expect(shardRead).toHaveBeenCalledTimes(1); + // Identity, not deep equality: a Prisma client is too large to deep-compare. + expect(shardRead.mock.calls[0][0]).toBe(prisma17); + expect(throwingLegacy).not.toHaveBeenCalled(); + } + ); + + heteroPostgresTest( + "gen-2 id on an UNCONFIGURED shard key logs an error and returns not-found, never throws", + async ({ prisma14, prisma17 }) => { + const logger = collectingLogger(); + const throwingLegacy = throwingClient("the legacy replica"); + const newRead = vi.fn((c: PrismaReplicaClient) => realRead(c, true)); + + // Shard "z" is not configured. A 500 here would be inducible by any caller that + // guesses a shard char, so the layer must degrade rather than throw. + const result = await readThroughRun({ + id: SHARD_Z_RUN_ID, + idKind: "run", + environmentId: "env_1", + readNew: newRead, + readLegacy: throwingLegacy, + deps: { + splitEnabled: true, + newClient: prisma17 as unknown as PrismaReplicaClient, + legacyReplica: prisma14 as unknown as PrismaReplicaClient, + shardReplicas: new Map([["a", prisma17 as unknown as PrismaReplicaClient]]), + logger, + }, + }); + + expect(result.found).toBe(false); + expect(result.found === false && result.reason).toBe("not-found"); + expect(logger.errors).toHaveLength(1); + expect(logger.errors[0].meta).toMatchObject({ shardKey: "z", configured: ["a"] }); + // It must not silently fall back onto a gen-1 store. + expect(newRead).not.toHaveBeenCalled(); + expect(throwingLegacy).not.toHaveBeenCalled(); + } + ); + + heteroPostgresTest( + "gen-1 RUN id reads the legacy replica only and never probes the new store", + async ({ prisma14 }) => { + const throwingNew = throwingClient("the new store"); + const legacyRead = vi.fn((c: PrismaReplicaClient) => realRead(c, true)); + + const result = await readThroughRun({ + id: LEGACY_RUN_ID, + idKind: "run", + environmentId: "env_1", + readNew: throwingNew, + readLegacy: legacyRead, + deps: { + splitEnabled: true, + newClient: prisma14 as unknown as PrismaReplicaClient, + legacyReplica: prisma14 as unknown as PrismaReplicaClient, + }, + }); + + expect(result.found && result.source).toBe("legacy-replica"); + expect(throwingNew).not.toHaveBeenCalled(); + expect(legacyRead).toHaveBeenCalledTimes(1); + } + ); + + heteroPostgresTest( + "cuid WAITPOINT id keeps the new-FIRST pair probe (frozen: cuid waitpoints co-locate on new)", + async ({ prisma14, prisma17 }) => { + const calls: string[] = []; + const newRead = vi.fn(async (c: PrismaReplicaClient) => { + calls.push("new"); + return realRead(c, false); + }); + const legacyRead = vi.fn(async (c: PrismaReplicaClient) => { + calls.push("legacy"); + return realRead(c, true); + }); + + const result = await readThroughRun({ + id: LEGACY_WAITPOINT_ID, + idKind: "waitpoint", + environmentId: "env_1", + readNew: newRead, + readLegacy: legacyRead, + deps: { + splitEnabled: true, + newClient: prisma17 as unknown as PrismaReplicaClient, + legacyReplica: prisma14 as unknown as PrismaReplicaClient, + }, + }); + + expect(result.found && result.source).toBe("legacy-replica"); + expect(calls).toEqual(["new", "legacy"]); + } + ); + + heteroPostgresTest( + "a cuid waitpoint found on the new store returns it without touching legacy", + async ({ prisma14, prisma17 }) => { + const throwingLegacy = throwingClient("the legacy replica"); + + const result = await readThroughRun({ + id: LEGACY_WAITPOINT_ID, + idKind: "waitpoint", environmentId: "env_1", readNew: (c) => realRead(c, true), readLegacy: throwingLegacy, @@ -146,7 +305,7 @@ describe("readThroughRun (legacy replica + new DB)", () => { }, }); - expect(result.source).toBe("new"); + expect(result.found && result.source).toBe("new"); expect(throwingLegacy).not.toHaveBeenCalled(); } ); diff --git a/apps/webapp/app/v3/runOpsMigration/readThrough.server.ts b/apps/webapp/app/v3/runOpsMigration/readThrough.server.ts index f15230ec442..6e1beaae62c 100644 --- a/apps/webapp/app/v3/runOpsMigration/readThrough.server.ts +++ b/apps/webapp/app/v3/runOpsMigration/readThrough.server.ts @@ -3,12 +3,18 @@ * (which carries the read load we are shedding). Disabled entirely when isSplitEnabled() * is false (single-DB passthrough). * - * During the retention window, old run-ops rows are served off the legacy read replica. - * Residency is decided purely by id-shape: a run-ops id (NEW) id reads new only, a cuid - * (LEGACY) id reads legacy only. An unclassifiable id falls back to a new-then-legacy - * probe. After termination, past-retention runs return the normal not-found response. - * Patterned on `mollifier/resolveRunForMutation.server.ts` (`?? default` DI), but with - * the legacy-primary/writer fallback deliberately removed: this layer has NO legacy-writer + * Residency is decided purely by id-shape, via `resolveShard`: a gen-2 body names its own + * shard (ONE read there), a gen-1 v1 body reads new only, everything else is legacy and + * routes on `idKind`. + * + * `idKind` is required because a cuid gives no way to tell a run id from a waitpoint id, + * and the two must route differently: a legacy-classified RUN id is legacy-resident (there + * is no cuid run migration), while a cuid WAITPOINT can be co-located with its run on the + * new store, which is what makes the new-first probe load-bearing for it. No default — + * a default would pick one of those arms silently. + * + * Patterned on `mollifier/resolveRunForMutation.server.ts` (`?? default` DI), but with the + * legacy-primary/writer fallback deliberately removed: this layer has NO legacy-writer * handle at all (structural guarantee). */ import type { PrismaReplicaClient } from "~/db.server"; @@ -17,90 +23,118 @@ import { runOpsNewReplica as defaultNewClient, } from "~/db.server"; import { logger as defaultLogger } from "~/services/logger.server"; -import { ownerEngine, UnclassifiableRunId } from "@trigger.dev/core/v3/isomorphic"; +import { resolveShard, type ShardKey } from "@trigger.dev/core/v3/isomorphic"; import { isSplitEnabled } from "./splitMode.server"; +import { runOpsShardReplicas } from "./shardHandles.server"; + +type ShardSource = `shard:${string}`; -type ReadThroughSource = "new" | "legacy-replica"; +type ReadThroughSource = "new" | "legacy-replica" | ShardSource; +/** + * `found` carries hit/miss STRUCTURALLY. `source` is open-ended once shards exist, so a + * consumer testing found-ness by listing hit sources reads a gen-2 hit as a miss; + * discriminating on `found` makes that a compile error instead. + */ export type ReadThroughResult = - | { source: ReadThroughSource; value: T } - | { source: "not-found" } - | { source: "past-retention" }; + | { found: true; source: ReadThroughSource; value: T } + | { found: false; reason: "not-found" | "past-retention" }; type ReadThroughDeps = { newClient?: PrismaReplicaClient; legacyReplica?: PrismaReplicaClient; + /** + * Gen-2 shard replicas by shard char; empty (RUN_OPS_SHARDS unset) makes the gen-2 arm + * unreachable. Load-bearing only for callers whose closures read a client DIRECTLY: + * `RoutingRunStore` never forwards a caller's client, so for store-backed closures the + * client picked here is only a read-your-writes signal. Not dead weight. + */ + shardReplicas?: ReadonlyMap; /** Resolved boot constant; never `await`ed per-request when supplied. */ splitEnabled?: boolean; - isPastRetention?: (runId: string) => boolean; - logger?: { warn: (m: string, meta?: unknown) => void }; + isPastRetention?: (id: string) => boolean; + logger?: { error: (m: string, meta?: Record) => void }; /** Saturation-signal emit hook: called on each legacy-replica hit. */ - onLegacyReplicaRead?: (runId: string) => void; + onLegacyReplicaRead?: (id: string) => void; }; type ReadThroughRunInput = { - runId: string; + id: string; + idKind: "run" | "waitpoint"; environmentId: string; readNew: (client: PrismaReplicaClient) => Promise; readLegacy: (replica: PrismaReplicaClient) => Promise; deps?: ReadThroughDeps; }; +function hit(source: ReadThroughSource, value: T): ReadThroughResult { + return { found: true, source, value }; +} + +function miss(reason: "not-found" | "past-retention"): ReadThroughResult { + return { found: false, reason }; +} + export async function readThroughRun( input: ReadThroughRunInput ): Promise> { - const { runId, deps } = input; + const { id, idKind, deps } = input; const newClient = deps?.newClient ?? defaultNewClient; const legacyReplica = deps?.legacyReplica ?? defaultLegacyReplica; + const shardReplicas = deps?.shardReplicas ?? runOpsShardReplicas; const logger = deps?.logger ?? defaultLogger; const splitEnabled = deps?.splitEnabled ?? (await isSplitEnabled()); - // Passthrough: single plain read against the one collapsed store. No legacy read, - // no second connection. + // Passthrough: single plain read against the one collapsed store. if (!splitEnabled) { const v = await input.readNew(newClient); - return v != null ? { source: "new", value: v } : { source: "not-found" }; + return v != null ? hit("new", v) : miss("not-found"); } - // Split is on. Classify residency; an unclassifiable id is treated as LEGACY - // (conservative — probe rather than drop a real run). - let residency: "LEGACY" | "NEW"; - try { - residency = ownerEngine(runId); - } catch (e) { - if (e instanceof UnclassifiableRunId) { - logger.warn("readThroughRun: UnclassifiableRunId, treating as LEGACY", { - runId, - valueLength: e.valueLength, + // Total: an unclassifiable id resolves to "legacy" (probe rather than drop a real run). + const shardKey = resolveShard(id); + + if (shardKey !== "new" && shardKey !== "legacy") { + const shardReplica = shardReplicas.get(shardKey); + if (shardReplica === undefined) { + // Deliberately not a throw: this id arrives from the caller (a URL param on the + // waitpoint route) and any base32hex core + [a-z0-9] + "2" parses as gen-2, so a + // throw is a 500 any client can induce. An error-logged not-found is neither silent + // nor a misroute. Throwing stays correct on the router path, where ids are minted. + logger.error("readThroughRun: gen-2 id resolved to an unconfigured shard key", { + id, + shardKey, + configured: [...shardReplicas.keys()], }); - residency = "LEGACY"; - } else { - throw e; + return miss("not-found"); } + // A gen-2 shard is a dedicated-schema store, exactly like `new`, so `readNew` fits. + const v = await input.readNew(shardReplica); + return v != null ? hit(`shard:${shardKey}`, v) : miss("not-found"); } - // A run-ops id can only live on the new DB — skip the legacy replica entirely. - if (residency === "NEW") { + if (shardKey === "new") { const v = await input.readNew(newClient); - return v != null ? { source: "new", value: v } : { source: "not-found" }; + return v != null ? hit("new", v) : miss("not-found"); } - // LEGACY (or unclassifiable→LEGACY) fan-out: new first. - const v = await input.readNew(newClient); - if (v != null) { - return { source: "new", value: v }; + if (idKind === "waitpoint") { + const v = await input.readNew(newClient); + if (v != null) { + return hit("new", v); + } } // Legacy READ REPLICA only — never a legacy writer/primary (no such handle exists). const lv = await input.readLegacy(legacyReplica); if (lv != null) { - deps?.onLegacyReplicaRead?.(runId); - return { source: "legacy-replica", value: lv }; + deps?.onLegacyReplicaRead?.(id); + return hit("legacy-replica", lv); } - if (deps?.isPastRetention?.(runId)) { - return { source: "past-retention" }; + if (deps?.isPastRetention?.(id)) { + return miss("past-retention"); } - return { source: "not-found" }; + return miss("not-found"); } diff --git a/apps/webapp/app/v3/runOpsMigration/shardHandles.server.test.ts b/apps/webapp/app/v3/runOpsMigration/shardHandles.server.test.ts new file mode 100644 index 00000000000..b90c1dfe7c4 --- /dev/null +++ b/apps/webapp/app/v3/runOpsMigration/shardHandles.server.test.ts @@ -0,0 +1,37 @@ +import { describe, expect, it } from "vitest"; +import { buildShardHandleMaps } from "./shardHandles.server"; + +// Two distinct sentinels per shard: the maps must not cross writer and replica. +function handle(key: string) { + return { + key, + writer: { tag: `${key}-writer` } as never, + replica: { tag: `${key}-replica` } as never, + }; +} + +describe("buildShardHandleMaps", () => { + it("yields empty maps when no shard is configured", () => { + const { replicas, writers } = buildShardHandleMaps([]); + + expect(replicas.size).toBe(0); + expect(writers.size).toBe(0); + }); + + it("keys each shard's replica and writer under its shard char", () => { + const { replicas, writers } = buildShardHandleMaps([handle("a"), handle("b")]); + + expect([...replicas.keys()].sort()).toEqual(["a", "b"]); + expect([...writers.keys()].sort()).toEqual(["a", "b"]); + expect(replicas.get("a")).toEqual({ tag: "a-replica" }); + expect(writers.get("a")).toEqual({ tag: "a-writer" }); + expect(replicas.get("b")).toEqual({ tag: "b-replica" }); + expect(writers.get("b")).toEqual({ tag: "b-writer" }); + }); + + it("never places a writer in the replica map", () => { + const { replicas } = buildShardHandleMaps([handle("a")]); + + expect(replicas.get("a")).not.toEqual({ tag: "a-writer" }); + }); +}); diff --git a/apps/webapp/app/v3/runOpsMigration/shardHandles.server.ts b/apps/webapp/app/v3/runOpsMigration/shardHandles.server.ts new file mode 100644 index 00000000000..cbe827be4dc --- /dev/null +++ b/apps/webapp/app/v3/runOpsMigration/shardHandles.server.ts @@ -0,0 +1,46 @@ +/** + * Gen-2 shard client handles, keyed by shard char, for the consumers that route by + * `resolveShard` outside the run-store boundary: read-through and the two cross-seam + * batch hydration sites. Both maps are empty unless RUN_OPS_SHARDS is configured, which + * is what keeps every gen-2 arm unreachable today. + */ +import type { PrismaClient } from "@trigger.dev/database"; +import type { ShardKey } from "@trigger.dev/core/v3/isomorphic"; +import type { PrismaReplicaClient } from "~/db.server"; +import { runOpsShardHandles } from "~/db.server"; + +type ShardHandle = { + key: string; + writer: unknown; + replica: unknown; +}; + +export function buildShardHandleMaps(handles: ShardHandle[]): { + replicas: ReadonlyMap; + writers: ReadonlyMap; +} { + const replicas = new Map(); + const writers = new Map(); + for (const handle of handles) { + replicas.set(handle.key, handle.replica as PrismaReplicaClient); + writers.set(handle.key, handle.writer as PrismaClient); + } + return { replicas, writers }; +} + +// A gen-2 shard is the same dedicated subset schema as the gen-1 new store, so these casts +// carry exactly the precedent (and the same residual risk) as `runOpsNewPrisma`'s. +// The try/catch mirrors `runStore.server.ts`'s handle resolution: a minimal `db.server` mock +// does not define this export at all, and accessing an undefined mock export throws. +function resolveShardHandles(): ShardHandle[] { + try { + return runOpsShardHandles ?? []; + } catch { + return []; + } +} + +const maps = buildShardHandleMaps(resolveShardHandles()); + +export const runOpsShardReplicas = maps.replicas; +export const runOpsShardWriters = maps.writers; diff --git a/apps/webapp/app/v3/runOpsMigration/track1-baseline.json b/apps/webapp/app/v3/runOpsMigration/track1-baseline.json index 63d27dbadca..26d038678ea 100644 --- a/apps/webapp/app/v3/runOpsMigration/track1-baseline.json +++ b/apps/webapp/app/v3/runOpsMigration/track1-baseline.json @@ -68,19 +68,19 @@ "WaitpointTag.project" ], "totals": { - "violations": 4, - "detectorI": 4, + "violations": 5, + "detectorI": 5, "detectorII": 0, "detectorIII": 0, "write": 0, - "read": 4, + "read": 5, "files": 1, "legacyAnnotations": 0 }, "violations": [ { "file": "apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts", - "line": 89, + "line": 93, "model": "BatchTaskRun", "delegate": "batchTaskRun", "callKind": "read", @@ -89,7 +89,7 @@ }, { "file": "apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts", - "line": 150, + "line": 154, "model": "BatchTaskRun", "delegate": "batchTaskRun", "callKind": "read", @@ -98,16 +98,25 @@ }, { "file": "apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts", - "line": 184, + "line": 214, "model": "TaskRun", "delegate": "taskRun", "callKind": "read", "detector": "i", - "snippet": "const newRows = (await newClient.taskRun.findMany({" + "snippet": "? ((await newClient.taskRun.findMany({" }, { "file": "apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts", - "line": 196, + "line": 224, + "model": "TaskRun", + "delegate": "taskRun", + "callKind": "read", + "detector": "i", + "snippet": "(await shardReplicas.get(shardKey)!.taskRun.findMany({" + }, + { + "file": "apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts", + "line": 241, "model": "TaskRun", "delegate": "taskRun", "callKind": "read", diff --git a/apps/webapp/app/v3/runOpsMigration/waitpointTokenResolve.server.test.ts b/apps/webapp/app/v3/runOpsMigration/waitpointTokenResolve.server.test.ts index 9ea849c8058..42dca92e1ce 100644 --- a/apps/webapp/app/v3/runOpsMigration/waitpointTokenResolve.server.test.ts +++ b/apps/webapp/app/v3/runOpsMigration/waitpointTokenResolve.server.test.ts @@ -138,7 +138,8 @@ describe("public wait-token resolution across the split boundary", () => { expect(gated?.id).toBe(waitpointId); const passthrough = await readThroughRun({ - runId: waitpointId, + id: waitpointId, + idKind: "waitpoint", environmentId: environment.id, readNew: (c) => read(c), readLegacy: (r) => read(r), @@ -150,7 +151,7 @@ describe("public wait-token resolution across the split boundary", () => { }); expect(gated).not.toBeNull(); - expect(passthrough.source).toBe("not-found"); + expect(passthrough.found === false && passthrough.reason).toBe("not-found"); } ); }); diff --git a/apps/webapp/app/v3/services/bulk/BulkActionV2.batchReadThrough.server.test.ts b/apps/webapp/app/v3/services/bulk/BulkActionV2.batchReadThrough.server.test.ts index 99d4cfd2dd7..779a7234748 100644 --- a/apps/webapp/app/v3/services/bulk/BulkActionV2.batchReadThrough.server.test.ts +++ b/apps/webapp/app/v3/services/bulk/BulkActionV2.batchReadThrough.server.test.ts @@ -13,6 +13,8 @@ vi.setConfig({ testTimeout: 60_000 }); // 25-char cuid body → LEGACY residency. 26-char v1 body (version "1" at index 25) → NEW residency. const LEGACY_RUN_ID = "run_" + "a".repeat(25); const NEW_RUN_ID = "run_" + "b".repeat(24) + "01"; +// 26-char gen-2 body: shard char at index 24, version "2" at index 25. +const SHARD_A_RUN_ID = "run_" + "c".repeat(24) + "a2"; type Row = { id: string }; @@ -90,4 +92,109 @@ describe("hydrateRunsAcrossSeam (PG14 legacy replica + PG17 new)", () => { expect(throwingLegacy).not.toHaveBeenCalled(); } ); + + heteroPostgresTest( + "(c) a gen-2 id hydrates from its OWN shard and is never read from the gen-1 stores", + async ({ prisma14, prisma17 }) => { + // Before the shard arm existed a gen-2 id joined the `new` group, missed, and was + // never legacy-probed either — so it vanished from the page with no error. + const onShardA = new Set([SHARD_A_RUN_ID]); + + const readNew = vi.fn(async (client: PrismaReplicaClient, ids: string[]): Promise => { + if ( + ids.includes(SHARD_A_RUN_ID) && + client !== (prisma17 as unknown as PrismaReplicaClient) + ) { + throw new Error("a gen-2 id must only be read on its own shard"); + } + return realReadFiltered(client, ids, onShardA); + }); + const readLegacyReplica = vi.fn( + async (_replica: PrismaReplicaClient, ids: string[]): Promise => { + if (ids.includes(SHARD_A_RUN_ID)) { + throw new Error("a gen-2 id must never reach the legacy probe"); + } + return []; + } + ); + + const rows = await hydrateRunsAcrossSeam({ + runIds: [SHARD_A_RUN_ID], + readNew, + readLegacyReplica, + deps: { + splitEnabled: true, + newClient: prisma14 as unknown as PrismaReplicaClient, + legacyReplica: prisma14 as unknown as PrismaReplicaClient, + shardReplicas: new Map([["a", prisma17 as unknown as PrismaReplicaClient]]), + }, + }); + + expect(rows.map((r) => r.id)).toEqual([SHARD_A_RUN_ID]); + expect(readLegacyReplica).not.toHaveBeenCalled(); + } + ); + + heteroPostgresTest( + "(d) a mixed gen-1 and gen-2 page hydrates every member", + async ({ prisma14, prisma17 }) => { + const onGenOneNew = new Set([NEW_RUN_ID]); + const onLegacy = new Set([LEGACY_RUN_ID]); + const onShardA = new Set([SHARD_A_RUN_ID]); + + const readNew = vi.fn(async (client: PrismaReplicaClient, ids: string[]): Promise => { + const present = ids.includes(SHARD_A_RUN_ID) ? onShardA : onGenOneNew; + return realReadFiltered(client, ids, present); + }); + const readLegacyReplica = vi.fn( + async (replica: PrismaReplicaClient, ids: string[]): Promise => + realReadFiltered(replica, ids, onLegacy) + ); + + const rows = await hydrateRunsAcrossSeam({ + runIds: [NEW_RUN_ID, LEGACY_RUN_ID, SHARD_A_RUN_ID], + readNew, + readLegacyReplica, + deps: { + splitEnabled: true, + newClient: prisma17 as unknown as PrismaReplicaClient, + legacyReplica: prisma14 as unknown as PrismaReplicaClient, + shardReplicas: new Map([["a", prisma17 as unknown as PrismaReplicaClient]]), + }, + }); + + expect(rows.map((r) => r.id).sort()).toEqual( + [NEW_RUN_ID, LEGACY_RUN_ID, SHARD_A_RUN_ID].sort() + ); + } + ); + + heteroPostgresTest( + "(e) a gen-2 id on an unconfigured shard is dropped with a logged error, not read elsewhere", + async ({ prisma14, prisma17 }) => { + const errors: unknown[] = []; + const readNew = vi.fn(async (client: PrismaReplicaClient, ids: string[]): Promise => { + if (ids.includes(SHARD_A_RUN_ID)) { + throw new Error("an unconfigured gen-2 id must not fall back to a gen-1 store"); + } + return realReadFiltered(client, ids, new Set([NEW_RUN_ID])); + }); + + const rows = await hydrateRunsAcrossSeam({ + runIds: [NEW_RUN_ID, SHARD_A_RUN_ID], + readNew, + readLegacyReplica: async () => [], + deps: { + splitEnabled: true, + newClient: prisma17 as unknown as PrismaReplicaClient, + legacyReplica: prisma14 as unknown as PrismaReplicaClient, + shardReplicas: new Map(), + logger: { error: (_m, meta) => errors.push(meta) }, + }, + }); + + expect(rows.map((r) => r.id)).toEqual([NEW_RUN_ID]); + expect(errors).toHaveLength(1); + } + ); }); diff --git a/apps/webapp/app/v3/services/bulk/BulkActionV2.batchReadThrough.server.ts b/apps/webapp/app/v3/services/bulk/BulkActionV2.batchReadThrough.server.ts index c7a0dc735e8..bc476cebf31 100644 --- a/apps/webapp/app/v3/services/bulk/BulkActionV2.batchReadThrough.server.ts +++ b/apps/webapp/app/v3/services/bulk/BulkActionV2.batchReadThrough.server.ts @@ -20,7 +20,8 @@ import { runOpsLegacyReplica as defaultLegacyReplica, runOpsNewReplica as defaultNewClient, } from "~/db.server"; -import { ownerEngine, UnclassifiableRunId } from "@trigger.dev/core/v3/isomorphic"; +import { resolveShard, type ShardKey } from "@trigger.dev/core/v3/isomorphic"; +import { runOpsShardReplicas as defaultShardReplicas } from "~/v3/runOpsMigration/shardHandles.server"; type SeamReadDeps = { /** @@ -30,7 +31,9 @@ type SeamReadDeps = { splitEnabled: boolean; newClient?: PrismaReplicaClient; legacyReplica?: PrismaReplicaClient; - logger?: { warn: (m: string, meta?: unknown) => void }; + /** Gen-2 shard replicas by shard char; empty (RUN_OPS_SHARDS unset) keeps today's behaviour. */ + shardReplicas?: ReadonlyMap; + logger?: { error: (m: string, meta?: Record) => void }; }; type HydrateRunsAcrossSeamInput = { @@ -61,28 +64,30 @@ export async function hydrateRunsAcrossSeam(input: HydrateRunsAcrossSeamInput return input.readNew(newClient, runIds); } - // Split is on. Classify residency; unclassifiable → LEGACY (probe rather than drop). + // Split is on. Partition by shard key; `resolveShard` is total, so an unclassifiable id + // resolves to "legacy" (probe rather than drop). A gen-2 id goes to its OWN shard and to + // no other store: it is directly routable, so it joins neither gen-1 group. + const shardReplicas = deps.shardReplicas ?? defaultShardReplicas; const newIds: string[] = []; const legacyCandidateIds: string[] = []; + const idsByShard = new Map(); for (const runId of runIds) { - let residency: "LEGACY" | "NEW"; - try { - residency = ownerEngine(runId); - } catch (e) { - if (e instanceof UnclassifiableRunId) { - deps.logger?.warn("hydrateRunsAcrossSeam: UnclassifiableRunId, treating as LEGACY", { - runId, - valueLength: e.valueLength, - }); - residency = "LEGACY"; - } else { - throw e; - } - } - if (residency === "NEW") { + const shardKey = resolveShard(runId); + if (shardKey === "new") { newIds.push(runId); - } else { + } else if (shardKey === "legacy") { legacyCandidateIds.push(runId); + } else if (shardReplicas.has(shardKey)) { + const group = idsByShard.get(shardKey); + group ? group.push(runId) : idsByShard.set(shardKey, [runId]); + } else { + // Not routable and not a gen-1 shape. Reading a gen-1 store would query the wrong + // database, so the id is dropped from the page — loudly, never silently. + deps.logger?.error("hydrateRunsAcrossSeam: gen-2 id on an unconfigured shard key", { + runId, + shardKey, + configured: [...shardReplicas.keys()], + }); } } @@ -103,6 +108,16 @@ export async function hydrateRunsAcrossSeam(input: HydrateRunsAcrossSeamInput legacyRows = await input.readLegacyReplica(legacyReplica, legacyToProbe); } + // Each configured shard is read once, in parallel: the groups are disjoint by id, so the + // results need no dedupe. + const shardRows = ( + await Promise.all( + [...idsByShard.entries()].map(([shardKey, ids]) => + input.readNew(shardReplicas.get(shardKey)!, ids) + ) + ) + ).flat(); + // Order within the page is irrelevant (downstream pMap does not depend on it). - return [...newRows, ...legacyRows]; + return [...newRows, ...legacyRows, ...shardRows]; } diff --git a/apps/webapp/app/v3/waitpointMigration/waitpointMintKind.server.ts b/apps/webapp/app/v3/waitpointMigration/waitpointMintKind.server.ts new file mode 100644 index 00000000000..75c00d328f3 --- /dev/null +++ b/apps/webapp/app/v3/waitpointMigration/waitpointMintKind.server.ts @@ -0,0 +1,66 @@ +import { $replica } from "~/db.server"; +import { env } from "~/env.server"; +import { BoundedTtlCache } from "~/services/realtime/boundedTtlCache"; +import { singleton } from "~/utils/singleton"; +import { logger } from "~/services/logger.server"; +import { FEATURE_FLAG } from "~/v3/featureFlags"; +import { computeWaitpointMintKind, type WaitpointMintKind } from "./waitpointMintKind.js"; + +export type { WaitpointMintKind }; + +// The two unions are declared separately, because the engine never imports from the +// webapp. Nothing pins them together here on purpose: every call site passes this value +// into an engine method, so a drift fails at those call sites, where the error is local +// to the code that actually broke. + +type WaitpointSystemFlag = "legacy" | "redis"; + +const mintCache = singleton( + "waitpointMintCache", + () => + new BoundedTtlCache( + env.WAITPOINT_MINT_FLAG_CACHE_TTL_MS, + env.WAITPOINT_MINT_FLAG_CACHE_MAX_ENTRIES + ) +); + +// ENV-BOUND wrapper — the only place env and $replica are read. +export async function resolveWaitpointMintKind(environment: { + organizationId: string; + id: string; + /** Pass environment.organization.featureFlags from the call site. */ + orgFeatureFlags?: unknown; +}): Promise { + return computeWaitpointMintKind(environment, { + globalDefault: env.WAITPOINT_SYSTEM_DEFAULT, + onError: (error) => + logger.error("[waitpointMintKind] flag read failed; minting legacy (fail-safe)", { error }), + flag: async (orgId, orgFeatureFlags) => { + // null is a cached "this org has no override", which must stay distinct from a miss: + // BoundedTtlCache reports a stored undefined as a miss, so never store undefined. + const cached = mintCache.get(orgId); + if (cached !== undefined) { + return cached ?? undefined; + } + + // Hot-path pass-through: only read the replica when the caller passed no org flags. + const overrides = + orgFeatureFlags !== undefined + ? orgFeatureFlags + : ( + await $replica.organization.findFirst({ + where: { id: orgId }, + select: { featureFlags: true }, + }) + )?.featureFlags; + + const value = (overrides as Record | null | undefined)?.[ + FEATURE_FLAG.waitpointSystem + ]; + const resolved = value === "redis" || value === "legacy" ? value : null; + + mintCache.set(orgId, resolved); + return resolved ?? undefined; + }, + }); +} diff --git a/apps/webapp/app/v3/waitpointMigration/waitpointMintKind.test.ts b/apps/webapp/app/v3/waitpointMigration/waitpointMintKind.test.ts new file mode 100644 index 00000000000..79c0dc37fe0 --- /dev/null +++ b/apps/webapp/app/v3/waitpointMigration/waitpointMintKind.test.ts @@ -0,0 +1,64 @@ +import { describe, expect, it, vi } from "vitest"; +import { computeWaitpointMintKind } from "./waitpointMintKind"; + +const environment = { organizationId: "org_1", id: "env_1" }; + +describe("computeWaitpointMintKind", () => { + it("returns legacy when the org has no override and the default is legacy", async () => { + const kind = await computeWaitpointMintKind(environment, { + globalDefault: "legacy", + flag: async () => undefined, + }); + + expect(kind).toBe("legacy"); + }); + + it("returns store when the org override is redis", async () => { + const kind = await computeWaitpointMintKind(environment, { + globalDefault: "legacy", + flag: async () => "redis", + }); + + expect(kind).toBe("store"); + }); + + it("lets an explicit org legacy override beat a redis global default", async () => { + const kind = await computeWaitpointMintKind(environment, { + globalDefault: "redis", + flag: async () => "legacy", + }); + + expect(kind).toBe("legacy"); + }); + + it("falls back to the global default when the org has no override", async () => { + const kind = await computeWaitpointMintKind(environment, { + globalDefault: "redis", + flag: async () => undefined, + }); + + expect(kind).toBe("store"); + }); + + it("fails safe to legacy when the flag read throws", async () => { + const kind = await computeWaitpointMintKind(environment, { + globalDefault: "redis", + flag: async () => { + throw new Error("replica down"); + }, + }); + + expect(kind).toBe("legacy"); + }); + + it("hands the pre-loaded org flags to the flag reader", async () => { + const flag = vi.fn(async () => "redis" as const); + + await computeWaitpointMintKind( + { ...environment, orgFeatureFlags: { waitpointSystem: "redis" } }, + { globalDefault: "legacy", flag } + ); + + expect(flag).toHaveBeenCalledWith("org_1", { waitpointSystem: "redis" }); + }); +}); diff --git a/apps/webapp/app/v3/waitpointMigration/waitpointMintKind.ts b/apps/webapp/app/v3/waitpointMigration/waitpointMintKind.ts new file mode 100644 index 00000000000..3b500f5f65c --- /dev/null +++ b/apps/webapp/app/v3/waitpointMigration/waitpointMintKind.ts @@ -0,0 +1,37 @@ +// Pure: no server-only imports, so a test can drive this without loading env.server. +/** + * Which coordinator mints a NEW waitpoint. Consulted at the mint and never again: every + * later operation routes by id shape. A flip therefore changes only where the NEXT + * waitpoint is born, which is why this needs no flip-grace machinery. + */ +export type WaitpointMintKind = "legacy" | "store"; + +/** The flag's vocabulary, deliberately not the coordinator's. */ +type WaitpointSystemFlag = "legacy" | "redis"; + +type MintKindDeps = { + globalDefault: WaitpointSystemFlag; + /** Undefined when the org has no override. Must not hit the DB when given org flags. */ + flag: ( + orgId: string, + orgFeatureFlags: unknown | undefined + ) => Promise; + /** Surfaced instead of logged, so this module pulls in no server-only import. */ + onError?: (error: unknown) => void; +}; + +// PURE CORE — no env import; the tests drive this directly. +export async function computeWaitpointMintKind( + environment: { organizationId: string; id: string; orgFeatureFlags?: unknown }, + deps: MintKindDeps +): Promise { + try { + const perOrg = await deps.flag(environment.organizationId, environment.orgFeatureFlags); + return (perOrg ?? deps.globalDefault) === "redis" ? "store" : "legacy"; + } catch (error) { + // Fail safe, as computeRunIdMintKind does: a flag-read failure degrades to the old + // path rather than becoming a trigger-path outage. + deps.onError?.(error); + return "legacy"; + } +} diff --git a/apps/webapp/test/apiBatchResultsPresenter.dedicatedSeam.test.ts b/apps/webapp/test/apiBatchResultsPresenter.dedicatedSeam.test.ts index eb322c48a1c..ade925f239c 100644 --- a/apps/webapp/test/apiBatchResultsPresenter.dedicatedSeam.test.ts +++ b/apps/webapp/test/apiBatchResultsPresenter.dedicatedSeam.test.ts @@ -3,10 +3,10 @@ // RESULTS READ assembles correctly when one batch's members are genuinely split across the real // dedicated run-ops subset schema (prisma17 / RunOpsPrismaClient) and the full control-plane // schema (prisma14) — not a mirrored full schema on both sides. No mocks. -import { heteroRunOpsPostgresTest } from "@internal/testcontainers"; +import { heteroRunOpsPostgresTest, makeNShardRunOpsPostgresTest } from "@internal/testcontainers"; import type { RunOpsPrismaClient } from "@internal/run-ops-database"; import type { PrismaClient } from "@trigger.dev/database"; -import { generateRunOpsId } from "@trigger.dev/core/v3/isomorphic"; +import { generateRunOpsId, generateRunOpsIdV2 } from "@trigger.dev/core/v3/isomorphic"; import { describe, expect, vi } from "vitest"; import type { PrismaReplicaClient } from "~/db.server"; import type { AuthenticatedEnvironment } from "~/services/apiAuth.server"; @@ -209,6 +209,10 @@ async function seedBatchOnNew( return batch; } +// One real gen-2 shard on its OWN database, so a member seeded there is genuinely absent +// from the gen-1 `new` store rather than merely routed away from it. +const oneShardTest = makeNShardRunOpsPostgresTest(1); + const env = (ctx: SeedCtx) => ({ id: ctx.environment.id, @@ -334,4 +338,141 @@ describe("ApiBatchResultsPresenter split mode — real run-ops dedicated schema expect(result!.items[0]).toMatchObject({ ok: true, id: "run_present" }); } ); + + // A gen-2 member is directly routable to its own shard. Before the shard arm existed it + // joined the gen-1 `new` read, missed there, and — classifying dedicated-family — never + // reached the legacy probe either, so it vanished from the batch results with no error. + oneShardTest( + "a gen-2 member is hydrated from its own shard database alongside a legacy-resident member", + async ({ legacyPrisma, newPrisma, shardPrismas }) => { + const shardPrisma = shardPrismas[0]!; + const ctx = await seedLegacyEnv(legacyPrisma, "gen2-shard"); + await relaxLegacyAttemptFk(legacyPrisma); + await relaxNewBatchItemFk(newPrisma); + + const shardMemberId = generateRunOpsIdV2("a"); + const legacyMemberId = generateLegacyCuid(); + + // The gen-2 member exists ONLY on the shard database. The gen-1 `new` store below is a + // different database, so routing this id there would genuinely miss. + await seedNewMember( + shardPrisma, + { envId: ctx.environment.id, orgId: ctx.organization.id, projectId: ctx.project.id }, + { + id: shardMemberId, + friendlyId: "run_shard_member", + status: "COMPLETED_SUCCESSFULLY", + output: JSON.stringify({ from: "shard-a" }), + } + ); + await seedLegacyMember(legacyPrisma, ctx, { + id: legacyMemberId, + friendlyId: "run_legacy_member", + status: "COMPLETED_WITH_ERRORS", + error: { type: "BUILT_IN_ERROR", name: "Err", message: "boom", stackTrace: "" }, + }); + + const batchFriendlyId = "batch_gen2_shard"; + await seedBatchOnNew(newPrisma, ctx.environment.id, batchFriendlyId, [ + shardMemberId, + legacyMemberId, + ]); + + const presenter = new ApiBatchResultsPresenter(throwingPrisma, throwingPrisma, { + splitEnabled: true, + newClient: newPrisma as unknown as PrismaReplicaClient, + legacyReplica: legacyPrisma as unknown as PrismaReplicaClient, + shardReplicas: new Map([["a", shardPrisma as unknown as PrismaReplicaClient]]), + }); + + const result = await presenter.call(batchFriendlyId, env(ctx)); + + expect(result).toBeDefined(); + expect(result!.items).toHaveLength(2); + const [first, second] = result!.items; + expect(first).toEqual({ + ok: true, + id: "run_shard_member", + taskIdentifier: "my-task", + output: JSON.stringify({ from: "shard-a" }), + outputType: "application/json", + }); + expect(second).toMatchObject({ ok: false, id: "run_legacy_member" }); + }, + 180_000 + ); + + // A gen-2 id naming a shard that is NOT configured must not fall back onto a gen-1 store: + // that reads the wrong database, misses, and (being dedicated-family) never reaches the + // legacy probe, so the member disappears with no error. Drop it, but loudly. + oneShardTest( + "a gen-2 member on an unconfigured shard is dropped without being read from a gen-1 store", + async ({ legacyPrisma, newPrisma, shardPrismas }) => { + const shardPrisma = shardPrismas[0]!; + const ctx = await seedLegacyEnv(legacyPrisma, "gen2-unconfigured"); + await relaxLegacyAttemptFk(legacyPrisma); + await relaxNewBatchItemFk(newPrisma); + + // Shard "z" is not in the configured map; shard "a" is. + const unconfiguredId = generateRunOpsIdV2("z"); + const legacyMemberId = generateLegacyCuid(); + + await seedNewMember( + shardPrisma, + { envId: ctx.environment.id, orgId: ctx.organization.id, projectId: ctx.project.id }, + { id: unconfiguredId, friendlyId: "run_unconfigured", status: "COMPLETED_SUCCESSFULLY" } + ); + await seedLegacyMember(legacyPrisma, ctx, { + id: legacyMemberId, + friendlyId: "run_legacy_member", + status: "COMPLETED_SUCCESSFULLY", + }); + + const batchFriendlyId = "batch_gen2_unconfigured"; + await seedBatchOnNew(newPrisma, ctx.environment.id, batchFriendlyId, [ + unconfiguredId, + legacyMemberId, + ]); + + // A closure-based recorder, not a mock: it records the id sets each store is asked for, + // so the assertion is about real reads rather than about a test double's behaviour. + const askedOf = (label: string, target: RunOpsPrismaClient | PrismaClient) => { + const asked: string[][] = []; + const handle = { + ...target, + taskRun: { + findMany: (args: { where?: { id?: { in?: string[] } } }) => { + asked.push(args.where?.id?.in ?? []); + return (target as unknown as PrismaReplicaClient).taskRun.findMany(args as never); + }, + }, + } as unknown as PrismaReplicaClient; + return { label, asked, handle }; + }; + const genOneNew = askedOf("new", newPrisma); + const legacy = askedOf("legacy", legacyPrisma); + + const presenter = new ApiBatchResultsPresenter(throwingPrisma, throwingPrisma, { + splitEnabled: true, + newClient: genOneNew.handle, + legacyReplica: legacy.handle, + shardReplicas: new Map([["a", shardPrisma as unknown as PrismaReplicaClient]]), + }); + + const result = await presenter.call(batchFriendlyId, env(ctx)); + + // The legacy member still resolves; the unconfigured gen-2 member is dropped. + expect(result).toBeDefined(); + expect(result!.items).toHaveLength(1); + expect(result!.items[0]).toMatchObject({ ok: true, id: "run_legacy_member" }); + + // The unconfigured id was never asked of a gen-1 store. + for (const store of [genOneNew, legacy]) { + for (const ids of store.asked) { + expect(ids).not.toContain(unconfiguredId); + } + } + }, + 180_000 + ); }); diff --git a/apps/webapp/test/readRunForEvent.replicaLag.test.ts b/apps/webapp/test/readRunForEvent.replicaLag.test.ts index 877f920e40f..9817a1b6e55 100644 --- a/apps/webapp/test/readRunForEvent.replicaLag.test.ts +++ b/apps/webapp/test/readRunForEvent.replicaLag.test.ts @@ -195,4 +195,76 @@ describe("readRunForEvent tolerates replica lag on its event-enrichment read", ( expect(onPrimary.friendlyId).toBe("run_rrfe_missing"); } ); + + // (c) SPLIT MODE, the gen-1 run fast path. A cuid run id classifies legacy, and there is no cuid + // run migration, so the new-store probe cannot find it. readRunForEvent declares idKind "run", + // which reads the legacy replica ONLY. The observable difference is the number of reads: one on + // the fast path, two on the old new-then-legacy pair probe. Counted by delegating through the + // real store rather than replacing it. + containerTest( + "readRunForEvent takes ONE read for a cuid run id under split, not a new-then-legacy pair", + async ({ prisma }) => { + const { organization, project, environment } = await seedEnvironment(prisma, "rrfe_split"); + + const runId = "d".repeat(25); // cuid-shaped -> classifies legacy + const friendlyId = "run_rrfe_split"; + + await prisma.taskRun.create({ + data: { + id: runId, + engine: "V2", + status: "COMPLETED_SUCCESSFULLY", + friendlyId, + taskIdentifier: "my-task", + payload: "{}", + payloadType: "application/json", + traceId: "trace_split", + spanId: "span_split", + queue: "task/my-task", + runtimeEnvironmentId: environment.id, + projectId: project.id, + organizationId: organization.id, + environmentType: "DEVELOPMENT", + isTest: false, + taskEventStore: "taskEvent", + }, + }); + + const realStore = new PostgresRunStore({ prisma, readOnlyPrisma: prisma as never }); + let findRunCalls = 0; + const countingStore = new Proxy(realStore, { + get(target, prop, receiver) { + if (prop === "findRun") { + return (...args: unknown[]) => { + findRunCalls += 1; + return (target.findRun as (...a: unknown[]) => unknown)(...args); + }; + } + return Reflect.get(target, prop, receiver); + }, + }); + + // The new side MUST miss for the two arms to be distinguishable: a pair probe that finds the + // row on its first read short-circuits and looks identical to the fast path. `missing` makes + // the new-store read return nothing, exactly as it would for a legacy-resident run. + const missingOnNew = laggingReplica(prisma, [{ model: "taskRun", mode: "missing" }]); + + const deps: EventReadDeps = { + store: countingStore as never, + newReplica: missingOnNew.client as never, + legacyReplica: prisma as never, + splitEnabled: true, + }; + + const run = await readRunForEvent(runId, environment.id, EVENT_SELECT, deps); + + // The run still resolves — the fast path must not cost the read. + expect(run).not.toBeNull(); + expect(run!.id).toBe(runId); + expect(run!.friendlyId).toBe(friendlyId); + + // ONE read. Two would mean the new store was probed first, which is the arm this removes. + expect(findRunCalls).toBe(1); + } + ); }); diff --git a/apps/webapp/test/resolveWaitpointThroughReadThrough.readthrough.test.ts b/apps/webapp/test/resolveWaitpointThroughReadThrough.readthrough.test.ts index c0b627262f7..09c3327ab57 100644 --- a/apps/webapp/test/resolveWaitpointThroughReadThrough.readthrough.test.ts +++ b/apps/webapp/test/resolveWaitpointThroughReadThrough.readthrough.test.ts @@ -1,7 +1,7 @@ import { heteroRunOpsPostgresTest, postgresTest } from "@internal/testcontainers"; import type { RunOpsPrismaClient } from "@internal/run-ops-database"; import type { PrismaClient } from "@trigger.dev/database"; -import { generateRunOpsId } from "@trigger.dev/core/v3/isomorphic"; +import { generateRunOpsId, generateRunOpsIdV2 } from "@trigger.dev/core/v3/isomorphic"; import { describe, expect, vi } from "vitest"; import type { PrismaReplicaClient } from "~/db.server"; import { resolveWaitpointThroughReadThrough } from "~/runEngine/concerns/resolveWaitpointThroughReadThrough.server"; @@ -286,4 +286,105 @@ describe("resolveWaitpointThroughReadThrough (hetero PG14 legacy + dedicated run expect(legacy.calls.length).toBe(0); } ); + + heteroRunOpsPostgresTest( + "gen-2 waitpoint resolves on its OWN shard replica; the gen-1 new store is never read", + async ({ prisma17, prisma14 }) => { + const id = generateRunOpsIdV2("a"); + const environmentId = generateRunOpsId(); + const projectId = generateRunOpsId(); + const seeded = await seedWaitpoint(prisma17, id, { id: environmentId, projectId }); + + // The gen-1 new store and the legacy replica are both forbidden: a gen-2 id must + // take one read on its shard and probe nothing else. + const newClient = recording(prisma14, { forbidden: true }); + const legacyReplica = recording(prisma14, { forbidden: true }); + const shardReplica = recording(prisma17); + + const result = await resolveWaitpointThroughReadThrough({ + waitpointId: id, + environmentId, + read: read(id, environmentId), + deps: { + splitEnabled: true, + newClient: newClient.handle, + legacyReplica: legacyReplica.handle, + newPrimary: newClient.handle, + shardReplicas: new Map([["a", shardReplica.handle]]), + }, + }); + + expect(result).not.toBeNull(); + expect(result!.id).toBe(seeded.id); + expect(shardReplica.calls.length).toBe(1); + expect(newClient.calls.length).toBe(0); + expect(legacyReplica.calls.length).toBe(0); + } + ); + + heteroRunOpsPostgresTest( + "a gen-2 waitpoint missing its shard REPLICA falls back to that shard's WRITER, not the gen-1 new writer", + async ({ prisma17, prisma14 }) => { + // Read-your-writes: a token completed immediately after mint may not have replicated. + // The fallback must read the shard's own primary. Reading the gen-1 new writer would + // query the wrong database and return null. + const id = generateRunOpsIdV2("a"); + const environmentId = generateRunOpsId(); + const projectId = generateRunOpsId(); + const seeded = await seedWaitpoint(prisma17, id, { id: environmentId, projectId }); + + const shardReplica = recording(prisma14); // lags: does not have the row + const shardWriter = recording(prisma17); // has the row + const forbiddenNewPrimary = recording(prisma17, { forbidden: true }); + + const result = await resolveWaitpointThroughReadThrough({ + waitpointId: id, + environmentId, + read: read(id, environmentId), + deps: { + splitEnabled: true, + newClient: recording(prisma14, { forbidden: true }).handle, + legacyReplica: recording(prisma14, { forbidden: true }).handle, + newPrimary: forbiddenNewPrimary.handle, + shardReplicas: new Map([["a", shardReplica.handle]]), + shardWriters: new Map([["a", shardWriter.handle]]), + }, + }); + + expect(result).not.toBeNull(); + expect(result!.id).toBe(seeded.id); + expect(shardReplica.calls.length).toBe(1); + expect(shardWriter.calls.length).toBe(1); + expect(forbiddenNewPrimary.calls.length).toBe(0); + } + ); + + heteroRunOpsPostgresTest( + "a gen-2 waitpoint with no configured shard writer returns null instead of reading a wrong database", + async ({ prisma17, prisma14 }) => { + const id = generateRunOpsIdV2("a"); + const environmentId = generateRunOpsId(); + const projectId = generateRunOpsId(); + await seedWaitpoint(prisma17, id, { id: environmentId, projectId }); + + const forbiddenNewPrimary = recording(prisma17, { forbidden: true }); + + const result = await resolveWaitpointThroughReadThrough({ + waitpointId: id, + environmentId, + read: read(id, environmentId), + deps: { + splitEnabled: true, + newClient: recording(prisma14, { forbidden: true }).handle, + legacyReplica: recording(prisma14, { forbidden: true }).handle, + newPrimary: forbiddenNewPrimary.handle, + shardReplicas: new Map([["a", recording(prisma14).handle]]), + shardWriters: new Map(), + }, + }); + + expect(result).toBeNull(); + expect(forbiddenNewPrimary.calls.length).toBe(0); + } + ); }); diff --git a/apps/webapp/test/unroutableIdStatus.test.ts b/apps/webapp/test/unroutableIdStatus.test.ts new file mode 100644 index 00000000000..89135b44e15 --- /dev/null +++ b/apps/webapp/test/unroutableIdStatus.test.ts @@ -0,0 +1,40 @@ +// `resolveShard` is pure id-shape, so any base32hex core plus `[a-z0-9]` plus "2" parses as gen-2 +// and names a shard — including one a caller invents. The routing store throws for a key it has no +// store for, which is correct and deliberately loud, but a read route that lets it reach the +// boundary answered 500 for caller-supplied input. These tests pin the boundary status. +import { describe, expect, it } from "vitest"; +import { json } from "@remix-run/server-runtime"; +import { UnknownShardKey } from "@internal/run-store"; +import { unroutableIdResponse } from "~/services/routeBuilders/unroutableId.server"; + +describe("unroutableIdResponse", () => { + it("answers 404 for an id naming a shard with no configured store", async () => { + const response = unroutableIdResponse(new UnknownShardKey("z", ["legacy", "new"])); + + expect(response).toBeDefined(); + expect(response!.status).toBe(404); + // Not retryable: no number of retries makes a topology grow a store. Contrast the + // waitpoint wait route, whose 404 IS retryable because a miss there can be replica lag. + expect(response!.headers.get("x-should-retry")).toBe("false"); + await expect(response!.json()).resolves.toEqual({ error: "Not Found" }); + }); + + it("declines an unrelated error so it still reaches the 500 path", () => { + expect(unroutableIdResponse(new Error("db down"))).toBeUndefined(); + expect(unroutableIdResponse(undefined)).toBeUndefined(); + expect(unroutableIdResponse("a string")).toBeUndefined(); + }); + + it("declines a deliberately thrown Response, which carries its own status", () => { + expect(unroutableIdResponse(json({ error: "nope" }, { status: 422 }))).toBeUndefined(); + }); + + it("keeps the key and the configured set on the error for the operator", () => { + // A 404 to the caller must not cost the operator what separates a forged id from a shard + // key dropped out of a config that is meant to be append-only. + const error = new UnknownShardKey("z", ["legacy", "new", "a"]); + + expect(error.shardKey).toBe("z"); + expect(error.configured).toContain("a"); + }); +}); diff --git a/apps/webapp/vitest.config.ts b/apps/webapp/vitest.config.ts index dabe517bf4f..c43e4213ac9 100644 --- a/apps/webapp/vitest.config.ts +++ b/apps/webapp/vitest.config.ts @@ -12,6 +12,7 @@ export default defineConfig({ include: [ "test/**/*.test.ts", "app/v3/runOpsMigration/**/*.test.ts", + "app/v3/waitpointMigration/**/*.test.ts", "app/v3/runStore.server.test.ts", "app/v3/utils/**/*.test.ts", "app/v3/services/bulk/**/*.test.ts", diff --git a/internal-packages/run-engine/src/engine/index.ts b/internal-packages/run-engine/src/engine/index.ts index ccfb60ca4d6..fff506c3aba 100644 --- a/internal-packages/run-engine/src/engine/index.ts +++ b/internal-packages/run-engine/src/engine/index.ts @@ -24,9 +24,10 @@ import { import type { TaskRunError } from "@trigger.dev/core/v3/schemas"; import { generateInternalId, + deriveWaitpointIdFromAnchor, parseNaturalLanguageDurationInMs, + parseWaitpointId, RunId, - WaitpointId, } from "@trigger.dev/core/v3/isomorphic"; import { type PrismaClient, @@ -92,6 +93,11 @@ import { } from "./controlPlaneResolver.js"; import { TtlSystem } from "./systems/ttlSystem.js"; import { WaitpointSystem } from "./systems/waitpointSystem.js"; +import { LegacyPostgresWaitpointCoordinator } from "./waitpointCoordinator/legacyPostgresCoordinator.js"; +import { WaitpointRouterCoordinator } from "./waitpointCoordinator/routerCoordinator.js"; +import { StoreWaitpointCoordinatorArm } from "./waitpointCoordinator/storeArm.js"; +import { WaitpointStoreCoordinator } from "./waitpointCoordinator/storeCoordinator.js"; +import type { WaitpointMintKind } from "./waitpointCoordinator/types.js"; import type { EngineWorker, HeartbeatTimeouts, @@ -130,6 +136,7 @@ export class RunEngine { runAttemptSystem: RunAttemptSystem; dequeueSystem: DequeueSystem; waitpointSystem: WaitpointSystem; + private waitpointStoreCoordinator?: WaitpointStoreCoordinator; batchSystem: BatchSystem; enqueueSystem: EnqueueSystem; checkpointSystem: CheckpointSystem; @@ -412,10 +419,34 @@ export class RunEngine { externalDeploymentParkDeadlineMs: options.externalDeploymentParkDeadlineMs, }); + this.waitpointStoreCoordinator = this.options.waitpointStore + ? new WaitpointStoreCoordinator({ + redisOptions: this.options.waitpointStore.redis, + logger: this.logger, + }) + : undefined; + this.waitpointSystem = new WaitpointSystem({ resources, executionSnapshotSystem: this.executionSnapshotSystem, enqueueSystem: this.enqueueSystem, + coordinator: new WaitpointRouterCoordinator({ + meter: this.meter, + legacy: new LegacyPostgresWaitpointCoordinator({ + runStore: this.runStore, + prisma: this.prisma, + logger: this.logger, + }), + store: this.waitpointStoreCoordinator + ? new StoreWaitpointCoordinatorArm({ + store: this.waitpointStoreCoordinator, + runStore: this.runStore, + logger: this.logger, + meter: this.meter, + }) + : undefined, + logger: this.logger, + }), }); this.ttlSystem = new TtlSystem({ @@ -849,6 +880,7 @@ export class RunEngine { replayedFromTaskRunFriendlyId, batch, resumeParentOnCompletion, + waitpointMintKind, depth, metadata, metadataType, @@ -971,6 +1003,23 @@ export class RunEngine { let taskRun: TaskRun & { associatedWaitpoint: Waitpoint | null }; const taskRunId = RunId.fromFriendlyId(friendlyId); + + // Mint the RUN waitpoint's identity BEFORE the run is created, so the decision to + // block the parent never depends on a Postgres relation that the store path does + // not write. Keying the block step off the row instead would silently stop + // suspending parents the moment a waitpoint stopped living in Postgres. + const associatedWaitpointData = + resumeParentOnCompletion && parentTaskRunId + ? this.waitpointSystem.buildRunAssociatedWaitpoint({ + projectId: environment.project.id, + environmentId: environment.id, + anchorRunId: taskRunId, + mintKind: waitpointMintKind, + }) + : undefined; + const associatedWaitpointRidesTheCreate = + associatedWaitpointData !== undefined && + parseWaitpointId(associatedWaitpointData.id).format === "legacy"; const initialSnapshotId = generateInternalId(); // App-level replacement for the dropped TaskRun env/project Cascade FKs. @@ -1080,15 +1129,14 @@ export class RunEngine { workerId, runnerId, }, - // Only create waitpoint if parent is waiting for this run to complete - // For standalone triggers (no waiting parent), waitpoint is created lazily if needed later - associatedWaitpoint: - resumeParentOnCompletion && parentTaskRunId - ? this.waitpointSystem.buildRunAssociatedWaitpoint({ - projectId: environment.project.id, - environmentId: environment.id, - }) - : undefined, + // Only create the waitpoint if a parent is waiting for this run. A standalone + // trigger gets one lazily later, if anything ever needs it. + // + // The store path deliberately passes nothing here: its waitpoint is created + // after the run commits, so the run's own insert carries no waitpoint row. + associatedWaitpoint: associatedWaitpointRidesTheCreate + ? associatedWaitpointData + : undefined, }, tx ); @@ -1131,8 +1179,18 @@ export class RunEngine { span.setAttribute("runId", taskRun.id); + // The store path's waitpoint is created here, after the run commits. Create-if-absent + // on an id derived from the run means a retry recomputes the same id, so this is + // idempotent and needs no lock. + if (associatedWaitpointData && !associatedWaitpointRidesTheCreate) { + await this.waitpointSystem.createRunAssociatedWaitpoint({ + runId: taskRun.id, + data: associatedWaitpointData, + }); + } + //triggerAndWait or batchTriggerAndWait - if (resumeParentOnCompletion && parentTaskRunId && taskRun.associatedWaitpoint) { + if (resumeParentOnCompletion && parentTaskRunId && associatedWaitpointData) { if (batch) { // Batch path: lockless insert. The parent is already EXECUTING_WITH_WAITPOINTS // from blockRunWithCreatedBatch, so we only need to insert the TaskRunWaitpoint @@ -1140,17 +1198,21 @@ export class RunEngine { // processing large batches with high concurrency. await this.waitpointSystem.blockRunWithWaitpointLockless({ runId: parentTaskRunId, - waitpoints: taskRun.associatedWaitpoint.id, - projectId: taskRun.associatedWaitpoint.projectId, + waitpoints: associatedWaitpointData.id, + projectId: associatedWaitpointData.projectId, batch, + // Derived, not looked up: the parent's BATCH waitpoint id is a pure function + // of the batch id, and the store arm needs it to assert the parent's pending + // set stays open for the whole absorb. + batchWaitpointId: deriveWaitpointIdFromAnchor(batch.id, "BATCH"), }); } else { // Single triggerAndWait: acquire the parent run lock to safely transition // the snapshot and insert the waitpoint await this.waitpointSystem.blockRunWithWaitpoint({ runId: parentTaskRunId, - waitpoints: taskRun.associatedWaitpoint.id, - projectId: taskRun.associatedWaitpoint.projectId, + waitpoints: associatedWaitpointData.id, + projectId: associatedWaitpointData.projectId, organizationId: environment.organization.id, batch, workerId, @@ -1309,6 +1371,7 @@ export class RunEngine { rootTaskRunId, depth, resumeParentOnCompletion, + waitpointMintKind, batch, traceId, spanId, @@ -1335,6 +1398,8 @@ export class RunEngine { /** Depth in the task tree (0 for root, parentDepth+1 for children). */ depth?: number; resumeParentOnCompletion?: boolean; + /** Which coordinator mints the associated waitpoint. Absent means legacy. */ + waitpointMintKind?: WaitpointMintKind; batch?: { id: string; index: number }; traceId?: string; spanId?: string; @@ -1367,14 +1432,19 @@ export class RunEngine { // App-level replacement for the dropped TaskRun env/project Cascade FKs. await this.controlPlaneResolver.assertEnvExists(environment.id); - // Build associated waitpoint data if parent is waiting for this run + // Minted before the create, for the same reason as the trigger path: the decision to + // block the parent must not depend on a row the store path never writes. const waitpointData = resumeParentOnCompletion && parentTaskRunId ? this.waitpointSystem.buildRunAssociatedWaitpoint({ projectId: environment.project.id, environmentId: environment.id, + anchorRunId: taskRunId, + mintKind: waitpointMintKind, }) : undefined; + const waitpointRidesTheCreate = + waitpointData !== undefined && parseWaitpointId(waitpointData.id).format === "legacy"; // No execution snapshot is needed: this run never gets dequeued, executed, // or heartbeated, so nothing will call getLatestExecutionSnapshot on it. @@ -1408,19 +1478,26 @@ export class RunEngine { resumeParentOnCompletion, taskEventStore, }, - associatedWaitpoint: waitpointData, + associatedWaitpoint: waitpointRidesTheCreate ? waitpointData : undefined, }, undefined ); span.setAttribute("runId", taskRun.id); + if (waitpointData && !waitpointRidesTheCreate) { + await this.waitpointSystem.createRunAssociatedWaitpoint({ + runId: taskRun.id, + data: waitpointData, + }); + } + // If parent is waiting, block it with the waitpoint then immediately // complete it with the error output so the parent can resume. - if (resumeParentOnCompletion && parentTaskRunId && taskRun.associatedWaitpoint) { + if (resumeParentOnCompletion && parentTaskRunId && waitpointData) { await this.waitpointSystem.blockRunAndCompleteWaitpoint({ runId: parentTaskRunId, - waitpointId: taskRun.associatedWaitpoint.id, + waitpointId: waitpointData.id, output: { value: JSON.stringify(error), isError: true }, projectId: environment.project.id, organizationId: environment.organization.id, @@ -1776,6 +1853,7 @@ export class RunEngine { completedAfter, idempotencyKey, idempotencyKeyExpiresAt, + waitpointMintKind, }: { /** The run that will block on this waitpoint. Co-locates the waitpoint with the run's DB. */ runId?: string; @@ -1784,6 +1862,8 @@ export class RunEngine { completedAfter: Date; idempotencyKey?: string; idempotencyKeyExpiresAt?: Date; + /** Which coordinator mints this waitpoint. Resolved from the org flag by the caller. */ + waitpointMintKind?: WaitpointMintKind; }) { return this.waitpointSystem.createDateTimeWaitpoint({ runId, @@ -1792,6 +1872,7 @@ export class RunEngine { completedAfter, idempotencyKey, idempotencyKeyExpiresAt, + waitpointMintKind, }); } @@ -1807,6 +1888,7 @@ export class RunEngine { timeout, tags, standaloneResidency, + waitpointMintKind, }: { /** The run that will block on this waitpoint. Co-locates the waitpoint with the run's DB. */ runId?: string; @@ -1818,6 +1900,8 @@ export class RunEngine { tags?: string[]; /** Standalone-token residency (no owning run) from the env mint kind; ignored when `runId` is set. */ standaloneResidency?: "NEW" | "LEGACY"; + /** Which coordinator mints this waitpoint. Resolved from the org flag by the caller. */ + waitpointMintKind?: WaitpointMintKind; }): Promise<{ waitpoint: Waitpoint; isCached: boolean }> { return this.waitpointSystem.createManualWaitpoint({ runId, @@ -1828,6 +1912,7 @@ export class RunEngine { timeout, tags, standaloneResidency, + waitpointMintKind, }); } @@ -1840,6 +1925,7 @@ export class RunEngine { environmentId, projectId, organizationId, + waitpointMintKind, tx, }: { runId: string; @@ -1847,24 +1933,23 @@ export class RunEngine { environmentId: string; projectId: string; organizationId: string; + waitpointMintKind?: WaitpointMintKind; tx?: PrismaClientOrTransaction; }): Promise { - try { - const waitpoint = await this.runStore.createWaitpoint( - { - data: { - ...WaitpointId.generate(), - type: "BATCH", - idempotencyKey: batchId, - userProvidedIdempotencyKey: false, - completedByBatchId: batchId, - environmentId, - projectId, - }, - }, - tx - ); + const waitpoint = await this.waitpointSystem.createBatchWaitpoint({ + batchId, + environmentId, + projectId, + mintKind: waitpointMintKind, + tx, + }); + + // Duplicate batch: the coordinator already reported it. + if (!waitpoint) { + return null; + } + try { await this.blockRunWithWaitpoint({ runId, waitpoints: waitpoint.id, @@ -1873,19 +1958,17 @@ export class RunEngine { batch: { id: batchId }, // No tx: the block edge routes to the run's owning DB, not the control-plane tx. }); - - return waitpoint; } catch (error) { - if (error instanceof Prisma.PrismaClientKnownRequestError) { - // duplicate idempotency key - if (error.code === "P2002") { - return null; - } else { - throw error; - } + // The previous shape wrapped the create AND the block in one catch, so a P2002 from + // the block step also returned null. Kept deliberately: narrowing it here would be a + // behaviour change smuggled into an extraction. + if (error instanceof Prisma.PrismaClientKnownRequestError && error.code === "P2002") { + return null; } throw error; } + + return waitpoint; } async tryCompleteBatch({ batchId }: { batchId: string }): Promise { @@ -2394,8 +2477,12 @@ export class RunEngine { const supportResults = await Promise.allSettled([ this.runLock.quit(), this.debounceSystem.quit(), + this.waitpointStoreCoordinator?.quit(), ]); - this.#logShutdownFailures(["runLock.quit", "debounceSystem.quit"], supportResults); + this.#logShutdownFailures( + ["runLock.quit", "debounceSystem.quit", "waitpointStore.quit"], + supportResults + ); // RunLocker/Redlock owns this client and normally closes it. Do not send a second QUIT, // but force-disconnect if Redlock failed to leave the connection in its terminal state. diff --git a/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts b/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts index 36738cfa983..791dc37e795 100644 --- a/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts @@ -11,9 +11,13 @@ import type { import { assertNever } from "assert-never"; import { sendNotificationToWorker } from "../eventBus.js"; import { isFinalRunStatus } from "../statuses.js"; -import { LegacyPostgresWaitpointCoordinator } from "../waitpointCoordinator/legacyPostgresCoordinator.js"; import { buildCompletedWaitpointRecords } from "../waitpointCoordinator/completedWaitpointRecords.js"; -import type { RunBlockEdge, WaitpointCoordinator } from "../waitpointCoordinator/types.js"; +import type { + AssociatedWaitpointData, + RunBlockEdge, + WaitpointCoordinator, + WaitpointMintKind, +} from "../waitpointCoordinator/types.js"; import type { EnqueueSystem } from "./enqueueSystem.js"; import type { ExecutionSnapshotSystem } from "./executionSnapshotSystem.js"; import { getLatestExecutionSnapshot } from "./executionSnapshotSystem.js"; @@ -23,6 +27,8 @@ export type WaitpointSystemOptions = { resources: SystemResources; executionSnapshotSystem: ExecutionSnapshotSystem; enqueueSystem: EnqueueSystem; + /** Which coordinator owns waitpoint state. The engine supplies a router over both arms. */ + coordinator: WaitpointCoordinator; }; type WaitpointContinuationWaitpoint = Pick; @@ -51,11 +57,7 @@ export class WaitpointSystem { this.$ = options.resources; this.executionSnapshotSystem = options.executionSnapshotSystem; this.enqueueSystem = options.enqueueSystem; - this.coordinator = new LegacyPostgresWaitpointCoordinator({ - runStore: this.$.runStore, - prisma: this.$.prisma, - logger: this.$.logger, - }); + this.coordinator = options.coordinator; } public async clearBlockingWaitpoints({ @@ -144,6 +146,7 @@ export class WaitpointSystem { completedAfter, idempotencyKey, idempotencyKeyExpiresAt, + waitpointMintKind, }: { runId?: string; projectId: string; @@ -151,8 +154,10 @@ export class WaitpointSystem { completedAfter: Date; idempotencyKey?: string; idempotencyKeyExpiresAt?: Date; + waitpointMintKind?: WaitpointMintKind; }) { const result = await this.coordinator.createDateTimeWaitpoint({ + mintKind: waitpointMintKind ?? "legacy", runId, projectId, environmentId, @@ -187,10 +192,12 @@ export class WaitpointSystem { timeout, tags, standaloneResidency, + waitpointMintKind, }: { runId?: string; environmentId: string; projectId: string; + waitpointMintKind?: WaitpointMintKind; idempotencyKey?: string; idempotencyKeyExpiresAt?: Date; timeout?: Date; @@ -201,6 +208,7 @@ export class WaitpointSystem { standaloneResidency?: "NEW" | "LEGACY"; }): Promise<{ waitpoint: Waitpoint; isCached: boolean }> { const result = await this.coordinator.createManualWaitpoint({ + mintKind: waitpointMintKind ?? "legacy", runId, environmentId, projectId, @@ -385,6 +393,7 @@ export class WaitpointSystem { timeout, spanIdToComplete, batch, + batchWaitpointId, }: { runId: string; waitpoints: string | string[]; @@ -392,6 +401,8 @@ export class WaitpointSystem { timeout?: Date; spanIdToComplete?: string; batch: { id: string; index?: number }; + /** The parent's BATCH waitpoint, so the store arm can assert it is still pending. */ + batchWaitpointId?: string; }): Promise { const $waitpoints = typeof waitpoints === "string" ? [waitpoints] : waitpoints; @@ -405,6 +416,7 @@ export class WaitpointSystem { spanIdToComplete, batchId: batch.id, batchIndex: batch.index, + batchWaitpointId, }); // Schedule timeout jobs if needed @@ -735,14 +747,59 @@ export class WaitpointSystem { }); // end of runlock } + /** The BATCH waitpoint for a batch. Returns null when the batch already has one. */ + public async createBatchWaitpoint(params: { + batchId: string; + environmentId: string; + projectId: string; + mintKind?: WaitpointMintKind; + tx?: PrismaClientOrTransaction; + }): Promise { + return this.coordinator.createBatchWaitpoint({ + ...params, + mintKind: params.mintKind ?? "legacy", + }); + } + + /** + * Mint the RUN waitpoint's data for a run that a parent will block on. + * + * A store mint derives the id from the anchor run's own id body, so the id is a pure + * function of the run id and create-if-absent needs no lock. Derivation only works when + * the run itself carries a run-ops id, so a legacy-shaped run keeps a legacy waitpoint + * even in a flipped organization, which is the coexistence rule the id routing relies on. + */ public buildRunAssociatedWaitpoint({ projectId, environmentId, + anchorRunId, + mintKind, }: { projectId: string; environmentId: string; + anchorRunId?: string; + mintKind?: WaitpointMintKind; }) { - return this.coordinator.mintAssociatedWaitpointData({ projectId, environmentId }); + return this.coordinator.mintAssociatedWaitpointData({ + projectId, + environmentId, + anchorRunId, + mintKind, + }); + } + + /** + * Create the RUN waitpoint that `buildRunAssociatedWaitpoint` minted. + * + * Only the store path calls this: the legacy path writes the row inside the run's own + * create. A crash between the run commit and this call leaves the waitpoint absent, and + * the parent's register step then fails loud rather than resuming without it. + */ + public async createRunAssociatedWaitpoint(params: { + runId: string; + data: AssociatedWaitpointData; + }): Promise { + return this.coordinator.createAssociatedWaitpoint(params); } /** diff --git a/internal-packages/run-engine/src/engine/tests/waitpointBatchCreate.test.ts b/internal-packages/run-engine/src/engine/tests/waitpointBatchCreate.test.ts new file mode 100644 index 00000000000..dbcb0edb74a --- /dev/null +++ b/internal-packages/run-engine/src/engine/tests/waitpointBatchCreate.test.ts @@ -0,0 +1,203 @@ +import type { RedisOptions } from "@internal/redis"; +import { assertNonNullable, containerTest } from "@internal/testcontainers"; +import { trace } from "@internal/tracing"; +import { + BatchId, + generateRunOpsId, + parseWaitpointId, + RunId, +} from "@trigger.dev/core/v3/isomorphic"; +import type { PrismaClient } from "@trigger.dev/database"; +import { describe, expect } from "vitest"; +import { RunEngine } from "../index.js"; +import { setupAuthenticatedEnvironment, setupBackgroundWorker } from "./setup.js"; + +vi.setConfig({ testTimeout: 60_000 }); + +type Arm = "legacy" | "store"; + +function engineFor(arm: Arm, prisma: PrismaClient, redisOptions: RedisOptions) { + return new RunEngine({ + prisma, + worker: { redis: redisOptions, workers: 1, tasksPerWorker: 10, pollIntervalMs: 100 }, + queue: { redis: redisOptions }, + runLock: { redis: redisOptions }, + waitpointStore: arm === "store" ? { redis: redisOptions } : undefined, + machines: { + defaultMachine: "small-1x", + machines: { + "small-1x": { name: "small-1x" as const, cpu: 0.5, memory: 0.5, centsPerMs: 0.0001 }, + }, + baseCostInCents: 0.0001, + }, + tracer: trace.getTracer("test", "0.0.0"), + }); +} + +function freshRunFriendlyId(arm: Arm) { + return arm === "store" ? RunId.toFriendlyId(generateRunOpsId()) : RunId.generate().friendlyId; +} + +function triggerParams(friendlyId: string, environment: any, taskIdentifier: string) { + return { + number: 1, + friendlyId, + environment, + taskIdentifier, + payload: "{}", + payloadType: "application/json", + context: {}, + traceContext: {}, + traceId: "t12345", + spanId: "s12345", + workerQueue: "main", + queue: `task/${taskIdentifier}`, + isTest: false, + tags: [], + }; +} + +async function seedBatch(prisma: PrismaClient, environment: any, arm: Arm) { + // Mirrors batchIdForMintKind: a run-ops batch carries a run-ops ROW id, and the BATCH + // waitpoint derives from that id, not from the friendly id. + const { id, friendlyId } = + arm === "store" + ? (() => { + const core = generateRunOpsId(); + return { id: core, friendlyId: BatchId.toFriendlyId(core) }; + })() + : BatchId.generate(); + + return prisma.batchTaskRun.create({ + data: { id, friendlyId, runtimeEnvironmentId: environment.id, runCount: 1 }, + }); +} + +describe.each(["legacy", "store"])("BATCH waitpoint create (%s arm)", (arm) => { + containerTest( + "blockRunWithCreatedBatch suspends the parent", + async ({ prisma, redisOptions }) => { + const environment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + const engine = engineFor(arm, prisma, redisOptions); + + try { + const taskIdentifier = "test-task"; + await setupBackgroundWorker(engine, environment, taskIdentifier); + + const parent = await engine.trigger( + triggerParams(freshRunFriendlyId(arm), environment, taskIdentifier), + prisma + ); + const batch = await seedBatch(prisma, environment, arm); + + const waitpoint = await engine.blockRunWithCreatedBatch({ + runId: parent.id, + batchId: batch.id, + environmentId: environment.id, + projectId: environment.project.id, + organizationId: environment.organization.id, + waitpointMintKind: arm, + }); + + assertNonNullable(waitpoint); + expect(waitpoint.type).toBe("BATCH"); + expect(waitpoint.completedByBatchId).toBe(batch.id); + expect(parseWaitpointId(waitpoint.id).format).toBe(arm === "store" ? "b32hexW" : "legacy"); + + // A parent that was never blocked stays QUEUED. + const snapshot = await engine.getRunExecutionData({ runId: parent.id }); + assertNonNullable(snapshot); + expect(snapshot.snapshot.executionStatus).toBe("SUSPENDED"); + } finally { + await engine.quit(); + } + } + ); + + containerTest("a duplicate batch returns null", async ({ prisma, redisOptions }) => { + const environment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + const engine = engineFor(arm, prisma, redisOptions); + + try { + const taskIdentifier = "test-task"; + await setupBackgroundWorker(engine, environment, taskIdentifier); + + const parent = await engine.trigger( + triggerParams(freshRunFriendlyId(arm), environment, taskIdentifier), + prisma + ); + const batch = await seedBatch(prisma, environment, arm); + const args = { + runId: parent.id, + batchId: batch.id, + environmentId: environment.id, + projectId: environment.project.id, + organizationId: environment.organization.id, + waitpointMintKind: arm, + } as const; + + expect(await engine.blockRunWithCreatedBatch(args)).not.toBeNull(); + // The legacy arm reports this through a unique-index violation, the store arm + // through its create-if-absent. Same contract either way. + expect(await engine.blockRunWithCreatedBatch(args)).toBeNull(); + } finally { + await engine.quit(); + } + }); +}); + +describe("BATCH waitpoint, the lockless absorb guard", () => { + // The invariant: the parent's BATCH waitpoint holds the pending set open for the whole + // absorb, so a completion arriving mid-absorb can never see an empty set and resume the + // parent before its items are registered. + containerTest( + "keeps the parent BATCH waitpoint pending while items absorb", + async ({ prisma, redisOptions }) => { + const environment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + const engine = engineFor("store", prisma, redisOptions); + + try { + const taskIdentifier = "test-task"; + await setupBackgroundWorker(engine, environment, taskIdentifier); + + const parent = await engine.trigger( + triggerParams(freshRunFriendlyId("store"), environment, taskIdentifier), + prisma + ); + const batch = await seedBatch(prisma, environment, "store"); + + const batchWaitpoint = await engine.blockRunWithCreatedBatch({ + runId: parent.id, + batchId: batch.id, + environmentId: environment.id, + projectId: environment.project.id, + organizationId: environment.organization.id, + waitpointMintKind: "store", + }); + assertNonNullable(batchWaitpoint); + + for (let index = 0; index < 3; index++) { + await engine.trigger( + { + ...triggerParams(freshRunFriendlyId("store"), environment, taskIdentifier), + parentTaskRunId: parent.id, + rootTaskRunId: parent.id, + resumeParentOnCompletion: true, + depth: 1, + batch: { id: batch.id, index }, + waitpointMintKind: "store", + }, + prisma + ); + + // After every item, the parent is still blocked by its BATCH waitpoint. + const snapshot = await engine.getRunExecutionData({ runId: parent.id }); + assertNonNullable(snapshot); + expect(snapshot.snapshot.executionStatus).toBe("SUSPENDED"); + } + } finally { + await engine.quit(); + } + } + ); +}); diff --git a/internal-packages/run-engine/src/engine/tests/waitpointRunCreate.test.ts b/internal-packages/run-engine/src/engine/tests/waitpointRunCreate.test.ts new file mode 100644 index 00000000000..cfff01eb430 --- /dev/null +++ b/internal-packages/run-engine/src/engine/tests/waitpointRunCreate.test.ts @@ -0,0 +1,236 @@ +import { createRedisClient, type RedisOptions } from "@internal/redis"; +import { assertNonNullable, containerTest } from "@internal/testcontainers"; +import { trace } from "@internal/tracing"; +import { + generateRunOpsId, + parseWaitpointId, + RunId, + deriveWaitpointIdFromAnchor, +} from "@trigger.dev/core/v3/isomorphic"; +import type { PrismaClient } from "@trigger.dev/database"; +import { describe, expect } from "vitest"; +import { RunEngine } from "../index.js"; +import { waitpointKeys } from "../waitpointCoordinator/keys.js"; +import { setupAuthenticatedEnvironment, setupBackgroundWorker } from "./setup.js"; + +vi.setConfig({ testTimeout: 60_000 }); + +type Arm = "legacy" | "store"; + +function engineFor(arm: Arm, prisma: PrismaClient, redisOptions: RedisOptions) { + return new RunEngine({ + prisma, + worker: { redis: redisOptions, workers: 1, tasksPerWorker: 10, pollIntervalMs: 100 }, + queue: { redis: redisOptions }, + runLock: { redis: redisOptions }, + waitpointStore: arm === "store" ? { redis: redisOptions } : undefined, + machines: { + defaultMachine: "small-1x", + machines: { + "small-1x": { name: "small-1x" as const, cpu: 0.5, memory: 0.5, centsPerMs: 0.0001 }, + }, + baseCostInCents: 0.0001, + }, + tracer: trace.getTracer("test", "0.0.0"), + }); +} + +/** + * A store RUN waitpoint derives its id from the anchor run's own id body, so a store-arm + * run has to be triggered with a run-ops friendly id. A legacy-shaped run in a flipped + * organization keeps a legacy waitpoint, which is a case worth its own test below. + */ +function freshRunFriendlyId(arm: Arm) { + return arm === "store" ? RunId.toFriendlyId(generateRunOpsId()) : RunId.generate().friendlyId; +} + +function triggerParams(friendlyId: string, environment: any, taskIdentifier: string) { + return { + number: 1, + friendlyId, + environment, + taskIdentifier, + payload: "{}", + payloadType: "application/json", + context: {}, + traceContext: {}, + traceId: "t12345", + spanId: "s12345", + workerQueue: "main", + queue: `task/${taskIdentifier}`, + isTest: false, + tags: [], + }; +} + +describe.each(["legacy", "store"])("trigger-time RUN waitpoint (%s arm)", (arm) => { + containerTest("triggerAndWait suspends the parent", async ({ prisma, redisOptions }) => { + const environment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + const engine = engineFor(arm, prisma, redisOptions); + + try { + const taskIdentifier = "test-task"; + await setupBackgroundWorker(engine, environment, taskIdentifier); + + const parent = await engine.trigger( + triggerParams(freshRunFriendlyId(arm), environment, taskIdentifier), + prisma + ); + + await engine.trigger( + { + ...triggerParams(freshRunFriendlyId(arm), environment, taskIdentifier), + parentTaskRunId: parent.id, + rootTaskRunId: parent.id, + resumeParentOnCompletion: true, + depth: 1, + waitpointMintKind: arm, + }, + prisma + ); + + // A parent that was never blocked stays QUEUED, so QUEUED must NOT be acceptable + // here. This is the assertion that catches the block step being skipped entirely. + const snapshot = await engine.getRunExecutionData({ runId: parent.id }); + assertNonNullable(snapshot); + expect(snapshot.snapshot.executionStatus).toBe("SUSPENDED"); + } finally { + await engine.quit(); + } + }); +}); + +describe("trigger-time RUN waitpoint, store specifics", () => { + containerTest( + "derives the waitpoint id from the child run's id", + async ({ prisma, redisOptions }) => { + const environment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + const engine = engineFor("store", prisma, redisOptions); + + try { + const taskIdentifier = "test-task"; + await setupBackgroundWorker(engine, environment, taskIdentifier); + + const parent = await engine.trigger( + triggerParams(freshRunFriendlyId("store"), environment, taskIdentifier), + prisma + ); + const childFriendlyId = freshRunFriendlyId("store"); + const child = await engine.trigger( + { + ...triggerParams(childFriendlyId, environment, taskIdentifier), + parentTaskRunId: parent.id, + rootTaskRunId: parent.id, + resumeParentOnCompletion: true, + depth: 1, + waitpointMintKind: "store", + }, + prisma + ); + + const expected = deriveWaitpointIdFromAnchor(child.id, "RUN"); + assertNonNullable(expected); + expect(parseWaitpointId(expected).format).toBe("b32hexW"); + + // No Postgres row: the store owns this waitpoint entirely. + const row = await prisma.waitpoint.findFirst({ where: { id: expected } }); + expect(row).toBeNull(); + } finally { + await engine.quit(); + } + } + ); + + // The Frozen-list rule: a crash between the run commit and the waitpoint create must fail + // loud at the parent's register step, never resume the parent as though nothing was owed. + containerTest( + "fails loud when the store waitpoint is missing at register", + async ({ prisma, redisOptions }) => { + const environment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + const engine = engineFor("store", prisma, redisOptions); + const redis = createRedisClient(redisOptions); + + try { + const taskIdentifier = "test-task"; + await setupBackgroundWorker(engine, environment, taskIdentifier); + + const parent = await engine.trigger( + triggerParams(freshRunFriendlyId("store"), environment, taskIdentifier), + prisma + ); + + const childFriendlyId = freshRunFriendlyId("store"); + const childId = RunId.fromFriendlyId(childFriendlyId); + const waitpointId = deriveWaitpointIdFromAnchor(childId, "RUN"); + assertNonNullable(waitpointId); + + // Stand in for the crash: the run commits, the waitpoint never reaches the store. + // Deleting the record before the parent registers reproduces that window exactly. + const failing = engine + .trigger( + { + ...triggerParams(childFriendlyId, environment, taskIdentifier), + parentTaskRunId: parent.id, + rootTaskRunId: parent.id, + resumeParentOnCompletion: true, + depth: 1, + waitpointMintKind: "store", + }, + prisma + ) + .then(async (run) => { + await redis.del(waitpointKeys(waitpointId).record); + await engine.blockRunWithWaitpoint({ + runId: parent.id, + waitpoints: waitpointId, + projectId: environment.project.id, + organizationId: environment.organization.id, + }); + return run; + }); + + await expect(failing).rejects.toThrow(); + } finally { + await redis.quit(); + await engine.quit(); + } + } + ); + + // Coexistence: a flipped organization still on legacy run ids keeps legacy waitpoints, + // because the derivation needs a run-ops anchor to work from. + containerTest( + "keeps a legacy waitpoint when the run id is legacy shaped", + async ({ prisma, redisOptions }) => { + const environment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + const engine = engineFor("store", prisma, redisOptions); + + try { + const taskIdentifier = "test-task"; + await setupBackgroundWorker(engine, environment, taskIdentifier); + + const parent = await engine.trigger( + triggerParams(freshRunFriendlyId("legacy"), environment, taskIdentifier), + prisma + ); + const child = await engine.trigger( + { + ...triggerParams(freshRunFriendlyId("legacy"), environment, taskIdentifier), + parentTaskRunId: parent.id, + rootTaskRunId: parent.id, + resumeParentOnCompletion: true, + depth: 1, + waitpointMintKind: "store", + }, + prisma + ); + + const row = await prisma.waitpoint.findFirst({ where: { completedByTaskRunId: child.id } }); + assertNonNullable(row); + expect(parseWaitpointId(row.id).format).toBe("legacy"); + } finally { + await engine.quit(); + } + } + ); +}); diff --git a/internal-packages/run-engine/src/engine/tests/waitpointStandaloneCreates.test.ts b/internal-packages/run-engine/src/engine/tests/waitpointStandaloneCreates.test.ts new file mode 100644 index 00000000000..1e06ba063f0 --- /dev/null +++ b/internal-packages/run-engine/src/engine/tests/waitpointStandaloneCreates.test.ts @@ -0,0 +1,155 @@ +import { containerTest } from "@internal/testcontainers"; +import { trace } from "@internal/tracing"; +import { parseWaitpointId } from "@trigger.dev/core/v3/isomorphic"; +import type { PrismaClient } from "@trigger.dev/database"; +import type { RedisOptions } from "@internal/redis"; +import { describe, expect } from "vitest"; +import { RunEngine } from "../index.js"; +import { setupAuthenticatedEnvironment } from "./setup.js"; + +vi.setConfig({ testTimeout: 60_000 }); + +type Arm = "legacy" | "store"; + +function engineFor(arm: Arm, prisma: PrismaClient, redisOptions: RedisOptions) { + return new RunEngine({ + prisma, + worker: { redis: redisOptions, workers: 1, tasksPerWorker: 10, pollIntervalMs: 100 }, + queue: { redis: redisOptions }, + runLock: { redis: redisOptions }, + // The arm under test is selected by whether a store is configured AT ALL, plus the mint + // kind each call passes. Both together are what a flipped organization looks like. + waitpointStore: arm === "store" ? { redis: redisOptions } : undefined, + machines: { + defaultMachine: "small-1x", + machines: { + "small-1x": { name: "small-1x" as const, cpu: 0.5, memory: 0.5, centsPerMs: 0.0001 }, + }, + baseCostInCents: 0.0001, + }, + tracer: trace.getTracer("test", "0.0.0"), + }); +} + +const expectedFormat: Record = { legacy: "legacy", store: "b32hexW" }; + +describe.each(["legacy", "store"])("standalone waitpoint creates (%s arm)", (arm) => { + containerTest( + "createManualWaitpoint mints into the expected system", + async ({ prisma, redisOptions }) => { + const environment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + const engine = engineFor(arm, prisma, redisOptions); + + try { + const { waitpoint } = await engine.createManualWaitpoint({ + environmentId: environment.id, + projectId: environment.project.id, + waitpointMintKind: arm, + }); + + expect(parseWaitpointId(waitpoint.id).format).toBe(expectedFormat[arm]); + expect(waitpoint.status).toBe("PENDING"); + expect(waitpoint.type).toBe("MANUAL"); + // Read unconditionally by the debounce path, so it must never be undefined. + expect(waitpoint.outputIsError).toBe(false); + } finally { + await engine.quit(); + } + } + ); + + containerTest( + "a repeated idempotency key returns the cached waitpoint", + async ({ prisma, redisOptions }) => { + const environment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + const engine = engineFor(arm, prisma, redisOptions); + + try { + const args = { + environmentId: environment.id, + projectId: environment.project.id, + idempotencyKey: "same-key", + waitpointMintKind: arm, + } as const; + + const first = await engine.createManualWaitpoint(args); + const second = await engine.createManualWaitpoint(args); + + expect(first.isCached).toBe(false); + expect(second.isCached).toBe(true); + expect(second.waitpoint.id).toBe(first.waitpoint.id); + } finally { + await engine.quit(); + } + } + ); + + containerTest( + "createDateTimeWaitpoint mints into the expected system", + async ({ prisma, redisOptions }) => { + const environment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + const engine = engineFor(arm, prisma, redisOptions); + + try { + const { waitpoint } = await engine.createDateTimeWaitpoint({ + environmentId: environment.id, + projectId: environment.project.id, + completedAfter: new Date(Date.now() + 60_000), + waitpointMintKind: arm, + }); + + expect(parseWaitpointId(waitpoint.id).format).toBe(expectedFormat[arm]); + expect(waitpoint.type).toBe("DATETIME"); + expect(waitpoint.completedAfter).not.toBeNull(); + } finally { + await engine.quit(); + } + } + ); +}); + +describe("standalone waitpoint creates, mint-kind fallback", () => { + // Reversibility: clearing the flag must revert the NEXT mint with no deploy, and an + // engine that has a store configured must still mint legacy when told to. + containerTest( + "a legacy mint stays legacy even with a store configured", + async ({ prisma, redisOptions }) => { + const environment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + const engine = engineFor("store", prisma, redisOptions); + + try { + const { waitpoint } = await engine.createManualWaitpoint({ + environmentId: environment.id, + projectId: environment.project.id, + waitpointMintKind: "legacy", + }); + + expect(parseWaitpointId(waitpoint.id).format).toBe("legacy"); + } finally { + await engine.quit(); + } + } + ); + + // Fail safe, not fail loud: a store mint on a process with no store configured must not + // turn every trigger for a flipped organization into an error. + containerTest( + "a store mint falls back to legacy when no store is configured", + async ({ prisma, redisOptions }) => { + const environment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + const engine = engineFor("legacy", prisma, redisOptions); + + try { + const { waitpoint } = await engine.createManualWaitpoint({ + environmentId: environment.id, + projectId: environment.project.id, + waitpointMintKind: "store", + }); + + expect(parseWaitpointId(waitpoint.id).format).toBe("legacy"); + } finally { + await engine.quit(); + } + } + ); +}); diff --git a/internal-packages/run-engine/src/engine/types.ts b/internal-packages/run-engine/src/engine/types.ts index 2516776373d..4bdae6b3df1 100644 --- a/internal-packages/run-engine/src/engine/types.ts +++ b/internal-packages/run-engine/src/engine/types.ts @@ -1,3 +1,4 @@ +import type { WaitpointMintKind } from "./waitpointCoordinator/types.js"; import { type RedisOptions } from "@internal/redis"; import type { Meter, Tracer } from "@internal/tracing"; import type { Logger, LogLevel } from "@trigger.dev/core/logger"; @@ -136,6 +137,13 @@ export type RunEngineOptions = { cache?: { redis: RedisOptions; }; + /** + * The waitpoint store. Absent means the store arm is unreachable and every waitpoint + * operation routes to Postgres, whatever an organization's mint flag says. + */ + waitpointStore?: { + redis: RedisOptions; + }; batchQueue?: { redis: RedisOptions; drr?: Partial; @@ -301,6 +309,11 @@ export type HeartbeatTimeouts = { }; export type TriggerParams = { + /** + * Which coordinator mints this run's associated waitpoint, when a parent waits on it. + * Resolved from the organization's flag by the caller; absent means legacy. + */ + waitpointMintKind?: WaitpointMintKind; number?: number; friendlyId: string; environment: MinimalAuthenticatedEnvironment; diff --git a/internal-packages/run-engine/src/engine/waitpointCoordinator/legacyPostgresCoordinator.ts b/internal-packages/run-engine/src/engine/waitpointCoordinator/legacyPostgresCoordinator.ts index 46eea8d740c..b21dfa68efa 100644 --- a/internal-packages/run-engine/src/engine/waitpointCoordinator/legacyPostgresCoordinator.ts +++ b/internal-packages/run-engine/src/engine/waitpointCoordinator/legacyPostgresCoordinator.ts @@ -10,6 +10,7 @@ import { fetchWaitpointsInChunks } from "../systems/executionSnapshotSystem.js"; import { envelopeSourceFromWaitpointRow } from "./completionEnvelopeSource.js"; import type { AssociatedWaitpointData, + CreateBatchWaitpointParams, ClearRunBlockStateParams, CompleteParams, CompleteResult, @@ -296,6 +297,43 @@ export class LegacyPostgresWaitpointCoordinator implements WaitpointCoordinator return { kind: "created", waitpoint }; } + /** + * The BATCH waitpoint for a batch, keyed on the batch id as its idempotency key. + * + * The P2002 catch IS the duplicate-batch contract: a second call for the same batch + * collides on the idempotencyKey unique index, and null is the caller's "this batch + * already has one" signal rather than an error. It stays on this arm because the code + * is dead against a non-Postgres store, where NX reports the duplicate instead. + */ + async createBatchWaitpoint({ + batchId, + environmentId, + projectId, + tx, + }: CreateBatchWaitpointParams): Promise { + try { + return await this.runStore.createWaitpoint( + { + data: { + ...WaitpointId.generate(), + type: "BATCH", + idempotencyKey: batchId, + userProvidedIdempotencyKey: false, + completedByBatchId: batchId, + environmentId, + projectId, + }, + }, + tx + ); + } catch (error) { + if (error instanceof Prisma.PrismaClientKnownRequestError && error.code === "P2002") { + return null; + } + throw error; + } + } + async createManualWaitpoint({ runId, environmentId, diff --git a/internal-packages/run-engine/src/engine/waitpointCoordinator/routerCoordinator.test.ts b/internal-packages/run-engine/src/engine/waitpointCoordinator/routerCoordinator.test.ts new file mode 100644 index 00000000000..5d5d577b30a --- /dev/null +++ b/internal-packages/run-engine/src/engine/waitpointCoordinator/routerCoordinator.test.ts @@ -0,0 +1,279 @@ +import { getMeter } from "@internal/tracing"; +import { Logger } from "@trigger.dev/core/logger"; +import { generateWaitpointId } from "@trigger.dev/core/v3/isomorphic"; +import type { Waitpoint } from "@trigger.dev/database"; +import { describe, expect, it } from "vitest"; +import { UnclassifiableWaitpointId } from "../errors.js"; +import { WaitpointRouterCoordinator } from "./routerCoordinator.js"; +import type { CompletionEnvelopeSource, RunBlockEdge, WaitpointCoordinator } from "./types.js"; + +const LEGACY_ID = "waitpoint_ckabc123def456ghi789jkl"; +const logger = new Logger("routerCoordinator.test", "error"); + +function storeId() { + return generateWaitpointId("MANUAL"); +} + +/** + * A recording double, not a mock: a real object satisfying the seam that remembers what it + * was asked. The router's whole job is dispatch, so what each arm receives IS the assertion. + */ +function arm(name: string, calls: string[], overrides: Partial = {}) { + const base: WaitpointCoordinator = { + async clearRunBlockState(params) { + calls.push(`${name}.clearRunBlockState:${JSON.stringify(params.edgeIds ?? null)}`); + return { count: params.edgeIds?.length ?? 0 }; + }, + async readRunBlockState(runId) { + calls.push(`${name}.readRunBlockState`); + return []; + }, + async readCompletionEnvelopes(params) { + calls.push(`${name}.readCompletionEnvelopes:${params.waitpointIds.length}`); + return []; + }, + async registerBlocks(params) { + calls.push(`${name}.registerBlocks:${params.waitpointIds.length}`); + return { pendingCount: 0 }; + }, + async registerBlocksLockless(params) { + calls.push(`${name}.registerBlocksLockless:${params.waitpointIds.length}`); + }, + async complete(params) { + calls.push(`${name}.complete`); + return { waitpoint: { id: params.waitpointId } as Waitpoint, blockedRuns: [] }; + }, + async createDateTimeWaitpoint() { + calls.push(`${name}.createDateTimeWaitpoint`); + return { kind: "created", waitpoint: {} as Waitpoint }; + }, + async createManualWaitpoint() { + calls.push(`${name}.createManualWaitpoint`); + return { kind: "created", waitpoint: {} as Waitpoint }; + }, + async createBatchWaitpoint() { + calls.push(`${name}.createBatchWaitpoint`); + return {} as Waitpoint; + }, + mintAssociatedWaitpointData() { + calls.push(`${name}.mintAssociatedWaitpointData`); + return {} as never; + }, + async createAssociatedWaitpoint(params) { + calls.push(`${name}.createAssociatedWaitpoint`); + return { id: params.data.id } as Waitpoint; + }, + }; + + return { ...base, ...overrides }; +} + +function router(calls: string[], opts: { withStore?: boolean } = { withStore: true }) { + return new WaitpointRouterCoordinator({ + legacy: arm("legacy", calls), + store: opts.withStore ? arm("store", calls) : undefined, + logger, + meter: getMeter("routerCoordinator.test"), + }); +} + +describe("WaitpointRouterCoordinator", () => { + describe("routing an operation by id shape", () => { + it("sends a legacy id to the legacy arm", async () => { + const calls: string[] = []; + await router(calls).complete({ waitpointId: LEGACY_ID }); + expect(calls).toEqual(["legacy.complete"]); + }); + + it("sends a store id to the store arm", async () => { + const calls: string[] = []; + await router(calls).complete({ waitpointId: storeId() }); + expect(calls).toEqual(["store.complete"]); + }); + + it("throws on a store id when no store arm is configured", async () => { + const calls: string[] = []; + await expect( + router(calls, { withStore: false }).complete({ waitpointId: storeId() }) + ).rejects.toBeInstanceOf(UnclassifiableWaitpointId); + expect(calls).toEqual([]); + }); + }); + + describe("fanning a mixed run across both arms", () => { + it("concatenates readRunBlockState from both", async () => { + const calls: string[] = []; + const legacyEdge = { id: "edge_legacy" } as RunBlockEdge; + const storeEdge = { id: "edge_store" } as RunBlockEdge; + const coordinator = new WaitpointRouterCoordinator({ + legacy: arm("legacy", calls, { readRunBlockState: async () => [legacyEdge] }), + store: arm("store", calls, { readRunBlockState: async () => [storeEdge] }), + logger, + meter: getMeter("routerCoordinator.test"), + }); + + const edges = await coordinator.readRunBlockState("run_1"); + + expect(edges.map((e) => e.id)).toEqual(["edge_legacy", "edge_store"]); + }); + + it("sums the pending count across both arms", async () => { + const calls: string[] = []; + const coordinator = new WaitpointRouterCoordinator({ + legacy: arm("legacy", calls, { registerBlocks: async () => ({ pendingCount: 1 }) }), + store: arm("store", calls, { registerBlocks: async () => ({ pendingCount: 2 }) }), + logger, + meter: getMeter("routerCoordinator.test"), + }); + + const { pendingCount } = await coordinator.registerBlocks({ + runId: "run_1", + waitpointIds: [LEGACY_ID, storeId()], + projectId: "proj_1", + client: {} as never, + }); + + expect(pendingCount).toBe(3); + }); + + it("gives each arm only the ids it owns", async () => { + const calls: string[] = []; + await router(calls).registerBlocks({ + runId: "run_1", + waitpointIds: [LEGACY_ID, storeId(), storeId()], + projectId: "proj_1", + client: {} as never, + }); + + expect(calls.sort()).toEqual(["legacy.registerBlocks:1", "store.registerBlocks:2"]); + }); + + it("concatenates completion envelopes from both arms", async () => { + const calls: string[] = []; + const coordinator = new WaitpointRouterCoordinator({ + legacy: arm("legacy", calls, { + readCompletionEnvelopes: async () => [{ id: "a" } as CompletionEnvelopeSource], + }), + store: arm("store", calls, { + readCompletionEnvelopes: async () => [{ id: "b" } as CompletionEnvelopeSource], + }), + logger, + meter: getMeter("routerCoordinator.test"), + }); + + const sources = await coordinator.readCompletionEnvelopes({ + runId: "run_1", + waitpointIds: [LEGACY_ID, storeId()], + }); + + expect(sources.map((s) => s.id)).toEqual(["a", "b"]); + }); + + it("skips an arm that owns none of the requested ids", async () => { + const calls: string[] = []; + await router(calls).registerBlocks({ + runId: "run_1", + waitpointIds: [LEGACY_ID], + projectId: "proj_1", + client: {} as never, + }); + + expect(calls).toEqual(["legacy.registerBlocks:1"]); + }); + }); + + describe("clearing block state", () => { + // The trap this pins: an omitted edgeIds means "clear the whole run", so a partition + // that comes out empty must send [] and never omit, or it wipes the other arm's edges. + it("sends an empty array, never an omission, to the arm with no edges", async () => { + const calls: string[] = []; + await router(calls).clearRunBlockState({ runId: "run_1", edgeIds: ["ckLegacyEdgeId"] }); + + expect(calls.sort()).toEqual([ + 'legacy.clearRunBlockState:["ckLegacyEdgeId"]', + "store.clearRunBlockState:[]", + ]); + }); + + it("routes a store edge id by the waitpoint id it carries", async () => { + const calls: string[] = []; + const edgeId = `${storeId()}#0`; + await router(calls).clearRunBlockState({ runId: "run_1", edgeIds: [edgeId] }); + + expect(calls.sort()).toEqual([ + "legacy.clearRunBlockState:[]", + `store.clearRunBlockState:["${edgeId}"]`, + ]); + }); + + it("forwards a full clear to both arms with edgeIds omitted", async () => { + const calls: string[] = []; + await router(calls).clearRunBlockState({ runId: "run_1" }); + + expect(calls.sort()).toEqual([ + "legacy.clearRunBlockState:null", + "store.clearRunBlockState:null", + ]); + }); + + it("sums the cleared counts", async () => { + const calls: string[] = []; + const { count } = await router(calls).clearRunBlockState({ + runId: "run_1", + edgeIds: ["ckLegacyEdgeId", `${storeId()}#0`], + }); + + expect(count).toBe(2); + }); + }); + + describe("routing a create by mint kind", () => { + it("sends a legacy mint to the legacy arm", async () => { + const calls: string[] = []; + await router(calls).createManualWaitpoint({ + mintKind: "legacy", + environmentId: "env_1", + projectId: "proj_1", + }); + + expect(calls).toEqual(["legacy.createManualWaitpoint"]); + }); + + it("sends a store mint to the store arm", async () => { + const calls: string[] = []; + await router(calls).createManualWaitpoint({ + mintKind: "store", + environmentId: "env_1", + projectId: "proj_1", + }); + + expect(calls).toEqual(["store.createManualWaitpoint"]); + }); + + // Fail safe at the mint, unlike an operation on an existing id: a misconfigured deploy + // must not fail every trigger for a flipped organization. + it("falls back to legacy when a store mint finds no store arm", async () => { + const calls: string[] = []; + await router(calls, { withStore: false }).createManualWaitpoint({ + mintKind: "store", + environmentId: "env_1", + projectId: "proj_1", + }); + + expect(calls).toEqual(["legacy.createManualWaitpoint"]); + }); + }); + + describe("routing an associated waitpoint", () => { + it("routes createAssociatedWaitpoint by the shape of the minted id", async () => { + const calls: string[] = []; + const id = storeId(); + await router(calls).createAssociatedWaitpoint({ + runId: "run_1", + data: { id } as never, + }); + + expect(calls).toEqual(["store.createAssociatedWaitpoint"]); + }); + }); +}); diff --git a/internal-packages/run-engine/src/engine/waitpointCoordinator/routerCoordinator.ts b/internal-packages/run-engine/src/engine/waitpointCoordinator/routerCoordinator.ts new file mode 100644 index 00000000000..fad34c1a9ee --- /dev/null +++ b/internal-packages/run-engine/src/engine/waitpointCoordinator/routerCoordinator.ts @@ -0,0 +1,281 @@ +import type { Counter, Meter } from "@internal/tracing"; +import type { Logger } from "@trigger.dev/core/logger"; +import { deriveWaitpointIdFromAnchor, parseWaitpointId } from "@trigger.dev/core/v3/isomorphic"; +import type { Waitpoint } from "@trigger.dev/database"; +import { UnclassifiableWaitpointId } from "../errors.js"; +import { waitpointIdFromEdgeField } from "./keys.js"; +import type { + AssociatedWaitpointData, + ClearRunBlockStateParams, + CompleteParams, + CompleteResult, + CompletionEnvelopeSource, + CreateBatchWaitpointParams, + CreateDateTimeWaitpointParams, + CreateManualWaitpointParams, + CreateWaitpointResult, + ReadCompletionEnvelopesParams, + RegisterBlocksLocklessParams, + RegisterBlocksParams, + RunBlockEdge, + WaitpointCoordinator, + WaitpointMintKind, +} from "./types.js"; + +export type WaitpointRouterCoordinatorOptions = { + legacy: WaitpointCoordinator; + /** Absent when no waitpoint store is configured, which makes the store path unreachable. */ + store?: WaitpointCoordinator; + logger: Logger; + meter: Meter; +}; + +/** + * Chooses which arm owns a waitpoint, and nothing else. + * + * Every method here is a partition followed by delegation. It holds no store client and no + * Prisma client of its own, so a branch that is not about ownership does not belong here. + * + * Two different rules, deliberately: + * + * - An OPERATION routes on the id's shape. The id already exists, so its residency is a + * fact. A store-shaped id with no store arm configured throws, because guessing would + * silently operate on the wrong system. + * - A CREATE routes on the caller's mint kind. There is no id yet, so nothing can be + * misrouted. A store mint with no store arm falls back to legacy and says so: refusing + * would turn one process with a bad configuration into a trigger outage for every + * organization that has the flag set. + */ +export class WaitpointRouterCoordinator implements WaitpointCoordinator { + private readonly legacy: WaitpointCoordinator; + private readonly store?: WaitpointCoordinator; + private readonly logger: Logger; + private readonly legacyAnchorDowngrades: Counter; + + constructor(options: WaitpointRouterCoordinatorOptions) { + this.legacy = options.legacy; + this.store = options.store; + this.logger = options.logger; + this.legacyAnchorDowngrades = options.meter.createCounter( + "waitpoint.legacy_anchor_downgrades", + { + description: + "Store mints that fell back to legacy because the anchor run carried a legacy id", + } + ); + } + + async clearRunBlockState(params: ClearRunBlockStateParams): Promise<{ count: number }> { + // An omitted edgeIds is the terminal "clear the whole run", so it must reach both arms + // as an omission. A partition, by contrast, must send [] to the arm with nothing to + // drain: omitting there would clear that arm's remaining edges for the run. + if (!params.edgeIds) { + const [legacy, store] = await Promise.all([ + this.legacy.clearRunBlockState(params), + this.store?.clearRunBlockState(params), + ]); + + return { count: legacy.count + (store?.count ?? 0) }; + } + + const split = this.#partitionEdgeIds(params.edgeIds); + const [legacy, store] = await Promise.all([ + this.legacy.clearRunBlockState({ ...params, edgeIds: split.legacy }), + this.store?.clearRunBlockState({ ...params, edgeIds: split.store }), + ]); + + return { count: legacy.count + (store?.count ?? 0) }; + } + + /** + * Both arms, always, because a run can be blocked by one of each and the pending set is + * only correct as the union. The store read is one round trip against possibly-absent + * keys, which answers empty for a run that never touched the store. + */ + async readRunBlockState(runId: string): Promise { + const [legacy, store] = await Promise.all([ + this.legacy.readRunBlockState(runId), + this.store?.readRunBlockState(runId), + ]); + + return [...legacy, ...(store ?? [])]; + } + + async readCompletionEnvelopes( + params: ReadCompletionEnvelopesParams + ): Promise { + const split = this.#partitionWaitpointIds(params.waitpointIds); + + const [legacy, store] = await Promise.all([ + split.legacy.length + ? this.legacy.readCompletionEnvelopes({ ...params, waitpointIds: split.legacy }) + : [], + split.store.length + ? this.#requireStore(split.store[0]!).readCompletionEnvelopes({ + ...params, + waitpointIds: split.store, + }) + : [], + ]); + + return [...legacy, ...store]; + } + + /** + * The dual pending check. Each arm counts only the ids it owns, and the sum is the run's + * whole pending set, so a run blocked by one waitpoint of each kind stays blocked until + * both complete. + */ + async registerBlocks(params: RegisterBlocksParams): Promise<{ pendingCount: number }> { + const split = this.#partitionWaitpointIds(params.waitpointIds); + + const [legacy, store] = await Promise.all([ + split.legacy.length + ? this.legacy.registerBlocks({ ...params, waitpointIds: split.legacy }) + : undefined, + split.store.length + ? this.#requireStore(split.store[0]!).registerBlocks({ + ...params, + waitpointIds: split.store, + }) + : undefined, + ]); + + return { pendingCount: (legacy?.pendingCount ?? 0) + (store?.pendingCount ?? 0) }; + } + + async registerBlocksLockless(params: RegisterBlocksLocklessParams): Promise { + const split = this.#partitionWaitpointIds(params.waitpointIds); + + await Promise.all([ + split.legacy.length + ? this.legacy.registerBlocksLockless({ ...params, waitpointIds: split.legacy }) + : undefined, + split.store.length + ? this.#requireStore(split.store[0]!).registerBlocksLockless({ + ...params, + waitpointIds: split.store, + }) + : undefined, + ]); + } + + async complete(params: CompleteParams): Promise { + return this.#armFor(params.waitpointId).complete(params); + } + + async createDateTimeWaitpoint( + params: CreateDateTimeWaitpointParams + ): Promise { + return this.#armForMint(params.mintKind).createDateTimeWaitpoint(params); + } + + async createManualWaitpoint(params: CreateManualWaitpointParams): Promise { + return this.#armForMint(params.mintKind).createManualWaitpoint(params); + } + + async createBatchWaitpoint(params: CreateBatchWaitpointParams): Promise { + return this.#armForMint(params.mintKind).createBatchWaitpoint(params); + } + + /** + * A RUN waitpoint's store id is derived from its anchor run's id body, so an anchor that + * is not itself a run-ops id has nothing to derive from. That run keeps a legacy + * waitpoint even in a flipped organization, which is the coexistence rule. + * + * Counted, not just logged: an organization whose runs are all legacy-shaped mints zero + * store waitpoints, and a wave gate that reads "no store problems" off an empty sample + * is measuring nothing. + */ + mintAssociatedWaitpointData(params: { + projectId: string; + environmentId: string; + anchorRunId?: string; + mintKind?: WaitpointMintKind; + }): AssociatedWaitpointData { + const mintKind = params.mintKind ?? "legacy"; + + if (mintKind === "store" && !this.#canDeriveFromAnchor(params.anchorRunId)) { + this.legacyAnchorDowngrades.add(1); + this.logger.info("waitpoint mint fell back to legacy: the anchor run is not a run-ops id", { + anchorRunId: params.anchorRunId, + }); + return this.legacy.mintAssociatedWaitpointData(params); + } + + return this.#armForMint(mintKind).mintAssociatedWaitpointData(params); + } + + #canDeriveFromAnchor(anchorRunId: string | undefined): boolean { + return ( + anchorRunId !== undefined && deriveWaitpointIdFromAnchor(anchorRunId, "RUN") !== undefined + ); + } + + /** Routes on the minted id, so it lands wherever mintAssociatedWaitpointData put it. */ + async createAssociatedWaitpoint(params: { + runId: string; + data: AssociatedWaitpointData; + }): Promise { + return this.#armFor(params.data.id).createAssociatedWaitpoint(params); + } + + #armFor(waitpointId: string): WaitpointCoordinator { + return parseWaitpointId(waitpointId).format === "b32hexW" + ? this.#requireStore(waitpointId) + : this.legacy; + } + + #armForMint(mintKind: WaitpointMintKind): WaitpointCoordinator { + if (mintKind !== "store") { + return this.legacy; + } + + if (!this.store) { + this.logger.error( + "waitpoint mint asked for the store with no store configured; minting legacy", + { mintKind } + ); + return this.legacy; + } + + return this.store; + } + + #requireStore(waitpointId: string): WaitpointCoordinator { + if (!this.store) { + throw new UnclassifiableWaitpointId(waitpointId); + } + + return this.store; + } + + #partitionWaitpointIds(waitpointIds: string[]): { legacy: string[]; store: string[] } { + const legacy: string[] = []; + const store: string[] = []; + + for (const waitpointId of waitpointIds) { + (parseWaitpointId(waitpointId).format === "b32hexW" ? store : legacy).push(waitpointId); + } + + return { legacy, store }; + } + + /** + * A store edge id is `#`; a legacy edge id is a Postgres row id + * with no separator, so the helper reports undefined for it and it partitions legacy. + */ + #partitionEdgeIds(edgeIds: string[]): { legacy: string[]; store: string[] } { + const legacy: string[] = []; + const store: string[] = []; + + for (const edgeId of edgeIds) { + const waitpointId = waitpointIdFromEdgeField(edgeId); + const isStore = + waitpointId !== undefined && parseWaitpointId(waitpointId).format === "b32hexW"; + (isStore ? store : legacy).push(edgeId); + } + + return { legacy, store }; + } +} diff --git a/internal-packages/run-engine/src/engine/waitpointCoordinator/storeArm.test.ts b/internal-packages/run-engine/src/engine/waitpointCoordinator/storeArm.test.ts new file mode 100644 index 00000000000..c7f17a19dba --- /dev/null +++ b/internal-packages/run-engine/src/engine/waitpointCoordinator/storeArm.test.ts @@ -0,0 +1,439 @@ +import { createRedisClient, type RedisOptions } from "@internal/redis"; +import { containerTest } from "@internal/testcontainers"; +import { getMeter } from "@internal/tracing"; +import { Logger } from "@trigger.dev/core/logger"; +import { generateRunOpsId, generateWaitpointId } from "@trigger.dev/core/v3/isomorphic"; +import type { PrismaClient } from "@trigger.dev/database"; +import { describe, expect } from "vitest"; +import { PostgresRunStore } from "@internal/run-store"; +import { setupAuthenticatedEnvironment } from "../tests/setup.js"; +import { runBlockKeys } from "./keys.js"; +import { StoreWaitpointCoordinatorArm } from "./storeArm.js"; +import { WaitpointStoreCoordinator, type WaitpointRecordInput } from "./storeCoordinator.js"; + +const RUN_ID = "run_blocked"; +const NOW = "2026-08-26T12:00:00.000Z"; + +function setup(redisOptions: RedisOptions, prisma: PrismaClient) { + const store = new WaitpointStoreCoordinator({ redisOptions }); + const arm = new StoreWaitpointCoordinatorArm({ + store, + runStore: new PostgresRunStore({ prisma, readOnlyPrisma: prisma }), + logger: new Logger("storeArm.test", "error"), + meter: getMeter("storeArm.test"), + }); + + return { store, arm }; +} + +function record( + id: string, + environmentId: string, + projectId: string, + overrides: Partial = {} +): WaitpointRecordInput { + return { + id, + friendlyId: `waitpoint_${id}`, + type: "MANUAL", + environmentId, + projectId, + createdAt: NOW, + updatedAt: NOW, + userProvidedIdempotencyKey: false, + tags: [], + idempotencyKey: `idem_${id}`, + ...overrides, + }; +} + +describe("StoreWaitpointCoordinatorArm", () => { + containerTest( + "reports COMPLETED once a blocked waitpoint is delivered", + async ({ prisma, redisOptions }) => { + const environment = await setupEnvironment(prisma); + const { store, arm } = setup(redisOptions, prisma); + + try { + const waitpointId = generateWaitpointId("MANUAL"); + await store.createIfAbsent({ + record: record(waitpointId, environment.id, environment.projectId), + status: "PENDING", + }); + + const { pendingCount } = await arm.registerBlocks({ + runId: RUN_ID, + waitpointIds: [waitpointId], + projectId: environment.projectId, + client: prisma, + }); + expect(pendingCount).toBe(1); + + const beforeComplete = await arm.readRunBlockState(RUN_ID); + expect(beforeComplete[0]!.waitpoint.status).toBe("PENDING"); + + await arm.complete({ waitpointId, output: { value: "42", isError: false } }); + + const afterComplete = await arm.readRunBlockState(RUN_ID); + expect(afterComplete).toHaveLength(1); + expect(afterComplete[0]!.waitpoint.status).toBe("COMPLETED"); + expect(afterComplete[0]!.waitpoint.type).toBe("MANUAL"); + } finally { + await store.quit(); + } + } + ); + + // I10, and the only premature-resume counterexample either TLA+ campaign produced. A + // run-shard loss removes the pending entry while the edge survives; "not pending, + // therefore complete" would resume a run whose waitpoint never completed. + containerTest( + "reports PENDING for an edge that is in neither the pending nor the delivered set", + async ({ prisma, redisOptions }) => { + const environment = await setupEnvironment(prisma); + const { store, arm } = setup(redisOptions, prisma); + const redis = createRedisClient(redisOptions); + + try { + const waitpointId = generateWaitpointId("MANUAL"); + await store.createIfAbsent({ + record: record(waitpointId, environment.id, environment.projectId), + status: "PENDING", + }); + await arm.registerBlocks({ + runId: RUN_ID, + waitpointIds: [waitpointId], + projectId: environment.projectId, + client: prisma, + }); + + await redis.srem(runBlockKeys(RUN_ID).pend, waitpointId); + + const edges = await arm.readRunBlockState(RUN_ID); + expect(edges).toHaveLength(1); + expect(edges[0]!.waitpoint.status).toBe("PENDING"); + } finally { + await redis.quit(); + await store.quit(); + } + } + ); + + // The case a "has a completion envelope" rule would wedge forever: a waitpoint may be + // COMPLETED with no envelope, which the reported box models on purpose. + containerTest( + "reports COMPLETED for a waitpoint completed before the run ever blocked on it", + async ({ prisma, redisOptions }) => { + const environment = await setupEnvironment(prisma); + const { store, arm } = setup(redisOptions, prisma); + + try { + const waitpointId = generateWaitpointId("MANUAL"); + await store.createIfAbsent({ + record: record(waitpointId, environment.id, environment.projectId), + status: "COMPLETED", + }); + + const { pendingCount } = await arm.registerBlocks({ + runId: RUN_ID, + waitpointIds: [waitpointId], + projectId: environment.projectId, + client: prisma, + }); + + expect(pendingCount).toBe(0); + + const edges = await arm.readRunBlockState(RUN_ID); + expect(edges[0]!.waitpoint.status).toBe("COMPLETED"); + } finally { + await store.quit(); + } + } + ); + + // §5.4's guard. Unmodeled in both campaigns, so this assertion is its only protection. + containerTest( + "refuses a lockless absorb when the parent BATCH waitpoint is absent", + async ({ prisma, redisOptions }) => { + const environment = await setupEnvironment(prisma); + const { store, arm } = setup(redisOptions, prisma); + + try { + const itemWaitpointId = generateWaitpointId("RUN"); + const batchWaitpointId = generateWaitpointId("BATCH"); + + await expect( + arm.registerBlocksLockless({ + runId: RUN_ID, + waitpointIds: [itemWaitpointId], + projectId: environment.projectId, + batchId: "batch_1", + batchIndex: 0, + batchWaitpointId, + }) + ).rejects.toThrow(/BATCH waitpoint/); + } finally { + await store.quit(); + } + } + ); + + // Present-but-not-pending is the half of the guard a presence-only check would miss. + containerTest( + "refuses a lockless absorb when the parent BATCH waitpoint is already complete", + async ({ prisma, redisOptions }) => { + const environment = await setupEnvironment(prisma); + const { store, arm } = setup(redisOptions, prisma); + + try { + const batchWaitpointId = generateWaitpointId("BATCH"); + await store.createIfAbsent({ + record: record(batchWaitpointId, environment.id, environment.projectId, { + type: "BATCH", + }), + status: "PENDING", + }); + await arm.registerBlocks({ + runId: RUN_ID, + waitpointIds: [batchWaitpointId], + projectId: environment.projectId, + client: prisma, + }); + await arm.complete({ waitpointId: batchWaitpointId, output: undefined }); + + await expect( + arm.registerBlocksLockless({ + runId: RUN_ID, + waitpointIds: [generateWaitpointId("RUN")], + projectId: environment.projectId, + batchId: "batch_1", + batchIndex: 0, + batchWaitpointId, + }) + ).rejects.toThrow(/BATCH waitpoint/); + } finally { + await store.quit(); + } + } + ); + + containerTest( + "allows a lockless absorb while the parent BATCH waitpoint is pending", + async ({ prisma, redisOptions }) => { + const environment = await setupEnvironment(prisma); + const { store, arm } = setup(redisOptions, prisma); + + try { + const batchWaitpointId = generateWaitpointId("BATCH"); + await store.createIfAbsent({ + record: record(batchWaitpointId, environment.id, environment.projectId, { + type: "BATCH", + }), + status: "PENDING", + }); + await arm.registerBlocks({ + runId: RUN_ID, + waitpointIds: [batchWaitpointId], + projectId: environment.projectId, + client: prisma, + }); + + const itemWaitpointId = generateWaitpointId("RUN"); + await store.createIfAbsent({ + record: record(itemWaitpointId, environment.id, environment.projectId, { type: "RUN" }), + status: "PENDING", + }); + + await arm.registerBlocksLockless({ + runId: RUN_ID, + waitpointIds: [itemWaitpointId], + projectId: environment.projectId, + batchId: "batch_1", + batchIndex: 0, + batchWaitpointId, + }); + + // The parent's BATCH waitpoint is still pending after the item absorbed, which is + // the invariant: the pending set is never momentarily empty mid-absorb. + const edges = await arm.readRunBlockState(RUN_ID); + const stillPending = edges.filter((e) => e.waitpoint.status === "PENDING"); + expect(stillPending.map((e) => e.waitpoint.id).sort()).toEqual( + [batchWaitpointId, itemWaitpointId].sort() + ); + } finally { + await store.quit(); + } + } + ); + + containerTest( + "writes the MANUAL projection row after the store commit", + async ({ prisma, redisOptions }) => { + const environment = await setupEnvironment(prisma); + const { store, arm } = setup(redisOptions, prisma); + + try { + const result = await arm.createManualWaitpoint({ + mintKind: "store", + environmentId: environment.id, + projectId: environment.projectId, + tags: ["alpha"], + }); + + expect(result.kind).toBe("created"); + + const row = await prisma.waitpoint.findFirst({ where: { id: result.waitpoint.id } }); + expect(row?.type).toBe("MANUAL"); + expect(row?.tags).toEqual(["alpha"]); + + // The store is the system of record; the row is a projection of it. + const held = await store.readWaitpoint(result.waitpoint.id); + expect(held?.status).toBe("PENDING"); + } finally { + await store.quit(); + } + } + ); + + // The token API and dashboard read status, output and completedAt from the projection + // row, so a completion that never reaches it reports a finished token as still waiting. + containerTest( + "reflects a MANUAL completion onto the projection row", + async ({ prisma, redisOptions }) => { + const environment = await setupEnvironment(prisma); + const { store, arm } = setup(redisOptions, prisma); + + try { + const created = await arm.createManualWaitpoint({ + mintKind: "store", + environmentId: environment.id, + projectId: environment.projectId, + }); + + await arm.complete({ + waitpointId: created.waitpoint.id, + output: { value: '{"done":true}', type: "application/json", isError: false }, + }); + + const row = await prisma.waitpoint.findFirst({ where: { id: created.waitpoint.id } }); + expect(row?.status).toBe("COMPLETED"); + expect(row?.output).toBe('{"done":true}'); + expect(row?.outputIsError).toBe(false); + expect(row?.completedAt).not.toBeNull(); + } finally { + await store.quit(); + } + } + ); + + // An unwired caller must fail, never silently disable the guard. + containerTest( + "refuses a lockless absorb that arrives with no parent BATCH waitpoint id", + async ({ prisma, redisOptions }) => { + const environment = await setupEnvironment(prisma); + const { store, arm } = setup(redisOptions, prisma); + + try { + await expect( + arm.registerBlocksLockless({ + runId: RUN_ID, + waitpointIds: [generateWaitpointId("RUN")], + projectId: environment.projectId, + batchId: "batch_1", + batchIndex: 0, + }) + ).rejects.toThrow(/no parent .*BATCH waitpoint id/); + } finally { + await store.quit(); + } + } + ); + + containerTest( + "returns the cached waitpoint for a repeated idempotency key", + async ({ prisma, redisOptions }) => { + const environment = await setupEnvironment(prisma); + const { store, arm } = setup(redisOptions, prisma); + + try { + const args = { + mintKind: "store" as const, + environmentId: environment.id, + projectId: environment.projectId, + idempotencyKey: "same-key", + }; + + const first = await arm.createManualWaitpoint(args); + const second = await arm.createManualWaitpoint(args); + + expect(first.kind).toBe("created"); + expect(second.kind).toBe("cached"); + expect(second.waitpoint.id).toBe(first.waitpoint.id); + } finally { + await store.quit(); + } + } + ); + + containerTest( + "returns null when the batch already has a waitpoint", + async ({ prisma, redisOptions }) => { + const environment = await setupEnvironment(prisma); + const { store, arm } = setup(redisOptions, prisma); + + try { + const batchId = `batch_${generateRunOpsId()}`; + const args = { + batchId, + environmentId: environment.id, + projectId: environment.projectId, + mintKind: "store" as const, + }; + + const first = await arm.createBatchWaitpoint(args); + expect(first).not.toBeNull(); + expect(first!.type).toBe("BATCH"); + expect(first!.completedByBatchId).toBe(batchId); + + const second = await arm.createBatchWaitpoint(args); + expect(second).toBeNull(); + } finally { + await store.quit(); + } + } + ); + + containerTest( + "creates the RUN waitpoint at the anchor-derived id, idempotently", + async ({ prisma, redisOptions }) => { + const environment = await setupEnvironment(prisma); + const { store, arm } = setup(redisOptions, prisma); + + try { + const runId = generateRunOpsId(); + const data = arm.mintAssociatedWaitpointData({ + projectId: environment.projectId, + environmentId: environment.id, + anchorRunId: runId, + }); + + // Pure function of the run id, which is what removes the need for a lock. + expect(data.id.slice(0, 24)).toBe(runId.slice(0, 24)); + + const first = await arm.createAssociatedWaitpoint({ runId, data }); + const second = await arm.createAssociatedWaitpoint({ runId, data }); + + expect(first.id).toBe(data.id); + expect(second.id).toBe(data.id); + expect(second.status).toBe("PENDING"); + } finally { + await store.quit(); + } + } + ); +}); + +async function setupEnvironment(prisma: PrismaClient) { + const environment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + return { id: environment.id, projectId: environment.project.id }; +} diff --git a/internal-packages/run-engine/src/engine/waitpointCoordinator/storeArm.ts b/internal-packages/run-engine/src/engine/waitpointCoordinator/storeArm.ts new file mode 100644 index 00000000000..0a146ed1ec7 --- /dev/null +++ b/internal-packages/run-engine/src/engine/waitpointCoordinator/storeArm.ts @@ -0,0 +1,557 @@ +import type { Meter, Counter } from "@internal/tracing"; +import type { RunStore } from "@internal/run-store"; +import type { Logger } from "@trigger.dev/core/logger"; +import { tryCatch } from "@trigger.dev/core/v3"; +import { + deriveWaitpointIdFromAnchor, + generateWaitpointId, + parseWaitpointId, + WaitpointId, +} from "@trigger.dev/core/v3/isomorphic"; +import type { Waitpoint } from "@trigger.dev/database"; +import { nanoid } from "nanoid"; +import type { + BlockEdge, + WaitpointCompletion, + WaitpointRecordInput, + WaitpointStoreCoordinator, +} from "./storeCoordinator.js"; +import type { + AssociatedWaitpointData, + ClearRunBlockStateParams, + CompleteParams, + CompleteResult, + CompletionEnvelopeSource, + CreateBatchWaitpointParams, + CreateDateTimeWaitpointParams, + CreateManualWaitpointParams, + CreateWaitpointResult, + ReadCompletionEnvelopesParams, + RegisterBlocksLocklessParams, + RegisterBlocksParams, + RunBlockEdge, + WaitpointCoordinator, +} from "./types.js"; +import { toPrismaWaitpoint } from "./waitpointShape.js"; + +export type StoreWaitpointCoordinatorArmOptions = { + store: WaitpointStoreCoordinator; + /** MANUAL projection writes only. Never read for coordination (I6). */ + runStore: RunStore; + logger: Logger; + meter: Meter; +}; + +/** + * Waitpoint coordination against the Redis store. + * + * The store is the system of record. Postgres keeps one derived artefact — the MANUAL + * projection row, written after the store commit so the dashboard and token API keep + * working — and no coordination path ever reads it back. + */ +export class StoreWaitpointCoordinatorArm implements WaitpointCoordinator { + private readonly store: WaitpointStoreCoordinator; + private readonly runStore: RunStore; + private readonly logger: Logger; + + private readonly resumeCrossCheckViolations: Counter; + private readonly batchGuardViolations: Counter; + private readonly projectionWriteFailures: Counter; + + constructor(options: StoreWaitpointCoordinatorArmOptions) { + this.store = options.store; + this.runStore = options.runStore; + this.logger = options.logger; + + this.resumeCrossCheckViolations = options.meter.createCounter( + "waitpoint.resume_crosscheck_violations", + { description: "Block edges found in neither the pending nor the delivered set" } + ); + this.batchGuardViolations = options.meter.createCounter("waitpoint.batch_guard_violations", { + description: "Lockless absorbs attempted without a pending parent BATCH waitpoint", + }); + this.projectionWriteFailures = options.meter.createCounter( + "waitpoint.projection_write_failures", + { description: "MANUAL projection rows that failed to write after the store commit" } + ); + } + + /** + * The store reports an outcome, not a delete count, and the seam's only consumer of the + * count is a debug log in the run-completion path. So this reports what was asked to + * drain rather than paying a read to confirm it. + */ + async clearRunBlockState({ runId, edgeIds }: ClearRunBlockStateParams): Promise<{ + count: number; + }> { + await this.store.clearBlockState({ runId, edgeIds }); + return { count: edgeIds?.length ?? 0 }; + } + + async readRunBlockState(runId: string): Promise { + const state = await this.store.readBlockState(runId); + const pending = new Set(state.pendingIds); + const delivered = new Set(state.deliveredIds); + + return state.edges.map((edge) => ({ + id: edge.edgeId, + batchId: edge.batchId ?? null, + batchIndex: edge.batchIndex ?? null, + waitpoint: { + id: edge.waitpointId, + status: this.#deriveStatus(runId, edge.waitpointId, pending, delivered), + type: edge.type, + completedAfter: edge.completedAfter ? new Date(edge.completedAfter) : null, + }, + })); + } + + /** + * I10. `runAbsorbBlockers` keeps every edge in exactly one of the pending or delivered + * sets. A run-shard data loss breaks that: the edge survives while its pending entry is + * gone. Reading "not pending, therefore complete" then resumes a run whose waitpoint + * never completed, which is the only premature-resume counterexample either TLA+ + * campaign produced. So an edge in neither set reports PENDING and is counted; the run + * stays blocked and a later sweep heals it. + */ + #deriveStatus( + runId: string, + waitpointId: string, + pending: Set, + delivered: Set + ): "PENDING" | "COMPLETED" { + if (delivered.has(waitpointId)) { + return "COMPLETED"; + } + + if (!pending.has(waitpointId)) { + this.resumeCrossCheckViolations.add(1); + this.logger.error("waitpoint edge is in neither the pending nor the delivered set", { + runId, + waitpointId, + }); + } + + return "PENDING"; + } + + readCompletionEnvelopes( + params: ReadCompletionEnvelopesParams + ): Promise { + return this.store.readCompletionEnvelopes(params); + } + + async registerBlocks(params: RegisterBlocksParams): Promise<{ pendingCount: number }> { + const edges = await this.#buildEdges(params); + const { pendingOfRequested } = await this.store.registerBlocks({ + runId: params.runId, + edges, + }); + + return { pendingCount: pendingOfRequested }; + } + + async registerBlocksLockless(params: RegisterBlocksLocklessParams): Promise { + await this.#assertBatchWaitpointPending(params); + + const edges = await this.#buildEdges(params); + await this.store.registerBlocks({ runId: params.runId, edges }); + } + + /** + * §5.4's guard invariant. A lockless absorb writes item edges one at a time without the + * run lock, which is only safe while the parent's BATCH waitpoint holds the pending set + * open. If it is absent or already complete, a concurrent completion could see an empty + * pending set mid-absorb and resume the parent early. + * + * Scope, stated precisely: this is a PREFLIGHT DETECTOR, not a barrier. It reads the run + * shard, then the absorb writes in a separate operation, so a completion landing between + * the two is detected on the next call, not prevented. Closing that window means moving + * the pending-set assertion inside the absorb script, so check and write share one + * atomic action. + * + * Neither TLA+ campaign models this variant, so until the race harness covers it this + * detector plus the fail-loud on a missing id is the whole protection. + */ + async #assertBatchWaitpointPending(params: RegisterBlocksLocklessParams): Promise { + if (!params.batchWaitpointId) { + // Never silently skip. An unwired caller would disable the guard rather than fail, + // which is the failure mode the guard exists to prevent. + this.batchGuardViolations.add(1); + throw new Error( + `Lockless absorb for run ${params.runId} reached the store arm with no parent ` + + `BATCH waitpoint id, so the pending-set guard has nothing to assert on` + ); + } + + const state = await this.store.readBlockState(params.runId); + if (state.pendingIds.includes(params.batchWaitpointId)) { + return; + } + + this.batchGuardViolations.add(1); + throw new Error( + `Lockless absorb for run ${params.runId} requires the parent BATCH waitpoint ` + + `${params.batchWaitpointId} to be present and pending on the run shard` + ); + } + + /** + * The edge blobs the run shard stores. + * + * `type` comes free from the id, which is what the positional id layout buys. Only + * DATETIME needs a record read, because its `completedAfter` rides the edge so the + * block-state read never has to touch each waitpoint's own key. RUN, BATCH and MANUAL + * skip it, which keeps `triggerAndWait` at one round trip per waitpoint. + */ + async #buildEdges(params: RegisterBlocksLocklessParams): Promise { + const createdAt = new Date().toISOString(); + const dateTimeIds = params.waitpointIds.filter((id) => { + const parsed = parseWaitpointId(id); + return parsed.format === "b32hexW" && parsed.type === "DATETIME"; + }); + const completedAfterById = await this.#readCompletedAfter(dateTimeIds); + + return params.waitpointIds.map((waitpointId) => { + const parsed = parseWaitpointId(waitpointId); + if (parsed.format !== "b32hexW") { + throw new Error(`Waitpoint ${waitpointId} is not a store-format id`); + } + + return { + waitpointId, + batchIndex: params.batchIndex ?? null, + batchId: params.batchId, + spanIdToComplete: params.spanIdToComplete, + createdAt, + type: parsed.type, + completedAfter: completedAfterById.get(waitpointId), + }; + }); + } + + async #readCompletedAfter(waitpointIds: string[]): Promise> { + const found = new Map(); + + for (const waitpointId of waitpointIds) { + const held = await this.store.readWaitpoint(waitpointId); + if (held?.record.completedAfter) { + found.set(waitpointId, held.record.completedAfter); + } + } + + return found; + } + + async complete({ waitpointId, output }: CompleteParams): Promise { + const completion: WaitpointCompletion = { + completedAt: new Date().toISOString(), + outputType: output?.type ?? "application/json", + outputIsError: output?.isError ?? false, + output: output ? { inline: output.value } : null, + }; + + const result = await this.store.complete({ waitpointId, completion }); + + // Deliver onto each watcher's own shard. The complete script returned the watchers + // atomically, so a watcher registered before the flip is always in this list. + for (const watcher of result.watchers) { + await this.store.deliverCompletion({ + runId: watcher.runId, + waitpointId, + completion: result.completion ?? completion, + }); + } + + const held = await this.store.readWaitpoint(waitpointId); + if (!held) { + throw new Error(`Waitpoint ${waitpointId} is not present in the store`); + } + + if (held.record.type === "MANUAL") { + await this.#completeManualProjection(waitpointId, held.completion ?? completion); + } + + return { + waitpoint: toPrismaWaitpoint(held.record, held.status, held.completion), + blockedRuns: result.watchers.map((watcher) => ({ + taskRunId: watcher.runId, + spanIdToComplete: watcher.spanIdToComplete ?? null, + createdAt: new Date(watcher.createdAt), + })), + }; + } + + async createDateTimeWaitpoint( + params: CreateDateTimeWaitpointParams + ): Promise { + return this.#createStandalone({ + type: "DATETIME", + environmentId: params.environmentId, + projectId: params.projectId, + idempotencyKey: params.idempotencyKey, + idempotencyKeyExpiresAt: params.idempotencyKeyExpiresAt, + completedAfter: params.completedAfter, + }); + } + + async createManualWaitpoint(params: CreateManualWaitpointParams): Promise { + const result = await this.#createStandalone({ + type: "MANUAL", + environmentId: params.environmentId, + projectId: params.projectId, + idempotencyKey: params.idempotencyKey, + idempotencyKeyExpiresAt: params.idempotencyKeyExpiresAt, + completedAfter: params.timeout, + tags: params.tags, + }); + + await this.#writeManualProjection(result.waitpoint); + return result; + } + + async createBatchWaitpoint({ + batchId, + environmentId, + projectId, + }: CreateBatchWaitpointParams): Promise { + const waitpointId = deriveWaitpointIdFromAnchor(batchId, "BATCH"); + if (!waitpointId) { + throw new Error(`Batch ${batchId} is not a run-ops id, so no BATCH waitpoint derives`); + } + + const record = this.#record({ + id: waitpointId, + type: "BATCH", + environmentId, + projectId, + idempotencyKey: batchId, + completedByBatchId: batchId, + }); + + const created = await this.store.createIfAbsent({ record, status: "PENDING" }); + + // The duplicate-batch contract. NX reports the second call, where the legacy arm gets + // a unique-index violation. + if (created.outcome === "exists") { + return null; + } + + return toPrismaWaitpoint(record, "PENDING"); + } + + mintAssociatedWaitpointData({ + projectId, + environmentId, + anchorRunId, + }: { + projectId: string; + environmentId: string; + anchorRunId?: string; + }): AssociatedWaitpointData { + const derived = anchorRunId ? deriveWaitpointIdFromAnchor(anchorRunId, "RUN") : undefined; + if (!derived) { + throw new Error( + `Run ${anchorRunId ?? "(none)"} is not a run-ops id, so no RUN waitpoint derives` + ); + } + + return { + id: derived, + friendlyId: WaitpointId.toFriendlyId(derived), + type: "RUN", + status: "PENDING", + idempotencyKey: nanoid(24), + userProvidedIdempotencyKey: false, + projectId, + environmentId, + }; + } + + /** + * Create-if-absent on the anchor-derived id. + * + * The lock and double-check the legacy arm needs are gone: the id is a pure function of + * the run id, so two racing callers compute the same id and NX settles it. A caller that + * finds it already present takes the existing record, which is what makes the crash + * window between the run commit and this call recoverable by retry. + */ + async createAssociatedWaitpoint({ + runId, + data, + }: { + runId: string; + data: AssociatedWaitpointData; + }): Promise { + const record = this.#record({ + id: data.id, + friendlyId: data.friendlyId, + type: "RUN", + environmentId: data.environmentId, + projectId: data.projectId, + idempotencyKey: data.idempotencyKey, + completedByTaskRunId: runId, + }); + + const created = await this.store.createIfAbsent({ record, status: "PENDING" }); + if (created.outcome === "exists") { + return toPrismaWaitpoint(created.record, created.status, created.completion); + } + + return toPrismaWaitpoint(record, "PENDING"); + } + + async #createStandalone(params: { + type: "DATETIME" | "MANUAL"; + environmentId: string; + projectId: string; + idempotencyKey?: string; + idempotencyKeyExpiresAt?: Date; + completedAfter?: Date; + tags?: string[]; + }): Promise { + const userProvidedIdempotencyKey = params.idempotencyKey !== undefined; + const record = this.#record({ + id: generateWaitpointId(params.type), + type: params.type, + environmentId: params.environmentId, + projectId: params.projectId, + idempotencyKey: params.idempotencyKey ?? nanoid(24), + userProvidedIdempotencyKey, + idempotencyKeyExpiresAt: params.idempotencyKeyExpiresAt?.toISOString(), + completedAfter: params.completedAfter?.toISOString(), + tags: params.tags, + }); + + // Without a user key there is nothing to dedupe against, so the reservation round trip + // is skipped entirely rather than reserved against a random key nobody will present. + if (!userProvidedIdempotencyKey) { + await this.store.createIfAbsent({ record, status: "PENDING" }); + return { kind: "created", waitpoint: toPrismaWaitpoint(record, "PENDING") }; + } + + const reserved = await this.store.createWithIdempotencyKey({ + record, + environmentId: params.environmentId, + idempotencyKey: params.idempotencyKey!, + }); + + if (reserved.created) { + return { kind: "created", waitpoint: toPrismaWaitpoint(record, "PENDING") }; + } + + const held = await this.store.readWaitpoint(reserved.waitpointId); + if (!held) { + throw new Error(`Waitpoint ${reserved.waitpointId} won the reservation but is absent`); + } + + return { + kind: "cached", + waitpoint: toPrismaWaitpoint(held.record, held.status, held.completion), + }; + } + + /** + * The MANUAL projection (I6). Written after the store commit, read by the dashboard and + * the token API, and never consulted for coordination. + * + * A failure here must not fail the create: the waitpoint already exists in the store and + * is already coordinating, so throwing would report failure for work that succeeded. + */ + async #writeManualProjection(waitpoint: Waitpoint): Promise { + const [error] = await tryCatch( + this.runStore.createWaitpoint({ + data: { + id: waitpoint.id, + friendlyId: waitpoint.friendlyId, + type: "MANUAL", + status: waitpoint.status, + idempotencyKey: waitpoint.idempotencyKey, + userProvidedIdempotencyKey: waitpoint.userProvidedIdempotencyKey, + idempotencyKeyExpiresAt: waitpoint.idempotencyKeyExpiresAt ?? undefined, + completedAfter: waitpoint.completedAfter ?? undefined, + environmentId: waitpoint.environmentId, + projectId: waitpoint.projectId, + tags: waitpoint.tags, + }, + }) + ); + + if (error) { + this.projectionWriteFailures.add(1); + this.logger.error("failed to write the MANUAL waitpoint projection row", { + waitpointId: waitpoint.id, + error, + }); + } + } + + /** + * Reflect a MANUAL completion onto the projection row. + * + * The token API and the dashboard read status, output and completedAt from this row, so + * leaving it PENDING would report a completed token as still waiting. Best effort, for + * the same reason as the create-time write: the store already completed the waitpoint. + */ + async #completeManualProjection( + waitpointId: string, + completion: WaitpointCompletion + ): Promise { + const output = completion.output; + + const [error] = await tryCatch( + this.runStore.updateManyWaitpoints({ + where: { id: waitpointId }, + data: { + status: "COMPLETED", + completedAt: new Date(completion.completedAt), + output: output ? ("inline" in output ? output.inline : output.ref) : null, + outputType: completion.outputType, + outputIsError: completion.outputIsError, + }, + }) + ); + + if (error) { + this.projectionWriteFailures.add(1); + this.logger.error("failed to complete the MANUAL waitpoint projection row", { + waitpointId, + error, + }); + } + } + + #record(params: { + id: string; + friendlyId?: string; + type: WaitpointRecordInput["type"]; + environmentId: string; + projectId: string; + idempotencyKey: string; + userProvidedIdempotencyKey?: boolean; + idempotencyKeyExpiresAt?: string; + completedAfter?: string; + completedByTaskRunId?: string; + completedByBatchId?: string; + tags?: string[]; + }): WaitpointRecordInput { + const now = new Date().toISOString(); + + return { + id: params.id, + friendlyId: params.friendlyId ?? WaitpointId.toFriendlyId(params.id), + type: params.type, + environmentId: params.environmentId, + projectId: params.projectId, + createdAt: now, + updatedAt: now, + userProvidedIdempotencyKey: params.userProvidedIdempotencyKey ?? false, + tags: params.tags ?? [], + idempotencyKey: params.idempotencyKey, + idempotencyKeyExpiresAt: params.idempotencyKeyExpiresAt, + completedAfter: params.completedAfter, + completedByTaskRunId: params.completedByTaskRunId, + completedByBatchId: params.completedByBatchId, + }; + } +} diff --git a/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.ts b/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.ts index 42be33724f6..164d0caf15a 100644 --- a/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.ts +++ b/internal-packages/run-engine/src/engine/waitpointCoordinator/storeCoordinator.ts @@ -506,6 +506,35 @@ export class WaitpointStoreCoordinator { return { pendingIds, deliveredIds, edges }; } + /** + * Read one waitpoint's three parts, or undefined when the store does not hold it. + * + * Single key, so no script and no #call guard: nothing here can span two slots. The + * seam needs this because its return types are the Postgres row shape, and only the + * immutable record carries the columns that shape requires. + */ + async readWaitpoint(waitpointId: string): Promise< + | { + record: WaitpointRecordInput; + status: WaitpointStatus; + completion?: WaitpointCompletion; + } + | undefined + > { + const fields = await this.redis.hmget(waitpointKeys(waitpointId).record, "r", "status", "c"); + + const record = parseJson(fields[0] ?? undefined); + if (!record) { + return undefined; + } + + return { + record, + status: fields[1] === "COMPLETED" ? "COMPLETED" : "PENDING", + completion: parseJson(fields[2] ?? undefined), + }; + } + /** * Source the envelope fields for a run's COMPLETED waitpoints. * diff --git a/internal-packages/run-engine/src/engine/waitpointCoordinator/types.ts b/internal-packages/run-engine/src/engine/waitpointCoordinator/types.ts index 8611a361b42..224d205ab28 100644 --- a/internal-packages/run-engine/src/engine/waitpointCoordinator/types.ts +++ b/internal-packages/run-engine/src/engine/waitpointCoordinator/types.ts @@ -24,9 +24,18 @@ export type WaitpointCoordinator = { complete(params: CompleteParams): Promise; createDateTimeWaitpoint(params: CreateDateTimeWaitpointParams): Promise; createManualWaitpoint(params: CreateManualWaitpointParams): Promise; + createBatchWaitpoint(params: CreateBatchWaitpointParams): Promise; mintAssociatedWaitpointData(params: { projectId: string; environmentId: string; + /** + * The run this waitpoint belongs to. A store arm derives the waitpoint id from the + * run's own id body, so the derivation is a pure function of the anchor and needs no + * lock. A Postgres arm mints a fresh id and ignores this. + */ + anchorRunId?: string; + /** Which arm mints it. Absent means legacy, which is what every existing caller wants. */ + mintKind?: WaitpointMintKind; }): AssociatedWaitpointData; createAssociatedWaitpoint(params: { runId: string; @@ -34,6 +43,23 @@ export type WaitpointCoordinator = { }): Promise; }; +/** + * Which coordinator mints a NEW waitpoint. Structurally identical to the webapp's own + * WaitpointMintKind; re-declared because the engine never imports from the webapp. + * + * Read at the mint and never again — every later operation routes by the minted id's shape. + */ +export type WaitpointMintKind = "legacy" | "store"; + +export type CreateBatchWaitpointParams = { + batchId: string; + environmentId: string; + projectId: string; + mintKind: WaitpointMintKind; + /** Legacy arm only: the create may join a caller transaction. A store arm ignores it. */ + tx?: PrismaClientOrTransaction; +}; + export type ReadCompletionEnvelopesParams = { runId: string; /** The DISTINCT completed waitpoint ids to source. Result order is not meaningful. */ @@ -110,7 +136,15 @@ export type RegisterBlocksParams = { * The lockless variant writes the edge and does not count. Two methods rather than * one method with a flag, so "the batch path issues no extra query" is structural. */ -export type RegisterBlocksLocklessParams = Omit; +export type RegisterBlocksLocklessParams = Omit & { + /** + * The parent's BATCH waitpoint id. A store arm asserts it is present and PENDING on the + * run's shard before writing any item edge, so the run's pending set can never be + * momentarily empty mid-absorb. Neither TLA+ campaign models this, so the assertion is + * the only protection. A legacy arm ignores it. + */ + batchWaitpointId?: string; +}; export type CompleteParams = { waitpointId: string; @@ -143,6 +177,7 @@ export type CreateWaitpointResult = | { kind: "created"; waitpoint: Waitpoint }; export type CreateDateTimeWaitpointParams = { + mintKind: WaitpointMintKind; /** When set, the waitpoint co-locates with this run's DB and the dedup probe targets it. */ runId?: string; projectId: string; @@ -153,6 +188,7 @@ export type CreateDateTimeWaitpointParams = { }; export type CreateManualWaitpointParams = { + mintKind: WaitpointMintKind; runId?: string; environmentId: string; projectId: string; diff --git a/internal-packages/run-engine/src/engine/waitpointCoordinator/waitpointShape.test.ts b/internal-packages/run-engine/src/engine/waitpointCoordinator/waitpointShape.test.ts new file mode 100644 index 00000000000..225f5694ee7 --- /dev/null +++ b/internal-packages/run-engine/src/engine/waitpointCoordinator/waitpointShape.test.ts @@ -0,0 +1,117 @@ +import { describe, expect, it } from "vitest"; +import type { WaitpointRecordInput } from "./storeCoordinator.js"; +import { toPrismaWaitpoint } from "./waitpointShape.js"; + +const record: WaitpointRecordInput = { + id: "abcdefghijklmnopqrstuvwxmw", + friendlyId: "waitpoint_abcdefghijklmnopqrstuvwxmw", + type: "MANUAL", + environmentId: "env_1", + projectId: "proj_1", + createdAt: "2026-08-26T10:00:00.000Z", + updatedAt: "2026-08-26T10:00:01.000Z", + userProvidedIdempotencyKey: true, + tags: ["alpha", "beta"], + idempotencyKey: "user-key", +}; + +describe("toPrismaWaitpoint", () => { + it("fills every non-null column on a PENDING waitpoint", () => { + const waitpoint = toPrismaWaitpoint(record, "PENDING"); + + expect(waitpoint.id).toBe(record.id); + expect(waitpoint.friendlyId).toBe(record.friendlyId); + expect(waitpoint.type).toBe("MANUAL"); + expect(waitpoint.status).toBe("PENDING"); + expect(waitpoint.idempotencyKey).toBe("user-key"); + expect(waitpoint.userProvidedIdempotencyKey).toBe(true); + expect(waitpoint.projectId).toBe("proj_1"); + expect(waitpoint.environmentId).toBe("env_1"); + expect(waitpoint.tags).toEqual(["alpha", "beta"]); + expect(waitpoint.createdAt).toEqual(new Date("2026-08-26T10:00:00.000Z")); + expect(waitpoint.updatedAt).toEqual(new Date("2026-08-26T10:00:01.000Z")); + + // The columns with database defaults, which a consumer reads unconditionally. + expect(waitpoint.outputType).toBe("application/json"); + expect(waitpoint.outputIsError).toBe(false); + + // Nullable columns that must be null rather than undefined: a consumer distinguishes + // "no value" from "field missing", and `inactiveIdempotencyKey` is not ported at all. + expect(waitpoint.completedAt).toBeNull(); + expect(waitpoint.output).toBeNull(); + expect(waitpoint.inactiveIdempotencyKey).toBeNull(); + expect(waitpoint.idempotencyKeyExpiresAt).toBeNull(); + expect(waitpoint.completedByTaskRunId).toBeNull(); + expect(waitpoint.completedByBatchId).toBeNull(); + expect(waitpoint.completedAfter).toBeNull(); + }); + + it("carries an inline completion onto a COMPLETED waitpoint", () => { + const waitpoint = toPrismaWaitpoint(record, "COMPLETED", { + completedAt: "2026-08-26T11:00:00.000Z", + outputType: "application/json", + outputIsError: true, + output: { inline: '{"boom":true}' }, + }); + + expect(waitpoint.status).toBe("COMPLETED"); + expect(waitpoint.completedAt).toEqual(new Date("2026-08-26T11:00:00.000Z")); + expect(waitpoint.output).toBe('{"boom":true}'); + expect(waitpoint.outputType).toBe("application/json"); + expect(waitpoint.outputIsError).toBe(true); + }); + + it("carries an offloaded reference in the output column, as the legacy row does", () => { + const waitpoint = toPrismaWaitpoint(record, "COMPLETED", { + completedAt: "2026-08-26T11:00:00.000Z", + outputType: "application/store", + outputIsError: false, + output: { ref: "waitpoints/abc/output.json" }, + }); + + expect(waitpoint.output).toBe("waitpoints/abc/output.json"); + expect(waitpoint.outputType).toBe("application/store"); + }); + + it("leaves output null when the completion carries none", () => { + // A BATCH completion, and the deriveFromRun case: the value is re-derived at read + // time and is never copied onto the row. + const waitpoint = toPrismaWaitpoint({ ...record, type: "BATCH" }, "COMPLETED", { + completedAt: "2026-08-26T11:00:00.000Z", + outputType: "application/json", + outputIsError: false, + output: null, + }); + + expect(waitpoint.status).toBe("COMPLETED"); + expect(waitpoint.output).toBeNull(); + }); + + it("maps the optional anchor and timing columns when the record carries them", () => { + const waitpoint = toPrismaWaitpoint( + { + ...record, + type: "RUN", + completedByTaskRunId: "run_1", + completedByBatchId: "batch_1", + completedAfter: "2026-08-27T00:00:00.000Z", + idempotencyKeyExpiresAt: "2026-08-28T00:00:00.000Z", + }, + "PENDING" + ); + + expect(waitpoint.completedByTaskRunId).toBe("run_1"); + expect(waitpoint.completedByBatchId).toBe("batch_1"); + expect(waitpoint.completedAfter).toEqual(new Date("2026-08-27T00:00:00.000Z")); + expect(waitpoint.idempotencyKeyExpiresAt).toEqual(new Date("2026-08-28T00:00:00.000Z")); + }); + + it("throws when the record carries no idempotency key", () => { + // The column is non-null and participates in the (environmentId, idempotencyKey) + // unique index, so inventing a value here could collide. The arm always mints one; + // an absent key means the arm has a defect, and it must surface as one. + const { idempotencyKey, ...withoutKey } = record; + + expect(() => toPrismaWaitpoint(withoutKey, "PENDING")).toThrow(/idempotency key/i); + }); +}); diff --git a/internal-packages/run-engine/src/engine/waitpointCoordinator/waitpointShape.ts b/internal-packages/run-engine/src/engine/waitpointCoordinator/waitpointShape.ts new file mode 100644 index 00000000000..3ab9da2ed38 --- /dev/null +++ b/internal-packages/run-engine/src/engine/waitpointCoordinator/waitpointShape.ts @@ -0,0 +1,59 @@ +import type { Waitpoint } from "@trigger.dev/database"; +import type { + WaitpointCompletion, + WaitpointRecordInput, + WaitpointStatus, +} from "./storeCoordinator.js"; + +/** + * Present a store-resident waitpoint as the Postgres row shape the seam returns. + * + * A store waitpoint has no row, but `WaitpointCoordinator`'s return types are the Prisma + * `Waitpoint`, and callers reach for its columns directly. Every column is listed + * explicitly rather than spread: a missing non-null column surfaces as `undefined` far + * from here, in a consumer that had no reason to guard. + */ +export function toPrismaWaitpoint( + record: WaitpointRecordInput, + status: WaitpointStatus, + completion?: WaitpointCompletion +): Waitpoint { + if (!record.idempotencyKey) { + // Non-null in the schema, and half of the (environmentId, idempotencyKey) unique + // index, so a synthesized value could collide with a real one. Every arm mints one. + throw new Error(`Waitpoint ${record.id} has no idempotency key`); + } + + const output = completion?.output; + + return { + id: record.id, + friendlyId: record.friendlyId, + type: record.type, + status, + completedAt: completion ? new Date(completion.completedAt) : null, + idempotencyKey: record.idempotencyKey, + userProvidedIdempotencyKey: record.userProvidedIdempotencyKey, + idempotencyKeyExpiresAt: optionalDate(record.idempotencyKeyExpiresAt), + // Not ported: clearing an idempotency key is a legacy debounce mechanism the store + // replaces with key expiry. + inactiveIdempotencyKey: null, + completedByTaskRunId: record.completedByTaskRunId ?? null, + completedAfter: optionalDate(record.completedAfter), + completedByBatchId: record.completedByBatchId ?? null, + // An offloaded reference rides the output column exactly as it does on a legacy row, + // with outputType naming it. A null output is re-derived at read time, never copied. + output: output ? ("inline" in output ? output.inline : output.ref) : null, + outputType: completion?.outputType ?? "application/json", + outputIsError: completion?.outputIsError ?? false, + projectId: record.projectId, + environmentId: record.environmentId, + createdAt: new Date(record.createdAt), + updatedAt: new Date(record.updatedAt), + tags: record.tags, + }; +} + +function optionalDate(value: string | undefined): Date | null { + return value ? new Date(value) : null; +} diff --git a/internal-packages/run-store/src/PostgresRunStore.ts b/internal-packages/run-store/src/PostgresRunStore.ts index 7af040eb99c..df718b4a1af 100644 --- a/internal-packages/run-store/src/PostgresRunStore.ts +++ b/internal-packages/run-store/src/PostgresRunStore.ts @@ -30,6 +30,7 @@ import type { TaskRunWithWaitpoint, } from "./types.js"; import type { TaskRunError } from "@trigger.dev/core/v3/schemas"; +import type { ShardKey } from "@trigger.dev/core/v3/isomorphic"; // Loose delegate method shape: each generated client types delegate methods as // `(args: PackageLocalArgs) => PrismaPromise<…>` against its own nominal @@ -2757,7 +2758,7 @@ export class PostgresRunStore implements RunStore { data: { environmentId: string; name: string; projectId: string; id?: string }, tx?: PrismaClientOrTransaction, // `residency` selects the store at the router; a single store has one client and ignores it. - _residency?: "NEW" | "LEGACY" + _residency?: ShardKey ): Promise { const prisma = tx ?? this.prisma; diff --git a/internal-packages/run-store/src/runOpsStore.shardMap.test.ts b/internal-packages/run-store/src/runOpsStore.shardMap.test.ts index b6ea71e9f60..8f2cc8c6485 100644 --- a/internal-packages/run-store/src/runOpsStore.shardMap.test.ts +++ b/internal-packages/run-store/src/runOpsStore.shardMap.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "vitest"; -import { RoutingRunStore } from "./runOpsStore.js"; +import { RoutingRunStore, UnknownShardKey } from "./runOpsStore.js"; import type { ReadClient, RunStore } from "./types.js"; // Pins the routing ALGEBRA: probe order, merge precedence, and the two id-less fallbacks that @@ -278,6 +278,65 @@ describe("RoutingRunStore id-to-shard-key seam", () => { ); expect(trace(log)).toEqual([]); }); + + // The case above injects a resolver. This one does NOT: it uses the real `resolveShard`, which + // the compat constructor defaults to. `resolveShard` is pure id-shape, so a gen-2 shaped id + // names its shard char whatever the topology holds — the two-store compat router therefore + // reaches this throw for any gen-2 id, with no shard configured anywhere. + // + // That matters beyond this class: these ids reach read routes as URL parameters, so whatever + // sits above the router must translate this throw into a 4xx rather than let it surface as a + // 5xx that any caller can induce. + it("reaches the unconfigured-shard throw for a real gen-2 id, even on the compat pair", () => { + const log: Call[] = []; + const router = new RoutingRunStore({ + new: fakeStore("new", log), + legacy: fakeStore("legacy", log), + }); + + const genTwoId = `${"0".repeat(24)}a2`; + + expect(() => router.findRun({ id: genTwoId })).toThrow( + 'no store is configured for shard key "a"' + ); + expect(trace(log)).toEqual([]); + }); + + // Typed, not a bare Error: the API boundary matches on it to answer 404 instead of 500, and + // the operator needs the key and the configured set to tell a forged id from a dropped shard. + it("throws a typed UnknownShardKey carrying the key and the configured set", () => { + const log: Call[] = []; + const router = new RoutingRunStore({ + new: fakeStore("new", log), + legacy: fakeStore("legacy", log), + }); + + let thrown: unknown; + try { + router.findRun({ id: `${"0".repeat(24)}a2` }); + } catch (error) { + thrown = error; + } + + expect(thrown).toBeInstanceOf(UnknownShardKey); + const error = thrown as UnknownShardKey; + expect(error.name).toBe("UnknownShardKey"); + expect(error.shardKey).toBe("a"); + expect([...error.configured].sort()).toEqual(["legacy", "new"]); + }); + + it("still routes gen-1 shapes on the compat pair with the real resolver", () => { + const log: Call[] = []; + const router = new RoutingRunStore({ + new: fakeStore("new", log), + legacy: fakeStore("legacy", log), + }); + + router.findRun({ id: `${"0".repeat(24)}01` }); + router.findRun({ id: "c".repeat(25) }); + + expect(trace(log)).toEqual(["new:findRun", "legacy:findRun"]); + }); }); function buildNShardRouter(shardKeys: string[], opts: { aliasOf?: Record } = {}) { @@ -301,6 +360,14 @@ function buildNShardRouter(shardKeys: string[], opts: { aliasOf?: Record { + // findRunsByIds reaches #fanOutPartitioned, the third unconfigured-shard guard. It must throw + // the typed error too, or this read path answers 500 where the boundary would give a 404. + it("throws a typed UnknownShardKey from the partitioned id fan-out", async () => { + const { router } = buildNShardRouter(["a"]); + + await expect(router.findRunsByIds(["a:r1", "z:r2"])).rejects.toBeInstanceOf(UnknownShardKey); + }); + it("routes an id to its gen-2 shard", async () => { const { router, log } = buildNShardRouter(["a", "b"]); await router.findRun({ id: "a:run_1" }); @@ -649,6 +716,17 @@ describe("RoutingRunStore countPendingWaitpoints — disjoint-sum partition", () ); }); + // The API boundary answers a non-retryable 404 by matching on the TYPE, so every + // unconfigured-shard guard has to throw the typed error and not a bare Error. Two other guards + // besides #shardStore reach an unconfigured key: this partition, and #fanOutPartitioned below. + it("throws a typed UnknownShardKey from the absent-id partition", async () => { + const { router } = partitionRouter({}); + + await expect( + router.countPendingWaitpoints(["c:w1"], undefined, "a:run") + ).rejects.toBeInstanceOf(UnknownShardKey); + }); + it("returns zero for an id absent everywhere", async () => { const { router } = partitionRouter({}); expect(await router.countPendingWaitpoints(["b:w9"], undefined, "a:run")).toBe(0); diff --git a/internal-packages/run-store/src/runOpsStore.ts b/internal-packages/run-store/src/runOpsStore.ts index 53089da21a5..7551e552b34 100644 --- a/internal-packages/run-store/src/runOpsStore.ts +++ b/internal-packages/run-store/src/runOpsStore.ts @@ -62,6 +62,28 @@ const LEGACY_SHARD: ShardKey = "legacy"; * and a merge lets a gen-2 shard win. A probe MUST iterate #probeOrder and a merge MUST iterate * #precedence. */ +/** + * An id resolved to a shard key the topology has no store for. Typed so a caller above the + * router can answer a 4xx instead of letting a routing failure surface as a 5xx: these ids + * arrive as URL parameters, and `resolveShard` is pure id-shape, so any gen-2 shaped id names + * a shard char whether or not one is configured. + */ +export class UnknownShardKey extends Error { + readonly shardKey: string; + readonly configured: string[]; + + constructor(shardKey: string, configured: string[], subject?: string) { + super( + subject === undefined + ? `RoutingRunStore: no store is configured for shard key "${shardKey}"` + : `RoutingRunStore: ${subject} resolves to unconfigured shard key "${shardKey}"` + ); + this.name = "UnknownShardKey"; + this.shardKey = shardKey; + this.configured = configured; + } +} + export class RoutingRunStore implements RunStore { readonly #shards: ReadonlyMap; // Sequential probe for a lookup with no routable id. The first non-null result wins, and the LAST @@ -173,12 +195,14 @@ export class RoutingRunStore implements RunStore { return client != null && !isReadReplicaClient(client) ? store.primaryReadClient : undefined; } - // The store for a shard key. Unreachable with the compat constructor — #shardKeyOfSafe yields only - // the two reserved keys — so this throw fires only if a caller wires a partial map. + // The store for a shard key. REACHABLE with the compat constructor: it defaults to the real + // `resolveShard`, which is pure id-shape, so any gen-2 shaped id names a shard char even when + // no shard is configured. Fails loud rather than reading the wrong database; the API boundary + // turns `UnknownShardKey` into a 404 so a caller-supplied id cannot induce a 5xx. #shardStore(key: ShardKey): RunStore { const store = this.#shards.get(key); if (store === undefined) { - throw new Error(`RoutingRunStore: no store is configured for shard key "${key}"`); + throw new UnknownShardKey(key, [...this.#shards.keys()]); } return store; } @@ -236,9 +260,7 @@ export class RoutingRunStore implements RunStore { // 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}"` - ); + throw new UnknownShardKey(key, [...this.#shards.keys()], `waitpoint "${id}"`); } const bucket = byKey.get(key); if (bucket) bucket.push(id); @@ -393,7 +415,7 @@ export class RoutingRunStore implements RunStore { // 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}"`); + throw new UnknownShardKey(key, [...this.#shards.keys()], `id "${id}"`); } const bucket = byShard.get(key); if (bucket) bucket.push(id);