diff --git a/apps/webapp/app/runEngine/concerns/resolveWaitpointThroughReadThrough.server.ts b/apps/webapp/app/runEngine/concerns/resolveWaitpointThroughReadThrough.server.ts index 284f842f56..fe17a66418 100644 --- a/apps/webapp/app/runEngine/concerns/resolveWaitpointThroughReadThrough.server.ts +++ b/apps/webapp/app/runEngine/concerns/resolveWaitpointThroughReadThrough.server.ts @@ -1,6 +1,8 @@ import type { PrismaReplicaClient } from "~/db.server"; import { $replica as defaultLegacyReplica, + runOpsLegacyPrisma as defaultLegacyPrimary, + runOpsNewPrisma as defaultNewPrimary, runOpsNewReplica as defaultNewClient, runOpsSplitReadEnabled as defaultSplitReadEnabled, } from "~/db.server"; @@ -9,6 +11,8 @@ import { readThroughRun } from "~/v3/runOpsMigration/readThrough.server"; type ResolveWaitpointDeps = { newClient?: PrismaReplicaClient; legacyReplica?: PrismaReplicaClient; + newPrimary?: PrismaReplicaClient; + legacyPrimary?: PrismaReplicaClient; splitEnabled?: boolean; isPastRetention?: (id: string) => boolean; }; @@ -18,12 +22,16 @@ type ResolveWaitpointDeps = { export type ResolveWaitpointReadThroughDefaults = { newClient: PrismaReplicaClient; legacyReplica: PrismaReplicaClient; + newPrimary: PrismaReplicaClient; + legacyPrimary: PrismaReplicaClient; splitEnabled: boolean; }; const productionDefaults: ResolveWaitpointReadThroughDefaults = { newClient: defaultNewClient, legacyReplica: defaultLegacyReplica, + newPrimary: defaultNewPrimary as unknown as PrismaReplicaClient, + legacyPrimary: defaultLegacyPrimary as unknown as PrismaReplicaClient, splitEnabled: defaultSplitReadEnabled, }; @@ -36,18 +44,42 @@ export async function resolveWaitpointThroughReadThrough(opts: { }): Promise { const defaults = opts.defaults ?? productionDefaults; + const splitEnabled = opts.deps?.splitEnabled ?? defaults.splitEnabled; + const result = await readThroughRun({ runId: opts.waitpointId, environmentId: opts.environmentId, readNew: (client) => opts.read(client), readLegacy: (replica) => opts.read(replica), deps: { - splitEnabled: opts.deps?.splitEnabled ?? defaults.splitEnabled, + splitEnabled, newClient: opts.deps?.newClient ?? defaults.newClient, legacyReplica: opts.deps?.legacyReplica ?? defaults.legacyReplica, isPastRetention: opts.deps?.isPastRetention, }, }); - return result.source === "new" || result.source === "legacy-replica" ? result.value : null; + if (result.source === "new" || result.source === "legacy-replica") { + return result.value; + } + // past-retention is an intentional not-found: the token is gone, don't hit the primary. + if (result.source === "past-retention") { + return null; + } + + // Read-your-writes fallback: readThroughRun is replica-only, so a token completed immediately after + // it was minted can miss on the replicas (not yet applied) and 404 a valid token - and the + // authoritative completeWaitpoint never runs. Re-read from the owning-store PRIMARY (new then legacy) + // before giving up. Bounded to this replica-miss; a genuinely-absent token still returns null. + const fromNewPrimary = await opts.read(opts.deps?.newPrimary ?? defaults.newPrimary); + if (fromNewPrimary != null) { + return fromNewPrimary; + } + if (splitEnabled) { + const fromLegacyPrimary = await opts.read(opts.deps?.legacyPrimary ?? defaults.legacyPrimary); + if (fromLegacyPrimary != null) { + return fromLegacyPrimary; + } + } + return null; } diff --git a/apps/webapp/app/v3/services/createCheckpoint.server.ts b/apps/webapp/app/v3/services/createCheckpoint.server.ts index 088001b826..f7359cb1d5 100644 --- a/apps/webapp/app/v3/services/createCheckpoint.server.ts +++ b/apps/webapp/app/v3/services/createCheckpoint.server.ts @@ -147,10 +147,14 @@ export class CreateCheckpointService extends BaseService { } case "WAIT_FOR_BATCH": { // Routed by friendlyId so a run-ops id (NEW-resident) batch is found on the owning DB; - // env-scoped to the dependent attempt's run (a batch shares its dependent's env). + // env-scoped to the dependent attempt's run (a batch shares its dependent's env). Read the + // primary: a batch that just resumed the parent may lag the replica, and a stale resumedAt + // (null) would checkpoint (suspend) an already-resumed run -> it stalls until a sweep. const batchRun = await this.runStore.findBatchTaskRunByFriendlyId( reason.batchFriendlyId, - attempt.taskRun.runtimeEnvironmentId + attempt.taskRun.runtimeEnvironmentId, + undefined, + this._prisma ); if (!batchRun) { @@ -361,11 +365,13 @@ export class CreateCheckpointService extends BaseService { }); await marqs?.cancelHeartbeat(attempt.taskRunId); - // Routed by friendlyId so a run-ops id (NEW-resident) batch is found on the owning DB; - // env-scoped to the dependent attempt's run (a batch shares its dependent's env). + // Routed by friendlyId; read the primary (this._prisma) so a just-resumed batch that still + // lags the replica doesn't leave a stale resumedAt and suspend an already-resumed run. const batchRun = await this.runStore.findBatchTaskRunByFriendlyId( reason.batchFriendlyId, - attempt.taskRun.runtimeEnvironmentId + attempt.taskRun.runtimeEnvironmentId, + undefined, + this._prisma ); if (!batchRun) { diff --git a/apps/webapp/test/createCheckpoint.batchReplicaLag.test.ts b/apps/webapp/test/createCheckpoint.batchReplicaLag.test.ts new file mode 100644 index 0000000000..39dbe7b3a0 --- /dev/null +++ b/apps/webapp/test/createCheckpoint.batchReplicaLag.test.ts @@ -0,0 +1,65 @@ +// Unit red-green for the checkpoint WAIT_FOR_BATCH replica-lag fix (createCheckpoint.server.ts). +// The service decides whether to suspend a run on `batchRun.resumedAt`; reading it from a lagging +// replica makes a just-resumed batch look unresumed -> it suspends an already-resumed run -> stall. +// The fix threads the primary (`this._prisma`) into `runStore.findBatchTaskRunByFriendlyId`. Here a +// spy runStore records which client the service passed and simulates the lag (only the primary read +// sees the fresh resumedAt): RED = no client -> stale null -> no early return; GREEN = primary -> kept alive. + +import { describe, expect, it, vi } from "vitest"; + +vi.mock("~/services/logger.server", () => ({ + logger: { debug: vi.fn(), info: vi.fn(), log: vi.fn(), error: vi.fn(), warn: vi.fn() }, +})); +vi.mock("~/v3/marqs/index.server", () => ({ + marqs: { replaceMessage: vi.fn(), cancelHeartbeat: vi.fn() }, +})); + +import { CreateCheckpointService } from "~/v3/services/createCheckpoint.server"; + +describe("checkpoint WAIT_FOR_BATCH reads the primary, not a lagging replica", () => { + it("threads the primary so an already-resumed batch keeps the run alive", async () => { + // A freezable attempt so control reaches the WAIT_FOR_BATCH arm. This object IS the primary the + // fix must thread into the batch read. + const prisma = { + taskRunAttempt: { + findFirst: async () => ({ + id: "attempt_1", + status: "EXECUTING", + taskRunId: "run_1", + taskRun: { id: "run_1", status: "EXECUTING", runtimeEnvironmentId: "env_1" }, + backgroundWorker: { id: "bw_1", deployment: { imageReference: "img:1" } }, + }), + }, + }; + + let seenClient: unknown = "NOT_CALLED"; + const runStore = { + findBatchTaskRunByFriendlyId: async ( + _friendlyId: string, + _environmentId: string, + _args: unknown, + client?: unknown + ) => { + seenClient = client; + // Lagging replica: only a read handed the primary sees the just-committed resumedAt. + return { resumedAt: client === prisma ? new Date() : null }; + }, + }; + + const service = new CreateCheckpointService(prisma as never, {} as never, runStore as never); + + let result: unknown; + try { + result = await service.call({ + attemptFriendlyId: "attempt_1", + reason: { type: "WAIT_FOR_BATCH", batchFriendlyId: "batch_1" }, + } as never); + } catch { + // Buggy path falls through the pre-check into checkpoint creation (unstubbed) and throws; the + // recorded client below is what distinguishes RED from GREEN. + } + + expect(seenClient).toBe(prisma); // the fix: primary threaded into the batch read + expect(result).toEqual({ success: false, keepRunAlive: true }); // early-return, run kept alive + }); +}); diff --git a/apps/webapp/test/resolveWaitpointThroughReadThrough.readthrough.test.ts b/apps/webapp/test/resolveWaitpointThroughReadThrough.readthrough.test.ts index 4bceb95152..cadb9e3fd7 100644 --- a/apps/webapp/test/resolveWaitpointThroughReadThrough.readthrough.test.ts +++ b/apps/webapp/test/resolveWaitpointThroughReadThrough.readthrough.test.ts @@ -151,6 +151,48 @@ describe("resolveWaitpointThroughReadThrough (hetero PG14 legacy + dedicated run } ); + // Read-your-writes: a token completed immediately after mint may not be on the replicas yet. The + // replica-only read-through misses, and without a primary fallback we 404 a valid token (the + // authoritative completeWaitpoint never runs). The primary fallback must find it. + heteroRunOpsPostgresTest( + "a token missing on both replicas is found on the primary (no spurious 404)", + async ({ prisma17, prisma14 }) => { + const id = generateLegacyCuid(); + const { project, environment } = await seedOrgProjectEnv(prisma14, "primary_fallback"); + const seeded = await seedWaitpoint(prisma14, id, { + id: environment.id, + projectId: project.id, + }); + + // Both replicas miss (not yet replicated): point them at the dedicated DB, which lacks this + // legacy cuid waitpoint. The LEGACY primary (control-plane writer, prisma14) has it. + const newClient = recording(prisma17); + const legacyReplica = recording(prisma17); + const newPrimary = recording(prisma17); + const legacyPrimary = recording(prisma14); + + const result = await resolveWaitpointThroughReadThrough({ + waitpointId: id, + environmentId: environment.id, + read: read(id, environment.id), + deps: { + splitEnabled: true, + newClient: newClient.handle, + legacyReplica: legacyReplica.handle, + newPrimary: newPrimary.handle, + legacyPrimary: legacyPrimary.handle, + }, + }); + + expect(result).not.toBeNull(); + expect(result!.id).toBe(seeded.id); + // Replicas probed and missed; the fallback read the primary. + expect(newClient.calls.length).toBe(1); + expect(legacyReplica.calls.length).toBe(1); + expect(legacyPrimary.calls.length).toBe(1); + } + ); + heteroRunOpsPostgresTest( "bare caller (no deps) resolves a NEW-resident waitpoint via the safe run-ops defaults", async ({ prisma17, prisma14 }) => { @@ -170,6 +212,9 @@ describe("resolveWaitpointThroughReadThrough (hetero PG14 legacy + dedicated run defaults: { newClient: prisma14 as unknown as PrismaReplicaClient, legacyReplica: prisma14 as unknown as PrismaReplicaClient, + // Primaries also miss the NEW-resident waitpoint, so the miss stays a miss. + newPrimary: prisma14 as unknown as PrismaReplicaClient, + legacyPrimary: prisma14 as unknown as PrismaReplicaClient, splitEnabled: true, }, }); @@ -183,6 +228,8 @@ describe("resolveWaitpointThroughReadThrough (hetero PG14 legacy + dedicated run defaults: { newClient: prisma17 as unknown as PrismaReplicaClient, legacyReplica: prisma14 as unknown as PrismaReplicaClient, + newPrimary: prisma17 as unknown as PrismaReplicaClient, + legacyPrimary: prisma14 as unknown as PrismaReplicaClient, splitEnabled: true, }, }); diff --git a/internal-packages/run-engine/src/engine/index.ts b/internal-packages/run-engine/src/engine/index.ts index f3091c93b8..00edf64405 100644 --- a/internal-packages/run-engine/src/engine/index.ts +++ b/internal-packages/run-engine/src/engine/index.ts @@ -1655,7 +1655,6 @@ export class RunEngine { completedAfter, idempotencyKey, idempotencyKeyExpiresAt, - tx, }: { /** The run that will block on this waitpoint. Co-locates the waitpoint with the run's DB. */ runId?: string; @@ -1664,7 +1663,6 @@ export class RunEngine { completedAfter: Date; idempotencyKey?: string; idempotencyKeyExpiresAt?: Date; - tx?: PrismaClientOrTransaction; }) { return this.waitpointSystem.createDateTimeWaitpoint({ runId, @@ -1673,7 +1671,6 @@ export class RunEngine { completedAfter, idempotencyKey, idempotencyKeyExpiresAt, - tx, }); } @@ -2097,13 +2094,24 @@ export class RunEngine { this.readOnlyPrisma !== this.prisma; const prisma = tx ?? (useReplica ? this.readOnlyPrisma : this.prisma); - const query = async (client: PrismaClientOrTransaction) => { - const snapshots = await getExecutionSnapshotsSince(client, runId, snapshotId, this.runStore); + const query = async ( + client: PrismaClientOrTransaction, + repairClient?: PrismaClientOrTransaction + ) => { + const snapshots = await getExecutionSnapshotsSince( + client, + runId, + snapshotId, + this.runStore, + repairClient + ); return snapshots.map(executionDataFromSnapshot); }; try { - return await query(prisma); + // When reading the replica, pass the primary so a snapshot whose completed-waitpoint join rows + // have not replicated yet is repaired from the primary instead of resuming the run waitpoint-less. + return await query(prisma, useReplica ? this.prisma : undefined); } catch (e) { if (useReplica && e instanceof ExecutionSnapshotNotFoundError) { // Replica lag: the runner learned this snapshot id from the writer before the @@ -2115,7 +2123,7 @@ export class RunEngine { if (maxMs > 0) { await setTimeout(minMs + Math.random() * Math.max(0, maxMs - minMs)); try { - const result = await query(this.readOnlyPrisma); + const result = await query(this.readOnlyPrisma, this.prisma); this.snapshotsSinceReplicaMissCounter.add(1, { outcome: "replica_retry" }); return result; } catch (replicaRetryError) { diff --git a/internal-packages/run-engine/src/engine/systems/executionSnapshotSystem.ts b/internal-packages/run-engine/src/engine/systems/executionSnapshotSystem.ts index 52417c1687..a6e6199a49 100644 --- a/internal-packages/run-engine/src/engine/systems/executionSnapshotSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/executionSnapshotSystem.ts @@ -138,6 +138,30 @@ async function getSnapshotWaitpointIds( return result.map((r) => r.B); } +// Reads the snapshot's completed-waitpoint ids AND whether the snapshot is visible on this reader, in +// one query, so a multi-reader replica can't return the ids from a different point-in-time than the +// snapshot. `present=false` -> this reader lacks the snapshot; its empty id list is not authoritative. +async function getSnapshotWaitpointIdsWithPresence( + prisma: PrismaClientOrTransaction, + snapshotId: string, + runStore?: RunStore +): Promise<{ present: boolean; ids: string[] }> { + if (runStore) { + return runStore.findSnapshotCompletedWaitpointIdsWithPresence(snapshotId, prisma); + } + + const rows = await prisma.$queryRaw<{ id: string; B: string | null }[]>` + SELECT s."id", cw."B" + FROM "TaskRunExecutionSnapshot" s + LEFT JOIN "_completedWaitpoints" cw ON cw."A" = s."id" + WHERE s."id" = ${snapshotId} + `; + return { + present: rows.length > 0, + ids: rows.filter((r) => r.B !== null).map((r) => r.B as string), + }; +} + /** * Fetches waitpoints in chunks to avoid NAPI string conversion limits. * This is necessary because waitpoints can have large outputs (100KB+), @@ -260,7 +284,10 @@ export async function getExecutionSnapshotsSince( prisma: PrismaClientOrTransaction, runId: string, sinceSnapshotId: string, - runStore?: RunStore + runStore?: RunStore, + // The primary, for read-repair when `prisma` is a lagging read replica (see Step 3). Omit when + // `prisma` is already the primary. + repairClient?: PrismaClientOrTransaction ): Promise { // Step 1: Find the createdAt of the sinceSnapshotId const sinceSnapshot = runStore @@ -316,10 +343,41 @@ export async function getExecutionSnapshotsSince( // Step 3: Get waitpoint IDs for the LATEST snapshot only (first in desc order) const latestSnapshot = snapshots[0]; - const waitpointIds = await getSnapshotWaitpointIds(prisma, latestSnapshot.id, runStore); + let readClient = prisma; + const { present, ids } = await getSnapshotWaitpointIdsWithPresence( + prisma, + latestSnapshot.id, + runStore + ); + let waitpointIds = ids; + + // Read-repair (multi-reader): Step 2 saw this snapshot, but on a multi-reader replica the join read + // can hit a laggier reader that does not have it yet. present=false means its empty id list is not + // authoritative - re-read from the primary so the runner is not handed a waitpoint-less continue + // (which it silently drops, hanging the run). Single-reader replicas never hit this (present stays true). + if (repairClient && repairClient !== prisma && !present) { + waitpointIds = await getSnapshotWaitpointIds(repairClient, latestSnapshot.id, runStore); + readClient = repairClient; + } + + // Read-repair (batch): completedWaitpointOrder is written on the snapshot row in the same statement as the + // snapshot, so it is authoritative. If it lists more waitpoints than the join read returned, the + // _completedWaitpoints rows are stale on this client (a lagging read replica) - re-read them from the + // primary so the runner is not handed a waitpoint-less continue, which it drops and the run hangs. + // Distinct: completedWaitpointOrder can list the same id twice (a run batched under one idempotency + // key) while the _completedWaitpoints join is deduped, so compare unique ids or the repair fires every + // poll even against a caught-up replica. + const expectedCount = new Set(latestSnapshot.completedWaitpointOrder ?? []).size; + if (repairClient && repairClient !== prisma && waitpointIds.length < expectedCount) { + const repaired = await getSnapshotWaitpointIds(repairClient, latestSnapshot.id, runStore); + if (repaired.length > waitpointIds.length) { + waitpointIds = repaired; + readClient = repairClient; + } + } // Step 4: Fetch waitpoints in chunks to avoid NAPI string conversion limits - const waitpoints = await fetchWaitpointsInChunks(prisma, waitpointIds, runStore); + const waitpoints = await fetchWaitpointsInChunks(readClient, waitpointIds, runStore); // Step 5: Build enhanced snapshots - only latest gets waitpoints, others get empty arrays // The runner only uses completedWaitpoints from the latest snapshot anyway diff --git a/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts b/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts index 948346fca4..d3523eaaf7 100644 --- a/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts @@ -211,7 +211,6 @@ export class WaitpointSystem { completedAfter, idempotencyKey, idempotencyKeyExpiresAt, - tx, }: { runId?: string; projectId: string; @@ -219,7 +218,6 @@ export class WaitpointSystem { completedAfter: Date; idempotencyKey?: string; idempotencyKeyExpiresAt?: Date; - tx?: PrismaClientOrTransaction; }) { // Co-location invariant: a DATETIME wait waitpoint lives on the same run-ops DB as the run that // blocks on it (so the block edge's local `Waitpoint` join resolves and completion/resume stay @@ -227,27 +225,20 @@ export class WaitpointSystem { // would always route to LEGACY and a run-ops run on NEW would hang. The (env,idempotencyKey) dedup // is within the owning run/tree (co-resident on one DB), so the dedup probe + rotation target the // SAME store. With no run id (a standalone token has no owning run yet) the lookup falls back to - // a cross-DB NEW-then-LEGACY scan and the upsert routes by id-shape. A caller-supplied tx pins a - // client (same physical DB as the control-plane tx → LEGACY), so it stays on direct prisma. + // a cross-DB NEW-then-LEGACY scan and the upsert routes by id-shape. Always routed through the + // run store (never a caller tx) so it can never bypass residency onto the wrong DB. const colocate = runId ? { coLocateWithRunId: runId } : undefined; const existingWaitpoint = idempotencyKey - ? tx - ? await tx.waitpoint.findFirst({ + ? await this.$.runStore.findWaitpoint( + { where: { environmentId, idempotencyKey, }, - }) - : await this.$.runStore.findWaitpoint( - { - where: { - environmentId, - idempotencyKey, - }, - }, - undefined, - colocate - ) + }, + undefined, + colocate + ) : undefined; if (existingWaitpoint) { @@ -266,11 +257,7 @@ export class WaitpointSystem { inactiveIdempotencyKey: existingWaitpoint.idempotencyKey, }, }; - if (tx) { - await tx.waitpoint.update(rotateArgs); - } else { - await this.$.runStore.updateWaitpoint(rotateArgs, undefined, colocate); - } + await this.$.runStore.updateWaitpoint(rotateArgs, undefined, colocate); //let it fall through to create a new waitpoint } else { @@ -297,9 +284,7 @@ export class WaitpointSystem { }, update: {}, }; - const waitpoint = tx - ? await tx.waitpoint.upsert(upsertArgs) - : await this.$.runStore.upsertWaitpoint(upsertArgs, undefined, colocate); + const waitpoint = await this.$.runStore.upsertWaitpoint(upsertArgs, undefined, colocate); await this.$.worker.enqueue({ id: `finishWaitpoint.${waitpoint.id}`, diff --git a/internal-packages/run-engine/src/engine/tests/getSnapshotsSince.test.ts b/internal-packages/run-engine/src/engine/tests/getSnapshotsSince.test.ts index 5a21bd784a..61ff3bff75 100644 --- a/internal-packages/run-engine/src/engine/tests/getSnapshotsSince.test.ts +++ b/internal-packages/run-engine/src/engine/tests/getSnapshotsSince.test.ts @@ -3,7 +3,9 @@ import { trace } from "@internal/tracing"; import { generateFriendlyId } from "@trigger.dev/core/v3/isomorphic"; import { setTimeout } from "node:timers/promises"; import { describe, expect } from "vitest"; +import type { PrismaClient } from "@trigger.dev/database"; import { RunEngine } from "../index.js"; +import { getExecutionSnapshotsSince } from "../systems/executionSnapshotSystem.js"; import { copySnapshotsToReplica, createTestMetricsMeter } from "./helpers/replicaTestHelpers.js"; import { setupTestScenario } from "./helpers/snapshotTestHelpers.js"; import { setupAuthenticatedEnvironment, setupBackgroundWorker } from "./setup.js"; @@ -1398,4 +1400,183 @@ describe("RunEngine getSnapshotsSince", () => { } } ); + + // This models a BATCH resume: production only populates completedWaitpointOrder for waitpoints that + // carry a batch index, so the read-repair only ever engages for batch resumes. Single-waitpoint + // continues have an empty order (repair is a no-op) and are covered instead by the atomic write + // (snapshot + join commit in one tx, which a replica applies atomically). + containerTest( + "repairs a batch resume's completed-waitpoints from the primary when the replica lags the join rows", + async ({ prisma, schemaOnlyPrisma, redisOptions }) => { + const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + + const engine = new RunEngine({ + prisma, + // Replica has the snapshot rows but lags on the _completedWaitpoints join (see below). + readOnlyPrisma: schemaOnlyPrisma, + readReplicaSnapshotsSinceEnabled: true, + readReplicaSnapshotsSinceRetryDelay: { minMs: 1, maxMs: 2 }, + worker: { redis: redisOptions, workers: 1, tasksPerWorker: 10, pollIntervalMs: 100 }, + queue: { redis: redisOptions }, + runLock: { redis: redisOptions }, + 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"), + }); + + try { + // Primary: a run whose latest snapshot is a warm-continue carrying one completed waitpoint - + // completedWaitpointOrder is set on the row AND the _completedWaitpoints join + waitpoint exist. + const scenario = await setupTestScenario(prisma, authenticatedEnvironment, { + totalWaitpoints: 1, + outputSizeKB: 1, + snapshotConfigs: [ + { status: "RUN_CREATED", completedWaitpointCount: 0 }, + { status: "EXECUTING", completedWaitpointCount: 0 }, + { status: "EXECUTING_WITH_WAITPOINTS", completedWaitpointCount: 0 }, + { status: "EXECUTING", completedWaitpointCount: 1 }, + ], + }); + const waitpointId = scenario.waitpoints[0].id; + + // Replica lag: it receives the snapshot ROWS (incl. completedWaitpointOrder) but NOT the + // _completedWaitpoints join rows or the waitpoint - the exact state that drops the resume. + await copySnapshotsToReplica(prisma, schemaOnlyPrisma, scenario.run.id); + + const result = await engine.getSnapshotsSince({ + runId: scenario.run.id, + snapshotId: scenario.snapshots[0].id, + }); + + expect(result).not.toBeNull(); + const latest = result![result!.length - 1]; + // The runner must receive the completed waitpoint (repaired from the primary), not an empty set. + expect(latest.completedWaitpoints.map((w) => w.id)).toEqual([waitpointId]); + } finally { + await engine.quit(); + } + } + ); + + // completedWaitpointOrder can contain the same waitpoint id twice (a run batched under one idempotency + // key), while the _completedWaitpoints join is deduped. The read-repair must compare DISTINCT ids, or it + // fires on every poll for such a snapshot even against a caught-up replica - silently defeating offload. + containerTest( + "does not repair from the primary when completedWaitpointOrder has duplicates but the join is complete", + async ({ prisma }) => { + const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + // Drives getExecutionSnapshotsSince directly, so it needs no RunEngine (Redis/worker) at all. + { + const scenario = await setupTestScenario(prisma, authenticatedEnvironment, { + totalWaitpoints: 1, + outputSizeKB: 1, + snapshotConfigs: [ + { status: "RUN_CREATED", completedWaitpointCount: 0 }, + { status: "EXECUTING", completedWaitpointCount: 1 }, + ], + }); + const waitpointId = scenario.waitpoints[0].id; + const latestSnapshot = scenario.snapshots[scenario.snapshots.length - 1]; + // The join stays a single row; the order column lists the same id twice (the batch-dup case). + await prisma.taskRunExecutionSnapshot.update({ + where: { id: latestSnapshot.id }, + data: { completedWaitpointOrder: [waitpointId, waitpointId] }, + }); + + // Count join reads on the repair client; the replica read (`prisma`) is already complete, so a + // correct repair must NEVER touch the repair client for this snapshot. + const repairCalls = { queryRaw: 0 }; + const repairClient = new Proxy(prisma, { + get(target, prop) { + if (prop === "$queryRaw") { + return (...args: unknown[]) => { + repairCalls.queryRaw++; + return (target as unknown as { $queryRaw: (...a: unknown[]) => unknown }).$queryRaw( + ...args + ); + }; + } + const value = (target as Record)[prop]; + return typeof value === "function" ? value.bind(target) : value; + }, + }) as unknown as PrismaClient; + + const result = await getExecutionSnapshotsSince( + prisma, + scenario.run.id, + scenario.snapshots[0].id, + undefined, + repairClient + ); + + const latest = result[result.length - 1]; + // The duplicate is intentional (batch dup) - enhance expands by completedWaitpointOrder. + expect(latest.completedWaitpoints.map((w) => w.id)).toEqual([waitpointId, waitpointId]); + // No spurious primary read: distinct(order) === join count, so the repair must not fire. + expect(repairCalls.queryRaw).toBe(0); + } + } + ); + + // A single triggerAndWait resume has an EMPTY completedWaitpointOrder (only batch-indexed waitpoints + // populate it) but a real join row. On a multi-reader replica the join read can hit a laggier reader + // that does not yet have the snapshot, returning 0 - and the order-based repair can't see it (order is + // empty). The presence check must detect the reader lacks the snapshot and repair from the primary, + // or the runner is handed a waitpoint-less continue and the run hangs. + containerTest( + "repairs a single-wait resume from the primary when the replica reader lacks the snapshot", + async ({ prisma }) => { + const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + const scenario = await setupTestScenario(prisma, authenticatedEnvironment, { + totalWaitpoints: 1, + outputSizeKB: 1, + snapshotConfigs: [ + { status: "RUN_CREATED", completedWaitpointCount: 0 }, + { status: "EXECUTING_WITH_WAITPOINTS", completedWaitpointCount: 0 }, + { status: "EXECUTING", completedWaitpointCount: 1 }, + ], + }); + const waitpointId = scenario.waitpoints[0].id; + const latestSnapshot = scenario.snapshots[scenario.snapshots.length - 1]; + // Single-wait: keep the join row, clear the order so it's the non-batch case. + await prisma.taskRunExecutionSnapshot.update({ + where: { id: latestSnapshot.id }, + data: { completedWaitpointOrder: [] }, + }); + + // A multi-reader replica reader that has NOT yet applied the snapshot: any raw read (the + // presence+join query) returns empty, while Step 2's findMany still returns the snapshot. + const laggyReader = new Proxy(prisma, { + get(target, prop) { + if (prop === "$queryRaw") { + return async () => []; + } + const value = (target as Record)[prop]; + return typeof value === "function" ? value.bind(target) : value; + }, + }) as unknown as PrismaClient; + + const result = await getExecutionSnapshotsSince( + laggyReader, + scenario.run.id, + scenario.snapshots[0].id, + undefined, + prisma + ); + + const latest = result[result.length - 1]; + // The runner must receive the completed waitpoint (repaired from the primary), not an empty set. + expect(latest.completedWaitpoints.map((w) => w.id)).toEqual([waitpointId]); + } + ); }); diff --git a/internal-packages/run-store/src/PostgresRunStore.dedicatedSelect.test.ts b/internal-packages/run-store/src/PostgresRunStore.dedicatedSelect.test.ts index 0c5fcbf21d..73a9bdd205 100644 --- a/internal-packages/run-store/src/PostgresRunStore.dedicatedSelect.test.ts +++ b/internal-packages/run-store/src/PostgresRunStore.dedicatedSelect.test.ts @@ -405,4 +405,36 @@ describe("PostgresRunStore dedicated caller-select adapter (P2-store-bodies-2)", expect(blocking[0].taskRun.friendlyId).toBe(r.friendlyId); } ); + + // The dedicated finders forward include/select to a subset client typed as the full schema, so a + // control-plane-only relation would throw an opaque Prisma 500 for NEW data (invisible to tsc). The + // boundary guard rejects the known keys with a clear message; the legacy (full-schema) store is a no-op. + heteroRunOpsPostgresTest( + "dedicated batch/attempt finders reject control-plane-only includes with a clear error", + async ({ prisma14, prisma17 }) => { + const dedicated = makeStore(prisma17, "dedicated"); + const legacy = makeStore(prisma14, "legacy"); + + await expect( + dedicated.findBatchTaskRunById("batch_x", { include: { runsBlocked: true } as never }) + ).rejects.toThrow(/not available on the dedicated run-ops subset/); + await expect( + dedicated.findTaskRunAttempt({ + where: { id: "att_x" }, + include: { backgroundWorker: true }, + } as never) + ).rejects.toThrow(/not available on the dedicated run-ops subset/); + + // A subset-present include passes the guard and resolves normally (null for a missing batch). + expect( + await dedicated.findBatchTaskRunById("batch_missing", { include: { items: true } as never }) + ).toBeNull(); + // Legacy store: guard is a no-op; the full schema has the relation, so no guard throw. + expect( + await legacy.findBatchTaskRunById("batch_missing", { + include: { runsBlocked: true } as never, + }) + ).toBeNull(); + } + ); }); diff --git a/internal-packages/run-store/src/PostgresRunStore.test.ts b/internal-packages/run-store/src/PostgresRunStore.test.ts index c0712bc3e3..b0423f24bd 100644 --- a/internal-packages/run-store/src/PostgresRunStore.test.ts +++ b/internal-packages/run-store/src/PostgresRunStore.test.ts @@ -888,6 +888,124 @@ describe("PostgresRunStore", () => { } ); + // lockRunToWorker creates a PENDING_EXECUTING snapshot (which can inherit completed waitpoints) and + // then connects the _completedWaitpoints join as a separate statement. These must commit together, or + // a replica-served /snapshots/since read can return the snapshot waitpoint-less and drop the resume. + postgresTest( + "lockRunToWorker writes the snapshot and its completed-waitpoint links atomically", + async ({ prisma }) => { + const { organization, project, environment } = await seedEnvironment(prisma); + const store = new PostgresRunStore({ prisma, readOnlyPrisma: prisma }); + const runId = "run_lock_atomic"; + + await store.createRun( + buildCreateRunInput({ + runId, + organizationId: organization.id, + projectId: project.id, + runtimeEnvironmentId: environment.id, + }) + ); + + const backgroundWorker = await prisma.backgroundWorker.create({ + data: { + friendlyId: "worker_atomic", + version: "20260601.1", + runtimeEnvironmentId: environment.id, + projectId: project.id, + contentHash: "abc", + sdkVersion: "3.0.0", + cliVersion: "3.0.0", + metadata: {}, + }, + }); + const workerTask = await prisma.backgroundWorkerTask.create({ + data: { + friendlyId: "task_atomic", + slug: "my-task", + filePath: "src/my-task.ts", + exportName: "myTask", + workerId: backgroundWorker.id, + runtimeEnvironmentId: environment.id, + projectId: project.id, + }, + }); + const queue = await prisma.taskQueue.create({ + data: { + friendlyId: "queue_atomic", + name: "task/my-task", + runtimeEnvironmentId: environment.id, + projectId: project.id, + }, + }); + const priorSnapshot = await prisma.taskRunExecutionSnapshot.create({ + data: { + engine: "V2", + executionStatus: "RUN_CREATED", + description: "prior", + runStatus: "PENDING", + environmentId: environment.id, + environmentType: "DEVELOPMENT", + projectId: project.id, + organizationId: organization.id, + runId, + }, + }); + const waitpoint = await prisma.waitpoint.create({ + data: { + friendlyId: "wp_lock_atomic", + type: "MANUAL", + status: "COMPLETED", + idempotencyKey: "idem-lock-atomic", + userProvidedIdempotencyKey: false, + projectId: project.id, + environmentId: environment.id, + }, + }); + + // Force the completed-waitpoint join insert to fail mid-write. + await prisma.$executeRawUnsafe('DROP TABLE "_completedWaitpoints"'); + + const snapshotId = "snap_lock_atomic"; + await expect( + // Base client as `tx` = how the dequeue path calls it; the store must still open its own tx. + store.lockRunToWorker( + runId, + { + lockedAt: new Date(), + lockedById: workerTask.id, + lockedToVersionId: backgroundWorker.id, + lockedQueueId: queue.id, + startedAt: new Date(), + baseCostInCents: 5, + machinePreset: "small-1x", + taskVersion: "20260601.1", + sdkVersion: "3.0.0", + cliVersion: "3.0.0", + maxDurationInSeconds: null, + snapshot: { + id: snapshotId, + previousSnapshotId: priorSnapshot.id, + environmentId: environment.id, + environmentType: "DEVELOPMENT", + projectId: project.id, + organizationId: organization.id, + completedWaitpointIds: [waitpoint.id], + completedWaitpointOrder: [waitpoint.id], + }, + }, + prisma + ) + ).rejects.toThrow(); + + // Atomic: neither the snapshot nor the run lock persists without the links. + const snap = await prisma.taskRunExecutionSnapshot.findUnique({ where: { id: snapshotId } }); + expect(snap).toBeNull(); + const run = await prisma.taskRun.findUniqueOrThrow({ where: { id: runId } }); + expect(run.status).not.toBe("DEQUEUED"); + } + ); + postgresTest( "parkPendingVersion sets status to PENDING_VERSION and stores statusReason", async ({ prisma }) => { diff --git a/internal-packages/run-store/src/PostgresRunStore.ts b/internal-packages/run-store/src/PostgresRunStore.ts index 94630cdb2c..16cb6be3e2 100644 --- a/internal-packages/run-store/src/PostgresRunStore.ts +++ b/internal-packages/run-store/src/PostgresRunStore.ts @@ -527,6 +527,25 @@ export class PostgresRunStore implements RunStore { ); } + // Run `fn` atomically: reuse the caller's interactive transaction if it gave us a real one, else open + // our own on this store's writer. A real interactive tx has no `$transaction` method; a base client + // (which callers, e.g. the engine's dequeue/resume paths, thread through for routing) DOES - so a base + // client still gets a fresh transaction. Used by write methods that create a row plus dependent rows + // (snapshot + completed-waitpoints, run + associated-waitpoint) which must commit together. + #withOptionalTransaction( + tx: PrismaClientOrTransaction | undefined, + fn: (client: PrismaClientOrTransaction) => Promise + ): Promise { + const alreadyInTransaction = + tx !== undefined && typeof (tx as { $transaction?: unknown }).$transaction !== "function"; + if (alreadyInTransaction) { + return fn(tx); + } + return (this.prisma as RunOpsTransactionalClient).$transaction((t) => + fn(t as unknown as PrismaClientOrTransaction) + ); + } + async createRun( params: CreateRunInput, tx?: PrismaClientOrTransaction @@ -547,15 +566,19 @@ export class PostgresRunStore implements RunStore { }; if (this.schemaVariant === "dedicated") { - const run = (await client.taskRun.create({ - data: { - ...params.data, - executionSnapshots: { create: snapshotCreate }, - }, - })) as TaskRun; + // The run + its associated RUN-type waitpoint are two writes here (the legacy branch below nests + // them). Commit them together so a crash / lagging read never leaves a run without its waitpoint. + return this.#withOptionalTransaction(tx, async (c) => { + const run = (await c.taskRun.create({ + data: { + ...params.data, + executionSnapshots: { create: snapshotCreate }, + }, + })) as TaskRun; - const associatedWaitpoint = await this.#createAssociatedWaitpoint(client, run.id, params); - return { ...run, associatedWaitpoint }; + const associatedWaitpoint = await this.#createAssociatedWaitpoint(c, run.id, params); + return { ...run, associatedWaitpoint }; + }); } return client.taskRun.create({ @@ -633,12 +656,15 @@ export class PostgresRunStore implements RunStore { const client = tx ?? this.prisma; if (this.schemaVariant === "dedicated") { - const run = (await client.taskRun.create({ - data: { ...params.data }, - })) as TaskRun; - - const associatedWaitpoint = await this.#createAssociatedWaitpoint(client, run.id, params); - return { ...run, associatedWaitpoint }; + // Run + associated RUN-type waitpoint are two writes here; commit them together (see createRun). + return this.#withOptionalTransaction(tx, async (c) => { + const run = (await c.taskRun.create({ + data: { ...params.data }, + })) as TaskRun; + + const associatedWaitpoint = await this.#createAssociatedWaitpoint(c, run.id, params); + return { ...run, associatedWaitpoint }; + }); } return client.taskRun.create({ @@ -954,8 +980,17 @@ export class PostgresRunStore implements RunStore { data: LockRunData, tx?: PrismaClientOrTransaction ): Promise> { - const prisma = tx ?? this.prisma; + // The run-lock update (with its nested PENDING_EXECUTING snapshot) and the completed-waitpoint + // connect must commit together, or a replica-served resume read can see the snapshot without its + // links and drop the resume. + return this.#withOptionalTransaction(tx, (c) => this.#lockRunToWorker(runId, data, c)); + } + async #lockRunToWorker( + runId: string, + data: LockRunData, + prisma: PrismaClientOrTransaction + ): Promise> { const dedicated = this.schemaVariant === "dedicated"; const result = await prisma.taskRun.update({ @@ -1533,8 +1568,17 @@ export class PostgresRunStore implements RunStore { input: CreateExecutionSnapshotInput, tx?: PrismaClientOrTransaction ): Promise> { - const prisma = tx ?? this.prisma; + // The snapshot row and its completed-waitpoint join rows MUST commit together. `/snapshots/since` + // can be served from a lagging read replica, so a snapshot that commits before its links can be + // read back waitpoint-less and the runner's resume is lost (the run hangs). This is the warm-continue + // path: the engine threads its base prisma through as `tx`, which is not a real transaction. + return this.#withOptionalTransaction(tx, (c) => this.#createExecutionSnapshot(input, c)); + } + async #createExecutionSnapshot( + input: CreateExecutionSnapshotInput, + prisma: PrismaClientOrTransaction + ): Promise> { const { run, snapshot, @@ -1618,6 +1662,41 @@ export class PostgresRunStore implements RunStore { return result.map((r) => r.B); } + // One query: LEFT JOIN the snapshot to its completed-waitpoint links so `present` (snapshot visible + // on this reader) and `ids` come from the SAME point-in-time. A multi-reader replica can otherwise + // return the snapshot (via a separate read) while a different, laggier reader returns 0 links. + async findSnapshotCompletedWaitpointIdsWithPresence( + snapshotId: string, + client?: ReadClient + ): Promise<{ present: boolean; ids: string[] }> { + const prisma = client ?? this.readOnlyPrisma; + + const joinDelegate = (prisma as RunOpsCapableClient).completedWaitpoint; + if (this.schemaVariant === "dedicated" && joinDelegate) { + const rows = await prisma.$queryRaw<{ id: string; waitpointId: string | null }[]>` + SELECT s."id", cw."waitpointId" + FROM "TaskRunExecutionSnapshot" s + LEFT JOIN "CompletedWaitpoint" cw ON cw."snapshotId" = s."id" + WHERE s."id" = ${snapshotId} + `; + return { + present: rows.length > 0, + ids: rows.filter((r) => r.waitpointId !== null).map((r) => r.waitpointId as string), + }; + } + + const rows = await prisma.$queryRaw<{ id: string; B: string | null }[]>` + SELECT s."id", cw."B" + FROM "TaskRunExecutionSnapshot" s + LEFT JOIN "_completedWaitpoints" cw ON cw."A" = s."id" + WHERE s."id" = ${snapshotId} + `; + return { + present: rows.length > 0, + ids: rows.filter((r) => r.B !== null).map((r) => r.B as string), + }; + } + // Reverse of `connectedRuns`: the run ids linked to a waitpoint. Co-resident with the RUN (the join // is written on the run's DB in blockRunWithWaitpointEdges), so the waitpoint's own store can MISS a // cross-DB run — the router fans this across BOTH DBs. @@ -1950,12 +2029,44 @@ export class PostgresRunStore implements RunStore { return prisma.taskRunWaitpoint.deleteMany(args); } + // The dedicated subset schema lacks control-plane relations; a pass-through include/select of one + // throws an opaque Prisma "Unknown field" 500 for NEW-resident data - invisible to tsc, since the + // run-ops client is typed as the full schema. Reject the known control-plane-only keys with a clear + // message at the boundary. No-op on the legacy (full-schema) store. + #assertSubsetSelectable( + fields: Record | null | undefined, + forbidden: readonly string[], + method: string + ): void { + if (this.schemaVariant !== "dedicated" || !fields) return; + for (const key of forbidden) { + if (fields[key]) { + throw new Error( + `${method}: "${key}" is not available on the dedicated run-ops subset schema ` + + `(control-plane-only); resolve it via the control-plane instead of selecting it here.` + ); + } + } + } + async findTaskRunAttempt( args: Prisma.SelectSubset, client?: ReadClient ): Promise | null> { const prisma = client ?? this.readOnlyPrisma; + const forbidden = ["backgroundWorker", "backgroundWorkerTask", "runtimeEnvironment", "queue"]; + this.#assertSubsetSelectable( + (args as { include?: Record }).include, + forbidden, + "findTaskRunAttempt" + ); + this.#assertSubsetSelectable( + (args as { select?: Record }).select, + forbidden, + "findTaskRunAttempt" + ); + return prisma.taskRunAttempt.findFirst( args ) as Promise | null>; @@ -2007,6 +2118,12 @@ export class PostgresRunStore implements RunStore { ): Promise | null> { const prisma = client ?? this.prisma; + this.#assertSubsetSelectable( + args?.include as Record | undefined, + ["runsBlocked", "waitpoints", "runtimeEnvironment"], + "findBatchTaskRunById" + ); + return prisma.batchTaskRun.findFirst({ where: { id }, ...(args?.include ? { include: args.include } : {}), @@ -2021,6 +2138,12 @@ export class PostgresRunStore implements RunStore { ): Promise | null> { const prisma = client ?? this.readOnlyPrisma; + this.#assertSubsetSelectable( + args?.include as Record | undefined, + ["runsBlocked", "waitpoints", "runtimeEnvironment"], + "findBatchTaskRunByFriendlyId" + ); + return prisma.batchTaskRun.findFirst({ where: { friendlyId, runtimeEnvironmentId: environmentId }, ...(args?.include ? { include: args.include } : {}), @@ -2037,6 +2160,12 @@ export class PostgresRunStore implements RunStore { ): Promise | null> { const prisma = client ?? this.prisma; + this.#assertSubsetSelectable( + args?.include as Record | undefined, + ["runsBlocked", "waitpoints", "runtimeEnvironment"], + "findBatchTaskRunByIdempotencyKey" + ); + return prisma.batchTaskRun.findFirst({ where: { runtimeEnvironmentId: environmentId, idempotencyKey }, ...(args?.include ? { include: args.include } : {}), diff --git a/internal-packages/run-store/src/PostgresRunStore.writeAtomicity.test.ts b/internal-packages/run-store/src/PostgresRunStore.writeAtomicity.test.ts index fb81eb90c1..37eca9e9df 100644 --- a/internal-packages/run-store/src/PostgresRunStore.writeAtomicity.test.ts +++ b/internal-packages/run-store/src/PostgresRunStore.writeAtomicity.test.ts @@ -396,6 +396,71 @@ describe("fan-out deleteManyTaskRunWaitpoints honors the caller's tx on the #leg ); }); +// createExecutionSnapshot writes the snapshot row and its completed-waitpoint join rows. These MUST +// commit together: with the flag off, `/snapshots/since` is served from a lagging read replica, so a +// snapshot that commits before its `_completedWaitpoints` rows can be read waitpoint-less, and the +// runner's EXECUTING branch no-ops on an empty completedWaitpoints -> the resume is lost -> hang. +describe("createExecutionSnapshot writes the snapshot and its completed-waitpoint links atomically", () => { + heteroRunOpsPostgresTest( + "rolls the snapshot back if the completed-waitpoint insert fails (no waitpoint-less snapshot persists)", + async ({ prisma14 }) => { + const legacy = makeLegacyStore(prisma14); + const env = await seedEnvironment(prisma14, "legacy", "ces_atomic"); + const runId = `run_${CUID_25}`; + await legacy.createRun( + buildCreateRunInput({ + runId, + friendlyId: "run_ces_atomic", + organizationId: env.organization.id, + projectId: env.project.id, + runtimeEnvironmentId: env.environment.id, + }) + ); + const waitpoint = await prisma14.waitpoint.create({ + data: { + friendlyId: "wp_ces_atomic", + type: "MANUAL", + status: "COMPLETED", + idempotencyKey: "idem-ces_atomic", + userProvidedIdempotencyKey: false, + projectId: env.project.id, + environmentId: env.environment.id, + }, + }); + + // Force the completed-waitpoint join insert to fail mid-write. + await prisma14.$executeRawUnsafe('DROP TABLE "_completedWaitpoints"'); + + await expect( + // Pass the base client as `tx` - exactly how the engine threads its base prisma through + // (continueRunIfUnblocked -> executionSnapshotSystem.createExecutionSnapshot(prisma, ...)). + // It is NOT an interactive transaction, so the store must still open its own to stay atomic. + legacy.createExecutionSnapshot( + { + run: { id: runId, status: "EXECUTING", attemptNumber: 1 }, + snapshot: { + executionStatus: "EXECUTING_WITH_WAITPOINTS", + description: "Run was blocked by a waitpoint.", + }, + environmentId: env.environment.id, + environmentType: "DEVELOPMENT", + projectId: env.project.id, + organizationId: env.project.id, + completedWaitpoints: [{ id: waitpoint.id, index: 0 }], + }, + prisma14 + ) + ).rejects.toThrow(); + + // The snapshot must NOT persist without its links, or a replica can serve it waitpoint-less. + const snap = await prisma14.taskRunExecutionSnapshot.findFirst({ + where: { runId, executionStatus: "EXECUTING_WITH_WAITPOINTS" }, + }); + expect(snap).toBeNull(); + } + ); +}); + // RoutingRunStore.createExecutionSnapshot accepts a caller tx but must forward it to the OWNING store // only when that store is #legacy: a control-plane tx can't wrap a #new (cross-DB) write, but it can // (and should) wrap a legacy-resident snapshot so it stays atomic with the caller's operation. @@ -457,3 +522,194 @@ describe("createExecutionSnapshot honors the caller's tx on the #legacy owning s } ); }); + +// On the dedicated subset schema the associated (RUN-type) waitpoint is created as a SEPARATE +// waitpoint.create after taskRun.create (the legacy schema nests it atomically). The pair must commit +// together, or a crash / lagging read leaves a run with no completion waitpoint and its parent never resumes. +function assocWaitpoint( + env: { project: { id: string }; environment: { id: string } }, + suffix: string +) { + return { + id: `wp_${suffix}`, + friendlyId: `waitpoint_${suffix}`, + type: "RUN" as const, + status: "PENDING" as const, + idempotencyKey: `idem_${suffix}`, + userProvidedIdempotencyKey: false, + projectId: env.project.id, + environmentId: env.environment.id, + }; +} + +describe("createRun / createFailedRun write the run and its associated waitpoint atomically (dedicated)", () => { + heteroRunOpsPostgresTest( + "createRun rolls the run back if the associated-waitpoint create fails", + async ({ prisma17 }) => { + const newStore = makeDedicatedStore(prisma17); + const env = await seedEnvironment(prisma17, "dedicated", "cr_atomic"); + const runId = `run_${NEW_ID_26}`; + const input = { + ...buildCreateRunInput({ + runId, + friendlyId: "run_cr_atomic", + organizationId: env.organization.id, + projectId: env.project.id, + runtimeEnvironmentId: env.environment.id, + }), + associatedWaitpoint: assocWaitpoint(env, "cr_atomic"), + }; + + // Force #createAssociatedWaitpoint (waitpoint.create) to fail after taskRun.create. + await prisma17.$executeRawUnsafe('DROP TABLE "Waitpoint"'); + + await expect(newStore.createRun(input)).rejects.toThrow(); + + const run = await prisma17.taskRun.findFirst({ where: { id: runId } }); + expect(run).toBeNull(); + } + ); + + heteroRunOpsPostgresTest( + "createFailedRun rolls the run back if the associated-waitpoint create fails", + async ({ prisma17 }) => { + const newStore = makeDedicatedStore(prisma17); + const env = await seedEnvironment(prisma17, "dedicated", "cf_atomic"); + const runId = `run_${NEW_ID_26}`; + const base = buildCreateRunInput({ + runId, + friendlyId: "run_cf_atomic", + organizationId: env.organization.id, + projectId: env.project.id, + runtimeEnvironmentId: env.environment.id, + }); + const input = { data: base.data, associatedWaitpoint: assocWaitpoint(env, "cf_atomic") }; + + await prisma17.$executeRawUnsafe('DROP TABLE "Waitpoint"'); + + await expect(newStore.createFailedRun(input)).rejects.toThrow(); + + const run = await prisma17.taskRun.findFirst({ where: { id: runId } }); + expect(run).toBeNull(); + } + ); +}); + +// The dedicated (#new) leg connects completed waitpoints through the `CompletedWaitpoint` join table +// (createMany), where the legacy leg uses the implicit `_completedWaitpoints` M2M. Both must commit the +// snapshot and its links together: a snapshot that commits before its links can be read waitpoint-less +// from a lagging replica, and the runner's EXECUTING branch no-ops on an empty set -> the resume hangs. +describe("createExecutionSnapshot / lockRunToWorker write the snapshot and its links atomically (dedicated)", () => { + heteroRunOpsPostgresTest( + "createExecutionSnapshot rolls the snapshot back if the CompletedWaitpoint insert fails", + async ({ prisma17 }) => { + const newStore = makeDedicatedStore(prisma17); + const env = await seedEnvironment(prisma17, "dedicated", "ces_ded"); + const runId = `run_${NEW_ID_26}`; + await newStore.createRun( + buildCreateRunInput({ + runId, + friendlyId: "run_ces_ded", + organizationId: env.organization.id, + projectId: env.project.id, + runtimeEnvironmentId: env.environment.id, + }) + ); + + // Force the dedicated join insert (completedWaitpoint.createMany) to fail mid-write. + await prisma17.$executeRawUnsafe('DROP TABLE "CompletedWaitpoint"'); + + await expect( + // Base client as `tx` = how the engine threads its base prisma through + // (continueRunIfUnblocked -> executionSnapshotSystem.createExecutionSnapshot(prisma, ...)). + // It is NOT an interactive transaction, so the store must still open its own to stay atomic. + newStore.createExecutionSnapshot( + { + run: { id: runId, status: "EXECUTING", attemptNumber: 1 }, + snapshot: { + executionStatus: "EXECUTING_WITH_WAITPOINTS", + description: "Run was blocked by a waitpoint.", + }, + environmentId: env.environment.id, + environmentType: "DEVELOPMENT", + projectId: env.project.id, + organizationId: env.project.id, + completedWaitpoints: [{ id: `wp_${NEW_ID_26}`, index: 0 }], + }, + prisma17 as never + ) + ).rejects.toThrow(); + + const snap = await prisma17.taskRunExecutionSnapshot.findFirst({ + where: { runId, executionStatus: "EXECUTING_WITH_WAITPOINTS" }, + }); + expect(snap).toBeNull(); + } + ); + + heteroRunOpsPostgresTest( + "lockRunToWorker rolls the snapshot and run lock back if the CompletedWaitpoint insert fails", + async ({ prisma17 }) => { + const newStore = makeDedicatedStore(prisma17); + const env = await seedEnvironment(prisma17, "dedicated", "lock_ded"); + const runId = `run_${NEW_ID_26}`; + await newStore.createRun( + buildCreateRunInput({ + runId, + friendlyId: "run_lock_ded", + organizationId: env.organization.id, + projectId: env.project.id, + runtimeEnvironmentId: env.environment.id, + }) + ); + const prior = await prisma17.taskRunExecutionSnapshot.findFirstOrThrow({ where: { runId } }); + + await prisma17.$executeRawUnsafe('DROP TABLE "CompletedWaitpoint"'); + + const snapshotId = `snap_${NEW_ID_26}`; + await expect( + // lockedById/lockedToVersionId/lockedQueueId are FK-free scalars on the dedicated subset, so + // synthetic ids are fine; the base client as `tx` mirrors the dequeue path (no interactive tx). + newStore.lockRunToWorker( + runId, + { + lockedAt: new Date(), + lockedById: `bwt_${NEW_ID_26}`, + lockedToVersionId: `bw_${NEW_ID_26}`, + lockedQueueId: `queue_${NEW_ID_26}`, + startedAt: new Date(), + baseCostInCents: 5, + machinePreset: "small-1x", + taskVersion: "20260601.1", + sdkVersion: "3.0.0", + cliVersion: "3.0.0", + maxDurationInSeconds: null, + snapshot: { + id: snapshotId, + previousSnapshotId: prior.id, + environmentId: env.environment.id, + environmentType: "DEVELOPMENT", + projectId: env.project.id, + organizationId: env.project.id, + completedWaitpointIds: [`wp_${NEW_ID_26}`], + completedWaitpointOrder: [`wp_${NEW_ID_26}`], + }, + }, + prisma17 as never + ) + ).rejects.toThrow(); + + const snap = await prisma17.taskRunExecutionSnapshot.findUnique({ + where: { id: snapshotId }, + }); + expect(snap).toBeNull(); + // The whole lock write must roll back, not just the status: no lock columns may leak through. + const run = await prisma17.taskRun.findUniqueOrThrow({ where: { id: runId } }); + expect(run.status).not.toBe("DEQUEUED"); + expect(run.lockedAt).toBeNull(); + expect(run.lockedById).toBeNull(); + expect(run.lockedToVersionId).toBeNull(); + expect(run.lockedQueueId).toBeNull(); + } + ); +}); diff --git a/internal-packages/run-store/src/checkpointBatchReplicaLag.test.ts b/internal-packages/run-store/src/checkpointBatchReplicaLag.test.ts new file mode 100644 index 0000000000..f9e19ad607 --- /dev/null +++ b/internal-packages/run-store/src/checkpointBatchReplicaLag.test.ts @@ -0,0 +1,109 @@ +// Repro for the checkpoint WAIT_FOR_BATCH replica-lag stall (createCheckpoint.server.ts:148-181). +// +// The service resolves the batch via `runStore.findBatchTaskRunByFriendlyId(friendlyId, envId)` with +// NO client, so the read is served from the REPLICA. Its decision hinges on `batchRun.resumedAt`: +// resumedAt set -> return keepRunAlive:true (batch already resumed; the run must keep executing) +// resumedAt null -> fall through -> create the checkpoint -> SUSPEND the run +// If the batch just resumed (primary has resumedAt), but the replica still lags (resumedAt null), the +// service suspends a run whose batch already completed -> it stalls until a sweep. The sibling +// WAIT_FOR_TASK arm reads the primary (this._prisma); only the batch arm defaults to the replica. +// +// This is invisible to the normal single-DB harness (no lag). We reintroduce the lag with the shared +// `laggingReplica` primitive: the store's replica is frozen at the pre-resume snapshot while the +// primary advances. RED = the current (client-less, replica) read -> SUSPEND. GREEN = threading the +// primary (the one-line fix) -> KEEP_ALIVE. + +import { laggingReplica, postgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import { describe, expect } from "vitest"; +import { PostgresRunStore } from "./PostgresRunStore.js"; +import type { ReadClient } from "./types.js"; + +type BatchRow = { resumedAt: Date | null } | null; + +// A line-for-line mirror of the service's WAIT_FOR_BATCH pre-check. `readClient` models the fix: the +// buggy code passes nothing (replica default); the fix threads the primary. +async function precheckWaitForBatch( + store: PostgresRunStore, + batchFriendlyId: string, + environmentId: string, + readClient?: ReadClient +): Promise<"DROP_RUN" | "KEEP_ALIVE" | "SUSPEND"> { + const batchRun = (await store.findBatchTaskRunByFriendlyId( + batchFriendlyId, + environmentId, + undefined, + readClient + )) as BatchRow; + if (!batchRun) return "DROP_RUN"; // keepRunAlive:false + if (batchRun.resumedAt) return "KEEP_ALIVE"; // batch already resumed -> run continues + return "SUSPEND"; // falls through -> checkpoint created -> run suspended +} + +async function seedEnvironment(prisma: PrismaClient, suffix: string) { + const organization = await prisma.organization.create({ + data: { title: `Org ${suffix}`, slug: `org-${suffix}` }, + }); + const project = await prisma.project.create({ + data: { + name: `Project ${suffix}`, + slug: `project-${suffix}`, + externalRef: `proj_${suffix}`, + organizationId: organization.id, + }, + }); + const environment = await prisma.runtimeEnvironment.create({ + data: { + type: "DEVELOPMENT", + slug: "dev", + projectId: project.id, + organizationId: organization.id, + apiKey: `tr_dev_${suffix}`, + pkApiKey: `pk_dev_${suffix}`, + shortcode: `short_${suffix}`, + }, + }); + return { organization, project, environment }; +} + +describe("checkpoint WAIT_FOR_BATCH under replica lag", () => { + postgresTest( + "a batch resumed on the primary but stale on the replica suspends an already-resumed run", + async ({ prisma }) => { + const { environment } = await seedEnvironment(prisma, "ckpt_lag"); + const friendlyId = "batch_ckpt_lag"; + const batch = await prisma.batchTaskRun.create({ + data: { friendlyId, runtimeEnvironmentId: environment.id }, + }); + + // Snapshot the batch as the replica still sees it (pre-resume: resumedAt = null). + const staleBatch = await prisma.batchTaskRun.findFirstOrThrow({ where: { id: batch.id } }); + expect(staleBatch.resumedAt).toBeNull(); + + // The batch completes and resumes the parent: primary now has resumedAt set... + await prisma.batchTaskRun.update({ + where: { id: batch.id }, + data: { resumedAt: new Date() }, + }); + + // ...but the replica lags, frozen at the pre-resume snapshot. + const replica = laggingReplica(prisma, [ + { model: "batchTaskRun", mode: "frozen", rows: [staleBatch] }, + ]); + const store = new PostgresRunStore({ + prisma, + readOnlyPrisma: replica.client, + schemaVariant: "legacy", + }); + + // RED - the current service call (no client -> replica): stale null -> SUSPEND an already-resumed run. + const buggy = await precheckWaitForBatch(store, friendlyId, environment.id); + expect(replica.wasHit("batchTaskRun")).toBe(true); // proves it read the (stale) replica + expect(buggy).toBe("SUSPEND"); // the stall bug + + // GREEN - the fix (thread the primary): sees resumedAt -> keep the run alive. + const fixed = await precheckWaitForBatch(store, friendlyId, environment.id, prisma); + expect(fixed).toBe("KEEP_ALIVE"); + } + ); +}); diff --git a/internal-packages/run-store/src/runOpsStore.batchItemMisroute.test.ts b/internal-packages/run-store/src/runOpsStore.batchItemMisroute.test.ts index b59b711bd8..68565b8ee3 100644 --- a/internal-packages/run-store/src/runOpsStore.batchItemMisroute.test.ts +++ b/internal-packages/run-store/src/runOpsStore.batchItemMisroute.test.ts @@ -92,4 +92,40 @@ describe("run-ops split — completing a batch item routes by the batch id, not expect(await prisma14.batchTaskRunItem.count({ where: { batchTaskRunId: batchId } })).toBe(0); } ); + + // createBatchTaskRunItem must co-locate the item with its BATCH (so the completion count, routed by + // batchTaskRunId, finds it), not with the child run. Routing by taskRunId would place a divergent- + // residency item on the child's DB -> invisible to the count -> the batch never completes -> parent hangs. + heteroRunOpsPostgresTest( + "createBatchTaskRunItem places the item on the batch's DB, routing by batchTaskRunId not taskRunId", + async ({ prisma14, prisma17 }: { prisma14: PrismaClient; prisma17: RunOpsPrismaClient }) => { + const router = makeSplitRouter(prisma14, prisma17); + const envId = "env_batchitem"; + const batchId = `batch_${NEW_ID_26}`; // run-ops id -> #new + const runId = "c".repeat(25); // cuid -> classifies LEGACY (the divergent routing key) + + await prisma17.batchTaskRun.create({ + data: { + id: batchId, + friendlyId: "batch_create_new", + runtimeEnvironmentId: envId, + runCount: 1, + status: "PROCESSING", + }, + }); + // The run physically lives on #new (the batch's DB) so the item's FKs resolve there. + await seedDedicatedRun(prisma17, envId, runId); + + await router.createBatchTaskRunItem({ + batchTaskRunId: batchId, + taskRunId: runId, + status: "PENDING", + }); + + // GREEN: routed by batchTaskRunId (NEW) -> item on #new, visible to the batch-completion count. + // RED: routed by the cuid taskRunId -> #legacy -> the create FK-fails / the count never sees it. + expect(await prisma17.batchTaskRunItem.count({ where: { batchTaskRunId: batchId } })).toBe(1); + expect(await prisma14.batchTaskRunItem.count({ where: { batchTaskRunId: batchId } })).toBe(0); + } + ); }); diff --git a/internal-packages/run-store/src/runOpsStore.snapshotCompletedWaitpoints.test.ts b/internal-packages/run-store/src/runOpsStore.snapshotCompletedWaitpoints.test.ts index b0a5a99e2f..0e300714d4 100644 --- a/internal-packages/run-store/src/runOpsStore.snapshotCompletedWaitpoints.test.ts +++ b/internal-packages/run-store/src/runOpsStore.snapshotCompletedWaitpoints.test.ts @@ -44,4 +44,60 @@ describe("run-ops split — completed waitpoints for a cuid snapshot are found o expect(ids).toEqual([WAITPOINT_ID]); } ); + + // WithPresence reports, in one read, whether the snapshot is visible on the reader (so a multi-reader + // replica can distinguish "no waitpoints" from "reader has not applied the snapshot yet"). + heteroRunOpsPostgresTest( + "findSnapshotCompletedWaitpointIdsWithPresence reports present+ids for a dedicated snapshot, absent otherwise", + async ({ prisma17 }: { prisma17: RunOpsPrismaClient }) => { + const newStore = new PostgresRunStore({ + prisma: prisma17 as never, + readOnlyPrisma: prisma17 as never, + schemaVariant: "dedicated", + }); + const runId = "run_" + "k".repeat(24) + "01"; + await prisma17.taskRun.create({ + data: { + id: runId, + friendlyId: "run_pres", + engine: "V2", + status: "PENDING", + taskIdentifier: "t", + payload: "{}", + payloadType: "application/json", + traceId: "tr", + spanId: "sp", + queue: "q", + runtimeEnvironmentId: "env_pres", + projectId: "proj_pres", + }, + }); + const snap = await prisma17.taskRunExecutionSnapshot.create({ + data: { + id: "c".repeat(25), + engine: "V2", + executionStatus: "EXECUTING", + description: "continue", + runStatus: "PENDING", + runId, + environmentId: "env_pres", + environmentType: "DEVELOPMENT", + projectId: "proj_pres", + organizationId: "org_pres", + }, + }); + await prisma17.completedWaitpoint.create({ + data: { snapshotId: snap.id, waitpointId: WAITPOINT_ID }, + }); + + expect(await newStore.findSnapshotCompletedWaitpointIdsWithPresence(snap.id)).toEqual({ + present: true, + ids: [WAITPOINT_ID], + }); + // A snapshot the reader does not have -> present:false, so its empty ids are not trusted as authoritative. + expect( + await newStore.findSnapshotCompletedWaitpointIdsWithPresence("c".repeat(24) + "zz") + ).toEqual({ present: false, ids: [] }); + } + ); }); diff --git a/internal-packages/run-store/src/runOpsStore.ts b/internal-packages/run-store/src/runOpsStore.ts index 102f56ea86..27cab57d1a 100644 --- a/internal-packages/run-store/src/runOpsStore.ts +++ b/internal-packages/run-store/src/runOpsStore.ts @@ -711,12 +711,14 @@ export class RoutingRunStore implements RunStore { // checkpoints). Mechanical residency-routing delegates so `implements RunStore` is satisfied. // --------------------------------------------------------------------------- - // Membership row lives on the run's residency — route by taskRunId. + // Route by batchTaskRunId so the item co-resides with its BatchTaskRun and is visible to the batch + // completion count/update (which route by batchTaskRunId). Routing by taskRunId would place the item + // on the child's DB if child and batch residency ever diverge, and the batch would never complete. async createBatchTaskRunItem( data: { batchTaskRunId: string; taskRunId: string; status: BatchTaskRunItemStatus }, tx?: PrismaClientOrTransaction ): Promise { - return (await this.#routeForWrite(data.taskRunId)).createBatchTaskRunItem(data); + return (await this.#routeForWrite(data.batchTaskRunId)).createBatchTaskRunItem(data); } // Snapshot reads route by OWNING run id (a SnapshotId is a cuid, NOT classifiable). The owning @@ -853,6 +855,28 @@ export class RoutingRunStore implements RunStore { return uniqueStrings([...fromNew, ...fromLegacy]); } + // Snapshot-id has no residency to route on, so fan out; the snapshot lives on exactly one store, so + // `present` is the OR and `ids` the union. + async findSnapshotCompletedWaitpointIdsWithPresence( + snapshotId: string, + client?: ReadClient + ): Promise<{ present: boolean; ids: string[] }> { + const [fromNew, fromLegacy] = await Promise.all([ + this.#new.findSnapshotCompletedWaitpointIdsWithPresence( + snapshotId, + RoutingRunStore.#ownPrimary(this.#new, client) + ), + this.#legacy.findSnapshotCompletedWaitpointIdsWithPresence( + snapshotId, + RoutingRunStore.#ownPrimary(this.#legacy, client) + ), + ]); + return { + present: fromNew.present || fromLegacy.present, + ids: uniqueStrings([...fromNew.ids, ...fromLegacy.ids]), + }; + } + // Keyed by waitpointId, but the WaitpointRunConnection / CompletedWaitpoint join co-locates with the // RUN/snapshot — which can be on the OTHER DB from a cross-DB token — so fan out to BOTH stores and // merge. Dedup by value: a token mirrored onto both DBs during drain can carry the same join diff --git a/internal-packages/run-store/src/types.ts b/internal-packages/run-store/src/types.ts index b1729347c4..1d39498f39 100644 --- a/internal-packages/run-store/src/types.ts +++ b/internal-packages/run-store/src/types.ts @@ -607,6 +607,12 @@ export interface RunStore { // Implicit-join group findSnapshotCompletedWaitpointIds(snapshotId: string, client?: ReadClient): Promise; + /** As above, but reports in the SAME read whether the snapshot is visible on the reader: `present=false` + * means this reader lacks the snapshot, so its empty id list is not authoritative (repair from primary). */ + findSnapshotCompletedWaitpointIdsWithPresence( + snapshotId: string, + client?: ReadClient + ): Promise<{ present: boolean; ids: string[] }>; /** Run ids connected to a waitpoint (WaitpointRunConnection / `_WaitpointRunConnections`), this DB only. */ findWaitpointConnectedRunIds(waitpointId: string, client?: ReadClient): Promise; /** Snapshot ids that completed a waitpoint (CompletedWaitpoint / `_completedWaitpoints`), this DB only. */ diff --git a/internal-packages/testcontainers/src/index.ts b/internal-packages/testcontainers/src/index.ts index 1bb9b3706b..ffa0b987a4 100644 --- a/internal-packages/testcontainers/src/index.ts +++ b/internal-packages/testcontainers/src/index.ts @@ -29,6 +29,7 @@ import { } from "./utils"; export { assertNonNullable, createPostgresContainer } from "./utils"; +export { laggingReplica, type LaggingModel } from "./laggingReplica"; export { logCleanup }; export type { MinIOConnectionConfig }; diff --git a/internal-packages/testcontainers/src/laggingReplica.ts b/internal-packages/testcontainers/src/laggingReplica.ts new file mode 100644 index 0000000000..d437aa8705 --- /dev/null +++ b/internal-packages/testcontainers/src/laggingReplica.ts @@ -0,0 +1,86 @@ +// A read-replica that has NOT caught up to its primary, for reproducing read-your-writes hazards the +// zero-lag single-DB test harness can't surface. Wraps a real Prisma client; for the configured +// models it serves STALE reads instead of live ones. Other models and all writes forward untouched. +// Modes: "missing" (row not replicated yet -> null/[]/0, *OrThrow throws); "frozen" (row out of date +// -> returns the provided pre-write snapshot rows, matched on top-level scalar `where` equality). +// `wasHit` asserts the stale replica was actually consulted. + +const READ_METHODS = new Set([ + "findFirst", + "findUnique", + "findFirstOrThrow", + "findUniqueOrThrow", + "findMany", + "count", +]); + +export type LaggingModel = + | { model: string; mode: "missing" } + | { model: string; mode: "frozen"; rows: readonly Record[] }; + +// Match a frozen row against a Prisma `where` using top-level scalar equality only (enough for the +// friendlyId / id / environmentId lookups these reads use); nested filters are treated as "matches". +function whereMatches(row: Record, where: unknown): boolean { + if (where == null || typeof where !== "object") return true; + return Object.entries(where as Record).every(([key, val]) => { + if (val !== null && typeof val === "object") return true; + return row[key] === val; + }); +} + +export function laggingReplica( + real: C, + configs: readonly LaggingModel[] +): { client: C; wasHit: (model?: string) => boolean } { + const byModel = new Map(configs.map((c) => [c.model, c])); + const hits = new Set(); + + const makeModelProxy = (modelName: string, realModel: object, cfg: LaggingModel) => + new Proxy(realModel, { + get(target, prop) { + if (typeof prop === "string" && READ_METHODS.has(prop)) { + return async (args?: { where?: unknown }) => { + hits.add(modelName); + if (cfg.mode === "missing") { + if (prop === "findMany") return []; + if (prop === "count") return 0; + if (prop === "findFirstOrThrow" || prop === "findUniqueOrThrow") { + throw new Error( + `laggingReplica: ${modelName}.${prop} - row not visible on replica yet` + ); + } + return null; + } + const matched = cfg.rows.filter((r) => whereMatches(r, args?.where)); + if (prop === "findMany") return matched; + if (prop === "count") return matched.length; + const first = matched[0] ?? null; + if ((prop === "findFirstOrThrow" || prop === "findUniqueOrThrow") && !first) { + throw new Error(`laggingReplica: ${modelName}.${prop} - no frozen row matched`); + } + return first; + }; + } + const value = (target as Record)[prop]; + return typeof value === "function" + ? (value as (...a: unknown[]) => unknown).bind(target) + : value; + }, + }); + + const client = new Proxy(real, { + get(target, prop) { + const cfg = typeof prop === "string" ? byModel.get(prop) : undefined; + const delegate = cfg ? (target as Record)[prop as string] : undefined; + if (cfg && delegate && typeof delegate === "object") { + return makeModelProxy(prop as string, delegate, cfg); + } + const value = (target as Record)[prop]; + return typeof value === "function" + ? (value as (...a: unknown[]) => unknown).bind(target) + : value; + }, + }) as C; + + return { client, wasHit: (model) => (model ? hits.has(model) : hits.size > 0) }; +}