From 55d1f8c67790639b8ff7a26f86028d2e850a123e Mon Sep 17 00:00:00 2001 From: nicktrn <55853254+nicktrn@users.noreply.github.com> Date: Thu, 20 Jun 2024 10:10:39 +0100 Subject: [PATCH] v3: fix dependency checkpoint race (#1171) * consolidate task statuses and utils * check dependency completion when creating checkpoints * add changeset --- .changeset/silly-buses-obey.md | 5 ++ apps/coordinator/src/index.ts | 57 +++++++++------ .../presenters/v3/RunListPresenter.server.ts | 4 +- apps/webapp/app/v3/handleSocketIo.server.ts | 15 +++- .../v3/marqs/sharedQueueConsumer.server.ts | 16 +---- .../app/v3/services/cancelAttempt.server.ts | 4 +- .../app/v3/services/cancelTaskRun.server.ts | 25 ++----- .../app/v3/services/crashTaskRun.server.ts | 23 ++---- .../v3/services/createCheckpoint.server.ts | 70 +++++++++++++++---- .../v3/services/restoreCheckpoint.server.ts | 13 ++-- apps/webapp/app/v3/taskStatus.ts | 69 ++++++++++++++++++ packages/core/src/v3/schemas/messages.ts | 4 ++ 12 files changed, 205 insertions(+), 100 deletions(-) create mode 100644 .changeset/silly-buses-obey.md create mode 100644 apps/webapp/app/v3/taskStatus.ts diff --git a/.changeset/silly-buses-obey.md b/.changeset/silly-buses-obey.md new file mode 100644 index 000000000..a3045ed0d --- /dev/null +++ b/.changeset/silly-buses-obey.md @@ -0,0 +1,5 @@ +--- +"@trigger.dev/core": patch +--- + +Add callback to checkpoint created message diff --git a/apps/coordinator/src/index.ts b/apps/coordinator/src/index.ts index 90ea13b84..b7cab90c5 100644 --- a/apps/coordinator/src/index.ts +++ b/apps/coordinator/src/index.ts @@ -1162,13 +1162,7 @@ class TaskCoordinator { return; } - if (!checkpoint.docker || !willSimulate) { - socket.emit("REQUEST_EXIT", { - version: "v1", - }); - } - - this.#platformSocket?.send("CHECKPOINT_CREATED", { + const ack = await this.#platformSocket?.sendWithAck("CHECKPOINT_CREATED", { version: "v1", attemptFriendlyId: message.attemptFriendlyId, docker: checkpoint.docker, @@ -1179,6 +1173,17 @@ class TaskCoordinator { now: message.now, }, }); + + if (ack?.keepRunAlive) { + logger.log("keeping run alive after duration checkpoint", { runId: socket.data.runId }); + return; + } + + if (!checkpoint.docker || !willSimulate) { + socket.emit("REQUEST_EXIT", { + version: "v1", + }); + } }); socket.on("WAIT_FOR_TASK", async (message, callback) => { @@ -1205,13 +1210,7 @@ class TaskCoordinator { return; } - if (!checkpoint.docker || !willSimulate) { - socket.emit("REQUEST_EXIT", { - version: "v1", - }); - } - - this.#platformSocket?.send("CHECKPOINT_CREATED", { + const ack = await this.#platformSocket?.sendWithAck("CHECKPOINT_CREATED", { version: "v1", attemptFriendlyId: message.attemptFriendlyId, docker: checkpoint.docker, @@ -1221,6 +1220,17 @@ class TaskCoordinator { friendlyId: message.friendlyId, }, }); + + if (ack?.keepRunAlive) { + logger.log("keeping run alive after task checkpoint", { runId: socket.data.runId }); + return; + } + + if (!checkpoint.docker || !willSimulate) { + socket.emit("REQUEST_EXIT", { + version: "v1", + }); + } }); socket.on("WAIT_FOR_BATCH", async (message, callback) => { @@ -1247,13 +1257,7 @@ class TaskCoordinator { return; } - if (!checkpoint.docker || !willSimulate) { - socket.emit("REQUEST_EXIT", { - version: "v1", - }); - } - - this.#platformSocket?.send("CHECKPOINT_CREATED", { + const ack = await this.#platformSocket?.sendWithAck("CHECKPOINT_CREATED", { version: "v1", attemptFriendlyId: message.attemptFriendlyId, docker: checkpoint.docker, @@ -1264,6 +1268,17 @@ class TaskCoordinator { runFriendlyIds: message.runFriendlyIds, }, }); + + if (ack?.keepRunAlive) { + logger.log("keeping run alive after batch checkpoint", { runId: socket.data.runId }); + return; + } + + if (!checkpoint.docker || !willSimulate) { + socket.emit("REQUEST_EXIT", { + version: "v1", + }); + } }); socket.on("INDEX_TASKS", async (message, callback) => { diff --git a/apps/webapp/app/presenters/v3/RunListPresenter.server.ts b/apps/webapp/app/presenters/v3/RunListPresenter.server.ts index f0e5d2303..e041f2a4b 100644 --- a/apps/webapp/app/presenters/v3/RunListPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/RunListPresenter.server.ts @@ -4,8 +4,8 @@ import { Direction } from "~/components/runs/RunStatuses"; import { FINISHED_STATUSES } from "~/components/runs/v3/TaskRunStatus"; import { sqlDatabaseSchema } from "~/db.server"; import { displayableEnvironment } from "~/models/runtimeEnvironment.server"; -import { CANCELLABLE_STATUSES } from "~/v3/services/cancelTaskRun.server"; import { BasePresenter } from "./basePresenter.server"; +import { isCancellableRunStatus } from "~/v3/taskStatus"; export type RunListOptions = { userId?: string; @@ -291,7 +291,7 @@ export class RunListPresenter extends BasePresenter { taskIdentifier: run.taskIdentifier, spanId: run.spanId, isReplayable: true, - isCancellable: CANCELLABLE_STATUSES.includes(run.status), + isCancellable: isCancellableRunStatus(run.status), environment: displayableEnvironment(environment, userId), idempotencyKey: run.idempotencyKey ? run.idempotencyKey : undefined, }; diff --git a/apps/webapp/app/v3/handleSocketIo.server.ts b/apps/webapp/app/v3/handleSocketIo.server.ts index 417dab007..ede5eb9b4 100644 --- a/apps/webapp/app/v3/handleSocketIo.server.ts +++ b/apps/webapp/app/v3/handleSocketIo.server.ts @@ -138,8 +138,19 @@ function createCoordinatorNamespace(io: Server) { await sharedQueueTasks.taskRunHeartbeat(message.runId); }, CHECKPOINT_CREATED: async (message) => { - const createCheckpoint = new CreateCheckpointService(); - await createCheckpoint.call(message); + try { + const createCheckpoint = new CreateCheckpointService(); + const result = await createCheckpoint.call(message); + + return { keepRunAlive: result?.keepRunAlive ?? false }; + } catch (error) { + logger.error("Error while creating checkpoint", { + rawMessage: message, + error: error instanceof Error ? error.message : error, + }); + + return { keepRunAlive: false }; + } }, CREATE_WORKER: async (message) => { try { diff --git a/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts b/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts index b93234e1e..9eb1140d2 100644 --- a/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts +++ b/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts @@ -17,7 +17,6 @@ import { BackgroundWorkerTask, RuntimeEnvironment, TaskRun, - TaskRunAttemptStatus, TaskRunStatus, } from "@trigger.dev/database"; import { z } from "zod"; @@ -43,6 +42,7 @@ import { generateJWTTokenForEnvironment } from "~/services/apiAuth.server"; import { EnvironmentVariable } from "../environmentVariables/repository"; import { machinePresetFromConfig } from "../machinePresets.server"; import { env } from "~/env.server"; +import { isFinalAttemptStatus, isFinalRunStatus } from "../taskStatus"; const WithTraceContext = z.object({ traceparent: z.string().optional(), @@ -962,19 +962,7 @@ class SharedQueueTasks { } if (setToExecuting) { - const FINAL_RUN_STATUSES: TaskRunStatus[] = [ - "CANCELED", - "COMPLETED_SUCCESSFULLY", - "COMPLETED_WITH_ERRORS", - "INTERRUPTED", - "SYSTEM_FAILURE", - ]; - const FINAL_ATTEMPT_STATUSES: TaskRunAttemptStatus[] = ["CANCELED", "COMPLETED", "FAILED"]; - - if ( - FINAL_ATTEMPT_STATUSES.includes(attempt.status) || - FINAL_RUN_STATUSES.includes(attempt.taskRun.status) - ) { + if (isFinalAttemptStatus(attempt.status) || isFinalRunStatus(attempt.taskRun.status)) { logger.error("Status already in final state", { attempt: { id: attempt.id, diff --git a/apps/webapp/app/v3/services/cancelAttempt.server.ts b/apps/webapp/app/v3/services/cancelAttempt.server.ts index 2413afac5..cbcb215e9 100644 --- a/apps/webapp/app/v3/services/cancelAttempt.server.ts +++ b/apps/webapp/app/v3/services/cancelAttempt.server.ts @@ -6,7 +6,7 @@ import { logger } from "~/services/logger.server"; import { PrismaClientOrTransaction, prisma } from "~/db.server"; import { ResumeTaskRunDependenciesService } from "./resumeTaskRunDependencies.server"; -import { CANCELLABLE_STATUSES } from "./cancelTaskRun.server"; +import { isCancellableRunStatus } from "../taskStatus"; export class CancelAttemptService extends BaseService { public async call( @@ -55,7 +55,7 @@ export class CancelAttemptService extends BaseService { taskRun: { update: { data: { - status: CANCELLABLE_STATUSES.includes(taskRunAttempt.taskRun.status) + status: isCancellableRunStatus(taskRunAttempt.taskRun.status) ? "INTERRUPTED" : undefined, }, diff --git a/apps/webapp/app/v3/services/cancelTaskRun.server.ts b/apps/webapp/app/v3/services/cancelTaskRun.server.ts index a4822d854..fffd2785d 100644 --- a/apps/webapp/app/v3/services/cancelTaskRun.server.ts +++ b/apps/webapp/app/v3/services/cancelTaskRun.server.ts @@ -1,4 +1,4 @@ -import { Prisma, TaskRun, TaskRunAttemptStatus, TaskRunStatus } from "@trigger.dev/database"; +import { Prisma, TaskRun } from "@trigger.dev/database"; import assertNever from "assert-never"; import { logger } from "~/services/logger.server"; import { marqs } from "~/v3/marqs/index.server"; @@ -7,22 +7,7 @@ import { socketIo } from "../handleSocketIo.server"; import { devPubSub } from "../marqs/devPubSub.server"; import { BaseService } from "./baseService.server"; import { CancelAttemptService } from "./cancelAttempt.server"; - -export const CANCELLABLE_STATUSES: Array = [ - "PENDING", - "WAITING_FOR_DEPLOY", - "EXECUTING", - "PAUSED", - "WAITING_TO_RESUME", - "PAUSED", - "RETRYING_AFTER_FAILURE", -]; - -const CANCELLABLE_ATTEMPT_STATUSES: Array = [ - "EXECUTING", - "PAUSED", - "PENDING", -]; +import { CANCELLABLE_ATTEMPT_STATUSES, isCancellableRunStatus } from "../taskStatus"; type ExtendedTaskRun = Prisma.TaskRunGetPayload<{ include: { @@ -53,7 +38,11 @@ export class CancelTaskRunService extends BaseService { }; // Make sure the task run is in a cancellable state - if (!CANCELLABLE_STATUSES.includes(taskRun.status)) { + if (!isCancellableRunStatus(taskRun.status)) { + logger.error("Task run is not in a cancellable state", { + runId: taskRun.id, + status: taskRun.status, + }); return; } diff --git a/apps/webapp/app/v3/services/crashTaskRun.server.ts b/apps/webapp/app/v3/services/crashTaskRun.server.ts index 774ec83d8..d7241f526 100644 --- a/apps/webapp/app/v3/services/crashTaskRun.server.ts +++ b/apps/webapp/app/v3/services/crashTaskRun.server.ts @@ -1,27 +1,11 @@ -import { - TaskRun, - TaskRunAttempt, - TaskRunAttemptStatus, - TaskRunStatus, -} from "@trigger.dev/database"; +import { TaskRun, TaskRunAttempt } from "@trigger.dev/database"; import { eventRepository } from "../eventRepository.server"; import { marqs } from "~/v3/marqs/index.server"; import { BaseService } from "./baseService.server"; import { logger } from "~/services/logger.server"; import { AuthenticatedEnvironment } from "~/services/apiAuth.server"; import { ResumeTaskRunDependenciesService } from "./resumeTaskRunDependencies.server"; - -export const CRASHABLE_RUN_STATUSES: Array = [ - "PENDING", - "WAITING_FOR_DEPLOY", - "EXECUTING", - "PAUSED", - "WAITING_TO_RESUME", - "PAUSED", - "RETRYING_AFTER_FAILURE", -]; - -const CRASHABLE_ATTEMPT_STATUSES: Array = ["EXECUTING", "PAUSED", "PENDING"]; +import { CRASHABLE_ATTEMPT_STATUSES, isCrashableRunStatus } from "../taskStatus"; export type CrashTaskRunServiceOptions = { reason?: string; @@ -52,7 +36,8 @@ export class CrashTaskRunService extends BaseService { } // Make sure the task run is in a crashable state - if (!CRASHABLE_RUN_STATUSES.includes(taskRun.status)) { + if (!isCrashableRunStatus(taskRun.status)) { + logger.error("Task run is not in a crashable state", { runId, status: taskRun.status }); return; } diff --git a/apps/webapp/app/v3/services/createCheckpoint.server.ts b/apps/webapp/app/v3/services/createCheckpoint.server.ts index 249796474..0ed9acab5 100644 --- a/apps/webapp/app/v3/services/createCheckpoint.server.ts +++ b/apps/webapp/app/v3/services/createCheckpoint.server.ts @@ -1,19 +1,13 @@ import { CoordinatorToPlatformMessages } from "@trigger.dev/core/v3"; import type { InferSocketMessageSchema } from "@trigger.dev/core/v3/zodSocket"; -import type { - CheckpointRestoreEvent, - TaskRunAttemptStatus, - TaskRunStatus, -} from "@trigger.dev/database"; +import type { Checkpoint, CheckpointRestoreEvent } from "@trigger.dev/database"; import { logger } from "~/services/logger.server"; import { generateFriendlyId } from "../friendlyIdentifiers"; import { marqs } from "~/v3/marqs/index.server"; import { CreateCheckpointRestoreEventService } from "./createCheckpointRestoreEvent.server"; import { BaseService } from "./baseService.server"; import { CrashTaskRunService } from "./crashTaskRun.server"; - -const FREEZABLE_RUN_STATUSES: TaskRunStatus[] = ["EXECUTING", "RETRYING_AFTER_FAILURE"]; -const FREEZABLE_ATTEMPT_STATUSES: TaskRunAttemptStatus[] = ["EXECUTING", "FAILED"]; +import { isFinalRunStatus, isFreezableAttemptStatus, isFreezableRunStatus } from "../taskStatus"; export class CreateCheckpointService extends BaseService { public async call( @@ -21,7 +15,14 @@ export class CreateCheckpointService extends BaseService { InferSocketMessageSchema, "version" > - ) { + ): Promise< + | { + checkpoint: Checkpoint; + event: CheckpointRestoreEvent; + keepRunAlive: boolean; + } + | undefined + > { logger.debug(`Creating checkpoint`, params); const attempt = await this._prisma.taskRunAttempt.findUnique({ @@ -49,8 +50,8 @@ export class CreateCheckpointService extends BaseService { } if ( - !FREEZABLE_ATTEMPT_STATUSES.includes(attempt.status) || - !FREEZABLE_RUN_STATUSES.includes(attempt.taskRun.status) + !isFreezableAttemptStatus(attempt.status) || + !isFreezableRunStatus(attempt.taskRun.status) ) { logger.error("Unfreezable state", { attempt: { @@ -115,7 +116,9 @@ export class CreateCheckpointService extends BaseService { }); const { reason } = params; + let checkpointEvent: CheckpointRestoreEvent | undefined; + let keepRunAlive = false; switch (reason.type) { case "WAIT_FOR_DURATION": { @@ -131,7 +134,12 @@ export class CreateCheckpointService extends BaseService { dependencyFriendlyRunId: reason.friendlyId, }); - await marqs?.acknowledgeMessage(attempt.taskRunId); + keepRunAlive = await this.#isRunCompleted(reason.friendlyId); + + if (!keepRunAlive) { + await marqs?.acknowledgeMessage(attempt.taskRunId); + } + break; } case "WAIT_FOR_BATCH": { @@ -140,7 +148,12 @@ export class CreateCheckpointService extends BaseService { batchDependencyFriendlyId: reason.batchFriendlyId, }); - await marqs?.acknowledgeMessage(attempt.taskRunId); + keepRunAlive = await this.#isBatchCompleted(reason.batchFriendlyId); + + if (!keepRunAlive) { + await marqs?.acknowledgeMessage(attempt.taskRunId); + } + break; } case "RETRYING_AFTER_FAILURE": { @@ -180,6 +193,37 @@ export class CreateCheckpointService extends BaseService { return { checkpoint, event: checkpointEvent, + keepRunAlive, }; } + + async #isBatchCompleted(friendlyId: string): Promise { + const batch = await this._prisma.batchTaskRun.findUnique({ + where: { + friendlyId, + }, + }); + + if (!batch) { + logger.error("Batch not found", { friendlyId }); + return false; + } + + return batch.status === "COMPLETED"; + } + + async #isRunCompleted(friendlyId: string): Promise { + const run = await this._prisma.taskRun.findUnique({ + where: { + friendlyId, + }, + }); + + if (!run) { + logger.error("Run not found", { friendlyId }); + return false; + } + + return isFinalRunStatus(run.status); + } } diff --git a/apps/webapp/app/v3/services/restoreCheckpoint.server.ts b/apps/webapp/app/v3/services/restoreCheckpoint.server.ts index a0017565e..dc66c3b00 100644 --- a/apps/webapp/app/v3/services/restoreCheckpoint.server.ts +++ b/apps/webapp/app/v3/services/restoreCheckpoint.server.ts @@ -1,12 +1,10 @@ -import { TaskRunAttemptStatus, TaskRunStatus, type Checkpoint } from "@trigger.dev/database"; +import { type Checkpoint } from "@trigger.dev/database"; import { logger } from "~/services/logger.server"; import { socketIo } from "../handleSocketIo.server"; import { machinePresetFromConfig } from "../machinePresets.server"; import { BaseService } from "./baseService.server"; import { CreateCheckpointRestoreEventService } from "./createCheckpointRestoreEvent.server"; - -const RESTORABLE_RUN_STATUSES: TaskRunStatus[] = ["WAITING_TO_RESUME"]; -const RESTORABLE_ATTEMPT_STATUSES: TaskRunAttemptStatus[] = ["PAUSED"]; +import { isRestorableAttemptStatus, isRestorableRunStatus } from "../taskStatus"; export class RestoreCheckpointService extends BaseService { public async call(params: { @@ -51,10 +49,7 @@ export class RestoreCheckpointService extends BaseService { const checkpoint = checkpointEvent.checkpoint; - const runIsRestorable = RESTORABLE_RUN_STATUSES.includes(checkpoint.run.status); - const attemptIsRestorable = RESTORABLE_ATTEMPT_STATUSES.includes(checkpoint.attempt.status); - - if (!runIsRestorable) { + if (!isRestorableRunStatus(checkpoint.run.status)) { logger.error("Run is unrestorable", { eventId: params.eventId, runId: checkpoint.runId, @@ -64,7 +59,7 @@ export class RestoreCheckpointService extends BaseService { return; } - if (!attemptIsRestorable && !params.isRetry) { + if (!isRestorableAttemptStatus(checkpoint.attempt.status) && !params.isRetry) { logger.error("Attempt is unrestorable", { eventId: params.eventId, runId: checkpoint.runId, diff --git a/apps/webapp/app/v3/taskStatus.ts b/apps/webapp/app/v3/taskStatus.ts new file mode 100644 index 000000000..3488f2d94 --- /dev/null +++ b/apps/webapp/app/v3/taskStatus.ts @@ -0,0 +1,69 @@ +import { TaskRunAttemptStatus, TaskRunStatus } from "@trigger.dev/database"; + +export const CANCELLABLE_RUN_STATUSES: TaskRunStatus[] = [ + "PENDING", + "WAITING_FOR_DEPLOY", + "EXECUTING", + "PAUSED", + "WAITING_TO_RESUME", + "PAUSED", + "RETRYING_AFTER_FAILURE", +]; +export const CANCELLABLE_ATTEMPT_STATUSES: TaskRunAttemptStatus[] = [ + "EXECUTING", + "PAUSED", + "PENDING", +]; + +export function isCancellableRunStatus(status: TaskRunStatus): boolean { + return CANCELLABLE_RUN_STATUSES.includes(status); +} +export function isCancellableAttemptStatus(status: TaskRunAttemptStatus): boolean { + return CANCELLABLE_ATTEMPT_STATUSES.includes(status); +} + +export const CRASHABLE_RUN_STATUSES: TaskRunStatus[] = CANCELLABLE_RUN_STATUSES; +export const CRASHABLE_ATTEMPT_STATUSES: TaskRunAttemptStatus[] = CANCELLABLE_ATTEMPT_STATUSES; + +export function isCrashableRunStatus(status: TaskRunStatus): boolean { + return CRASHABLE_RUN_STATUSES.includes(status); +} +export function isCrashableAttemptStatus(status: TaskRunAttemptStatus): boolean { + return CRASHABLE_ATTEMPT_STATUSES.includes(status); +} + +export const FINAL_RUN_STATUSES: TaskRunStatus[] = [ + "CANCELED", + "COMPLETED_SUCCESSFULLY", + "COMPLETED_WITH_ERRORS", + "INTERRUPTED", + "SYSTEM_FAILURE", +]; +export const FINAL_ATTEMPT_STATUSES: TaskRunAttemptStatus[] = ["CANCELED", "COMPLETED", "FAILED"]; + +export const FREEZABLE_RUN_STATUSES: TaskRunStatus[] = ["EXECUTING", "RETRYING_AFTER_FAILURE"]; +export const FREEZABLE_ATTEMPT_STATUSES: TaskRunAttemptStatus[] = ["EXECUTING", "FAILED"]; + +export function isFreezableRunStatus(status: TaskRunStatus): boolean { + return FREEZABLE_RUN_STATUSES.includes(status); +} +export function isFreezableAttemptStatus(status: TaskRunAttemptStatus): boolean { + return FREEZABLE_ATTEMPT_STATUSES.includes(status); +} + +export function isFinalRunStatus(status: TaskRunStatus): boolean { + return FINAL_RUN_STATUSES.includes(status); +} +export function isFinalAttemptStatus(status: TaskRunAttemptStatus): boolean { + return FINAL_ATTEMPT_STATUSES.includes(status); +} + +export const RESTORABLE_RUN_STATUSES: TaskRunStatus[] = ["WAITING_TO_RESUME"]; +export const RESTORABLE_ATTEMPT_STATUSES: TaskRunAttemptStatus[] = ["PAUSED"]; + +export function isRestorableRunStatus(status: TaskRunStatus): boolean { + return RESTORABLE_RUN_STATUSES.includes(status); +} +export function isRestorableAttemptStatus(status: TaskRunAttemptStatus): boolean { + return RESTORABLE_ATTEMPT_STATUSES.includes(status); +} diff --git a/packages/core/src/v3/schemas/messages.ts b/packages/core/src/v3/schemas/messages.ts index 86e0163f8..ed696edaf 100644 --- a/packages/core/src/v3/schemas/messages.ts +++ b/packages/core/src/v3/schemas/messages.ts @@ -554,6 +554,10 @@ export const CoordinatorToPlatformMessages = { }), ]), }), + callback: z.object({ + version: z.literal("v1").default("v1"), + keepRunAlive: z.boolean(), + }), }, INDEXING_FAILED: { message: z.object({