From ff290dfe2f6de3f54c524d058a91fa93f7901640 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 16 Apr 2026 15:43:16 +0100 Subject: [PATCH] perf(run-engine): merge dequeue snapshot creation into taskRun.update transaction [TRI-8450] (#3395) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Nests the `TaskRunExecutionSnapshot` creation inside the `taskRun.update()` Prisma call in the dequeue flow, reducing **2 DB commits → 1** per dequeue operation. This is the highest-volume of the five unmerged flows identified in TRI-8450 (~9,200 commits/sec on the engine service). **Pattern**: Follows the same nested-write approach already used in the completion path (`runAttemptSystem.ts:735`) and trigger path (`engine/index.ts:674`). **Changes**: - `dequeueSystem.ts`: Moved snapshot creation into `executionSnapshots: { create: {...} }` within the existing `taskRun.update()`. Pre-generates the snapshot ID via `generateInternalId()` (plain cuid, matching what Prisma's `@default(cuid())` produces) so the event emission, heartbeat enqueue, and return value can all be constructed from data already in scope — **no extra DB read needed** after the merged write. `SnapshotId.toFriendlyId()` is used only for the return value's `friendlyId` field, matching the original `createExecutionSnapshot` behavior. - `executionSnapshotSystem.ts`: Added public `enqueueHeartbeatIfNeeded()` method that exposes the heartbeat scheduling logic (previously only available internally via `createExecutionSnapshot`). This is needed because `PENDING_EXECUTING` requires a heartbeat, unlike the `FINISHED` status in the completion reference pattern. This method is reusable by future merge targets (retry-immediate, checkpoint, cancel, requeue). **Net DB change per dequeue**: eliminates 1 write transaction (the separate `TaskRunExecutionSnapshot.create`). No extra reads added — the snapshot ID is pre-generated and the `executionSnapshotCreated` event payload is constructed inline from values already available in the closure. ## Review & Testing Checklist for Human - [ ] **Verify manually-constructed event payload matches DB state**: The `executionSnapshotCreated` event is now built inline (not read back from DB). Confirm the field values (`runStatus: "PENDING"`, `attemptNumber`, `checkpointId`, `workerId`, `runnerId`, `completedWaitpointIds`) match what Prisma actually writes. A mismatch here would be silent — event consumers would get stale/wrong data. - [ ] **Verify `attemptNumber` source is equivalent**: Old code used `lockedTaskRun.attemptNumber` (post-update result). New code uses `result.run.attemptNumber` (pre-update). The `taskRun.update()` data payload does NOT include `attemptNumber`, so they should be identical — but confirm this assumption holds for all dequeue scenarios (e.g. retried runs). - [ ] **Verify `isValid` defaults to `true` in schema**: The old `createExecutionSnapshot` explicitly set `isValid: error ? false : true`. The nested create omits `isValid` (no error in the dequeue happy path). Confirm the Prisma schema default for `TaskRunExecutionSnapshot.isValid` is `true`. - [ ] **Verify `runStatus: "PENDING"` hardcoding matches the mapping**: The old code passed `lockedTaskRun.status` ("DEQUEUED") to `createExecutionSnapshot`, which mapped it to "PENDING" via `run.status === "DEQUEUED" ? "PENDING" : run.status`. The new code hardcodes `"PENDING"` directly. This is correct but brittle if `status` ever changes from "DEQUEUED" to something else upstream. - [ ] **Spot-check `completedWaitpoints` connect + order logic**: The nested create replicates the connect/order logic from `createExecutionSnapshot` (lines 387-393). Verify the `snapshot.completedWaitpoints` type provides `id` and `index` fields compatible with this usage. - [ ] **Verify `checkpoint` in return value**: The return now uses `snapshot.checkpoint` (from the *previous* snapshot) instead of reading the newly-created snapshot's checkpoint relation. Since `checkpointId` is passed through unchanged, they should be identical — but worth a sanity check. **Recommended test plan**: deploy to staging, run the `sample_pg_activity.py` sampler for a 5-minute window, and verify the COMMIT count drop on the engine service + proportional `IO:XactSync` reduction. ### Notes - This only covers the **dequeue** flow (flow #1 from TRI-8450). The remaining four flows (retry-immediate, checkpoint, requeue, cancel) are separate follow-ups. - The new `enqueueHeartbeatIfNeeded` method is deliberately designed for reuse by those follow-up PRs. - CI note: the `priority.test.ts` failure in shard 7 is a flaky ordering assertion unrelated to this change (it compares `friendlyId` values in dequeue order). The `audit` check is also pre-existing/unrelated. Link to Devin session: https://app.devin.ai/sessions/034fe0e7224f49278a2de260203e1377 Requested by: @ericallam --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Eric Allam --- ...merge-dequeue-snapshot-into-transaction.md | 6 ++ .../src/engine/systems/dequeueSystem.ts | 98 +++++++++++++------ .../engine/systems/executionSnapshotSystem.ts | 21 ++++ 3 files changed, 94 insertions(+), 31 deletions(-) create mode 100644 .server-changes/merge-dequeue-snapshot-into-transaction.md diff --git a/.server-changes/merge-dequeue-snapshot-into-transaction.md b/.server-changes/merge-dequeue-snapshot-into-transaction.md new file mode 100644 index 000000000..62c9a0ec6 --- /dev/null +++ b/.server-changes/merge-dequeue-snapshot-into-transaction.md @@ -0,0 +1,6 @@ +--- +area: webapp +type: improvement +--- + +Merge execution snapshot creation into the dequeue taskRun.update transaction, reducing 2 DB commits to 1 per dequeue operation diff --git a/internal-packages/run-engine/src/engine/systems/dequeueSystem.ts b/internal-packages/run-engine/src/engine/systems/dequeueSystem.ts index 9476a081f..15d79e76b 100644 --- a/internal-packages/run-engine/src/engine/systems/dequeueSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/dequeueSystem.ts @@ -3,7 +3,7 @@ import { startSpan } from "@internal/tracing"; import { assertExhaustive, tryCatch } from "@trigger.dev/core"; import { DequeuedMessage, RetryOptions, RunAnnotations } from "@trigger.dev/core/v3"; import { placementTag } from "@trigger.dev/core/v3/serverOnly"; -import { getMaxDuration } from "@trigger.dev/core/v3/isomorphic"; +import { generateInternalId, getMaxDuration, SnapshotId } from "@trigger.dev/core/v3/isomorphic"; import { BackgroundWorker, BackgroundWorkerTask, @@ -416,6 +416,9 @@ export class DequeueSystem { ? undefined : result.task.retryConfig; + // Pre-generate snapshot ID so we can construct the result without an extra read + const snapshotId = generateInternalId(); + const lockedTaskRun = await prisma.taskRun.update({ where: { id: runId, @@ -435,6 +438,33 @@ export class DequeueSystem { cliVersion: result.worker.cliVersion, maxDurationInSeconds, maxAttempts: maxAttempts ?? undefined, + executionSnapshots: { + create: { + id: snapshotId, + engine: "V2", + executionStatus: "PENDING_EXECUTING", + description: "Run was dequeued for execution", + // Map DEQUEUED -> PENDING for backwards compatibility with older runners + runStatus: "PENDING", + attemptNumber: result.run.attemptNumber ?? undefined, + previousSnapshotId: snapshot.id, + environmentId: snapshot.environmentId, + environmentType: snapshot.environmentType, + projectId: snapshot.projectId, + organizationId: snapshot.organizationId, + checkpointId: snapshot.checkpointId ?? undefined, + batchId: snapshot.batchId ?? undefined, + completedWaitpoints: { + connect: snapshot.completedWaitpoints.map((w) => ({ id: w.id })), + }, + completedWaitpointOrder: snapshot.completedWaitpoints + .filter((c) => c.index !== undefined) + .sort((a, b) => a.index! - b.index!) + .map((w) => w.id), + workerId, + runnerId, + }, + }, }, include: { runtimeEnvironment: true, @@ -516,44 +546,50 @@ export class DequeueSystem { hasPrivateLink = billingResult.val.hasPrivateLink; } - const newSnapshot = await this.executionSnapshotSystem.createExecutionSnapshot( - prisma, - { - run: { - id: runId, - status: lockedTaskRun.status, - attemptNumber: lockedTaskRun.attemptNumber, - }, - snapshot: { - executionStatus: "PENDING_EXECUTING", - description: "Run was dequeued for execution", - }, - previousSnapshotId: snapshot.id, - environmentId: snapshot.environmentId, - environmentType: snapshot.environmentType, - projectId: snapshot.projectId, - organizationId: snapshot.organizationId, - checkpointId: snapshot.checkpointId ?? undefined, - batchId: snapshot.batchId ?? undefined, - completedWaitpoints: snapshot.completedWaitpoints, - workerId, - runnerId, - } - ); + // Snapshot was created as part of the taskRun.update above (single transaction). + // Construct the snapshot info from data we already have and handle side effects + // (heartbeat + event) manually — no extra DB read needed. + const snapshotCreatedAt = new Date(); + + this.$.eventBus.emit("executionSnapshotCreated", { + time: snapshotCreatedAt, + run: { + id: runId, + }, + snapshot: { + id: snapshotId, + executionStatus: "PENDING_EXECUTING", + description: "Run was dequeued for execution", + runStatus: "PENDING", + attemptNumber: result.run.attemptNumber ?? null, + checkpointId: snapshot.checkpointId ?? null, + workerId: workerId ?? null, + runnerId: runnerId ?? null, + isValid: true, + error: null, + completedWaitpointIds: snapshot.completedWaitpoints.map((wp) => wp.id), + }, + }); + + await this.executionSnapshotSystem.enqueueHeartbeatIfNeeded({ + id: snapshotId, + runId, + executionStatus: "PENDING_EXECUTING", + }); return { version: "1" as const, dequeuedAt: new Date(), workerQueueLength: message.workerQueueLength, snapshot: { - id: newSnapshot.id, - friendlyId: newSnapshot.friendlyId, - executionStatus: newSnapshot.executionStatus, - description: newSnapshot.description, - createdAt: newSnapshot.createdAt, + id: snapshotId, + friendlyId: SnapshotId.toFriendlyId(snapshotId), + executionStatus: "PENDING_EXECUTING" as const, + description: "Run was dequeued for execution", + createdAt: snapshotCreatedAt, }, image: result.deployment?.imageReference ?? undefined, - checkpoint: newSnapshot.checkpoint ?? undefined, + checkpoint: snapshot.checkpoint ?? undefined, completedWaitpoints: snapshot.completedWaitpoints, backgroundWorker: { id: result.worker.id, diff --git a/internal-packages/run-engine/src/engine/systems/executionSnapshotSystem.ts b/internal-packages/run-engine/src/engine/systems/executionSnapshotSystem.ts index a224e5a86..d615c066b 100644 --- a/internal-packages/run-engine/src/engine/systems/executionSnapshotSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/executionSnapshotSystem.ts @@ -518,6 +518,27 @@ export class ExecutionSnapshotSystem { return executionResultFromSnapshot(latestSnapshot); } + /** + * Enqueues a heartbeat job for a snapshot if the execution status requires one. + * Use this after nesting a snapshot create inside a taskRun.update() to replicate + * the heartbeat side effect that createExecutionSnapshot normally handles. + */ + public async enqueueHeartbeatIfNeeded(snapshot: { + id: string; + runId: string; + executionStatus: TaskRunExecutionStatus; + }) { + const intervalMs = this.#getHeartbeatIntervalMs(snapshot.executionStatus); + if (intervalMs !== null) { + await this.$.worker.enqueue({ + id: `heartbeatSnapshot.${snapshot.runId}`, + job: "heartbeatSnapshot", + payload: { snapshotId: snapshot.id, runId: snapshot.runId }, + availableAt: new Date(Date.now() + intervalMs), + }); + } + } + #getHeartbeatIntervalMs(status: TaskRunExecutionStatus): number | null { switch (status) { case "PENDING_EXECUTING": {