Waitpoint timeouts and idempotency expiry

Author: matt-aitken

internal-packages/run-engine/src/engine/index.ts

@@ -23,6 +23,7 @@ import {
   TaskRunInternalError,
   TaskRunSuccessfulExecutionResult,
   WaitForDurationResult,
+  WAITPOINT_TIMEOUT_ERROR_CODE,
 } from "@trigger.dev/core/v3";
 import {
   BatchId,
@@ -118,6 +119,13 @@ const workerCatalog = {
     }),
     visibilityTimeoutMs: 10_000,
   },
+  timeoutWaitpoint: {
+    schema: z.object({
+      waitpointId: z.string(),
+      timeout: z.coerce.date(),
+    }),
+    visibilityTimeoutMs: 10_000,
+  },
 };
 
 type EngineWorker = Worker<typeof workerCatalog>;
@@ -199,6 +207,12 @@ export class RunEngine {
             runId: payload.runId,
           });
         },
+        timeoutWaitpoint: async ({ payload }) => {
+          await this.#timeoutWaitpoint({
+            waitpointId: payload.waitpointId,
+            timeout: payload.timeout,
+          });
+        },
       },
     }).start();
 
@@ -1735,10 +1749,14 @@ export class RunEngine {
     environmentId,
     projectId,
     idempotencyKey,
+    idempotencyKeyExpiresAt,
+    timeout,
   }: {
     environmentId: string;
     projectId: string;
     idempotencyKey?: string;
+    idempotencyKeyExpiresAt?: Date;
+    timeout?: Date;
   }): Promise<Waitpoint> {
     const existingWaitpoint = idempotencyKey
       ? await this.prisma.waitpoint.findUnique({
@@ -1752,19 +1770,59 @@ export class RunEngine {
       : undefined;
 
     if (existingWaitpoint) {
-      return existingWaitpoint;
+      if (
+        existingWaitpoint.idempotencyKeyExpiresAt &&
+        new Date() > existingWaitpoint.idempotencyKeyExpiresAt
+      ) {
+        //the idempotency key has expired
+        //remove the waitpoint idempotencyKey
+        await this.prisma.waitpoint.update({
+          where: {
+            id: existingWaitpoint.id,
+          },
+          data: {
+            idempotencyKey: nanoid(24),
+            inactiveIdempotencyKey: existingWaitpoint.idempotencyKey,
+          },
+        });
+
+        //let it fall through to create a new waitpoint
+      } else {
+        return existingWaitpoint;
+      }
     }
 
-    return this.prisma.waitpoint.create({
-      data: {
+    const waitpoint = await this.prisma.waitpoint.upsert({
+      where: {
+        environmentId_idempotencyKey: {
+          environmentId,
+          idempotencyKey: idempotencyKey ?? nanoid(24),
+        },
+      },
+      create: {
         ...WaitpointId.generate(),
         type: "MANUAL",
         idempotencyKey: idempotencyKey ?? nanoid(24),
+        idempotencyKeyExpiresAt,
         userProvidedIdempotencyKey: !!idempotencyKey,
         environmentId,
         projectId,
+        completedAfter: timeout,
       },
+      update: {},
     });
+
+    //schedule the timeout
+    if (timeout) {
+      await this.worker.enqueue({
+        id: `timeoutWaitpoint.${waitpoint.id}`,
+        job: "timeoutWaitpoint",
+        payload: { waitpointId: waitpoint.id, timeout },
+        availableAt: timeout,
+      });
+    }
+
+    return waitpoint;
   }
 
   /** This block a run with a BATCH waitpoint.
@@ -2032,14 +2090,6 @@ export class RunEngine {
       isError: boolean;
     };
   }): Promise<Waitpoint> {
-    const waitpoint = await this.prisma.waitpoint.findUnique({
-      where: { id },
-    });
-
-    if (!waitpoint) {
-      throw new Error(`Waitpoint ${id} not found`);
-    }
-
     const result = await $transaction(
       this.prisma,
       async (tx) => {
@@ -2051,23 +2101,32 @@ export class RunEngine {
 
         if (affectedTaskRuns.length === 0) {
           this.logger.warn(`No TaskRunWaitpoints found for waitpoint`, {
-            waitpoint,
+            waitpointId: id,
           });
         }
 
-        // 2. Update the waitpoint to completed
-        const updatedWaitpoint = await tx.waitpoint.update({
-          where: { id },
-          data: {
-            status: "COMPLETED",
-            completedAt: new Date(),
-            output: output?.value,
-            outputType: output?.type,
-            outputIsError: output?.isError,
-          },
-        });
+        // 2. Update the waitpoint to completed (only if it's pending)
+        const waitpoint = await tx.waitpoint
+          .update({
+            where: { id, status: "PENDING" },
+            data: {
+              status: "COMPLETED",
+              completedAt: new Date(),
+              output: output?.value,
+              outputType: output?.type,
+              outputIsError: output?.isError,
+            },
+          })
+          .catch(async (error) => {
+            if (error instanceof Prisma.PrismaClientKnownRequestError && error.code === "P2025") {
+              return tx.waitpoint.findUnique({
+                where: { id },
+              });
+            }
+            throw error;
+          });
 
-        return { updatedWaitpoint, affectedTaskRuns };
+        return { waitpoint, affectedTaskRuns };
       },
       (error) => {
         this.logger.error(`Error completing waitpoint ${id}, retrying`, { error });
@@ -2079,6 +2138,10 @@ export class RunEngine {
       throw new Error(`Waitpoint couldn't be updated`);
     }
 
+    if (!result.waitpoint) {
+      throw new Error(`Waitpoint ${id} not found`);
+    }
+
     //schedule trying to continue the runs
     for (const run of result.affectedTaskRuns) {
       await this.worker.enqueue({
@@ -2100,7 +2163,7 @@ export class RunEngine {
       }
     }
 
-    return result.updatedWaitpoint;
+    return result.waitpoint;
   }
 
   async createCheckpoint({
@@ -3491,6 +3554,19 @@ export class RunEngine {
     };
   }
 
+  async #timeoutWaitpoint({ waitpointId, timeout }: { waitpointId: string; timeout: Date }) {
+    await this.completeWaitpoint({
+      id: waitpointId,
+      output: {
+        value: JSON.stringify({
+          code: WAITPOINT_TIMEOUT_ERROR_CODE,
+          message: `Waitpoint timed out at ${timeout.toISOString()}`,
+        }),
+        isError: true,
+      },
+    });
+  }
+
   async #clearBlockingWaitpoints({ runId, tx }: { runId: string; tx?: PrismaClientOrTransaction }) {
     const prisma = tx ?? this.prisma;
     const deleted = await prisma.taskRunWaitpoint.deleteMany({

internal-packages/run-engine/src/engine/tests/waitpoints.test.ts

@@ -9,6 +9,7 @@ import { expect } from "vitest";
 import { RunEngine } from "../index.js";
 import { setTimeout } from "timers/promises";
 import { EventBusEventArgs } from "../eventBus.js";
+import { isWaitpointOutputTimeout } from "@trigger.dev/core/v3";
 
 describe("RunEngine Waitpoints", () => {
   containerTest("waitForDuration", { timeout: 15_000 }, async ({ prisma, redisOptions }) => {
@@ -398,7 +399,7 @@ describe("RunEngine Waitpoints", () => {
   );
 
   containerTest(
-    "Manual waitpoint timeout",
+    "Manual waitpoint failAfter",
     { timeout: 15_000 },
     async ({ prisma, redisOptions }) => {
       //create environment
@@ -665,4 +666,155 @@ describe("RunEngine Waitpoints", () => {
       }
     }
   );
+
+  containerTest(
+    "Create a Manual waitpoint and let it timeout",
+    { timeout: 15_000 },
+    async ({ prisma, redisOptions }) => {
+      //create environment
+      const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
+
+      const engine = new RunEngine({
+        prisma,
+        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 {
+        const taskIdentifier = "test-task";
+
+        //create background worker
+        await setupBackgroundWorker(prisma, authenticatedEnvironment, taskIdentifier);
+
+        //trigger the run
+        const run = await engine.trigger(
+          {
+            number: 1,
+            friendlyId: "run_p1234",
+            environment: authenticatedEnvironment,
+            taskIdentifier,
+            payload: "{}",
+            payloadType: "application/json",
+            context: {},
+            traceContext: {},
+            traceId: "t12345",
+            spanId: "s12345",
+            masterQueue: "main",
+            queueName: "task/test-task",
+            isTest: false,
+            tags: [],
+          },
+          prisma
+        );
+
+        //dequeue the run
+        const dequeued = await engine.dequeueFromMasterQueue({
+          consumerId: "test_12345",
+          masterQueue: run.masterQueue,
+          maxRunCount: 10,
+        });
+
+        //create an attempt
+        const attemptResult = await engine.startRunAttempt({
+          runId: dequeued[0].run.id,
+          snapshotId: dequeued[0].snapshot.id,
+        });
+        expect(attemptResult.snapshot.executionStatus).toBe("EXECUTING");
+
+        //create a manual waitpoint with timeout
+        const timeout = new Date(Date.now() + 1_000);
+        const waitpoint = await engine.createManualWaitpoint({
+          environmentId: authenticatedEnvironment.id,
+          projectId: authenticatedEnvironment.projectId,
+          timeout,
+        });
+        expect(waitpoint.status).toBe("PENDING");
+        expect(waitpoint.completedAfter).toStrictEqual(timeout);
+
+        //block the run
+        await engine.blockRunWithWaitpoint({
+          runId: run.id,
+          waitpoints: waitpoint.id,
+          environmentId: authenticatedEnvironment.id,
+          projectId: authenticatedEnvironment.projectId,
+        });
+
+        const executionData = await engine.getRunExecutionData({ runId: run.id });
+        expect(executionData?.snapshot.executionStatus).toBe("EXECUTING_WITH_WAITPOINTS");
+
+        //check there is a waitpoint blocking the parent run
+        const runWaitpointBefore = await prisma.taskRunWaitpoint.findFirst({
+          where: {
+            taskRunId: run.id,
+          },
+          include: {
+            waitpoint: true,
+          },
+        });
+        expect(runWaitpointBefore?.waitpointId).toBe(waitpoint.id);
+
+        let event: EventBusEventArgs<"workerNotification">[0] | undefined = undefined;
+        engine.eventBus.on("workerNotification", (result) => {
+          event = result;
+        });
+
+        await setTimeout(1_250);
+
+        const executionData2 = await engine.getRunExecutionData({ runId: run.id });
+        expect(executionData2?.snapshot.executionStatus).toBe("EXECUTING");
+
+        assertNonNullable(event);
+        const notificationEvent = event as EventBusEventArgs<"workerNotification">[0];
+        expect(notificationEvent.run.id).toBe(run.id);
+
+        //check there are no waitpoints blocking the parent run
+        const runWaitpoint = await prisma.taskRunWaitpoint.findFirst({
+          where: {
+            taskRunId: run.id,
+          },
+          include: {
+            waitpoint: true,
+          },
+        });
+        expect(runWaitpoint).toBeNull();
+
+        const waitpoint2 = await prisma.waitpoint.findUnique({
+          where: {
+            id: waitpoint.id,
+          },
+        });
+        assertNonNullable(waitpoint2);
+        expect(waitpoint2.status).toBe("COMPLETED");
+        expect(waitpoint2.outputIsError).toBe(true);
+        assertNonNullable(waitpoint2.output);
+        const isTimeout = isWaitpointOutputTimeout(waitpoint2.output);
+        expect(isTimeout).toBe(true);
+      } finally {
+        engine.quit();
+      }
+    }
+  );
 });