chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,444 @@
|
||||
import type { CoordinatorToPlatformMessages, ManualCheckpointMetadata } from "@trigger.dev/core/v3";
|
||||
import type { InferSocketMessageSchema } from "@trigger.dev/core/v3/zodSocket";
|
||||
import type { Checkpoint, CheckpointRestoreEvent } from "@trigger.dev/database";
|
||||
import { logger } from "~/services/logger.server";
|
||||
import { marqs } from "~/v3/marqs/index.server";
|
||||
import { isFreezableAttemptStatus, isFreezableRunStatus } from "../taskStatus";
|
||||
import { BaseService } from "./baseService.server";
|
||||
import { CreateCheckpointRestoreEventService } from "./createCheckpointRestoreEvent.server";
|
||||
import { ResumeBatchRunService } from "./resumeBatchRun.server";
|
||||
import { ResumeDependentParentsService } from "./resumeDependentParents.server";
|
||||
import { CheckpointId } from "@trigger.dev/core/v3/isomorphic";
|
||||
|
||||
export class CreateCheckpointService extends BaseService {
|
||||
public async call(
|
||||
params: Omit<
|
||||
InferSocketMessageSchema<typeof CoordinatorToPlatformMessages, "CHECKPOINT_CREATED">,
|
||||
"version"
|
||||
>
|
||||
): Promise<
|
||||
| {
|
||||
success: true;
|
||||
checkpoint: Checkpoint;
|
||||
event: CheckpointRestoreEvent;
|
||||
keepRunAlive: boolean;
|
||||
}
|
||||
| {
|
||||
success: false;
|
||||
keepRunAlive?: boolean;
|
||||
}
|
||||
> {
|
||||
logger.debug(`Creating checkpoint`, params);
|
||||
|
||||
const attempt = await this._prisma.taskRunAttempt.findFirst({
|
||||
where: {
|
||||
friendlyId: params.attemptFriendlyId,
|
||||
},
|
||||
include: {
|
||||
taskRun: true,
|
||||
backgroundWorker: {
|
||||
select: {
|
||||
id: true,
|
||||
deployment: {
|
||||
select: {
|
||||
imageReference: true,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
if (!attempt) {
|
||||
logger.error("Attempt not found", params);
|
||||
|
||||
return {
|
||||
success: false,
|
||||
};
|
||||
}
|
||||
|
||||
if (
|
||||
!isFreezableAttemptStatus(attempt.status) ||
|
||||
!isFreezableRunStatus(attempt.taskRun.status)
|
||||
) {
|
||||
logger.error("Unfreezable state", {
|
||||
attempt: {
|
||||
id: attempt.id,
|
||||
status: attempt.status,
|
||||
},
|
||||
run: {
|
||||
id: attempt.taskRunId,
|
||||
status: attempt.taskRun.status,
|
||||
},
|
||||
params,
|
||||
});
|
||||
|
||||
return {
|
||||
success: false,
|
||||
keepRunAlive: true,
|
||||
};
|
||||
}
|
||||
|
||||
const imageRef = attempt.backgroundWorker.deployment?.imageReference;
|
||||
|
||||
if (!imageRef) {
|
||||
logger.error("Missing deployment or image ref", {
|
||||
attemptId: attempt.id,
|
||||
workerId: attempt.backgroundWorker.id,
|
||||
params,
|
||||
});
|
||||
|
||||
return {
|
||||
success: false,
|
||||
};
|
||||
}
|
||||
|
||||
const { reason } = params;
|
||||
|
||||
// Check if we should accept this checkpoint
|
||||
switch (reason.type) {
|
||||
case "MANUAL": {
|
||||
// Always accept manual checkpoints
|
||||
break;
|
||||
}
|
||||
case "WAIT_FOR_DURATION": {
|
||||
// Always accept duration checkpoints
|
||||
break;
|
||||
}
|
||||
case "WAIT_FOR_TASK": {
|
||||
const childRun = await this._prisma.taskRun.findFirst({
|
||||
where: {
|
||||
friendlyId: reason.friendlyId,
|
||||
},
|
||||
select: {
|
||||
dependency: {
|
||||
select: {
|
||||
resumedAt: true,
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
if (!childRun) {
|
||||
logger.error("CreateCheckpointService: Pre-check - WAIT_FOR_TASK child run not found", {
|
||||
friendlyId: reason.friendlyId,
|
||||
params,
|
||||
});
|
||||
|
||||
return {
|
||||
success: false,
|
||||
keepRunAlive: false,
|
||||
};
|
||||
}
|
||||
|
||||
if (childRun.dependency?.resumedAt) {
|
||||
logger.info("CreateCheckpointService: Child run already resumed", {
|
||||
childRun,
|
||||
params,
|
||||
});
|
||||
|
||||
return {
|
||||
success: false,
|
||||
keepRunAlive: true,
|
||||
};
|
||||
}
|
||||
|
||||
break;
|
||||
}
|
||||
case "WAIT_FOR_BATCH": {
|
||||
// Routed by friendlyId so a run-ops id (NEW-resident) batch is found on the owning DB;
|
||||
// env-scoped to the dependent attempt's run (a batch shares its dependent's env). Read the
|
||||
// primary: a batch that just resumed the parent may lag the replica, and a stale resumedAt
|
||||
// (null) would checkpoint (suspend) an already-resumed run -> it stalls until a sweep.
|
||||
const batchRun = await this.runStore.findBatchTaskRunByFriendlyId(
|
||||
reason.batchFriendlyId,
|
||||
attempt.taskRun.runtimeEnvironmentId,
|
||||
undefined,
|
||||
this._prisma
|
||||
);
|
||||
|
||||
if (!batchRun) {
|
||||
logger.error("CreateCheckpointService: Pre-check - Batch not found", {
|
||||
batchFriendlyId: reason.batchFriendlyId,
|
||||
params,
|
||||
});
|
||||
|
||||
return {
|
||||
success: false,
|
||||
keepRunAlive: false,
|
||||
};
|
||||
}
|
||||
|
||||
if (batchRun.resumedAt) {
|
||||
logger.info("CreateCheckpointService: Batch already resumed", {
|
||||
batchRun,
|
||||
params,
|
||||
});
|
||||
|
||||
return {
|
||||
success: false,
|
||||
keepRunAlive: true,
|
||||
};
|
||||
}
|
||||
|
||||
break;
|
||||
}
|
||||
default: {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
//sleep to test slow checkpoints
|
||||
// Sleep a random value between 4 and 30 seconds
|
||||
// await new Promise((resolve) => {
|
||||
// const waitSeconds = Math.floor(Math.random() * 26) + 4;
|
||||
// logger.log(`Sleep for ${waitSeconds} seconds`);
|
||||
// setTimeout(resolve, waitSeconds * 1000);
|
||||
// });
|
||||
|
||||
let metadata: string;
|
||||
|
||||
if (params.reason.type === "MANUAL") {
|
||||
metadata = JSON.stringify({
|
||||
...params.reason,
|
||||
attemptId: attempt.id,
|
||||
previousAttemptStatus: attempt.status,
|
||||
previousRunStatus: attempt.taskRun.status,
|
||||
} satisfies ManualCheckpointMetadata);
|
||||
} else {
|
||||
metadata = JSON.stringify(params.reason);
|
||||
}
|
||||
|
||||
const checkpoint = await this._prisma.checkpoint.create({
|
||||
data: {
|
||||
...CheckpointId.generate(),
|
||||
runtimeEnvironmentId: attempt.taskRun.runtimeEnvironmentId,
|
||||
projectId: attempt.taskRun.projectId,
|
||||
attemptId: attempt.id,
|
||||
attemptNumber: attempt.number,
|
||||
runId: attempt.taskRunId,
|
||||
location: params.location,
|
||||
type: params.docker ? "DOCKER" : "KUBERNETES",
|
||||
reason: params.reason.type,
|
||||
metadata,
|
||||
imageRef,
|
||||
},
|
||||
});
|
||||
|
||||
const eventService = new CreateCheckpointRestoreEventService(this._prisma);
|
||||
|
||||
await this._prisma.taskRunAttempt.update({
|
||||
where: {
|
||||
id: attempt.id,
|
||||
},
|
||||
data: {
|
||||
status: params.reason.type === "RETRYING_AFTER_FAILURE" ? undefined : "PAUSED",
|
||||
taskRun: {
|
||||
update: {
|
||||
status: "WAITING_TO_RESUME",
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
let checkpointEvent: CheckpointRestoreEvent | undefined;
|
||||
|
||||
switch (reason.type) {
|
||||
case "MANUAL":
|
||||
case "WAIT_FOR_DURATION": {
|
||||
let restoreAtUnixTimeMs: number;
|
||||
|
||||
if (reason.type === "MANUAL") {
|
||||
// Restore immediately if not specified, useful for live migration
|
||||
restoreAtUnixTimeMs = reason.restoreAtUnixTimeMs ?? Date.now();
|
||||
} else {
|
||||
restoreAtUnixTimeMs = reason.now + reason.ms;
|
||||
}
|
||||
|
||||
checkpointEvent = await eventService.checkpoint({
|
||||
checkpointId: checkpoint.id,
|
||||
});
|
||||
|
||||
if (checkpointEvent) {
|
||||
await marqs.requeueMessage(
|
||||
attempt.taskRunId,
|
||||
{
|
||||
type: "RESUME_AFTER_DURATION",
|
||||
resumableAttemptId: attempt.id,
|
||||
checkpointEventId: checkpointEvent.id,
|
||||
},
|
||||
restoreAtUnixTimeMs,
|
||||
"resume"
|
||||
);
|
||||
|
||||
return {
|
||||
success: true,
|
||||
checkpoint,
|
||||
event: checkpointEvent,
|
||||
keepRunAlive: false,
|
||||
};
|
||||
}
|
||||
|
||||
break;
|
||||
}
|
||||
case "WAIT_FOR_TASK": {
|
||||
checkpointEvent = await eventService.checkpoint({
|
||||
checkpointId: checkpoint.id,
|
||||
dependencyFriendlyRunId: reason.friendlyId,
|
||||
});
|
||||
|
||||
if (checkpointEvent) {
|
||||
//heartbeats will start again when the run resumes
|
||||
logger.log("CreateCheckpointService: Canceling heartbeat", {
|
||||
attemptId: attempt.id,
|
||||
taskRunId: attempt.taskRunId,
|
||||
type: "WAIT_FOR_TASK",
|
||||
reason,
|
||||
params,
|
||||
});
|
||||
await marqs?.cancelHeartbeat(attempt.taskRunId);
|
||||
|
||||
const childRun = await this._prisma.taskRun.findFirst({
|
||||
where: {
|
||||
friendlyId: reason.friendlyId,
|
||||
},
|
||||
});
|
||||
|
||||
if (!childRun) {
|
||||
logger.error("CreateCheckpointService: WAIT_FOR_TASK child run not found", {
|
||||
friendlyId: reason.friendlyId,
|
||||
params,
|
||||
});
|
||||
|
||||
return {
|
||||
success: true,
|
||||
checkpoint,
|
||||
event: checkpointEvent,
|
||||
keepRunAlive: false,
|
||||
};
|
||||
}
|
||||
|
||||
const resumeService = new ResumeDependentParentsService(this._prisma);
|
||||
const result = await resumeService.call({ id: childRun.id });
|
||||
|
||||
if (result.success) {
|
||||
logger.log("CreateCheckpointService: Resumed dependent parents", {
|
||||
result,
|
||||
childRun,
|
||||
attempt,
|
||||
checkpointEvent,
|
||||
params,
|
||||
});
|
||||
} else {
|
||||
logger.error("CreateCheckpointService: Failed to resume dependent parents", {
|
||||
result,
|
||||
childRun,
|
||||
attempt,
|
||||
checkpointEvent,
|
||||
params,
|
||||
});
|
||||
}
|
||||
|
||||
return {
|
||||
success: true,
|
||||
checkpoint,
|
||||
event: checkpointEvent,
|
||||
keepRunAlive: false,
|
||||
};
|
||||
}
|
||||
|
||||
break;
|
||||
}
|
||||
case "WAIT_FOR_BATCH": {
|
||||
checkpointEvent = await eventService.checkpoint({
|
||||
checkpointId: checkpoint.id,
|
||||
batchDependencyFriendlyId: reason.batchFriendlyId,
|
||||
});
|
||||
|
||||
if (checkpointEvent) {
|
||||
//heartbeats will start again when the run resumes
|
||||
logger.log("CreateCheckpointService: Canceling heartbeat", {
|
||||
attemptId: attempt.id,
|
||||
taskRunId: attempt.taskRunId,
|
||||
type: "WAIT_FOR_BATCH",
|
||||
params,
|
||||
});
|
||||
await marqs?.cancelHeartbeat(attempt.taskRunId);
|
||||
|
||||
// Routed by friendlyId; read the primary (this._prisma) so a just-resumed batch that still
|
||||
// lags the replica doesn't leave a stale resumedAt and suspend an already-resumed run.
|
||||
const batchRun = await this.runStore.findBatchTaskRunByFriendlyId(
|
||||
reason.batchFriendlyId,
|
||||
attempt.taskRun.runtimeEnvironmentId,
|
||||
undefined,
|
||||
this._prisma
|
||||
);
|
||||
|
||||
if (!batchRun) {
|
||||
logger.error("CreateCheckpointService: Batch not found", {
|
||||
friendlyId: reason.batchFriendlyId,
|
||||
params,
|
||||
});
|
||||
|
||||
return {
|
||||
success: true,
|
||||
checkpoint,
|
||||
event: checkpointEvent,
|
||||
keepRunAlive: false,
|
||||
};
|
||||
}
|
||||
|
||||
//if there's a message in the queue, we make sure the checkpoint event is on it
|
||||
await marqs.replaceMessage(attempt.taskRun.id, {
|
||||
checkpointEventId: checkpointEvent.id,
|
||||
});
|
||||
|
||||
await ResumeBatchRunService.enqueue(batchRun.id, batchRun.batchVersion === "v3");
|
||||
|
||||
return {
|
||||
success: true,
|
||||
checkpoint,
|
||||
event: checkpointEvent,
|
||||
keepRunAlive: false,
|
||||
};
|
||||
}
|
||||
|
||||
break;
|
||||
}
|
||||
case "RETRYING_AFTER_FAILURE": {
|
||||
checkpointEvent = await eventService.checkpoint({
|
||||
checkpointId: checkpoint.id,
|
||||
});
|
||||
|
||||
// ACK is already handled by attempt completion
|
||||
break;
|
||||
}
|
||||
default: {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if (!checkpointEvent) {
|
||||
logger.error("No checkpoint event", {
|
||||
attemptId: attempt.id,
|
||||
checkpointId: checkpoint.id,
|
||||
params,
|
||||
});
|
||||
await marqs?.acknowledgeMessage(
|
||||
attempt.taskRunId,
|
||||
"No checkpoint event in CreateCheckpointService"
|
||||
);
|
||||
|
||||
return {
|
||||
success: false,
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
success: true,
|
||||
checkpoint,
|
||||
event: checkpointEvent,
|
||||
keepRunAlive: false,
|
||||
};
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user