More changes to blocking to support continuing after idempotent completed runs

Author: matt-aitken

apps/webapp/app/v3/services/triggerTaskV2.server.ts

@@ -90,7 +90,7 @@ export class TriggerTaskServiceV2 extends WithRunEngine {
         } else {
           //We're using `andWait` so we need to block the parent run with a waitpoint
           if (
-            existingRun.associatedWaitpoint?.status === "PENDING" &&
+            existingRun.associatedWaitpoint &&
             body.options?.resumeParentOnCompletion &&
             body.options?.parentRunId
           ) {

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

@@ -1593,6 +1593,29 @@ export class RunEngine {
     return await this.runLock.lock([runId], 5000, async (signal) => {
       let snapshot: TaskRunExecutionSnapshot = await getLatestExecutionSnapshot(prisma, runId);
 
+      //block the run with the waitpoints, returning how many waitpoints are pending
+      const insert = await prisma.$queryRaw<{ pending_count: number }[]>`
+        WITH inserted AS (
+          INSERT INTO "TaskRunWaitpoint" ("id", "taskRunId", "waitpointId", "projectId", "createdAt", "updatedAt")
+          SELECT
+            gen_random_uuid(),
+            ${runId},
+            w.id,
+            ${projectId},
+            NOW(),
+            NOW()
+          FROM "Waitpoint" w
+          WHERE w.id IN (${Prisma.join(waitpointIds)})
+          ON CONFLICT DO NOTHING
+          RETURNING "waitpointId"
+        )
+        SELECT COUNT(*) as pending_count
+        FROM inserted i
+        JOIN "Waitpoint" w ON w.id = i."waitpointId"
+        WHERE w.status = 'PENDING';`;
+
+      const pendingCount = insert.at(0)?.pending_count ?? 0;
+
       let newStatus: TaskRunExecutionStatus = "BLOCKED_BY_WAITPOINTS";
       if (
         snapshot.executionStatus === "EXECUTING" ||
@@ -1601,25 +1624,6 @@ export class RunEngine {
         newStatus = "EXECUTING_WITH_WAITPOINTS";
       }
 
-      const insertedBlockers = await prisma.$executeRaw`
-        INSERT INTO "TaskRunWaitpoint" ("id", "taskRunId", "waitpointId", "projectId", "createdAt", "updatedAt")
-        SELECT
-          gen_random_uuid(),
-          ${runId},
-          w.id,
-          ${projectId},
-          NOW(),
-          NOW()
-        FROM "Waitpoint" w
-        WHERE w.id IN (${Prisma.join(waitpointIds)})
-          AND w.status = 'PENDING'
-        ON CONFLICT DO NOTHING;`;
-
-      //if no runs were blocked we don't want to do anything more
-      if (insertedBlockers === 0) {
-        return snapshot;
-      }
-
       //if the state has changed, create a new snapshot
       if (newStatus !== snapshot.executionStatus) {
         snapshot = await this.#createExecutionSnapshot(prisma, {
@@ -1646,6 +1650,19 @@ export class RunEngine {
         }
       }
 
+      //no pending waitpoint, schedule unblocking the run
+      //debounce if we're rapidly adding waitpoints
+      if (pendingCount === 0) {
+        await this.worker.enqueue({
+          //this will debounce the call
+          id: `continueRunIfUnblocked:${runId}`,
+          job: "continueRunIfUnblocked",
+          payload: { runId: runId },
+          //100ms in the future
+          availableAt: new Date(Date.now() + 100),
+        });
+      }
+
       return snapshot;
     });
   }

references/hello-world/src/trigger/example.ts

@@ -60,7 +60,11 @@ export const batchParentTask = task({
 export const childTask = task({
   id: "child",
   run: async (
-    { message, failureChance = 0.3 }: { message?: string; failureChance?: number },
+    {
+      message,
+      failureChance = 0.3,
+      duration = 3_000,
+    }: { message?: string; failureChance?: number; duration?: number },
     { ctx }
   ) => {
     logger.info("Hello, world from the child", { message, failureChance });
@@ -69,7 +73,7 @@ export const childTask = task({
       throw new Error("Random error at start");
     }
 
-    await setTimeout(3_000);
+    await setTimeout(duration);
 
     if (Math.random() < failureChance) {
       throw new Error("Random error at end");

references/hello-world/src/trigger/idempotency.ts

@@ -0,0 +1,76 @@
+import { batch, idempotencyKeys, logger, task, timeout, usage, wait } from "@trigger.dev/sdk/v3";
+import { setTimeout } from "timers/promises";
+import { childTask } from "./example.js";
+
+export const idempotency = task({
+  id: "idempotency",
+  run: async (payload: any, { ctx }) => {
+    logger.log("Hello, world from the parent", { payload });
+
+    const child1Key = await idempotencyKeys.create("a", { scope: "global" });
+
+    const child1 = await childTask.triggerAndWait(
+      { message: "Hello, world!", duration: 10_000 },
+      { idempotencyKey: child1Key, idempotencyKeyTTL: "60s" }
+    );
+    logger.log("Child 1", { child1 });
+
+    ctx.attempt.id;
+
+    const child2 = await childTask.triggerAndWait(
+      { message: "Hello, world!", duration: 10_000 },
+      { idempotencyKey: child1Key, idempotencyKeyTTL: "60s" }
+    );
+    logger.log("Child 2", { child2 });
+
+    // const results = await childTask.batchTriggerAndWait([
+    //   {
+    //     payload: { message: "Hello, world!" },
+    //     //@ts-ignore
+    //     options: { idempotencyKey: "1", idempotencyKeyTTL: "60s" },
+    //   },
+    //   {
+    //     payload: { message: "Hello, world 2!" },
+    //     //@ts-ignore
+    //     options: { idempotencyKey: "2", idempotencyKeyTTL: "60s" },
+    //   },
+    // ]);
+    // logger.log("Results", { results });
+
+    // const results2 = await batch.triggerAndWait<typeof childTask>([
+    //   {
+    //     id: "child",
+    //     payload: { message: "Hello, world !" },
+    //     //@ts-ignore
+    //     options: { idempotencyKey: "1", idempotencyKeyTTL: "60s" },
+    //   },
+    //   {
+    //     id: "child",
+    //     payload: { message: "Hello, world 2!" },
+    //     //@ts-ignore
+    //     options: { idempotencyKey: "2", idempotencyKeyTTL: "60s" },
+    //   },
+    // ]);
+    // logger.log("Results 2", { results2 });
+
+    // const results3 = await batch.triggerByTask([
+    //   {
+    //     task: childTask,
+    //     payload: { message: "Hello, world !" },
+    //     options: { idempotencyKey: "1", idempotencyKeyTTL: "60s" },
+    //   },
+    //   {
+    //     task: childTask,
+    //     payload: { message: "Hello, world 2!" },
+    //     options: { idempotencyKey: "2", idempotencyKeyTTL: "60s" },
+    //   },
+    // ]);
+    // logger.log("Results 3", { results3 });
+
+    // const results4 = await batch.triggerByTaskAndWait([
+    //   { task: childTask, payload: { message: "Hello, world !" } },
+    //   { task: childTask, payload: { message: "Hello, world 2!" } },
+    // ]);
+    // logger.log("Results 4", { results4 });
+  },
+});