diff --git a/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts b/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts index 5d5a80772a6..3dbed999445 100644 --- a/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts @@ -1,5 +1,4 @@ -import { timeoutError, tryCatch } from "@trigger.dev/core/v3"; -import { WaitpointId } from "@trigger.dev/core/v3/isomorphic"; +import { timeoutError } from "@trigger.dev/core/v3"; import type { PrismaClientOrTransaction, TaskRun, @@ -7,13 +6,11 @@ import type { TaskRunExecutionStatus, Waitpoint, } from "@trigger.dev/database"; -import { Prisma, boundedIn } from "@trigger.dev/database"; -import type { RunStore } from "@internal/run-store"; import { assertNever } from "assert-never"; -import { nanoid } from "nanoid"; -import { UnclassifiableWaitpointId } from "../errors.js"; import { sendNotificationToWorker } from "../eventBus.js"; import { isFinalRunStatus } from "../statuses.js"; +import { LegacyPostgresWaitpointCoordinator } from "../waitpointCoordinator/legacyPostgresCoordinator.js"; +import type { WaitpointCoordinator } from "../waitpointCoordinator/types.js"; import type { EnqueueSystem } from "./enqueueSystem.js"; import type { ExecutionSnapshotSystem } from "./executionSnapshotSystem.js"; import { getLatestExecutionSnapshot } from "./executionSnapshotSystem.js"; @@ -45,11 +42,17 @@ export class WaitpointSystem { private readonly $: SystemResources; private readonly executionSnapshotSystem: ExecutionSnapshotSystem; private readonly enqueueSystem: EnqueueSystem; + private readonly coordinator: WaitpointCoordinator; constructor(private readonly options: WaitpointSystemOptions) { this.$ = options.resources; this.executionSnapshotSystem = options.executionSnapshotSystem; this.enqueueSystem = options.enqueueSystem; + this.coordinator = new LegacyPostgresWaitpointCoordinator({ + runStore: this.$.runStore, + prisma: this.$.prisma, + logger: this.$.logger, + }); } public async clearBlockingWaitpoints({ @@ -59,14 +62,7 @@ export class WaitpointSystem { runId: string; tx?: PrismaClientOrTransaction; }) { - // A run's edges co-locate with the run (the edge write routes by runId), so the router routes this - // taskRunId-keyed delete to the run's store rather than fanning out. The caller's `tx` is not - // forwarded — the delete runs on the owning store's own client (the router never threads a - // control-plane tx into a routed write). - const deleted = await this.$.runStore.deleteManyTaskRunWaitpoints( - { where: { taskRunId: runId } }, - tx - ); + const deleted = await this.coordinator.clearRunBlockState({ runId, tx }); return deleted.count; } @@ -84,86 +80,19 @@ export class WaitpointSystem { isError: boolean; }; }): Promise { - // Residency store-selection guard. completeWaitpoint arrives with only - // (waitpointId, output) — no run id — so the owning run-ops store is selected - // by the waitpoint's own residency. In single-DB this is the one store - // (no classification). An unclassifiable id throws loud — never default-routes. - let store: RunStore; - try { - store = await this.$.runStore.forWaitpointCompletion(id, { routeKind: "MANUAL" }); - } catch (error) { - this.$.logger.error("completeWaitpoint: unclassifiable waitpointId", { - waitpointId: id, - error, - }); - throw new UnclassifiableWaitpointId(id, { cause: error }); - } - - // 1. Complete the Waitpoint (if not completed) - const [updateError, updateResult] = await tryCatch( - store.updateManyWaitpoints({ - where: { id, status: "PENDING" }, - data: { - status: "COMPLETED", - completedAt: new Date(), - output: output?.value, - outputType: output?.type, - outputIsError: output?.isError, - }, - }) - ); - - if (updateError) { - this.$.logger.error("completeWaitpoint: error updating waitpoint:", { updateError }); - throw updateError; - } - - if (updateResult.count === 0) { - this.$.logger.info( - "completeWaitpoint: attempted to complete a waitpoint that is not PENDING", - { waitpointId: id } - ); - } - - // Re-read the just-written row from the RESOLVED store's PRIMARY: the replica (findWaitpoint's - // default) can miss it under lag → false "not found" → the parent hangs; this.$.prisma would - // instead hit the wrong DB. findWaitpointOnPrimary reads the owning store's primary. - const waitpoint = await store.findWaitpointOnPrimary({ - where: { id }, + const { waitpoint, blockedRuns } = await this.coordinator.complete({ + waitpointId: id, + output, }); - if (!waitpoint) { - this.$.logger.error("completeWaitpoint: waitpoint not found", { waitpointId: id }); - throw new Error("Waitpoint not found"); - } - - if (waitpoint.status !== "COMPLETED") { - this.$.logger.error(`completeWaitpoint: waitpoint is not completed`, { - waitpointId: id, - }); - throw new Error("Waitpoint not completed"); - } - - // 2. Find the TaskRuns blocked by this waitpoint. The edge (TaskRunWaitpoint) co-locates - // with its RUN, not this token, so it can live on the OTHER run-ops DB: read via the router - // (which fans the waitpointId lookup across both DBs) rather than the token's own `store`, - // or a cross-DB blocked run is never found and hangs forever. - const affectedTaskRuns = await this.$.runStore.findManyTaskRunWaitpoints( - { - where: { waitpointId: id }, - select: { taskRunId: true, spanIdToComplete: true, createdAt: true }, - }, - this.$.prisma - ); - - if (affectedTaskRuns.length === 0) { + if (blockedRuns.length === 0) { this.$.logger.debug(`completeWaitpoint: no TaskRunWaitpoints found for waitpoint`, { waitpointId: id, }); } // 3. Schedule trying to continue the runs - for (const run of affectedTaskRuns) { + for (const run of blockedRuns) { const jobId = `continueRunIfUnblocked:${run.taskRunId}`; //50ms in the future const availableAt = new Date(Date.now() + 50); @@ -220,81 +149,27 @@ export class WaitpointSystem { idempotencyKey?: string; idempotencyKeyExpiresAt?: Date; }) { - // 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 - // local). The minted waitpoint id is always a cuid, so without `coLocateWithRunId` the upsert - // 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. 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 - ? await this.$.runStore.findWaitpoint( - { - where: { - environmentId, - idempotencyKey, - }, - }, - undefined, - colocate - ) - : undefined; - - if (existingWaitpoint) { - if ( - existingWaitpoint.idempotencyKeyExpiresAt && - new Date() > existingWaitpoint.idempotencyKeyExpiresAt - ) { - //the idempotency key has expired - //remove the waitpoint idempotencyKey - const rotateArgs = { - where: { - id: existingWaitpoint.id, - }, - data: { - idempotencyKey: nanoid(24), - inactiveIdempotencyKey: existingWaitpoint.idempotencyKey, - }, - }; - await this.$.runStore.updateWaitpoint(rotateArgs, undefined, colocate); + const result = await this.coordinator.createDateTimeWaitpoint({ + runId, + projectId, + environmentId, + completedAfter, + idempotencyKey, + idempotencyKeyExpiresAt, + }); - //let it fall through to create a new waitpoint - } else { - return { waitpoint: existingWaitpoint, isCached: true }; - } + if (result.kind === "cached") { + return { waitpoint: result.waitpoint, isCached: true }; } - const upsertArgs = { - where: { - environmentId_idempotencyKey: { - environmentId, - idempotencyKey: idempotencyKey ?? nanoid(24), - }, - }, - create: { - ...WaitpointId.generate(), - type: "DATETIME" as const, - idempotencyKey: idempotencyKey ?? nanoid(24), - idempotencyKeyExpiresAt, - userProvidedIdempotencyKey: !!idempotencyKey, - environmentId, - projectId, - completedAfter, - }, - update: {}, - }; - const waitpoint = await this.$.runStore.upsertWaitpoint(upsertArgs, undefined, colocate); - await this.$.worker.enqueue({ - id: `finishWaitpoint.${waitpoint.id}`, + id: `finishWaitpoint.${result.waitpoint.id}`, job: "finishWaitpoint", - payload: { waitpointId: waitpoint.id }, + payload: { waitpointId: result.waitpoint.id }, availableAt: completedAfter, }); - return { waitpoint, isCached: false }; + return { waitpoint: result.waitpoint, isCached: false }; } /** This creates a MANUAL waitpoint, that can be explicitly completed (or failed). @@ -322,117 +197,35 @@ export class WaitpointSystem { // to LEGACY by its cuid id-shape. Ignored when `runId` is set (co-location wins). standaloneResidency?: "NEW" | "LEGACY"; }): Promise<{ waitpoint: Waitpoint; isCached: boolean }> { - // Co-location invariant (see createDateTimeWaitpoint): when a `runId` is supplied the waitpoint - // co-locates with that run's DB and the (env,idempotencyKey) dedup is per-run (co-resident). A - // standalone token (api.v1.waitpoints.tokens.ts) passes no run id — it is created without an - // owner, blocked later by whichever run waits on it (possibly cross-DB, resolved by the - // run-co-resident block edge + completion fan-out). With no owner it reads the env mint kind via - // `standaloneResidency` so a minted-new env keeps its tokens on NEW; unset, it routes by id-shape. No tx here. - const colocate = runId - ? { coLocateWithRunId: runId } - : standaloneResidency - ? { residency: standaloneResidency } - : undefined; - const existingWaitpoint = idempotencyKey - ? await this.$.runStore.findWaitpoint( - { - where: { - environmentId, - idempotencyKey, - }, - }, - undefined, - colocate - ) - : undefined; - - if (existingWaitpoint) { - if ( - existingWaitpoint.idempotencyKeyExpiresAt && - new Date() > existingWaitpoint.idempotencyKeyExpiresAt - ) { - //the idempotency key has expired - //remove the waitpoint idempotencyKey - await this.$.runStore.updateWaitpoint( - { - where: { - id: existingWaitpoint.id, - }, - data: { - idempotencyKey: nanoid(24), - inactiveIdempotencyKey: existingWaitpoint.idempotencyKey, - }, - }, - undefined, - colocate - ); + const result = await this.coordinator.createManualWaitpoint({ + runId, + environmentId, + projectId, + idempotencyKey, + idempotencyKeyExpiresAt, + timeout, + tags, + standaloneResidency, + }); - //let it fall through to create a new waitpoint - } else { - return { waitpoint: existingWaitpoint, isCached: true }; - } + if (result.kind === "cached") { + return { waitpoint: result.waitpoint, isCached: true }; } - const maxRetries = 5; - let attempts = 0; - - while (attempts < maxRetries) { - try { - const waitpoint = await this.$.runStore.upsertWaitpoint( - { - where: { - environmentId_idempotencyKey: { - environmentId, - idempotencyKey: idempotencyKey ?? nanoid(24), - }, - }, - create: { - ...WaitpointId.generate(), - type: "MANUAL", - idempotencyKey: idempotencyKey ?? nanoid(24), - idempotencyKeyExpiresAt, - userProvidedIdempotencyKey: !!idempotencyKey, - environmentId, - projectId, - completedAfter: timeout, - tags, - }, - update: {}, - }, - undefined, - colocate - ); - - //schedule the timeout - if (timeout) { - await this.$.worker.enqueue({ - id: `finishWaitpoint.${waitpoint.id}`, - job: "finishWaitpoint", - payload: { - waitpointId: waitpoint.id, - error: JSON.stringify(timeoutError(timeout)), - }, - availableAt: timeout, - }); - } - - return { waitpoint, isCached: false }; - } catch (error) { - if (error instanceof Prisma.PrismaClientKnownRequestError && error.code === "P2002") { - // Handle unique constraint violation (conflict) - attempts++; - if (attempts >= maxRetries) { - throw new Error( - `Failed to create waitpoint after ${maxRetries} attempts due to conflicts.` - ); - } - } else { - throw error; // Re-throw other errors - } - } + //schedule the timeout + if (timeout) { + await this.$.worker.enqueue({ + id: `finishWaitpoint.${result.waitpoint.id}`, + job: "finishWaitpoint", + payload: { + waitpointId: result.waitpoint.id, + error: JSON.stringify(timeoutError(timeout)), + }, + availableAt: timeout, + }); } - throw new Error(`Failed to create waitpoint after ${maxRetries} attempts due to conflicts.`); + return { waitpoint: result.waitpoint, isCached: false }; } /** @@ -489,25 +282,19 @@ export class WaitpointSystem { this.$.runStore ); - // Insert the blocking + historical connections via the run-ops store, routed by the owning - // run id so the edge co-resides with the run. Never pinned to the caller's control-plane tx: - // that joined `Waitpoint` on the wrong DB and wrote 0 edges. The pending check stays a - // SEPARATE store call so it gets its own READ COMMITTED snapshot (see the doc comment above). - await this.$.runStore.blockRunWithWaitpointEdges({ + // Insert the blocking + historical connections and re-check the pending count. The + // coordinator keeps these as two separate store statements, in this order, for the READ + // COMMITTED reason documented on the method and in the doc comment above. + const { pendingCount } = await this.coordinator.registerBlocks({ runId, waitpointIds: $waitpoints, projectId, spanIdToComplete, batchId: batch?.id, batchIndex: batch?.index, + client: prisma, }); - // Check if the run is actually blocked using a separate query (see above). Pass the writer so the - // pending re-read is read-your-writes on the owning PRIMARY (a lagging replica can strand the run). - // Route by the blocked run id: its blocking waitpoints co-locate with the run, so the router - // counts on the run's store and only falls back to the other DB for a cross-tree token. - const pendingCount = await this.$.runStore.countPendingWaitpoints($waitpoints, prisma, runId); - const isRunBlocked = pendingCount > 0; let newStatus: TaskRunExecutionStatus = "SUSPENDED"; @@ -605,10 +392,10 @@ export class WaitpointSystem { }): Promise { const $waitpoints = typeof waitpoints === "string" ? [waitpoints] : waitpoints; - // Same routed edge write as blockRunWithWaitpoint, routed by the owning run id. No lock - // needed: ON CONFLICT DO NOTHING makes concurrent inserts safe, and the parent snapshot is - // already EXECUTING_WITH_WAITPOINTS from blockRunWithCreatedBatch. - await this.$.runStore.blockRunWithWaitpointEdges({ + // Same routed edge write as blockRunWithWaitpoint. No lock needed: ON CONFLICT DO NOTHING + // makes concurrent inserts safe, and the parent snapshot is already + // EXECUTING_WITH_WAITPOINTS from blockRunWithCreatedBatch. No pending count here. + await this.coordinator.registerBlocksLockless({ runId, waitpointIds: $waitpoints, projectId, @@ -682,20 +469,7 @@ export class WaitpointSystem { return await this.$.runLock.lock("continueRunIfUnblocked", [runId], async () => { // 1. Get the any blocking waitpoints - const blockingWaitpoints = await this.$.runStore.findManyTaskRunWaitpoints( - { - where: { taskRunId: runId }, - select: { - id: true, - batchId: true, - batchIndex: true, - waitpoint: { - select: { id: true, status: true, type: true, completedAfter: true }, - }, - }, - }, - this.$.prisma - ); + const blockingWaitpoints = await this.coordinator.readRunBlockState(runId); // 2. There are blockers still, so do nothing if (blockingWaitpoints.some((w) => w.waitpoint.status !== "COMPLETED")) { @@ -926,11 +700,9 @@ export class WaitpointSystem { if (blockingWaitpoints.length > 0) { //5. Remove the blocking waitpoints - await this.$.runStore.deleteManyTaskRunWaitpoints({ - where: { - taskRunId: runId, - id: { in: boundedIn(blockingWaitpoints.map((b) => b.id)) }, - }, + await this.coordinator.clearRunBlockState({ + runId, + edgeIds: blockingWaitpoints.map((b) => b.id), }); this.$.logger.debug(`continueRunIfUnblocked: removed blocking waitpoints`, { @@ -953,15 +725,7 @@ export class WaitpointSystem { projectId: string; environmentId: string; }) { - return { - ...WaitpointId.generate(), - type: "RUN" as const, - status: "PENDING" as const, - idempotencyKey: nanoid(24), - userProvidedIdempotencyKey: false, - projectId, - environmentId, - }; + return this.coordinator.mintAssociatedWaitpointData({ projectId, environmentId }); } /** @@ -1045,12 +809,9 @@ export class WaitpointSystem { // Create waitpoint and link to run atomically const waitpointData = this.buildRunAssociatedWaitpoint({ projectId, environmentId }); - // RUN-type within-tree waitpoint that belongs to runId; routes by owning run id. - const waitpoint = await this.$.runStore.createWaitpoint({ - data: { - ...waitpointData, - completedByTaskRunId: runId, - }, + const waitpoint = await this.coordinator.createAssociatedWaitpoint({ + runId, + data: waitpointData, }); // If run has already finished (per snapshot), complete the waitpoint immediately so the parent can resume diff --git a/internal-packages/run-engine/src/engine/waitpointCoordinator/legacyPostgresCoordinator.ts b/internal-packages/run-engine/src/engine/waitpointCoordinator/legacyPostgresCoordinator.ts new file mode 100644 index 00000000000..d1e48fa4f8d --- /dev/null +++ b/internal-packages/run-engine/src/engine/waitpointCoordinator/legacyPostgresCoordinator.ts @@ -0,0 +1,437 @@ +import type { RunStore } from "@internal/run-store"; +import { tryCatch } from "@trigger.dev/core/v3"; +import { WaitpointId } from "@trigger.dev/core/v3/isomorphic"; +import type { Logger } from "@trigger.dev/core/logger"; +import type { PrismaClient, Waitpoint } from "@trigger.dev/database"; +import { boundedIn, Prisma } from "@trigger.dev/database"; +import { nanoid } from "nanoid"; +import { UnclassifiableWaitpointId } from "../errors.js"; +import type { + AssociatedWaitpointData, + ClearRunBlockStateParams, + CompleteParams, + CompleteResult, + CreateDateTimeWaitpointParams, + CreateManualWaitpointParams, + CreateWaitpointResult, + RegisterBlocksLocklessParams, + RegisterBlocksParams, + RunBlockEdge, + WaitpointCoordinator, +} from "./types.js"; + +export type LegacyPostgresWaitpointCoordinatorOptions = { + runStore: RunStore; + prisma: PrismaClient; + logger: Logger; +}; + +/** + * Waitpoint coordination against Postgres, through the run-ops store. + * + * Dependencies are deliberately narrow: no run lock, no worker, no event bus. + * That makes "this owns waitpoint state only" structural rather than a convention. + */ +export class LegacyPostgresWaitpointCoordinator implements WaitpointCoordinator { + private readonly runStore: RunStore; + private readonly prisma: PrismaClient; + private readonly logger: Logger; + + constructor(options: LegacyPostgresWaitpointCoordinatorOptions) { + this.runStore = options.runStore; + this.prisma = options.prisma; + this.logger = options.logger; + } + + async clearRunBlockState({ + runId, + edgeIds, + tx, + }: ClearRunBlockStateParams): Promise<{ count: number }> { + if (edgeIds) { + // Bounded delete of named edges, on the unblock path. No tx: that path is not inside a + // caller transaction, and boundedIn caps the id-list arity for Prisma. + return this.runStore.deleteManyTaskRunWaitpoints({ + where: { + taskRunId: runId, + id: { in: boundedIn(edgeIds) }, + }, + }); + } + + // A run's edges co-locate with the run (the edge write routes by runId), so the router routes + // this taskRunId-keyed delete to the run's store rather than fanning out. The caller's `tx` is + // passed through: a routing store strips it, and a single store joins it. + return this.runStore.deleteManyTaskRunWaitpoints({ where: { taskRunId: runId } }, tx); + } + + async readRunBlockState(runId: string): Promise { + return this.runStore.findManyTaskRunWaitpoints( + { + where: { taskRunId: runId }, + select: { + id: true, + batchId: true, + batchIndex: true, + waitpoint: { + select: { id: true, status: true, type: true, completedAfter: true }, + }, + }, + }, + this.prisma + ); + } + + async registerBlocks({ + client, + ...edge + }: RegisterBlocksParams): Promise<{ pendingCount: number }> { + await this.#writeBlockEdges(edge); + + // Check if the run is actually blocked using a separate query. The separate statement is the + // point: under PostgreSQL READ COMMITTED each statement gets its own snapshot, so a + // concurrent completion that commits between the edge write and this check is still seen. + // It queries ALL requested ids, not just inserted ones: a row that already existed (ON + // CONFLICT skipped the insert) but is still PENDING must still block. Pass the caller's + // client so the re-read is read-your-writes on the owning PRIMARY, and pass the run id so + // the router counts on the run's store instead of fanning out to both DBs. + const pendingCount = await this.runStore.countPendingWaitpoints( + edge.waitpointIds, + client, + edge.runId + ); + + return { pendingCount }; + } + + async registerBlocksLockless(params: RegisterBlocksLocklessParams): Promise { + await this.#writeBlockEdges(params); + } + + async complete({ waitpointId, output }: CompleteParams): Promise { + // Residency store-selection guard. complete arrives with only (waitpointId, output) — no run + // id — so the owning run-ops store is selected by the waitpoint's own residency. In single-DB + // this is the one store (no classification). An unclassifiable id throws loud — never + // default-routes. The try wraps ONLY the resolve: widening it would swallow the + // "Waitpoint not found" path that a single store relies on. + let store: RunStore; + try { + store = await this.runStore.forWaitpointCompletion(waitpointId, { routeKind: "MANUAL" }); + } catch (error) { + this.logger.error("completeWaitpoint: unclassifiable waitpointId", { + waitpointId, + error, + }); + throw new UnclassifiableWaitpointId(waitpointId, { cause: error }); + } + + // 1. Complete the Waitpoint (if not completed) + const [updateError, updateResult] = await tryCatch( + store.updateManyWaitpoints({ + where: { id: waitpointId, status: "PENDING" }, + data: { + status: "COMPLETED", + completedAt: new Date(), + output: output?.value, + outputType: output?.type, + outputIsError: output?.isError, + }, + }) + ); + + if (updateError) { + this.logger.error("completeWaitpoint: error updating waitpoint:", { updateError }); + throw updateError; + } + + if (updateResult.count === 0) { + this.logger.info("completeWaitpoint: attempted to complete a waitpoint that is not PENDING", { + waitpointId, + }); + } + + // Re-read the just-written row from the RESOLVED store's PRIMARY: the replica (findWaitpoint's + // default) can miss it under lag → false "not found" → the parent hangs. Going back through + // the router would re-resolve the store and change the routing, so use the handle. + const waitpoint = await store.findWaitpointOnPrimary({ + where: { id: waitpointId }, + }); + + if (!waitpoint) { + this.logger.error("completeWaitpoint: waitpoint not found", { waitpointId }); + throw new Error("Waitpoint not found"); + } + + if (waitpoint.status !== "COMPLETED") { + this.logger.error(`completeWaitpoint: waitpoint is not completed`, { waitpointId }); + throw new Error("Waitpoint not completed"); + } + + // 2. Find the TaskRuns blocked by this waitpoint. The edge (TaskRunWaitpoint) co-locates + // with its RUN, not this token, so it can live on the OTHER run-ops DB: read via the router + // (which fans the waitpointId lookup across both DBs) rather than the token's own `store`, + // or a cross-DB blocked run is never found and hangs forever. + const blockedRuns = await this.runStore.findManyTaskRunWaitpoints( + { + where: { waitpointId }, + select: { taskRunId: true, spanIdToComplete: true, createdAt: true }, + }, + this.prisma + ); + + return { waitpoint, blockedRuns }; + } + + async createDateTimeWaitpoint({ + runId, + projectId, + environmentId, + completedAfter, + idempotencyKey, + idempotencyKeyExpiresAt, + }: CreateDateTimeWaitpointParams): Promise { + // 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 + // local). The minted waitpoint id is always a cuid, so without `coLocateWithRunId` the upsert + // 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. 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 + ? await this.runStore.findWaitpoint( + { + where: { + environmentId, + idempotencyKey, + }, + }, + undefined, + colocate + ) + : undefined; + + if (existingWaitpoint) { + if ( + existingWaitpoint.idempotencyKeyExpiresAt && + new Date() > existingWaitpoint.idempotencyKeyExpiresAt + ) { + //the idempotency key has expired + //remove the waitpoint idempotencyKey + const rotateArgs = { + where: { + id: existingWaitpoint.id, + }, + data: { + idempotencyKey: nanoid(24), + inactiveIdempotencyKey: existingWaitpoint.idempotencyKey, + }, + }; + await this.runStore.updateWaitpoint(rotateArgs, undefined, colocate); + + //let it fall through to create a new waitpoint + } else { + return { kind: "cached", waitpoint: existingWaitpoint }; + } + } + + // The two `nanoid(24)` calls below are deliberately separate and produce DIFFERENT values: + // the upsert `where` key must not match the `create` key, or a guaranteed-miss upsert becomes + // a possible update. Do not hoist either to a shared constant. + const upsertArgs = { + where: { + environmentId_idempotencyKey: { + environmentId, + idempotencyKey: idempotencyKey ?? nanoid(24), + }, + }, + create: { + ...WaitpointId.generate(), + type: "DATETIME" as const, + idempotencyKey: idempotencyKey ?? nanoid(24), + idempotencyKeyExpiresAt, + userProvidedIdempotencyKey: !!idempotencyKey, + environmentId, + projectId, + completedAfter, + }, + update: {}, + }; + const waitpoint = await this.runStore.upsertWaitpoint(upsertArgs, undefined, colocate); + + return { kind: "created", waitpoint }; + } + + async createManualWaitpoint({ + runId, + environmentId, + projectId, + idempotencyKey, + idempotencyKeyExpiresAt, + timeout, + tags, + standaloneResidency, + }: CreateManualWaitpointParams): Promise { + // Co-location invariant (see createDateTimeWaitpoint): when a `runId` is supplied the waitpoint + // co-locates with that run's DB and the (env,idempotencyKey) dedup is per-run (co-resident). A + // standalone token (api.v1.waitpoints.tokens.ts) passes no run id — it is created without an + // owner, blocked later by whichever run waits on it (possibly cross-DB, resolved by the + // run-co-resident block edge + completion fan-out). With no owner it reads the env mint kind via + // `standaloneResidency` so a minted-new env keeps its tokens on NEW; unset, it routes by id-shape. No tx here. + const colocate = runId + ? { coLocateWithRunId: runId } + : standaloneResidency + ? { residency: standaloneResidency } + : undefined; + const existingWaitpoint = idempotencyKey + ? await this.runStore.findWaitpoint( + { + where: { + environmentId, + idempotencyKey, + }, + }, + undefined, + colocate + ) + : undefined; + + if (existingWaitpoint) { + if ( + existingWaitpoint.idempotencyKeyExpiresAt && + new Date() > existingWaitpoint.idempotencyKeyExpiresAt + ) { + //the idempotency key has expired + //remove the waitpoint idempotencyKey + await this.runStore.updateWaitpoint( + { + where: { + id: existingWaitpoint.id, + }, + data: { + idempotencyKey: nanoid(24), + inactiveIdempotencyKey: existingWaitpoint.idempotencyKey, + }, + }, + undefined, + colocate + ); + + //let it fall through to create a new waitpoint + } else { + return { kind: "cached", waitpoint: existingWaitpoint }; + } + } + + const maxRetries = 5; + let attempts = 0; + + while (attempts < maxRetries) { + try { + // As in createDateTimeWaitpoint, the two `nanoid(24)` calls are deliberately separate and + // differ. Both, and `WaitpointId.generate()`, are re-evaluated on every attempt: that is + // what makes a retry after a unique-constraint conflict try a fresh key. + const waitpoint = await this.runStore.upsertWaitpoint( + { + where: { + environmentId_idempotencyKey: { + environmentId, + idempotencyKey: idempotencyKey ?? nanoid(24), + }, + }, + create: { + ...WaitpointId.generate(), + type: "MANUAL", + idempotencyKey: idempotencyKey ?? nanoid(24), + idempotencyKeyExpiresAt, + userProvidedIdempotencyKey: !!idempotencyKey, + environmentId, + projectId, + completedAfter: timeout, + tags, + }, + update: {}, + }, + undefined, + colocate + ); + + return { kind: "created", waitpoint }; + } catch (error) { + if (error instanceof Prisma.PrismaClientKnownRequestError && error.code === "P2002") { + // Handle unique constraint violation (conflict) + attempts++; + if (attempts >= maxRetries) { + throw new Error( + `Failed to create waitpoint after ${maxRetries} attempts due to conflicts.` + ); + } + } else { + throw error; // Re-throw other errors + } + } + } + + throw new Error(`Failed to create waitpoint after ${maxRetries} attempts due to conflicts.`); + } + + mintAssociatedWaitpointData({ + projectId, + environmentId, + }: { + projectId: string; + environmentId: string; + }): AssociatedWaitpointData { + return { + ...WaitpointId.generate(), + type: "RUN" as const, + status: "PENDING" as const, + idempotencyKey: nanoid(24), + userProvidedIdempotencyKey: false, + projectId, + environmentId, + }; + } + + async createAssociatedWaitpoint({ + runId, + data, + }: { + runId: string; + data: AssociatedWaitpointData; + }): Promise { + // RUN-type within-tree waitpoint that belongs to runId; routes by owning run id. + return this.runStore.createWaitpoint({ + data: { + ...data, + completedByTaskRunId: runId, + }, + }); + } + + /** + * The edge write, shared by both register paths so they cannot drift. + * + * Routed by the owning run id so the edge co-resides with the run. Never pinned to a caller + * transaction: that joined `Waitpoint` on the wrong DB, wrote 0 edges, and silently never + * suspended the parent. The write is idempotent (ON CONFLICT DO NOTHING). + */ + #writeBlockEdges({ + runId, + waitpointIds, + projectId, + spanIdToComplete, + batchId, + batchIndex, + }: RegisterBlocksLocklessParams): Promise { + return this.runStore.blockRunWithWaitpointEdges({ + runId, + waitpointIds, + projectId, + spanIdToComplete, + batchId, + batchIndex, + }); + } +} diff --git a/internal-packages/run-engine/src/engine/waitpointCoordinator/types.ts b/internal-packages/run-engine/src/engine/waitpointCoordinator/types.ts new file mode 100644 index 00000000000..8a50abb7d1c --- /dev/null +++ b/internal-packages/run-engine/src/engine/waitpointCoordinator/types.ts @@ -0,0 +1,145 @@ +import type { ReadClient } from "@internal/run-store"; +import type { PrismaClientOrTransaction, Waitpoint } from "@trigger.dev/database"; + +/** + * The waitpoint and edge state operations that `WaitpointSystem` delegates. + * + * Orchestration stays in `WaitpointSystem`: the run lock, snapshot transitions, + * worker-job enqueues, event emissions and racepoints. This owns waitpoint and + * edge state only, so a non-Postgres implementation can replace it without any + * caller learning that it changed. + * + * The residency hints and `tx` are opaque pass-throughs. Opaque does not mean + * type-free — a Prisma type appears here — it means a non-Postgres implementation + * never reads the value. + */ +export type WaitpointCoordinator = { + clearRunBlockState(params: ClearRunBlockStateParams): Promise<{ count: number }>; + readRunBlockState(runId: string): Promise; + registerBlocks(params: RegisterBlocksParams): Promise<{ pendingCount: number }>; + registerBlocksLockless(params: RegisterBlocksLocklessParams): Promise; + complete(params: CompleteParams): Promise; + createDateTimeWaitpoint(params: CreateDateTimeWaitpointParams): Promise; + createManualWaitpoint(params: CreateManualWaitpointParams): Promise; + mintAssociatedWaitpointData(params: { + projectId: string; + environmentId: string; + }): AssociatedWaitpointData; + createAssociatedWaitpoint(params: { + runId: string; + data: AssociatedWaitpointData; + }): Promise; +}; + +export type ClearRunBlockStateParams = { + runId: string; + /** Edge ids to delete. Omit to clear every edge for the run. */ + edgeIds?: string[]; + /** + * Forwarded verbatim on the full-clear leg only, and never on the bounded leg + * or an edge write. A routing store strips it; a single store joins it. + */ + tx?: PrismaClientOrTransaction; +}; + +/** + * One block edge, with the fields the unblock decision reads. + * + * `batchId` is read by no logic. It rides inside the two `logger.debug` payloads in + * `continueRunIfUnblocked`, so removing it changes log output. + */ +export type RunBlockEdge = { + id: string; + batchId: string | null; + batchIndex: number | null; + waitpoint: Pick; +}; + +export type RegisterBlocksParams = { + runId: string; + waitpointIds: string[]; + projectId: string; + spanIdToComplete?: string; + batchId?: string; + batchIndex?: number; + /** + * Read client for the pending count only. The caller resolves `tx ?? prisma` once + * and passes the result, so the writer is used when the caller is inside a + * transaction and the pending re-read is read-your-writes on the owning primary. + * Never forwarded to the edge write. + */ + client: ReadClient; +}; + +/** + * 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 CompleteParams = { + waitpointId: string; + output?: { + value: string; + type?: string; + isError: boolean; + }; +}; + +/** One run blocked by the completed waitpoint, with the fields the caller's fan-out loop reads. */ +type BlockedRun = { + taskRunId: string; + spanIdToComplete: string | null; + createdAt: Date; +}; + +export type CompleteResult = { + waitpoint: Waitpoint; + blockedRuns: BlockedRun[]; +}; + +/** + * Discriminated on purpose. The caller enqueues the `finishWaitpoint` job only in the + * `created` branch, because today's create methods return before their enqueue on the + * cached path. A boolean would let a later edit enqueue on both branches. + */ +export type CreateWaitpointResult = + | { kind: "cached"; waitpoint: Waitpoint } + | { kind: "created"; waitpoint: Waitpoint }; + +export type CreateDateTimeWaitpointParams = { + /** When set, the waitpoint co-locates with this run's DB and the dedup probe targets it. */ + runId?: string; + projectId: string; + environmentId: string; + completedAfter: Date; + idempotencyKey?: string; + idempotencyKeyExpiresAt?: Date; +}; + +export type CreateManualWaitpointParams = { + runId?: string; + environmentId: string; + projectId: string; + idempotencyKey?: string; + idempotencyKeyExpiresAt?: Date; + timeout?: Date; + tags?: string[]; + /** + * See the `standaloneResidency` param doc on `WaitpointSystem.createManualWaitpoint` for the + * full rationale. Only a Postgres implementation reads this. + */ + standaloneResidency?: "NEW" | "LEGACY"; +}; + +/** The RUN-waitpoint row data. Pure — no store touch — so the mint is coordinator-owned. */ +export type AssociatedWaitpointData = { + id: string; + friendlyId: string; + type: "RUN"; + status: "PENDING"; + idempotencyKey: string; + userProvidedIdempotencyKey: false; + projectId: string; + environmentId: string; +};