fix dependency resumes, properly this time

nicktrn · nicktrn · commit 8cc207e4171d · 2024-04-15T16:37:22.000+01:00
diff --git a/apps/webapp/app/v3/services/resumeAttempt.server.ts b/apps/webapp/app/v3/services/resumeAttempt.server.ts
@@ -4,12 +4,13 @@ import {
   TaskRunExecution,
   TaskRunExecutionResult,
 } from "@trigger.dev/core/v3";
-import { $transaction } from "~/db.server";
+import { $transaction, PrismaClientOrTransaction } from "~/db.server";
 import { logger } from "~/services/logger.server";
 import { marqs } from "~/v3/marqs/index.server";
 import { socketIo } from "../handleSocketIo.server";
 import { sharedQueueTasks } from "../marqs/sharedQueueConsumer.server";
 import { BaseService } from "./baseService.server";
+import { TaskRunAttempt } from "@trigger.dev/database";
 
 export class ResumeAttemptService extends BaseService {
   public async call(
@@ -24,7 +25,7 @@ export class ResumeAttemptService extends BaseService {
         },
         include: {
           taskRun: true,
-          taskRunDependency: {
+          dependencies: {
             include: {
               taskRun: {
                 include: {
@@ -40,8 +41,12 @@ export class ResumeAttemptService extends BaseService {
                 },
               },
             },
+            orderBy: {
+              createdAt: "desc",
+            },
+            take: 1,
           },
-          batchTaskRunDependency: {
+          batchDependencies: {
             include: {
               items: {
                 include: {
@@ -61,6 +66,10 @@ export class ResumeAttemptService extends BaseService {
                 },
               },
             },
+            orderBy: {
+              createdAt: "desc",
+            },
+            take: 1,
           },
         },
       });
@@ -78,6 +87,8 @@ export class ResumeAttemptService extends BaseService {
         return;
       }
 
+      let completedAttemptIds: string[] = [];
+
       switch (params.type) {
         case "WAIT_FOR_DURATION": {
           logger.error(
@@ -93,148 +104,140 @@ export class ResumeAttemptService extends BaseService {
           });
           break;
         }
-        case "WAIT_FOR_TASK":
-        case "WAIT_FOR_BATCH": {
-          let completedAttemptIds: string[] = [];
-
-          if (attempt.taskRunDependency) {
-            const dependentAttempt = attempt.taskRunDependency.taskRun.attempts[0];
+        case "WAIT_FOR_TASK": {
+          if (attempt.dependencies.length) {
+            // We only care about the latest dependency
+            const dependentAttempt = attempt.dependencies[0].taskRun.attempts[0];
 
             if (!dependentAttempt) {
               logger.error("No dependent attempt", { attemptId: attempt.id });
               return;
             }
 
             completedAttemptIds = [dependentAttempt.id];
-
-            await tx.taskRunAttempt.update({
-              where: {
-                id: attempt.id,
-              },
-              data: {
-                taskRunDependency: {
-                  disconnect: true,
-                },
-              },
-            });
-          } else if (attempt.batchTaskRunDependency) {
-            const dependentBatchItems = attempt.batchTaskRunDependency.items;
+          } else {
+            logger.error("No task dependency", { attemptId: attempt.id });
+            return;
+          }
+          break;
+        }
+        case "WAIT_FOR_BATCH": {
+          if (attempt.batchDependencies) {
+            // We only care about the latest batch dependency
+            const dependentBatchItems = attempt.batchDependencies[0].items;
 
             if (!dependentBatchItems) {
               logger.error("No dependent batch items", { attemptId: attempt.id });
               return;
             }
 
             completedAttemptIds = dependentBatchItems.map((item) => item.taskRun.attempts[0]?.id);
-
-            await tx.taskRunAttempt.update({
-              where: {
-                id: attempt.id,
-              },
-              data: {
-                batchTaskRunDependency: {
-                  disconnect: true,
-                },
-              },
-            });
           } else {
-            logger.error("No dependencies", { attemptId: attempt.id });
-            return;
-          }
-
-          if (completedAttemptIds.length === 0) {
-            logger.error("No completed attempt IDs", { attemptId: attempt.id });
+            logger.error("No batch dependency", { attemptId: attempt.id });
             return;
           }
+          break;
+        }
+        default: {
+          break;
+        }
+      }
 
-          const completions: TaskRunExecutionResult[] = [];
-          const executions: TaskRunExecution[] = [];
+      await this.#handleDependencyResume(attempt, completedAttemptIds, tx);
+    });
+  }
 
-          for (const completedAttemptId of completedAttemptIds) {
-            const completedAttempt = await tx.taskRunAttempt.findUnique({
-              where: {
-                id: completedAttemptId,
-                taskRun: {
-                  lockedAt: {
-                    not: null,
-                  },
-                  lockedById: {
-                    not: null,
-                  },
-                },
-              },
-            });
-
-            if (!completedAttempt) {
-              logger.error("Completed attempt not found", {
-                attemptId: attempt.id,
-                completedAttemptId,
-              });
-              await marqs?.acknowledgeMessage(attempt.taskRunId);
-              return;
-            }
+  async #handleDependencyResume(
+    attempt: TaskRunAttempt,
+    completedAttemptIds: string[],
+    tx: PrismaClientOrTransaction
+  ) {
+    if (completedAttemptIds.length === 0) {
+      logger.error("No completed attempt IDs", { attemptId: attempt.id });
+      return;
+    }
+
+    const completions: TaskRunExecutionResult[] = [];
+    const executions: TaskRunExecution[] = [];
+
+    for (const completedAttemptId of completedAttemptIds) {
+      const completedAttempt = await tx.taskRunAttempt.findUnique({
+        where: {
+          id: completedAttemptId,
+          taskRun: {
+            lockedAt: {
+              not: null,
+            },
+            lockedById: {
+              not: null,
+            },
+          },
+        },
+      });
 
-            const completion = await sharedQueueTasks.getCompletionPayloadFromAttempt(
-              completedAttempt.id
-            );
+      if (!completedAttempt) {
+        logger.error("Completed attempt not found", {
+          attemptId: attempt.id,
+          completedAttemptId,
+        });
+        await marqs?.acknowledgeMessage(attempt.taskRunId);
+        return;
+      }
 
-            if (!completion) {
-              logger.error("Failed to get completion payload", {
-                attemptId: attempt.id,
-                completedAttemptId,
-              });
-              await marqs?.acknowledgeMessage(attempt.taskRunId);
-              return;
-            }
+      const completion = await sharedQueueTasks.getCompletionPayloadFromAttempt(
+        completedAttempt.id
+      );
 
-            completions.push(completion);
+      if (!completion) {
+        logger.error("Failed to get completion payload", {
+          attemptId: attempt.id,
+          completedAttemptId,
+        });
+        await marqs?.acknowledgeMessage(attempt.taskRunId);
+        return;
+      }
 
-            const executionPayload = await sharedQueueTasks.getExecutionPayloadFromAttempt(
-              completedAttempt.id
-            );
+      completions.push(completion);
 
-            if (!executionPayload) {
-              logger.error("Failed to get execution payload", {
-                attemptId: attempt.id,
-                completedAttemptId,
-              });
-              await marqs?.acknowledgeMessage(attempt.taskRunId);
-              return;
-            }
+      const executionPayload = await sharedQueueTasks.getExecutionPayloadFromAttempt(
+        completedAttempt.id
+      );
 
-            executions.push(executionPayload.execution);
-          }
+      if (!executionPayload) {
+        logger.error("Failed to get execution payload", {
+          attemptId: attempt.id,
+          completedAttemptId,
+        });
+        await marqs?.acknowledgeMessage(attempt.taskRunId);
+        return;
+      }
 
-          const updated = await tx.taskRunAttempt.update({
-            where: {
-              id: attempt.id,
-            },
+      executions.push(executionPayload.execution);
+    }
+
+    const updated = await tx.taskRunAttempt.update({
+      where: {
+        id: attempt.id,
+      },
+      data: {
+        status: "EXECUTING",
+        taskRun: {
+          update: {
             data: {
-              status: "EXECUTING",
-              taskRun: {
-                update: {
-                  data: {
-                    status: attempt.number > 1 ? "RETRYING_AFTER_FAILURE" : "EXECUTING",
-                  },
-                },
-              },
+              status: attempt.number > 1 ? "RETRYING_AFTER_FAILURE" : "EXECUTING",
             },
-          });
+          },
+        },
+      },
+    });
 
-          socketIo.coordinatorNamespace.emit("RESUME_AFTER_DEPENDENCY", {
-            version: "v1",
-            runId: attempt.taskRunId,
-            attemptId: attempt.id,
-            attemptFriendlyId: attempt.friendlyId,
-            completions,
-            executions,
-          });
-          break;
-        }
-        default: {
-          break;
-        }
-      }
+    socketIo.coordinatorNamespace.emit("RESUME_AFTER_DEPENDENCY", {
+      version: "v1",
+      runId: attempt.taskRunId,
+      attemptId: attempt.id,
+      attemptFriendlyId: attempt.friendlyId,
+      completions,
+      executions,
     });
   }
 }
diff --git a/packages/database/prisma/migrations/20240415135132_make_attempt_dependencies_an_array/migration.sql b/packages/database/prisma/migrations/20240415135132_make_attempt_dependencies_an_array/migration.sql
@@ -0,0 +1,2 @@
+-- DropIndex
+DROP INDEX "TaskRunDependency_dependentAttemptId_key";
diff --git a/packages/database/prisma/migrations/20240415143325_remove_attempt_unique_constraint_from_batch_runs/migration.sql b/packages/database/prisma/migrations/20240415143325_remove_attempt_unique_constraint_from_batch_runs/migration.sql
@@ -0,0 +1,2 @@
+-- DropIndex
+DROP INDEX "BatchTaskRun_dependentTaskAttemptId_key";
diff --git a/packages/database/prisma/schema.prisma b/packages/database/prisma/schema.prisma
@@ -1667,8 +1667,8 @@ model TaskRunDependency {
   checkpointEventId String?                 @unique
 
   /// An attempt that is dependent on this task run. 
-  dependentAttempt   TaskRunAttempt? @relation("dependentAttempt", fields: [dependentAttemptId], references: [id])
-  dependentAttemptId String?         @unique
+  dependentAttempt   TaskRunAttempt? @relation(fields: [dependentAttemptId], references: [id])
+  dependentAttemptId String?
 
   /// A batch run that is dependent on this task run
   dependentBatchRun   BatchTaskRun? @relation("dependentBatchRun", fields: [dependentBatchRunId], references: [id])
@@ -1732,8 +1732,8 @@ model TaskRunAttempt {
   output     String?
   outputType String  @default("application/json")
 
-  taskRunDependency      TaskRunDependency? @relation("dependentAttempt")
-  batchTaskRunDependency BatchTaskRun?
+  dependencies      TaskRunDependency[]
+  batchDependencies BatchTaskRun[]
 
   checkpoints            Checkpoint[]
   batchTaskRunItems      BatchTaskRunItem[]
@@ -1904,7 +1904,7 @@ model BatchTaskRun {
   runtimeEnvironmentId String
 
   dependentTaskAttempt   TaskRunAttempt? @relation(fields: [dependentTaskAttemptId], references: [id], onDelete: Cascade, onUpdate: Cascade)
-  dependentTaskAttemptId String?         @unique
+  dependentTaskAttemptId String?
 
   items           BatchTaskRunItem[]
   runDependencies TaskRunDependency[] @relation("dependentBatchRun")

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+-- DropIndex`
	`2`	`+DROP INDEX "TaskRunDependency_dependentAttemptId_key";`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+-- DropIndex`
	`2`	`+DROP INDEX "BatchTaskRun_dependentTaskAttemptId_key";`