From e73577ae5f9bb69631503623992b318c4c02cd26 Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Fri, 21 Aug 2026 13:37:49 +0100 Subject: [PATCH] refactor(run-engine): move block-edge registration behind the coordinator --- .../src/engine/systems/waitpointSystem.ts | 24 +++----- .../legacyPostgresCoordinator.ts | 59 ++++++++++++++++++- .../src/engine/waitpointCoordinator/types.ts | 25 ++++++++ 3 files changed, 92 insertions(+), 16 deletions(-) diff --git a/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts b/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts index 752dbe90d..2a9b1d55a 100644 --- a/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts @@ -490,25 +490,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"; @@ -606,10 +600,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, diff --git a/internal-packages/run-engine/src/engine/waitpointCoordinator/legacyPostgresCoordinator.ts b/internal-packages/run-engine/src/engine/waitpointCoordinator/legacyPostgresCoordinator.ts index b7ca49f15..d310c4e2b 100644 --- a/internal-packages/run-engine/src/engine/waitpointCoordinator/legacyPostgresCoordinator.ts +++ b/internal-packages/run-engine/src/engine/waitpointCoordinator/legacyPostgresCoordinator.ts @@ -2,7 +2,13 @@ import type { RunStore } from "@internal/run-store"; import type { Logger } from "@trigger.dev/core/logger"; import type { PrismaClient } from "@trigger.dev/database"; import { boundedIn } from "@trigger.dev/database"; -import type { ClearRunBlockStateParams, RunBlockEdge, WaitpointCoordinator } from "./types.js"; +import type { + ClearRunBlockStateParams, + RegisterBlocksLocklessParams, + RegisterBlocksParams, + RunBlockEdge, + WaitpointCoordinator, +} from "./types.js"; export type LegacyPostgresWaitpointCoordinatorOptions = { runStore: RunStore; @@ -65,4 +71,55 @@ export class LegacyPostgresWaitpointCoordinator implements WaitpointCoordinator 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); + } + + /** + * 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 index f4c630638..faaa77f40 100644 --- a/internal-packages/run-engine/src/engine/waitpointCoordinator/types.ts +++ b/internal-packages/run-engine/src/engine/waitpointCoordinator/types.ts @@ -1,3 +1,4 @@ +import type { ReadClient } from "@internal/run-store"; import type { PrismaClientOrTransaction, Waitpoint } from "@trigger.dev/database"; /** @@ -15,6 +16,8 @@ import type { PrismaClientOrTransaction, Waitpoint } from "@trigger.dev/database export type WaitpointCoordinator = { clearRunBlockState(params: ClearRunBlockStateParams): Promise<{ count: number }>; readRunBlockState(runId: string): Promise; + registerBlocks(params: RegisterBlocksParams): Promise<{ pendingCount: number }>; + registerBlocksLockless(params: RegisterBlocksLocklessParams): Promise; }; export type ClearRunBlockStateParams = { @@ -40,3 +43,25 @@ export type RunBlockEdge = { 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;