fix(run-store): commit run snapshots and their completed-waitpoint links atomically

d-cs · d-cs · commit a8c63233bfcd · 2026-07-06T01:55:04.000+01:00
createExecutionSnapshot and lockRunToWorker wrote the execution snapshot and its
completed-waitpoint join rows as two separate statements, and createRun and
createFailedRun (dedicated schema) wrote the run and its associated waitpoint
separately. Served back from a read replica that applied the snapshot but not yet
the join, the runner gets a waitpoint-less continue and never resumes, so the run
hangs (or partial state persists on a crash between the writes). Wrap each primary
write and its dependent write in one transaction, reusing the caller's transaction
when it already has a real one, so a replica can never observe the partial state.
diff --git 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
@@ -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<R>(
+    tx: PrismaClientOrTransaction | undefined,
+    fn: (client: PrismaClientOrTransaction) => Promise<R>
+  ): Promise<R> {
+    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<Prisma.TaskRunGetPayload<{}>> {
-    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<Prisma.TaskRunGetPayload<{}>> {
     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<Prisma.TaskRunExecutionSnapshotGetPayload<{ include: { checkpoint: true } }>> {
-    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<Prisma.TaskRunExecutionSnapshotGetPayload<{ include: { checkpoint: true } }>> {
     const {
       run,
       snapshot,
diff --git 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,75 @@ 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();
+    }
+  );
+});
