unify failed and crashed run retries

nicktrn · nicktrn · commit 78ef8342602d · 2024-10-21T15:05:44.000+01:00
diff --git a/apps/webapp/app/v3/failedTaskRun.server.ts b/apps/webapp/app/v3/failedTaskRun.server.ts
@@ -36,12 +36,11 @@ export class FailedTaskRunService extends BaseService {
 
     const isFriendlyId = anyRunId.startsWith("run_");
 
-    const taskRun = await this._prisma.taskRun.findUnique({
+    const taskRun = await this._prisma.taskRun.findFirst({
       where: {
         friendlyId: isFriendlyId ? anyRunId : undefined,
         id: !isFriendlyId ? anyRunId : undefined,
       },
-      include: includeAttempts,
     });
 
     if (!taskRun) {
@@ -62,24 +61,13 @@ export class FailedTaskRunService extends BaseService {
       return;
     }
 
-    const retriableExecution = await this.#getRetriableAttemptExecution(taskRun, completion);
-
-    if (retriableExecution) {
-      logger.debug("[FailedTaskRunService] Completing attempt", { taskRun, completion });
-
-      const executionRetry =
-        completion.retry ?? (await this.#getExecutionRetry(taskRun, retriableExecution));
-
-      const completeAttempt = new CompleteAttemptService(this._prisma);
-      await completeAttempt.call({
-        completion: {
-          ...completion,
-          retry: executionRetry,
-        },
-        execution: retriableExecution,
-        isSystemFailure: true,
-      });
+    const retryHelper = new FailedTaskRunRetryHelper(this._prisma);
+    const retryResult = await retryHelper.call({
+      runId: taskRun.id,
+      completion,
+    });
 
+    if (retryResult !== undefined) {
       return;
     }
 
@@ -112,6 +100,58 @@ export class FailedTaskRunService extends BaseService {
       ],
     });
   }
+}
+
+export class FailedTaskRunRetryHelper extends BaseService {
+  async call({
+    runId,
+    completion,
+    isCrash,
+  }: {
+    runId: string;
+    completion: TaskRunFailedExecutionResult;
+    isCrash?: boolean;
+  }) {
+    const taskRun = await this._prisma.taskRun.findFirst({
+      where: {
+        id: runId,
+      },
+      include: includeAttempts,
+    });
+
+    if (!taskRun) {
+      logger.error("[FailedTaskRunRetryHelper] Task run not found", {
+        runId,
+        completion,
+      });
+
+      return;
+    }
+
+    const retriableExecution = await this.#getRetriableAttemptExecution(taskRun, completion);
+
+    if (!retriableExecution) {
+      return;
+    }
+
+    logger.debug("[FailedTaskRunRetryHelper] Completing attempt", { taskRun, completion });
+
+    const executionRetry =
+      completion.retry ?? (await this.#getExecutionRetry(taskRun, retriableExecution));
+
+    const completeAttempt = new CompleteAttemptService(this._prisma);
+    const completeResult = await completeAttempt.call({
+      completion: {
+        ...completion,
+        retry: executionRetry,
+      },
+      execution: retriableExecution,
+      isSystemFailure: !isCrash,
+      isCrash,
+    });
+
+    return completeResult;
+  }
 
   async #getRetriableAttemptExecution(
     run: TaskRunWithAttempts,
@@ -123,7 +163,7 @@ export class FailedTaskRunService extends BaseService {
     // - None exists yet
     // - The last attempt has a final status, e.g. we failed between attempts
     if (!attempt || isFinalAttemptStatus(attempt.status)) {
-      logger.error("[FailedTaskRunService] No attempts found", {
+      logger.error("[FailedTaskRunRetryHelper] No attempts found", {
         run,
         completion,
       });
@@ -134,7 +174,7 @@ export class FailedTaskRunService extends BaseService {
         const { execution } = await createAttempt.call(run.id);
         return execution;
       } catch (error) {
-        logger.error("[FailedTaskRunService] Failed to create attempt", {
+        logger.error("[FailedTaskRunRetryHelper] Failed to create attempt", {
           run,
           completion,
           error,
@@ -153,7 +193,7 @@ export class FailedTaskRunService extends BaseService {
 
       return executionPayload?.execution;
     } catch (error) {
-      logger.error("[FailedTaskRunService] Failed to get execution payload", {
+      logger.error("[FailedTaskRunRetryHelper] Failed to get execution payload", {
         run,
         completion,
         error,
@@ -170,7 +210,7 @@ export class FailedTaskRunService extends BaseService {
     const parsedRetryConfig = RetryOptions.safeParse(run.lockedBy?.retryConfig);
 
     if (!parsedRetryConfig.success) {
-      logger.error("[FailedTaskRunService] Invalid retry config", {
+      logger.error("[FailedTaskRunRetryHelper] Invalid retry config", {
         run,
         execution,
       });
@@ -181,7 +221,7 @@ export class FailedTaskRunService extends BaseService {
     const delay = calculateNextRetryDelay(parsedRetryConfig.data, execution.attempt.number);
 
     if (!delay) {
-      logger.debug("[FailedTaskRunService] No more retries", {
+      logger.debug("[FailedTaskRunRetryHelper] No more retries", {
         run,
         execution,
       });
diff --git a/apps/webapp/app/v3/services/completeAttempt.server.ts b/apps/webapp/app/v3/services/completeAttempt.server.ts
@@ -22,7 +22,7 @@ import { MAX_TASK_RUN_ATTEMPTS } from "~/consts";
 import { CreateCheckpointService } from "./createCheckpoint.server";
 import { TaskRun } from "@trigger.dev/database";
 import { RetryAttemptService } from "./retryAttempt.server";
-import { isFinalAttemptStatus, isFinalRunStatus } from "../taskStatus";
+import { FAILED_RUN_STATUSES, isFinalAttemptStatus, isFinalRunStatus } from "../taskStatus";
 import { FinalizeTaskRunService } from "./finalizeTaskRun.server";
 import { env } from "~/env.server";
 
@@ -41,13 +41,15 @@ export class CompleteAttemptService extends BaseService {
     checkpoint,
     supportsRetryCheckpoints,
     isSystemFailure,
+    isCrash,
   }: {
     completion: TaskRunExecutionResult;
     execution: TaskRunExecution;
     env?: AuthenticatedEnvironment;
     checkpoint?: CheckpointData;
     supportsRetryCheckpoints?: boolean;
     isSystemFailure?: boolean;
+    isCrash?: boolean;
   }): Promise<"COMPLETED" | "RETRIED"> {
     const taskRunAttempt = await findAttempt(this._prisma, execution.attempt.id);
 
@@ -114,6 +116,7 @@ export class CompleteAttemptService extends BaseService {
         checkpoint,
         supportsRetryCheckpoints,
         isSystemFailure,
+        isCrash,
       });
     }
   }
@@ -175,6 +178,7 @@ export class CompleteAttemptService extends BaseService {
     checkpoint,
     supportsRetryCheckpoints,
     isSystemFailure,
+    isCrash,
   }: {
     completion: TaskRunFailedExecutionResult;
     execution: TaskRunExecution;
@@ -183,6 +187,7 @@ export class CompleteAttemptService extends BaseService {
     checkpoint?: CheckpointData;
     supportsRetryCheckpoints?: boolean;
     isSystemFailure?: boolean;
+    isCrash?: boolean;
   }): Promise<"COMPLETED" | "RETRIED"> {
     if (
       completion.error.type === "INTERNAL_ERROR" &&
@@ -260,11 +265,20 @@ export class CompleteAttemptService extends BaseService {
       },
     });
 
-    const status = isSystemFailure
-      ? "SYSTEM_FAILURE"
-      : sanitizedError.type === "INTERNAL_ERROR" && sanitizedError.code === "MAX_DURATION_EXCEEDED"
-      ? "TIMED_OUT"
-      : "COMPLETED_WITH_ERRORS";
+    let status: FAILED_RUN_STATUSES;
+
+    if (isSystemFailure) {
+      status = "SYSTEM_FAILURE";
+    } else if (isCrash) {
+      status = "CRASHED";
+    } else if (
+      sanitizedError.type === "INTERNAL_ERROR" &&
+      sanitizedError.code === "MAX_DURATION_EXCEEDED"
+    ) {
+      status = "TIMED_OUT";
+    } else {
+      status = "COMPLETED_WITH_ERRORS";
+    }
 
     const finalizeService = new FinalizeTaskRunService();
     await finalizeService.call({
diff --git a/apps/webapp/app/v3/services/crashTaskRun.server.ts b/apps/webapp/app/v3/services/crashTaskRun.server.ts
@@ -6,6 +6,7 @@ import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { CRASHABLE_ATTEMPT_STATUSES, isCrashableRunStatus } from "../taskStatus";
 import { sanitizeError, TaskRunInternalError } from "@trigger.dev/core/v3";
 import { FinalizeTaskRunService } from "./finalizeTaskRun.server";
+import { FailedTaskRunRetryHelper } from "../failedTaskRun.server";
 
 export type CrashTaskRunServiceOptions = {
   reason?: string;
@@ -35,16 +36,43 @@ export class CrashTaskRunService extends BaseService {
     });
 
     if (!taskRun) {
-      logger.error("Task run not found", { runId });
+      logger.error("[CrashTaskRunService] Task run not found", { runId });
       return;
     }
 
     // Make sure the task run is in a crashable state
     if (!opts.overrideCompletion && !isCrashableRunStatus(taskRun.status)) {
-      logger.error("Task run is not in a crashable state", { runId, status: taskRun.status });
+      logger.error("[CrashTaskRunService] Task run is not in a crashable state", {
+        runId,
+        status: taskRun.status,
+      });
+      return;
+    }
+
+    logger.debug("[CrashTaskRunService] Completing attempt", { runId, options });
+
+    const retryHelper = new FailedTaskRunRetryHelper(this._prisma);
+    const retryResult = await retryHelper.call({
+      runId,
+      completion: {
+        ok: false,
+        id: runId,
+        error: {
+          type: "INTERNAL_ERROR",
+          code: opts.errorCode ?? "TASK_RUN_CRASHED",
+          message: opts.reason,
+          stackTrace: opts.logs,
+        },
+      },
+    });
+
+    if (retryResult === "RETRIED") {
+      logger.debug("[CrashTaskRunService] Retried task run", { runId });
       return;
     }
 
+    logger.debug("[CrashTaskRunService] Overriding completion", { runId, options });
+
     const finalizeService = new FinalizeTaskRunService();
     const crashedTaskRun = await finalizeService.call({
       id: taskRun.id,
@@ -86,7 +114,7 @@ export class CrashTaskRunService extends BaseService {
       options?.overrideCompletion
     );
 
-    logger.debug("Crashing in-progress events", {
+    logger.debug("[CrashTaskRunService] Crashing in-progress events", {
       inProgressEvents: inProgressEvents.map((event) => event.id),
     });
 
@@ -135,25 +163,29 @@ export class CrashTaskRunService extends BaseService {
       code?: TaskRunInternalError["code"];
     }
   ) {
-    return await this.traceWithEnv("failAttempt()", environment, async (span) => {
-      span.setAttribute("taskRunId", run.id);
-      span.setAttribute("attemptId", attempt.id);
+    return await this.traceWithEnv(
+      "[CrashTaskRunService] failAttempt()",
+      environment,
+      async (span) => {
+        span.setAttribute("taskRunId", run.id);
+        span.setAttribute("attemptId", attempt.id);
 
-      await this._prisma.taskRunAttempt.update({
-        where: {
-          id: attempt.id,
-        },
-        data: {
-          status: "FAILED",
-          completedAt: failedAt,
-          error: sanitizeError({
-            type: "INTERNAL_ERROR",
-            code: error.code ?? "TASK_RUN_CRASHED",
-            message: error.reason,
-            stackTrace: error.logs,
-          }),
-        },
-      });
-    });
+        await this._prisma.taskRunAttempt.update({
+          where: {
+            id: attempt.id,
+          },
+          data: {
+            status: "FAILED",
+            completedAt: failedAt,
+            error: sanitizeError({
+              type: "INTERNAL_ERROR",
+              code: error.code ?? "TASK_RUN_CRASHED",
+              message: error.reason,
+              stackTrace: error.logs,
+            }),
+          },
+        });
+      }
+    );
   }
 }
diff --git a/apps/webapp/app/v3/taskStatus.ts b/apps/webapp/app/v3/taskStatus.ts
@@ -49,6 +49,8 @@ export const FAILED_RUN_STATUSES = [
   "TIMED_OUT",
 ] satisfies TaskRunStatus[];
 
+export type FAILED_RUN_STATUSES = (typeof FAILED_RUN_STATUSES)[number];
+
 export const CANCELLABLE_RUN_STATUSES = NON_FINAL_RUN_STATUSES;
 export const CANCELLABLE_ATTEMPT_STATUSES = NON_FINAL_ATTEMPT_STATUSES;
 
