chore: import upstream snapshot with attribution

2026-07-13 13:32:57 +08:00
commit cd420f9332
4811 changed files with 884702 additions and 0 deletions
@@ -0,0 +1,399 @@
+import type { PrismaClientOrTransaction } from "~/db.server";
+import { commonWorker } from "../commonWorker.server";
+import { marqs } from "~/v3/marqs/index.server";
+import { BaseService } from "./baseService.server";
+import { logger } from "~/services/logger.server";
+import type { BatchTaskRun, Prisma } from "@trigger.dev/database";
+import { findEnvironmentById } from "~/models/runtimeEnvironment.server";
+import type { AuthenticatedEnvironment } from "~/services/apiAuth.server";
+import { workerQueue } from "~/services/worker.server";
+import { isV3Disabled } from "../engineDeprecation.server";
+
+const finishedBatchRunStatuses = ["COMPLETED", "FAILED", "CANCELED"];
+
+const BATCH_RUN_INCLUDE = {
+  items: {
+    select: {
+      status: true,
+      taskRunAttemptId: true,
+    },
+  },
+} satisfies Prisma.BatchTaskRunInclude;
+
+type RetrieveBatchRunResult = Prisma.BatchTaskRunGetPayload<{
+  include: typeof BATCH_RUN_INCLUDE;
+}>;
+
+export class ResumeBatchRunService extends BaseService {
+  public async call(batchRunId: string) {
+    const batchRun = await this.runStore.findBatchTaskRunById(batchRunId, {
+      include: BATCH_RUN_INCLUDE,
+    });
+
+    if (!batchRun) {
+      logger.error(
+        "ResumeBatchRunService: Batch run doesn't exist or doesn't have a dependent attempt",
+        {
+          batchRunId,
+        }
+      );
+
+      return "ERROR";
+    }
+
+    // BatchTaskRun -> RuntimeEnvironment FK is dropped; resolve the env from the scalar id.
+    const environment = await findEnvironmentById(batchRun.runtimeEnvironmentId);
+    if (!environment) {
+      logger.error("ResumeBatchRunService: Environment not found", {
+        batchRunId,
+        runtimeEnvironmentId: batchRun.runtimeEnvironmentId,
+      });
+
+      return "ERROR";
+    }
+
+    // v3 (engine V1) shutdown: don't resume batches for abandoned V1 projects. v4 is unaffected.
+    // The BatchTaskRun -> RuntimeEnvironment relation is dropped, so read the engine from the
+    // resolved environment's project rather than the unloaded batchRun.runtimeEnvironment relation.
+    if (isV3Disabled() && environment.project.engine === "V1") {
+      logger.debug("[ResumeBatchRunService] Skipping resume for shut-down v3 batch", {
+        batchRunId,
+      });
+      return "ERROR";
+    }
+
+    if (batchRun.batchVersion === "v3") {
+      return await this.#handleV3BatchRun(batchRun, environment);
+    } else {
+      return await this.#handleLegacyBatchRun(batchRun, environment);
+    }
+  }
+
+  async #handleV3BatchRun(batchRun: RetrieveBatchRunResult, environment: AuthenticatedEnvironment) {
+    // V3 batch runs should already be complete by the time this is called
+    if (batchRun.status !== "COMPLETED") {
+      logger.debug("ResumeBatchRunService: Batch run is already completed", {
+        batchRunId: batchRun.id,
+        batchRun: {
+          id: batchRun.id,
+          status: batchRun.status,
+        },
+      });
+
+      return "ERROR";
+    }
+
+    // Even though we are in v3, we still need to check if the batch run has a dependent attempt
+    if (!batchRun.dependentTaskAttemptId) {
+      logger.debug("ResumeBatchRunService: Batch run doesn't have a dependent attempt", {
+        batchRunId: batchRun.id,
+      });
+
+      return "ERROR";
+    }
+
+    return await this.#handleDependentTaskAttempt(
+      batchRun,
+      batchRun.dependentTaskAttemptId,
+      environment
+    );
+  }
+
+  async #handleLegacyBatchRun(
+    batchRun: RetrieveBatchRunResult,
+    environment: AuthenticatedEnvironment
+  ) {
+    if (batchRun.status === "COMPLETED") {
+      logger.debug("ResumeBatchRunService: Batch run is already completed", {
+        batchRunId: batchRun.id,
+        batchRun: {
+          id: batchRun.id,
+          status: batchRun.status,
+        },
+      });
+
+      return "ERROR";
+    }
+
+    if (batchRun.batchVersion === "v2") {
+      if (batchRun.items.length < batchRun.runCount) {
+        logger.debug("ResumeBatchRunService: All items aren't yet completed [v2]", {
+          batchRunId: batchRun.id,
+          batchRun: {
+            id: batchRun.id,
+            status: batchRun.status,
+            itemsLength: batchRun.items.length,
+            runCount: batchRun.runCount,
+          },
+        });
+
+        return "PENDING";
+      }
+    }
+
+    if (batchRun.items.some((item) => !finishedBatchRunStatuses.includes(item.status))) {
+      logger.debug("ResumeBatchRunService: All items aren't yet completed [v1]", {
+        batchRunId: batchRun.id,
+        batchRun: {
+          id: batchRun.id,
+          status: batchRun.status,
+        },
+      });
+
+      return "PENDING";
+    }
+
+    // If we are in development, or there is no dependent attempt, we can just mark the batch as completed and return
+    if (environment.type === "DEVELOPMENT" || !batchRun.dependentTaskAttemptId) {
+      // We need to update the batchRun status so we don't resume it again
+      await this.runStore.updateBatchTaskRun({
+        where: {
+          id: batchRun.id,
+        },
+        data: {
+          status: "COMPLETED",
+        },
+        select: { id: true },
+      });
+
+      return "COMPLETED";
+    }
+
+    return await this.#handleDependentTaskAttempt(
+      batchRun,
+      batchRun.dependentTaskAttemptId,
+      environment
+    );
+  }
+
+  async #handleDependentTaskAttempt(
+    batchRun: RetrieveBatchRunResult,
+    dependentTaskAttemptId: string,
+    environment: AuthenticatedEnvironment
+  ) {
+    const dependentTaskAttempt = await this._prisma.taskRunAttempt.findFirst({
+      where: {
+        id: dependentTaskAttemptId,
+      },
+      select: {
+        status: true,
+        id: true,
+        taskRun: {
+          select: {
+            id: true,
+            queue: true,
+            taskIdentifier: true,
+            concurrencyKey: true,
+            createdAt: true,
+            queueTimestamp: true,
+          },
+        },
+      },
+    });
+
+    if (!dependentTaskAttempt) {
+      logger.error("ResumeBatchRunService: Dependent attempt not found", {
+        batchRunId: batchRun.id,
+        dependentTaskAttemptId: batchRun.dependentTaskAttemptId,
+      });
+
+      return "ERROR";
+    }
+
+    // This batch has a dependent attempt and just finalized, we should resume that attempt
+    const dependentRun = dependentTaskAttempt.taskRun;
+
+    if (dependentTaskAttempt.status === "PAUSED" && batchRun.checkpointEventId) {
+      logger.debug("ResumeBatchRunService: Attempt is paused and has a checkpoint event", {
+        batchRunId: batchRun.id,
+        dependentTaskAttempt: dependentTaskAttempt,
+        checkpointEventId: batchRun.checkpointEventId,
+      });
+
+      // We need to update the batchRun status so we don't resume it again
+      const wasUpdated = await this.#setBatchToResumedOnce(batchRun);
+
+      if (wasUpdated) {
+        logger.debug("ResumeBatchRunService: Resuming dependent run with checkpoint", {
+          batchRunId: batchRun.id,
+          dependentTaskAttemptId: dependentTaskAttempt.id,
+        });
+
+        await marqs.enqueueMessage(
+          environment,
+          dependentRun.queue,
+          dependentRun.id,
+          {
+            type: "RESUME",
+            completedAttemptIds: [],
+            resumableAttemptId: dependentTaskAttempt.id,
+            checkpointEventId: batchRun.checkpointEventId,
+            taskIdentifier: dependentTaskAttempt.taskRun.taskIdentifier,
+            projectId: environment.projectId,
+            environmentId: environment.id,
+            environmentType: environment.type,
+          },
+          dependentRun.concurrencyKey ?? undefined,
+          dependentRun.queueTimestamp ?? dependentRun.createdAt,
+          undefined,
+          "resume"
+        );
+
+        return "COMPLETED";
+      } else {
+        logger.debug("ResumeBatchRunService: with checkpoint was already completed", {
+          batchRunId: batchRun.id,
+          dependentTaskAttempt: dependentTaskAttempt,
+          checkpointEventId: batchRun.checkpointEventId,
+          hasCheckpointEvent: !!batchRun.checkpointEventId,
+        });
+
+        return "ALREADY_COMPLETED";
+      }
+    } else {
+      logger.debug("ResumeBatchRunService: attempt is not paused or there's no checkpoint event", {
+        batchRunId: batchRun.id,
+        dependentTaskAttempt: dependentTaskAttempt,
+        checkpointEventId: batchRun.checkpointEventId,
+        hasCheckpointEvent: !!batchRun.checkpointEventId,
+      });
+
+      if (dependentTaskAttempt.status === "PAUSED" && !batchRun.checkpointEventId) {
+        // In case of race conditions the status can be PAUSED without a checkpoint event
+        // When the checkpoint is created, it will continue the run
+        logger.error("ResumeBatchRunService: attempt is paused but there's no checkpoint event", {
+          batchRunId: batchRun.id,
+          dependentTaskAttempt: dependentTaskAttempt,
+          checkpointEventId: batchRun.checkpointEventId,
+          hasCheckpointEvent: !!batchRun.checkpointEventId,
+        });
+
+        return "ERROR";
+      }
+
+      // We need to update the batchRun status so we don't resume it again
+      const wasUpdated = await this.#setBatchToResumedOnce(batchRun);
+
+      if (wasUpdated) {
+        logger.debug("ResumeBatchRunService: Resuming dependent run without checkpoint", {
+          batchRunId: batchRun.id,
+          dependentTaskAttempt: dependentTaskAttempt,
+          checkpointEventId: batchRun.checkpointEventId,
+          hasCheckpointEvent: !!batchRun.checkpointEventId,
+        });
+
+        await marqs.requeueMessage(
+          dependentRun.id,
+          {
+            type: "RESUME",
+            completedAttemptIds: batchRun.items
+              .map((item) => item.taskRunAttemptId)
+              .filter(Boolean),
+            resumableAttemptId: dependentTaskAttempt.id,
+            checkpointEventId: batchRun.checkpointEventId ?? undefined,
+            taskIdentifier: dependentTaskAttempt.taskRun.taskIdentifier,
+            projectId: environment.projectId,
+            environmentId: environment.id,
+            environmentType: environment.type,
+          },
+          (
+            dependentTaskAttempt.taskRun.queueTimestamp ?? dependentTaskAttempt.taskRun.createdAt
+          ).getTime(),
+          "resume"
+        );
+
+        return "COMPLETED";
+      } else {
+        logger.debug("ResumeBatchRunService: without checkpoint was already completed", {
+          batchRunId: batchRun.id,
+          dependentTaskAttempt: dependentTaskAttempt,
+          checkpointEventId: batchRun.checkpointEventId,
+          hasCheckpointEvent: !!batchRun.checkpointEventId,
+        });
+
+        return "ALREADY_COMPLETED";
+      }
+    }
+  }
+
+  async #setBatchToResumedOnce(batchRun: BatchTaskRun) {
+    // v3 batches don't use the status for deciding whether a batch has been resumed
+    if (batchRun.batchVersion === "v3") {
+      const result = await this.runStore.updateManyBatchTaskRun({
+        where: {
+          id: batchRun.id,
+          resumedAt: null,
+        },
+        data: {
+          resumedAt: new Date(),
+        },
+      });
+
+      if (result.count > 0) {
+        return true;
+      } else {
+        return false;
+      }
+    }
+
+    const result = await this.runStore.updateManyBatchTaskRun({
+      where: {
+        id: batchRun.id,
+        status: {
+          not: "COMPLETED", // Ensure the status is not already "COMPLETED"
+        },
+      },
+      data: {
+        status: "COMPLETED",
+      },
+    });
+
+    if (result.count > 0) {
+      return true;
+    } else {
+      return false;
+    }
+  }
+
+  static async enqueue(
+    batchRunId: string,
+    skipJobKey: boolean,
+    tx?: PrismaClientOrTransaction,
+    runAt?: Date
+  ) {
+    if (tx) {
+      logger.debug("ResumeBatchRunService: Enqueuing resume batch run using workerQueue", {
+        batchRunId,
+        skipJobKey,
+        runAt,
+      });
+
+      return await workerQueue.enqueue(
+        "v3.resumeBatchRun",
+        {
+          batchRunId,
+        },
+        {
+          jobKey: skipJobKey ? undefined : `resumeBatchRun-${batchRunId}`,
+          runAt,
+          tx,
+        }
+      );
+    } else {
+      logger.debug("ResumeBatchRunService: Enqueuing resume batch run using commonWorker", {
+        batchRunId,
+        skipJobKey,
+        runAt,
+      });
+
+      return await commonWorker.enqueue({
+        id: skipJobKey ? undefined : `resumeBatchRun-${batchRunId}`,
+        job: "v3.resumeBatchRun",
+        payload: {
+          batchRunId,
+        },
+        availableAt: runAt,
+      });
+    }
+  }
+}