From bb382619475a9be641b70491bfa5ccd95b683f13 Mon Sep 17 00:00:00 2001 From: Matt Aitken Date: Thu, 8 Aug 2024 15:32:06 +0100 Subject: [PATCH] FinalizeRunService (#1250) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * WIP notes on each location where we’ll use finalize * Initial FinalizeTaskRunService * ExpireEnqueuedRunService uses FinalizeTaskRunService * FailedTaskRunService uses FinalizeTaskRunService * Allow passing in an include when finalizing the run * CrashTaskRunService using FinalizeTaskRunService * Remove comments * Status is optional * CancelAttemptService using FinalizeTaskRunService * Import tidy * CancelTaskRunService using FinalizeTaskRunService * Import tidying * CompleteAttemptService system failure switched to FinalizeTaskRunService * Added more logging to Finalizing * CompleteAttemptStatus COMPLETED_SUCCESSFULLY * CompletedAttempt “SYSTEM_FAILURE” * CompletedService final pair * Use satisfies so we can derive types from the groups * Only allow final states to be used with this service * BaseService tx support, minor improvements --- .../app/components/runs/v3/TaskRunStatus.tsx | 16 ++- apps/webapp/app/v3/failedTaskRun.server.ts | 22 ++-- .../app/v3/services/cancelAttempt.server.ts | 48 ++++---- .../app/v3/services/cancelTaskRun.server.ts | 23 ++-- .../app/v3/services/completeAttempt.server.ts | 114 +++++++++--------- .../app/v3/services/crashTaskRun.server.ts | 18 +-- .../v3/services/expireEnqueuedRun.server.ts | 19 ++- .../app/v3/services/finalizeTaskRun.server.ts | 59 +++++++++ 8 files changed, 175 insertions(+), 144 deletions(-) create mode 100644 apps/webapp/app/v3/services/finalizeTaskRun.server.ts diff --git a/apps/webapp/app/components/runs/v3/TaskRunStatus.tsx b/apps/webapp/app/components/runs/v3/TaskRunStatus.tsx index b35771688..c3a384a5a 100644 --- a/apps/webapp/app/components/runs/v3/TaskRunStatus.tsx +++ b/apps/webapp/app/components/runs/v3/TaskRunStatus.tsx @@ -67,15 +67,19 @@ const taskRunStatusDescriptions: Record = { EXPIRED: "Task has surpassed its ttl and won't be executed", }; -export const QUEUED_STATUSES: TaskRunStatus[] = ["PENDING", "WAITING_FOR_DEPLOY", "DELAYED"]; +export const QUEUED_STATUSES = [ + "PENDING", + "WAITING_FOR_DEPLOY", + "DELAYED", +] satisfies TaskRunStatus[]; -export const RUNNING_STATUSES: TaskRunStatus[] = [ +export const RUNNING_STATUSES = [ "EXECUTING", "RETRYING_AFTER_FAILURE", "WAITING_TO_RESUME", -]; +] satisfies TaskRunStatus[]; -export const FINISHED_STATUSES: TaskRunStatus[] = [ +export const FINISHED_STATUSES = [ "COMPLETED_SUCCESSFULLY", "CANCELED", "COMPLETED_WITH_ERRORS", @@ -83,7 +87,9 @@ export const FINISHED_STATUSES: TaskRunStatus[] = [ "SYSTEM_FAILURE", "CRASHED", "EXPIRED", -]; +] satisfies TaskRunStatus[]; + +export type FINISHED_STATUSES = (typeof FINISHED_STATUSES)[number]; export function descriptionForTaskRunStatus(status: TaskRunStatus): string { return taskRunStatusDescriptions[status]; diff --git a/apps/webapp/app/v3/failedTaskRun.server.ts b/apps/webapp/app/v3/failedTaskRun.server.ts index 94c5ebae3..247e196bc 100644 --- a/apps/webapp/app/v3/failedTaskRun.server.ts +++ b/apps/webapp/app/v3/failedTaskRun.server.ts @@ -1,10 +1,9 @@ import { TaskRunFailedExecutionResult } from "@trigger.dev/core/v3"; -import { logger } from "~/services/logger.server"; -import { marqs } from "~/v3/marqs/index.server"; - import { TaskRunStatus } from "@trigger.dev/database"; +import { logger } from "~/services/logger.server"; import { createExceptionPropertiesFromError, eventRepository } from "./eventRepository.server"; import { BaseService } from "./services/baseService.server"; +import { FinalizeTaskRunService } from "./services/finalizeTaskRun.server"; const FAILABLE_TASK_RUN_STATUSES: TaskRunStatus[] = ["EXECUTING", "PENDING", "WAITING_FOR_DEPLOY"]; @@ -40,7 +39,12 @@ export class FailedTaskRunService extends BaseService { // No more retries, we need to fail the task run logger.debug("[FailedTaskRunService] Failing task run", { taskRun, completion }); - await marqs?.acknowledgeMessage(taskRun.id); + const finalizeService = new FinalizeTaskRunService(); + await finalizeService.call({ + id: taskRun.id, + status: "SYSTEM_FAILURE", + completedAt: new Date(), + }); // Now we need to "complete" the task run event/span await eventRepository.completeEvent(taskRun.spanId, { @@ -58,15 +62,5 @@ export class FailedTaskRunService extends BaseService { }, ], }); - - await this._prisma.taskRun.update({ - where: { - id: taskRun.id, - }, - data: { - status: "SYSTEM_FAILURE", - completedAt: new Date(), - }, - }); } } diff --git a/apps/webapp/app/v3/services/cancelAttempt.server.ts b/apps/webapp/app/v3/services/cancelAttempt.server.ts index 418e08037..bdbaa2f3f 100644 --- a/apps/webapp/app/v3/services/cancelAttempt.server.ts +++ b/apps/webapp/app/v3/services/cancelAttempt.server.ts @@ -1,11 +1,10 @@ -import { AuthenticatedEnvironment } from "~/services/apiAuth.server"; +import { $transaction, type PrismaClientOrTransaction, prisma } from "~/db.server"; +import { type AuthenticatedEnvironment } from "~/services/apiAuth.server"; import { logger } from "~/services/logger.server"; -import { marqs } from "~/v3/marqs/index.server"; import { eventRepository } from "../eventRepository.server"; -import { BaseService } from "./baseService.server"; - -import { PrismaClientOrTransaction, prisma } from "~/db.server"; import { isCancellableRunStatus } from "../taskStatus"; +import { BaseService } from "./baseService.server"; +import { FinalizeTaskRunService } from "./finalizeTaskRun.server"; import { ResumeTaskRunDependenciesService } from "./resumeTaskRunDependencies.server"; export class CancelAttemptService extends BaseService { @@ -51,28 +50,25 @@ export class CancelAttemptService extends BaseService { return; } - await marqs?.acknowledgeMessage(taskRunId); - - await this._prisma.taskRunAttempt.update({ - where: { - friendlyId: attemptId, - }, - data: { - status: "CANCELED", - completedAt: cancelledAt, - taskRun: { - update: { - data: { - status: isCancellableRunStatus(taskRunAttempt.taskRun.status) - ? "INTERRUPTED" - : undefined, - completedAt: isCancellableRunStatus(taskRunAttempt.taskRun.status) - ? cancelledAt - : undefined, - }, - }, + await $transaction(this._prisma, async (tx) => { + await tx.taskRunAttempt.update({ + where: { + friendlyId: attemptId, }, - }, + data: { + status: "CANCELED", + completedAt: cancelledAt, + }, + }); + + const finalizeService = new FinalizeTaskRunService(tx); + await finalizeService.call({ + id: taskRunId, + status: isCancellableRunStatus(taskRunAttempt.taskRun.status) ? "INTERRUPTED" : undefined, + completedAt: isCancellableRunStatus(taskRunAttempt.taskRun.status) + ? cancelledAt + : undefined, + }); }); const inProgressEvents = await eventRepository.queryIncompleteEvents({ diff --git a/apps/webapp/app/v3/services/cancelTaskRun.server.ts b/apps/webapp/app/v3/services/cancelTaskRun.server.ts index 6f06b6888..f6810cfc5 100644 --- a/apps/webapp/app/v3/services/cancelTaskRun.server.ts +++ b/apps/webapp/app/v3/services/cancelTaskRun.server.ts @@ -1,14 +1,14 @@ -import { Prisma, TaskRun } from "@trigger.dev/database"; +import { type Prisma, type TaskRun } from "@trigger.dev/database"; import assertNever from "assert-never"; import { logger } from "~/services/logger.server"; -import { marqs } from "~/v3/marqs/index.server"; import { eventRepository } from "../eventRepository.server"; import { socketIo } from "../handleSocketIo.server"; import { devPubSub } from "../marqs/devPubSub.server"; +import { CANCELLABLE_ATTEMPT_STATUSES, isCancellableRunStatus } from "../taskStatus"; import { BaseService } from "./baseService.server"; import { CancelAttemptService } from "./cancelAttempt.server"; -import { CANCELLABLE_ATTEMPT_STATUSES, isCancellableRunStatus } from "../taskStatus"; import { CancelTaskAttemptDependenciesService } from "./cancelTaskAttemptDependencies.server"; +import { FinalizeTaskRunService } from "./finalizeTaskRun.server"; type ExtendedTaskRun = Prisma.TaskRunGetPayload<{ include: { @@ -47,18 +47,11 @@ export class CancelTaskRunService extends BaseService { return; } - // Remove the task run from the queue if it's there for some reason - await marqs?.acknowledgeMessage(taskRun.id); - - // Set the task run status to cancelled - const cancelledTaskRun = await this._prisma.taskRun.update({ - where: { - id: taskRun.id, - }, - data: { - status: "CANCELED", - completedAt: opts.cancelledAt, - }, + const finalizeService = new FinalizeTaskRunService(); + const cancelledTaskRun = await finalizeService.call({ + id: taskRun.id, + status: "CANCELED", + completedAt: opts.cancelledAt, include: { attempts: { where: { diff --git a/apps/webapp/app/v3/services/completeAttempt.server.ts b/apps/webapp/app/v3/services/completeAttempt.server.ts index f121f37ab..e71e7bd96 100644 --- a/apps/webapp/app/v3/services/completeAttempt.server.ts +++ b/apps/webapp/app/v3/services/completeAttempt.server.ts @@ -8,7 +8,7 @@ import { flattenAttributes, sanitizeError, } from "@trigger.dev/core/v3"; -import { PrismaClientOrTransaction } from "~/db.server"; +import { $transaction, PrismaClientOrTransaction } from "~/db.server"; import { AuthenticatedEnvironment } from "~/services/apiAuth.server"; import { logger } from "~/services/logger.server"; import { safeJsonParse } from "~/utils/json"; @@ -23,6 +23,7 @@ import { TaskRun } from "@trigger.dev/database"; import { PerformTaskAttemptAlertsService } from "./alerts/performTaskAttemptAlerts.server"; import { RetryAttemptService } from "./retryAttempt.server"; import { isFinalAttemptStatus, isFinalRunStatus } from "../taskStatus"; +import { FinalizeTaskRunService } from "./finalizeTaskRun.server"; type FoundAttempt = Awaited>; @@ -50,17 +51,30 @@ export class CompleteAttemptService extends BaseService { id: execution.attempt.id, }); - // Update the task run to be failed - await this._prisma.taskRun.update({ + const run = await this._prisma.taskRun.findFirst({ where: { friendlyId: execution.run.id, }, - data: { - status: "SYSTEM_FAILURE", - completedAt: new Date(), + select: { + id: true, }, }); + if (!run) { + logger.error("[CompleteAttemptService] Task run not found", { + friendlyId: execution.run.id, + }); + + return "COMPLETED"; + } + + const finalizeService = new FinalizeTaskRunService(); + await finalizeService.call({ + id: run.id, + status: "SYSTEM_FAILURE", + completedAt: new Date(), + }); + // No attempt, so there's no message to ACK return "COMPLETED"; } @@ -96,29 +110,26 @@ export class CompleteAttemptService extends BaseService { taskRunAttempt: NonNullable, env?: AuthenticatedEnvironment ): Promise<"COMPLETED"> { - await this._prisma.taskRunAttempt.update({ - where: { id: taskRunAttempt.id }, - data: { - status: "COMPLETED", - completedAt: new Date(), - output: completion.output, - outputType: completion.outputType, - usageDurationMs: completion.usage?.durationMs, - taskRun: { - update: { - data: { - status: "COMPLETED_SUCCESSFULLY", - completedAt: new Date(), - }, - }, + await $transaction(this._prisma, async (tx) => { + await tx.taskRunAttempt.update({ + where: { id: taskRunAttempt.id }, + data: { + status: "COMPLETED", + completedAt: new Date(), + output: completion.output, + outputType: completion.outputType, + usageDurationMs: completion.usage?.durationMs, }, - }, + }); + + const finalizeService = new FinalizeTaskRunService(tx); + await finalizeService.call({ + id: taskRunAttempt.taskRunId, + status: "COMPLETED_SUCCESSFULLY", + completedAt: new Date(), + }); }); - logger.debug("Completed attempt successfully, ACKing message"); - - await marqs?.acknowledgeMessage(taskRunAttempt.taskRunId); - // Now we need to "complete" the task run event/span await eventRepository.completeEvent(taskRunAttempt.taskRun.spanId, { endTime: new Date(), @@ -255,19 +266,13 @@ export class CompleteAttemptService extends BaseService { if (!checkpointCreateResult) { logger.error("Failed to create checkpoint", { checkpoint, execution: execution.run.id }); - // Update the task run to be failed - await this._prisma.taskRun.update({ - where: { - friendlyId: execution.run.id, - }, - data: { - status: "SYSTEM_FAILURE", - completedAt: new Date(), - }, + const finalizeService = new FinalizeTaskRunService(); + await finalizeService.call({ + id: taskRunAttempt.taskRunId, + status: "SYSTEM_FAILURE", + completedAt: new Date(), }); - await marqs?.acknowledgeMessage(taskRunAttempt.taskRunId); - return "COMPLETED"; } @@ -279,11 +284,6 @@ export class CompleteAttemptService extends BaseService { return "RETRIED"; } else { - // No more retries, we need to fail the task run - logger.debug("Completed attempt, ACKing message", taskRunAttempt); - - await marqs?.acknowledgeMessage(taskRunAttempt.taskRunId); - // Now we need to "complete" the task run event/span await eventRepository.completeEvent(taskRunAttempt.taskRun.spanId, { endTime: new Date(), @@ -305,6 +305,13 @@ export class CompleteAttemptService extends BaseService { sanitizedError.type === "INTERNAL_ERROR" && sanitizedError.code === "GRACEFUL_EXIT_TIMEOUT" ) { + const finalizeService = new FinalizeTaskRunService(); + await finalizeService.call({ + id: taskRunAttempt.taskRunId, + status: "SYSTEM_FAILURE", + completedAt: new Date(), + }); + // We need to fail all incomplete spans const inProgressEvents = await eventRepository.queryIncompleteEvents({ attemptId: execution.attempt.id, @@ -328,25 +335,12 @@ export class CompleteAttemptService extends BaseService { }); }) ); - - await this._prisma.taskRun.update({ - where: { - id: taskRunAttempt.taskRunId, - }, - data: { - status: "SYSTEM_FAILURE", - completedAt: new Date(), - }, - }); } else { - await this._prisma.taskRun.update({ - where: { - id: taskRunAttempt.taskRunId, - }, - data: { - status: "COMPLETED_WITH_ERRORS", - completedAt: new Date(), - }, + const finalizeService = new FinalizeTaskRunService(); + await finalizeService.call({ + id: taskRunAttempt.taskRunId, + status: "COMPLETED_WITH_ERRORS", + completedAt: new Date(), }); } diff --git a/apps/webapp/app/v3/services/crashTaskRun.server.ts b/apps/webapp/app/v3/services/crashTaskRun.server.ts index 2775ff9ab..7820fc2b0 100644 --- a/apps/webapp/app/v3/services/crashTaskRun.server.ts +++ b/apps/webapp/app/v3/services/crashTaskRun.server.ts @@ -7,6 +7,7 @@ import { AuthenticatedEnvironment } from "~/services/apiAuth.server"; import { ResumeTaskRunDependenciesService } from "./resumeTaskRunDependencies.server"; import { CRASHABLE_ATTEMPT_STATUSES, isCrashableRunStatus } from "../taskStatus"; import { sanitizeError } from "@trigger.dev/core/v3"; +import { FinalizeTaskRunService } from "./finalizeTaskRun.server"; export type CrashTaskRunServiceOptions = { reason?: string; @@ -43,18 +44,11 @@ export class CrashTaskRunService extends BaseService { return; } - // Remove the task run from the queue if it's there for some reason - await marqs?.acknowledgeMessage(taskRun.id); - - // Set the task run status to crashed - const crashedTaskRun = await this._prisma.taskRun.update({ - where: { - id: taskRun.id, - }, - data: { - status: "CRASHED", - completedAt: new Date(), - }, + const finalizeService = new FinalizeTaskRunService(); + const crashedTaskRun = await finalizeService.call({ + id: taskRun.id, + status: "CRASHED", + completedAt: new Date(), include: { attempts: { where: { diff --git a/apps/webapp/app/v3/services/expireEnqueuedRun.server.ts b/apps/webapp/app/v3/services/expireEnqueuedRun.server.ts index 5dae2e20e..d1a3abb99 100644 --- a/apps/webapp/app/v3/services/expireEnqueuedRun.server.ts +++ b/apps/webapp/app/v3/services/expireEnqueuedRun.server.ts @@ -1,7 +1,7 @@ import { logger } from "~/services/logger.server"; -import { marqs } from "~/v3/marqs/index.server"; import { BaseService } from "./baseService.server"; import { eventRepository } from "../eventRepository.server"; +import { FinalizeTaskRunService } from "./finalizeTaskRun.server"; export class ExpireEnqueuedRunService extends BaseService { public async call(runId: string) { @@ -39,15 +39,12 @@ export class ExpireEnqueuedRunService extends BaseService { run, }); - await this._prisma.taskRun.update({ - where: { - id: run.id, - }, - data: { - status: "EXPIRED", - expiredAt: new Date(), - completedAt: new Date(), - }, + const finalizeService = new FinalizeTaskRunService(); + await finalizeService.call({ + id: run.id, + status: "EXPIRED", + expiredAt: new Date(), + completedAt: new Date(), }); await eventRepository.completeEvent(run.spanId, { @@ -67,7 +64,5 @@ export class ExpireEnqueuedRunService extends BaseService { }, ], }); - - await marqs?.acknowledgeMessage(run.id); } } diff --git a/apps/webapp/app/v3/services/finalizeTaskRun.server.ts b/apps/webapp/app/v3/services/finalizeTaskRun.server.ts new file mode 100644 index 000000000..68d40b5de --- /dev/null +++ b/apps/webapp/app/v3/services/finalizeTaskRun.server.ts @@ -0,0 +1,59 @@ +import { type Prisma, type TaskRun } from "@trigger.dev/database"; +import { type FINISHED_STATUSES } from "~/components/runs/v3/TaskRunStatus"; +import { logger } from "~/services/logger.server"; +import { marqs } from "~/v3/marqs/index.server"; +import { BaseService } from "./baseService.server"; + +type BaseInput = { + id: string; + status?: FINISHED_STATUSES; + expiredAt?: Date; + completedAt?: Date; +}; + +type InputWithInclude = BaseInput & { + include: T; +}; + +type InputWithoutInclude = BaseInput & { + include?: undefined; +}; + +type Output = T extends Prisma.TaskRunInclude + ? Prisma.TaskRunGetPayload<{ include: T }> + : TaskRun; + +export class FinalizeTaskRunService extends BaseService { + public async call({ + id, + status, + expiredAt, + completedAt, + include, + }: T extends Prisma.TaskRunInclude ? InputWithInclude : InputWithoutInclude): Promise< + Output + > { + logger.debug("Finalizing run marqs ack", { + id, + status, + expiredAt, + completedAt, + }); + await marqs?.acknowledgeMessage(id); + + logger.debug("Finalizing run updating run status", { + id, + status, + expiredAt, + completedAt, + }); + + const run = await this._prisma.taskRun.update({ + where: { id }, + data: { status, expiredAt, completedAt }, + ...(include ? { include } : {}), + }); + + return run as Output; + } +}