chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,278 @@
|
||||
import type { V3TaskRunExecution } from "@trigger.dev/core/v3";
|
||||
import { parsePacket } from "@trigger.dev/core/v3";
|
||||
import type { TaskRun, TaskRunAttempt } from "@trigger.dev/database";
|
||||
import { MAX_TASK_RUN_ATTEMPTS } from "~/consts";
|
||||
import type { PrismaClientOrTransaction } from "~/db.server";
|
||||
import { $transaction, prisma } from "~/db.server";
|
||||
import { findQueueInEnvironment } from "~/models/taskQueue.server";
|
||||
import type { AuthenticatedEnvironment } from "~/services/apiAuth.server";
|
||||
import { logger } from "~/services/logger.server";
|
||||
import { reportInvocationUsage } from "~/services/platform.v3.server";
|
||||
import { generateFriendlyId } from "../friendlyIdentifiers";
|
||||
import { machinePresetFromConfig, machinePresetFromRun } from "../machinePresets.server";
|
||||
import { FINAL_RUN_STATUSES } from "../taskStatus";
|
||||
import { BaseService, ServiceValidationError } from "./baseService.server";
|
||||
import { CrashTaskRunService } from "./crashTaskRun.server";
|
||||
import { ExpireEnqueuedRunService } from "./expireEnqueuedRun.server";
|
||||
import { runStore } from "../runStore.server";
|
||||
import { controlPlaneResolver } from "~/v3/runOpsMigration/controlPlaneResolver.server";
|
||||
|
||||
export class CreateTaskRunAttemptService extends BaseService {
|
||||
public async call({
|
||||
runId,
|
||||
authenticatedEnv,
|
||||
setToExecuting = true,
|
||||
startAtZero = false,
|
||||
}: {
|
||||
runId: string;
|
||||
authenticatedEnv?: AuthenticatedEnvironment;
|
||||
setToExecuting?: boolean;
|
||||
startAtZero?: boolean;
|
||||
}): Promise<{
|
||||
execution: V3TaskRunExecution;
|
||||
run: TaskRun;
|
||||
attempt: TaskRunAttempt;
|
||||
}> {
|
||||
const environment =
|
||||
authenticatedEnv ?? (await getAuthenticatedEnvironmentFromRun(runId, this._prisma));
|
||||
|
||||
if (!environment) {
|
||||
throw new ServiceValidationError("Environment not found", 404);
|
||||
}
|
||||
|
||||
const isFriendlyId = runId.startsWith("run_");
|
||||
|
||||
return await this.traceWithEnv("call()", environment, async (span) => {
|
||||
if (isFriendlyId) {
|
||||
span.setAttribute("taskRunFriendlyId", runId);
|
||||
} else {
|
||||
span.setAttribute("taskRunId", runId);
|
||||
}
|
||||
|
||||
const taskRun = await this.runStore.findRun(
|
||||
{
|
||||
id: !isFriendlyId ? runId : undefined,
|
||||
friendlyId: isFriendlyId ? runId : undefined,
|
||||
runtimeEnvironmentId: environment.id,
|
||||
},
|
||||
{
|
||||
include: {
|
||||
attempts: {
|
||||
take: 1,
|
||||
orderBy: {
|
||||
number: "desc",
|
||||
},
|
||||
},
|
||||
batchItems: {
|
||||
include: {
|
||||
batchTaskRun: {
|
||||
select: {
|
||||
friendlyId: true,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
this._prisma
|
||||
);
|
||||
|
||||
logger.debug("Creating a task run attempt", { taskRun });
|
||||
|
||||
if (!taskRun) {
|
||||
throw new ServiceValidationError("Task run not found", 404);
|
||||
}
|
||||
|
||||
span.setAttribute("taskRunId", taskRun.id);
|
||||
span.setAttribute("taskRunFriendlyId", taskRun.friendlyId);
|
||||
span.setAttribute("taskRunStatus", taskRun.status);
|
||||
|
||||
if (taskRun.status === "CANCELED") {
|
||||
throw new ServiceValidationError("Task run is cancelled", 400);
|
||||
}
|
||||
|
||||
// If the run is finalized, it's pointless to create another attempt
|
||||
if (FINAL_RUN_STATUSES.includes(taskRun.status)) {
|
||||
throw new ServiceValidationError("Task run is already finished", 400);
|
||||
}
|
||||
|
||||
const lockedWorker = await controlPlaneResolver.resolveRunLockedWorker({
|
||||
lockedById: taskRun.lockedById,
|
||||
});
|
||||
const lockedBy = lockedWorker?.lockedBy;
|
||||
|
||||
if (!lockedBy) {
|
||||
throw new ServiceValidationError("Task run is not locked", 400);
|
||||
}
|
||||
|
||||
const queue = await findQueueInEnvironment(taskRun.queue, environment.id, lockedBy.id);
|
||||
|
||||
if (!queue) {
|
||||
throw new ServiceValidationError("Queue not found", 404);
|
||||
}
|
||||
|
||||
const nextAttemptNumber = taskRun.attempts[0]
|
||||
? taskRun.attempts[0].number + 1
|
||||
: startAtZero
|
||||
? 0
|
||||
: 1;
|
||||
|
||||
if (nextAttemptNumber > MAX_TASK_RUN_ATTEMPTS) {
|
||||
const service = new CrashTaskRunService(this._prisma);
|
||||
await service.call(taskRun.id, {
|
||||
reason: lockedBy.worker.supportsLazyAttempts
|
||||
? "Max attempts reached."
|
||||
: "Max attempts reached. Please upgrade your CLI and SDK.",
|
||||
});
|
||||
|
||||
throw new ServiceValidationError("Max attempts reached", 400);
|
||||
}
|
||||
|
||||
const taskRunAttempt = await $transaction(this._prisma, "create attempt", async (tx) => {
|
||||
const taskRunAttempt = await tx.taskRunAttempt.create({
|
||||
data: {
|
||||
number: nextAttemptNumber,
|
||||
friendlyId: generateFriendlyId("attempt"),
|
||||
taskRunId: taskRun.id,
|
||||
startedAt: new Date(),
|
||||
backgroundWorkerId: lockedBy.worker.id,
|
||||
backgroundWorkerTaskId: lockedBy.id,
|
||||
status: setToExecuting ? "EXECUTING" : "PENDING",
|
||||
queueId: queue.id,
|
||||
runtimeEnvironmentId: environment.id,
|
||||
},
|
||||
});
|
||||
|
||||
await tx.taskRun.update({
|
||||
where: {
|
||||
id: taskRun.id,
|
||||
},
|
||||
data: {
|
||||
status: setToExecuting ? "EXECUTING" : undefined,
|
||||
executedAt: taskRun.executedAt ?? new Date(),
|
||||
attemptNumber: nextAttemptNumber,
|
||||
},
|
||||
});
|
||||
|
||||
if (taskRun.ttl) {
|
||||
await ExpireEnqueuedRunService.ack(taskRun.id, tx);
|
||||
}
|
||||
|
||||
return taskRunAttempt;
|
||||
});
|
||||
|
||||
if (!taskRunAttempt) {
|
||||
logger.error("Failed to create task run attempt", { runId: taskRun.id, nextAttemptNumber });
|
||||
throw new ServiceValidationError("Failed to create task run attempt", 500);
|
||||
}
|
||||
|
||||
if (taskRunAttempt.number === 1 && taskRun.baseCostInCents > 0) {
|
||||
await reportInvocationUsage(environment.organizationId, taskRun.baseCostInCents, {
|
||||
runId: taskRun.id,
|
||||
});
|
||||
}
|
||||
|
||||
const machinePreset =
|
||||
machinePresetFromRun(taskRun) ?? machinePresetFromConfig(lockedBy.machineConfig ?? {});
|
||||
|
||||
const metadata = await parsePacket({
|
||||
data: taskRun.metadata ?? undefined,
|
||||
dataType: taskRun.metadataType,
|
||||
});
|
||||
|
||||
const execution: V3TaskRunExecution = {
|
||||
task: {
|
||||
id: lockedBy.slug,
|
||||
filePath: lockedBy.filePath,
|
||||
exportName: lockedBy.exportName ?? "@deprecated",
|
||||
},
|
||||
attempt: {
|
||||
id: taskRunAttempt.friendlyId,
|
||||
number: taskRunAttempt.number,
|
||||
startedAt: taskRunAttempt.startedAt ?? taskRunAttempt.createdAt,
|
||||
backgroundWorkerId: lockedBy.worker.id,
|
||||
backgroundWorkerTaskId: lockedBy.id,
|
||||
status: "EXECUTING" as const,
|
||||
},
|
||||
run: {
|
||||
id: taskRun.friendlyId,
|
||||
payload: taskRun.payload,
|
||||
payloadType: taskRun.payloadType,
|
||||
context: taskRun.context,
|
||||
createdAt: taskRun.createdAt,
|
||||
tags: taskRun.runTags ?? [],
|
||||
isTest: taskRun.isTest,
|
||||
isReplay: !!taskRun.replayedFromTaskRunFriendlyId,
|
||||
idempotencyKey: taskRun.idempotencyKey ?? undefined,
|
||||
startedAt: taskRun.startedAt ?? taskRun.createdAt,
|
||||
durationMs: taskRun.usageDurationMs,
|
||||
costInCents: taskRun.costInCents,
|
||||
baseCostInCents: taskRun.baseCostInCents,
|
||||
maxAttempts: taskRun.maxAttempts ?? undefined,
|
||||
version: lockedBy.worker.version,
|
||||
metadata,
|
||||
maxDuration: taskRun.maxDurationInSeconds ?? undefined,
|
||||
},
|
||||
queue: {
|
||||
id: queue.friendlyId,
|
||||
name: queue.name,
|
||||
},
|
||||
environment: {
|
||||
id: environment.id,
|
||||
slug: environment.slug,
|
||||
type: environment.type,
|
||||
},
|
||||
organization: {
|
||||
id: environment.organization.id,
|
||||
slug: environment.organization.slug,
|
||||
name: environment.organization.title,
|
||||
},
|
||||
project: {
|
||||
id: environment.project.id,
|
||||
ref: environment.project.externalRef,
|
||||
slug: environment.project.slug,
|
||||
name: environment.project.name,
|
||||
},
|
||||
batch:
|
||||
taskRun.batchItems[0] && taskRun.batchItems[0].batchTaskRun
|
||||
? { id: taskRun.batchItems[0].batchTaskRun.friendlyId }
|
||||
: undefined,
|
||||
machine: machinePreset,
|
||||
};
|
||||
|
||||
return {
|
||||
execution,
|
||||
run: taskRun,
|
||||
attempt: taskRunAttempt,
|
||||
};
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
async function getAuthenticatedEnvironmentFromRun(
|
||||
friendlyId: string,
|
||||
prismaClient?: PrismaClientOrTransaction
|
||||
) {
|
||||
const isFriendlyId = friendlyId.startsWith("run_");
|
||||
|
||||
const taskRun = await runStore.findRun(
|
||||
{
|
||||
id: !isFriendlyId ? friendlyId : undefined,
|
||||
friendlyId: isFriendlyId ? friendlyId : undefined,
|
||||
},
|
||||
{
|
||||
select: {
|
||||
runtimeEnvironmentId: true,
|
||||
},
|
||||
},
|
||||
prismaClient ?? prisma
|
||||
);
|
||||
|
||||
if (!taskRun) {
|
||||
return;
|
||||
}
|
||||
|
||||
return (
|
||||
(await controlPlaneResolver.resolveAuthenticatedEnv(taskRun.runtimeEnvironmentId)) ?? undefined
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user