v3: checkpoint restore events (#951)

* add checkpoint restore events

* fix retries.enabledInDev

---------

Co-authored-by: Eric Allam <eric@trigger.dev>
This commit is contained in:
nicktrn
2024-03-19 09:02:56 +00:00
committed by GitHub
parent f26ba9a5cd
commit b585d1f2c5
7 changed files with 185 additions and 50 deletions
@@ -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);
@@ -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<typeof CoordinatorToPlatformMessages, "CHECKPOINT_CREATED">
): Promise<Checkpoint> {
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,
@@ -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<CheckpointRestoreEvent | undefined> {
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;
});
}
}
@@ -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<Checkpoint> {
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;
}
}
@@ -233,6 +233,7 @@ export class TaskExecutor {
const delay = calculateNextRetryDelay(retry, execution.attempt.number);
if (
execution.environment.type === "DEVELOPMENT" &&
typeof this._config.retries?.enabledInDev === "boolean" &&
!this._config.retries.enabledInDev
) {
@@ -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;
+58 -21
View File
@@ -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())