From 593b66bfed9c83c5674009332bb9191c685d4778 Mon Sep 17 00:00:00 2001 From: nicktrn <55853254+nicktrn@users.noreply.github.com> Date: Mon, 18 Mar 2024 14:57:09 +0000 Subject: [PATCH] add checkpoint restore events --- .../v3/marqs/sharedQueueConsumer.server.ts | 34 ++------ .../v3/services/createCheckpoint.server.ts | 8 +- .../createCheckpointRestoreEvent.server.ts | 41 ++++++++++ .../v3/services/restoreCheckpoint.server.ts | 38 +++++++++ .../migration.sql | 34 ++++++++ packages/database/prisma/schema.prisma | 79 ++++++++++++++----- 6 files changed, 184 insertions(+), 50 deletions(-) create mode 100644 apps/webapp/app/v3/services/createCheckpointRestoreEvent.server.ts create mode 100644 apps/webapp/app/v3/services/restoreCheckpoint.server.ts create mode 100644 packages/database/prisma/migrations/20240318135831_add_checkpoint_restore_event/migration.sql diff --git a/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts b/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts index 82d73db69..46c96ae82 100644 --- a/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts +++ b/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts @@ -20,6 +20,7 @@ import { EnvironmentVariablesRepository } from "../environmentVariables/environm import { CancelAttemptService } from "../services/cancelAttempt.server"; import { socketIo } from "../handleSocketIo.server"; import { singleton } from "~/utils/singleton"; +import { RestoreCheckpointService } from "../services/restoreCheckpoint.server"; const tracer = trace.getTracer("sharedQueueConsumer"); @@ -464,15 +465,8 @@ export class SharedQueueConsumer { return; } - socketIo.providerNamespace.emit("RESTORE", { - version: "v1", - checkpointId: latestCheckpoint.id, - runId: latestCheckpoint.runId, - attemptId: latestCheckpoint.attemptId, - type: latestCheckpoint.type, - location: latestCheckpoint.location, - reason: latestCheckpoint.reason ?? undefined, - }); + const restoreService = new RestoreCheckpointService(); + await restoreService.call({ checkpointId: latestCheckpoint.id }); } else { await this._sender.send("BACKGROUND_WORKER_MESSAGE", { backgroundWorkerId: deployment.worker.friendlyId, @@ -608,15 +602,8 @@ export class SharedQueueConsumer { return; } - socketIo.providerNamespace.emit("RESTORE", { - version: "v1", - checkpointId: latestCheckpoint.id, - runId: latestCheckpoint.runId, - attemptId: latestCheckpoint.attemptId, - type: latestCheckpoint.type, - location: latestCheckpoint.location, - reason: latestCheckpoint.reason ?? undefined, - }); + const restoreService = new RestoreCheckpointService(); + await restoreService.call({ checkpointId: latestCheckpoint.id }); setTimeout(() => this.#doWork(), this._options.interval); return; @@ -750,15 +737,8 @@ export class SharedQueueConsumer { } // The attempt will resume automatically after restore - socketIo.providerNamespace.emit("RESTORE", { - version: "v1", - checkpointId: latestCheckpoint.id, - runId: latestCheckpoint.runId, - attemptId: latestCheckpoint.attemptId, - type: latestCheckpoint.type, - location: latestCheckpoint.location, - reason: latestCheckpoint.reason ?? undefined, - }); + const restoreService = new RestoreCheckpointService(); + await restoreService.call({ checkpointId: latestCheckpoint.id }); } catch (e) { if (e instanceof Error) { this._currentSpan?.recordException(e); diff --git a/apps/webapp/app/v3/services/createCheckpoint.server.ts b/apps/webapp/app/v3/services/createCheckpoint.server.ts index 07e640967..f59092fdf 100644 --- a/apps/webapp/app/v3/services/createCheckpoint.server.ts +++ b/apps/webapp/app/v3/services/createCheckpoint.server.ts @@ -4,6 +4,7 @@ import { PrismaClient, prisma } from "~/db.server"; import { logger } from "~/services/logger.server"; import { generateFriendlyId } from "../friendlyIdentifiers"; import { marqs } from "../marqs.server"; +import { CreateCheckpointRestoreEventService } from "./createCheckpointRestoreEvent.server"; export class CreateCheckpointService { #prismaClient: PrismaClient; @@ -15,6 +16,8 @@ export class CreateCheckpointService { public async call( params: InferSocketMessageSchema ): Promise { + logger.debug(`Creating checkpoint`, params); + const attempt = await this.#prismaClient.taskRunAttempt.findUniqueOrThrow({ where: { id: params.attemptId, @@ -24,8 +27,6 @@ export class CreateCheckpointService { }, }); - logger.debug(`Creating checkpoint`, params); - const checkpoint = await this.#prismaClient.checkpoint.create({ data: { friendlyId: generateFriendlyId("checkpoint"), @@ -40,6 +41,9 @@ export class CreateCheckpointService { }, }); + const eventService = new CreateCheckpointRestoreEventService(this.#prismaClient); + await eventService.call({ checkpointId: checkpoint.id, type: "CHECKPOINT" }); + await this.#prismaClient.taskRunAttempt.update({ where: { id: params.attemptId, diff --git a/apps/webapp/app/v3/services/createCheckpointRestoreEvent.server.ts b/apps/webapp/app/v3/services/createCheckpointRestoreEvent.server.ts new file mode 100644 index 000000000..16a9f4a71 --- /dev/null +++ b/apps/webapp/app/v3/services/createCheckpointRestoreEvent.server.ts @@ -0,0 +1,41 @@ +import type { CheckpointRestoreEvent, CheckpointRestoreEventType } from "@trigger.dev/database"; +import { $transaction, PrismaClient, prisma } from "~/db.server"; +import { logger } from "~/services/logger.server"; + +export class CreateCheckpointRestoreEventService { + #prismaClient: PrismaClient; + + constructor(prismaClient: PrismaClient = prisma) { + this.#prismaClient = prismaClient; + } + + public async call(params: { + checkpointId: string; + type: CheckpointRestoreEventType; + }): Promise { + return await $transaction(this.#prismaClient, async (tx) => { + const checkpoint = await this.#prismaClient.checkpoint.findUniqueOrThrow({ + where: { + id: params.checkpointId, + }, + }); + + logger.debug(`Creating checkpoint/restore event`, params); + + const checkpointEvent = await this.#prismaClient.checkpointRestoreEvent.create({ + data: { + checkpointId: checkpoint.id, + runtimeEnvironmentId: checkpoint.runtimeEnvironmentId, + projectId: checkpoint.projectId, + attemptId: checkpoint.attemptId, + runId: checkpoint.runId, + type: params.type, + reason: checkpoint.reason, + metadata: checkpoint.metadata, + }, + }); + + return checkpointEvent; + }); + } +} diff --git a/apps/webapp/app/v3/services/restoreCheckpoint.server.ts b/apps/webapp/app/v3/services/restoreCheckpoint.server.ts new file mode 100644 index 000000000..3704cb4e5 --- /dev/null +++ b/apps/webapp/app/v3/services/restoreCheckpoint.server.ts @@ -0,0 +1,38 @@ +import { type Checkpoint } from "@trigger.dev/database"; +import { PrismaClient, prisma } from "~/db.server"; +import { logger } from "~/services/logger.server"; +import { socketIo } from "../handleSocketIo.server"; +import { CreateCheckpointRestoreEventService } from "./createCheckpointRestoreEvent.server"; + +export class RestoreCheckpointService { + #prismaClient: PrismaClient; + + constructor(prismaClient: PrismaClient = prisma) { + this.#prismaClient = prismaClient; + } + + public async call(params: { checkpointId: string }): Promise { + logger.debug(`Restoring checkpoint`, params); + + const checkpoint = await this.#prismaClient.checkpoint.findUniqueOrThrow({ + where: { + id: params.checkpointId, + }, + }); + + const eventService = new CreateCheckpointRestoreEventService(this.#prismaClient); + await eventService.call({ checkpointId: checkpoint.id, type: "RESTORE" }); + + socketIo.providerNamespace.emit("RESTORE", { + version: "v1", + checkpointId: checkpoint.id, + runId: checkpoint.runId, + attemptId: checkpoint.attemptId, + type: checkpoint.type, + location: checkpoint.location, + reason: checkpoint.reason ?? undefined, + }); + + return checkpoint; + } +} diff --git a/packages/database/prisma/migrations/20240318135831_add_checkpoint_restore_event/migration.sql b/packages/database/prisma/migrations/20240318135831_add_checkpoint_restore_event/migration.sql new file mode 100644 index 000000000..549802f42 --- /dev/null +++ b/packages/database/prisma/migrations/20240318135831_add_checkpoint_restore_event/migration.sql @@ -0,0 +1,34 @@ +-- CreateEnum +CREATE TYPE "CheckpointRestoreEventType" AS ENUM ('CHECKPOINT', 'RESTORE'); + +-- CreateTable +CREATE TABLE "CheckpointRestoreEvent" ( + "id" TEXT NOT NULL, + "type" "CheckpointRestoreEventType" NOT NULL, + "reason" TEXT, + "metadata" TEXT, + "checkpointId" TEXT NOT NULL, + "runId" TEXT NOT NULL, + "attemptId" TEXT NOT NULL, + "projectId" TEXT NOT NULL, + "runtimeEnvironmentId" TEXT NOT NULL, + "createdAt" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updatedAt" TIMESTAMP(3) NOT NULL, + + CONSTRAINT "CheckpointRestoreEvent_pkey" PRIMARY KEY ("id") +); + +-- AddForeignKey +ALTER TABLE "CheckpointRestoreEvent" ADD CONSTRAINT "CheckpointRestoreEvent_checkpointId_fkey" FOREIGN KEY ("checkpointId") REFERENCES "Checkpoint"("id") ON DELETE CASCADE ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "CheckpointRestoreEvent" ADD CONSTRAINT "CheckpointRestoreEvent_runId_fkey" FOREIGN KEY ("runId") REFERENCES "TaskRun"("id") ON DELETE CASCADE ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "CheckpointRestoreEvent" ADD CONSTRAINT "CheckpointRestoreEvent_attemptId_fkey" FOREIGN KEY ("attemptId") REFERENCES "TaskRunAttempt"("id") ON DELETE CASCADE ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "CheckpointRestoreEvent" ADD CONSTRAINT "CheckpointRestoreEvent_projectId_fkey" FOREIGN KEY ("projectId") REFERENCES "Project"("id") ON DELETE CASCADE ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "CheckpointRestoreEvent" ADD CONSTRAINT "CheckpointRestoreEvent_runtimeEnvironmentId_fkey" FOREIGN KEY ("runtimeEnvironmentId") REFERENCES "RuntimeEnvironment"("id") ON DELETE CASCADE ON UPDATE CASCADE; diff --git a/packages/database/prisma/schema.prisma b/packages/database/prisma/schema.prisma index 4b04de2ca..8a59cdf9c 100644 --- a/packages/database/prisma/schema.prisma +++ b/packages/database/prisma/schema.prisma @@ -394,6 +394,7 @@ model RuntimeEnvironment { workerDeployments WorkerDeployment[] workerDeploymentPromotions WorkerDeploymentPromotion[] taskRunAttempts TaskRunAttempt[] + CheckpointRestoreEvent CheckpointRestoreEvent[] @@unique([projectId, slug, orgMemberId]) @@unique([projectId, shortcode]) @@ -422,23 +423,24 @@ model Project { version ProjectVersion @default(V2) - environments RuntimeEnvironment[] - endpoints Endpoint[] - jobs Job[] - jobVersion JobVersion[] - events EventRecord[] - runs JobRun[] - sources TriggerSource[] - httpEndpoints TriggerHttpEndpoint[] - webhooks Webhook[] - backgroundWorkers BackgroundWorker[] - backgroundWorkerTasks BackgroundWorkerTask[] - taskRuns TaskRun[] - taskTags TaskTag[] - taskQueues TaskQueue[] - environmentVariables EnvironmentVariable[] - checkpoints Checkpoint[] - WorkerDeployment WorkerDeployment[] + environments RuntimeEnvironment[] + endpoints Endpoint[] + jobs Job[] + jobVersion JobVersion[] + events EventRecord[] + runs JobRun[] + sources TriggerSource[] + httpEndpoints TriggerHttpEndpoint[] + webhooks Webhook[] + backgroundWorkers BackgroundWorker[] + backgroundWorkerTasks BackgroundWorkerTask[] + taskRuns TaskRun[] + taskTags TaskTag[] + taskQueues TaskQueue[] + environmentVariables EnvironmentVariable[] + checkpoints Checkpoint[] + WorkerDeployment WorkerDeployment[] + CheckpointRestoreEvent CheckpointRestoreEvent[] } enum ProjectVersion { @@ -1604,8 +1606,9 @@ model TaskRun { concurrencyKey String? - batchItem BatchTaskRunItem? - dependency TaskRunDependency? + batchItem BatchTaskRunItem? + dependency TaskRunDependency? + CheckpointRestoreEvent CheckpointRestoreEvent[] @@unique([runtimeEnvironmentId, idempotencyKey]) } @@ -1717,8 +1720,9 @@ model TaskRunAttempt { taskRunDependency TaskRunDependency? @relation("dependentAttempt") batchTaskRunDependency BatchTaskRun? - checkpoints Checkpoint[] - batchTaskRunItems BatchTaskRunItem[] + checkpoints Checkpoint[] + batchTaskRunItems BatchTaskRunItem[] + CheckpointRestoreEvent CheckpointRestoreEvent[] @@unique([taskRunId, number]) } @@ -1955,6 +1959,8 @@ model Checkpoint { reason String? metadata String? + events CheckpointRestoreEvent[] + run TaskRun @relation(fields: [runId], references: [id], onDelete: Cascade, onUpdate: Cascade) runId String @@ -1976,6 +1982,37 @@ enum CheckpointType { KUBERNETES } +model CheckpointRestoreEvent { + id String @id @default(cuid()) + + type CheckpointRestoreEventType + reason String? + metadata String? + + checkpoint Checkpoint @relation(fields: [checkpointId], references: [id], onDelete: Cascade, onUpdate: Cascade) + checkpointId String + + run TaskRun @relation(fields: [runId], references: [id], onDelete: Cascade, onUpdate: Cascade) + runId String + + attempt TaskRunAttempt @relation(fields: [attemptId], references: [id], onDelete: Cascade, onUpdate: Cascade) + attemptId String + + project Project @relation(fields: [projectId], references: [id], onDelete: Cascade, onUpdate: Cascade) + projectId String + + runtimeEnvironment RuntimeEnvironment @relation(fields: [runtimeEnvironmentId], references: [id], onDelete: Cascade, onUpdate: Cascade) + runtimeEnvironmentId String + + createdAt DateTime @default(now()) + updatedAt DateTime @updatedAt +} + +enum CheckpointRestoreEventType { + CHECKPOINT + RESTORE +} + model WorkerDeployment { id String @id @default(cuid())