diff --git a/apps/coordinator/package.json b/apps/coordinator/package.json index 19f7e2d8e..4cf789c16 100644 --- a/apps/coordinator/package.json +++ b/apps/coordinator/package.json @@ -4,7 +4,6 @@ "version": "0.0.1", "description": "", "main": "dist/index.cjs", - "type": "module", "scripts": { "build": "npm run build:bundle", "build:bundle": "esbuild src/index.ts --bundle --outfile=dist/index.mjs --platform=node --format=esm --target=esnext --banner:js=\"const require = createRequire(import.meta.url);\"", @@ -31,4 +30,4 @@ "tsx": "^4.7.0", "typescript": "^5.3.3" } -} +} \ No newline at end of file diff --git a/apps/coordinator/src/index.ts b/apps/coordinator/src/index.ts index 185b4c84c..25c1f5a85 100644 --- a/apps/coordinator/src/index.ts +++ b/apps/coordinator/src/index.ts @@ -301,6 +301,16 @@ class TaskCoordinator { taskSocket.emit("RESUME_AFTER_DURATION", message); }, + REQUEST_ATTEMPT_CANCELLATION: async (message) => { + const taskSocket = await this.#getAttemptSocket(message.attemptId); + + if (!taskSocket) { + logger.log("Socket for attempt not found", { attemptId: message.attemptId }); + return; + } + + taskSocket.emit("REQUEST_ATTEMPT_CANCELLATION", message); + }, }, }); @@ -385,11 +395,21 @@ class TaskCoordinator { if (!executionAck) { logger.error("no execution ack", { attemptId: socket.data.attemptId }); + + socket.emit("REQUEST_EXIT", { + version: "v1", + }); + return; } if (!executionAck.success) { logger.error("execution unsuccessful", { attemptId: socket.data.attemptId }); + + socket.emit("REQUEST_EXIT", { + version: "v1", + }); + return; } diff --git a/apps/coordinator/tsconfig.json b/apps/coordinator/tsconfig.json index 4c3965549..967c7281e 100644 --- a/apps/coordinator/tsconfig.json +++ b/apps/coordinator/tsconfig.json @@ -9,10 +9,10 @@ "strict": true, "skipLibCheck": true, "paths": { - "@trigger.dev/core/v3": ["../core/src/v3"], - "@trigger.dev/core/v3/*": ["../core/src/v3/*"], - "@trigger.dev/core-apps": ["../core-apps/src"], - "@trigger.dev/core-apps/*": ["../core-apps/src/*"] + "@trigger.dev/core/v3": ["../../packages/core/src/v3"], + "@trigger.dev/core/v3/*": ["../../packages/core/src/v3/*"], + "@trigger.dev/core-apps": ["../../packages/core-apps/src"], + "@trigger.dev/core-apps/*": ["../../packages/core-apps/src/*"] } } } diff --git a/apps/docker-provider/package.json b/apps/docker-provider/package.json index a123fbeb4..1b39e0026 100644 --- a/apps/docker-provider/package.json +++ b/apps/docker-provider/package.json @@ -4,7 +4,6 @@ "version": "0.0.1", "description": "", "main": "dist/index.cjs", - "type": "module", "scripts": { "build": "npm run build:bundle", "build:bundle": "esbuild src/index.ts --bundle --outfile=dist/index.mjs --platform=node --format=esm --target=esnext --banner:js=\"const require = createRequire(import.meta.url);\"", @@ -29,4 +28,4 @@ "tsx": "^4.7.0", "typescript": "^5.3.3" } -} +} \ No newline at end of file diff --git a/apps/docker-provider/tsconfig.json b/apps/docker-provider/tsconfig.json index a491b740b..345326da8 100644 --- a/apps/docker-provider/tsconfig.json +++ b/apps/docker-provider/tsconfig.json @@ -7,10 +7,10 @@ "strict": true, "skipLibCheck": true, "paths": { - "@trigger.dev/core/v3": ["../core/src/v3"], - "@trigger.dev/core/v3/*": ["../core/src/v3/*"], - "@trigger.dev/core-apps": ["../core-apps/src"], - "@trigger.dev/core-apps/*": ["../core-apps/src/*"] + "@trigger.dev/core/v3": ["../../packages/core/src/v3"], + "@trigger.dev/core/v3/*": ["../../packages/core/src/v3/*"], + "@trigger.dev/core-apps": ["../../packages/core-apps/src"], + "@trigger.dev/core-apps/*": ["../../packages/core-apps/src/*"] } } } diff --git a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.v3.$projectParam.runs.$runParam/route.tsx b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.v3.$projectParam.runs.$runParam/route.tsx index 3a1aab99a..ab7bf168f 100644 --- a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.v3.$projectParam.runs.$runParam/route.tsx +++ b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.v3.$projectParam.runs.$runParam/route.tsx @@ -334,7 +334,6 @@ function TasksTreeView({ onClick={(e) => { e.stopPropagation(); toggleExpandNode(node.id); - selectNode(node.id); scrollToNode(node.id); }} > diff --git a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.v3.$projectParam.test/route.tsx b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.v3.$projectParam.test/route.tsx index 64ccd7c24..8f12c3cb7 100644 --- a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.v3.$projectParam.test/route.tsx +++ b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.v3.$projectParam.test/route.tsx @@ -52,7 +52,7 @@ export default function Page() {
-
+
{tasks.length === 0 ? ( ) : ( @@ -81,7 +81,7 @@ function TaskSelector({ tasks }: { tasks: TaskListItem[] }) { const project = useProject(); return ( -
+
{tasks.map((t) => ( {({ isActive, isPending }) => ( diff --git a/apps/webapp/app/routes/api.v2.runs.$runParam.cancel.ts b/apps/webapp/app/routes/api.v2.runs.$runParam.cancel.ts new file mode 100644 index 000000000..d23fad9e2 --- /dev/null +++ b/apps/webapp/app/routes/api.v2.runs.$runParam.cancel.ts @@ -0,0 +1,52 @@ +import type { ActionFunctionArgs } from "@remix-run/server-runtime"; +import { json } from "@remix-run/server-runtime"; +import { z } from "zod"; +import { prisma } from "~/db.server"; +import { authenticateApiRequest } from "~/services/apiAuth.server"; +import { CancelTaskRunService } from "~/v3/services/cancelTaskRun.server"; + +const ParamsSchema = z.object({ + runParam: z.string(), +}); + +export async function action({ request, params }: ActionFunctionArgs) { + // Ensure this is a POST request + if (request.method.toUpperCase() !== "POST") { + return { status: 405, body: "Method Not Allowed" }; + } + + // Authenticate the request + const authenticationResult = await authenticateApiRequest(request); + + if (!authenticationResult) { + return json({ error: "Invalid or Missing API Key" }, { status: 401 }); + } + + const parsed = ParamsSchema.safeParse(params); + + if (!parsed.success) { + return json({ error: "Invalid or Missing runId" }, { status: 400 }); + } + + const { runParam } = parsed.data; + + const taskRun = await prisma.taskRun.findUnique({ + where: { + friendlyId: runParam, + }, + }); + + if (!taskRun) { + return json({ error: "Run not found" }, { status: 404 }); + } + + const service = new CancelTaskRunService(); + + try { + await service.call(taskRun); + } catch (error) { + return json({ error: "Internal Server Error" }, { status: 500 }); + } + + return json({ message: "Run cancelled" }, { status: 200 }); +} diff --git a/apps/webapp/app/services/worker.server.ts b/apps/webapp/app/services/worker.server.ts index e3a269160..954ed20cb 100644 --- a/apps/webapp/app/services/worker.server.ts +++ b/apps/webapp/app/services/worker.server.ts @@ -30,6 +30,9 @@ import { DeliverWebhookRequestService } from "./sources/deliverWebhookRequest.se import { PerformTaskOperationService } from "./tasks/performTaskOperation.server"; import { ProcessCallbackTimeoutService } from "./tasks/processCallbackTimeout.server"; import { ResumeTaskService } from "./tasks/resumeTask.server"; +import { ResumeTaskRunDependenciesService } from "~/v3/services/resumeTaskRunDependencies.server"; +import { ResumeBatchRunService } from "~/v3/services/resumeBatchRun.server"; +import { ResumeTaskDependencyService } from "~/v3/services/resumeTaskDependency.server"; const workerCatalog = { indexEndpoint: z.object({ @@ -107,6 +110,17 @@ const workerCatalog = { "v3.indexDeployment": z.object({ id: z.string(), }), + "v3.resumeTaskRunDependencies": z.object({ + attemptId: z.string(), + }), + "v3.resumeBatchRun": z.object({ + batchRunId: z.string(), + sourceTaskAttemptId: z.string(), + }), + "v3.resumeTaskDependency": z.object({ + dependencyId: z.string(), + sourceTaskAttemptId: z.string(), + }), }; const executionWorkerCatalog = { @@ -443,6 +457,33 @@ function getWorkerQueue() { return await service.call(payload.id); }, }, + "v3.resumeTaskRunDependencies": { + priority: 0, + maxAttempts: 5, + handler: async (payload, job) => { + const service = new ResumeTaskRunDependenciesService(); + + return await service.call(payload.attemptId); + }, + }, + "v3.resumeBatchRun": { + priority: 0, + maxAttempts: 5, + handler: async (payload, job) => { + const service = new ResumeBatchRunService(); + + return await service.call(payload.batchRunId, payload.sourceTaskAttemptId); + }, + }, + "v3.resumeTaskDependency": { + priority: 0, + maxAttempts: 5, + handler: async (payload, job) => { + const service = new ResumeTaskDependencyService(); + + return await service.call(payload.dependencyId, payload.sourceTaskAttemptId); + }, + }, }, }); } diff --git a/apps/webapp/app/v3/authenticatedSocketConnection.server.ts b/apps/webapp/app/v3/authenticatedSocketConnection.server.ts index d88055916..d5d3daa5c 100644 --- a/apps/webapp/app/v3/authenticatedSocketConnection.server.ts +++ b/apps/webapp/app/v3/authenticatedSocketConnection.server.ts @@ -10,6 +10,7 @@ import { AuthenticatedEnvironment } from "~/services/apiAuth.server"; import { logger } from "~/services/logger.server"; import { DevQueueConsumer } from "./marqs/devQueueConsumer.server"; import type { WebSocket, MessageEvent, CloseEvent, ErrorEvent } from "ws"; +import { env } from "~/env.server"; export class AuthenticatedSocketConnection { public id: string; @@ -26,6 +27,10 @@ export class AuthenticatedSocketConnection { schema: serverWebsocketMessages, sender: async (message) => { return new Promise((resolve, reject) => { + if (!ws.OPEN) { + return reject(new Error("Websocket is not open")); + } + ws.send(JSON.stringify(message), {}, (err) => { if (err) { reject(err); @@ -84,6 +89,8 @@ export class AuthenticatedSocketConnection { } async #handleClose(ev: CloseEvent) { + logger.debug("[AuthenticatedSocketConnection] Websocket closed", { ev }); + await this._consumer.stop(); this.onClose.post(ev); diff --git a/apps/webapp/app/v3/eventRepository.server.ts b/apps/webapp/app/v3/eventRepository.server.ts index c7fd2cfd4..084481485 100644 --- a/apps/webapp/app/v3/eventRepository.server.ts +++ b/apps/webapp/app/v3/eventRepository.server.ts @@ -44,6 +44,7 @@ export type TraceAttributes = Partial< CreatableEvent, | "attemptId" | "isError" + | "isCancelled" | "runId" | "runIsTest" | "output" diff --git a/apps/webapp/app/v3/marqs/devPubSub.server.ts b/apps/webapp/app/v3/marqs/devPubSub.server.ts new file mode 100644 index 000000000..7d3de2ff6 --- /dev/null +++ b/apps/webapp/app/v3/marqs/devPubSub.server.ts @@ -0,0 +1,38 @@ +import { z } from "zod"; +import { singleton } from "~/utils/singleton"; +import { ZodPubSub, ZodSubscriber } from "../utils/zodPubSub.server"; +import { env } from "~/env.server"; + +const messageCatalog = { + CANCEL_ATTEMPT: z.object({ + version: z.literal("v1").default("v1"), + backgroundWorkerId: z.string(), + attemptId: z.string(), + taskRunId: z.string(), + }), +}; + +export type DevSubscriber = ZodSubscriber; + +export const devPubSub = singleton("devPubSub", initializeDevPubSub); + +function initializeDevPubSub() { + return new ZodPubSub({ + redis: { + port: env.REDIS_PORT, + host: env.REDIS_HOST, + username: env.REDIS_USERNAME, + password: env.REDIS_PASSWORD, + enableAutoPipelining: true, + ...(env.REDIS_TLS_DISABLED === "true" ? {} : { tls: {} }), + }, + schema: { + CANCEL_ATTEMPT: z.object({ + version: z.literal("v1").default("v1"), + backgroundWorkerId: z.string(), + attemptId: z.string(), + taskRunId: z.string(), + }), + }, + }); +} diff --git a/apps/webapp/app/v3/marqs/devQueueConsumer.server.ts b/apps/webapp/app/v3/marqs/devQueueConsumer.server.ts index 78621c38c..fa05a4f05 100644 --- a/apps/webapp/app/v3/marqs/devQueueConsumer.server.ts +++ b/apps/webapp/app/v3/marqs/devQueueConsumer.server.ts @@ -17,6 +17,7 @@ import { marqs } from "../marqs.server"; import { CancelAttemptService } from "../services/cancelAttempt.server"; import { CompleteAttemptService } from "../services/completeAttempt.server"; import { attributesFromAuthenticatedEnv } from "../tracer.server"; +import { DevSubscriber, devPubSub } from "./devPubSub.server"; const tracer = trace.getTracer("devQueueConsumer"); @@ -36,9 +37,11 @@ export type DevQueueConsumerOptions = { export class DevQueueConsumer { private _backgroundWorkers: Map = new Map(); + private _backgroundWorkerSubscriber: Map = new Map(); private _deprecatedWorkers: Map = new Map(); private _enabled = false; - private _options: Required; + private _maximumItemsPerTrace: number; + private _traceTimeoutSeconds: number; private _perTraceCountdown: number | undefined; private _lastNewTrace: Date | undefined; private _currentSpanContext: Context | undefined; @@ -51,12 +54,10 @@ export class DevQueueConsumer { constructor( public env: AuthenticatedEnvironment, private _sender: ZodMessageSender, - options: DevQueueConsumerOptions = {} + private _options: DevQueueConsumerOptions = {} ) { - this._options = { - maximumItemsPerTrace: options.maximumItemsPerTrace ?? 1_000, // 1k items per trace - traceTimeoutSeconds: options.traceTimeoutSeconds ?? 60, // 60 seconds - }; + this._traceTimeoutSeconds = _options.traceTimeoutSeconds ?? 60; + this._maximumItemsPerTrace = _options.maximumItemsPerTrace ?? 1_000; } // This method is called when a background worker is deprecated and will no longer be used unless a run is locked to it @@ -87,6 +88,21 @@ export class DevQueueConsumer { logger.debug("Registered background worker", { backgroundWorker: backgroundWorker.id }); + const subscriber = await devPubSub.subscribe(`backgroundWorker:${backgroundWorker.id}:*`); + + subscriber.on("CANCEL_ATTEMPT", async (message) => { + await this._sender.send("BACKGROUND_WORKER_MESSAGE", { + backgroundWorkerId: backgroundWorker.friendlyId, + data: { + type: "CANCEL_ATTEMPT", + taskAttemptId: message.attemptId, + taskRunId: message.taskRunId, + }, + }); + }); + + this._backgroundWorkerSubscriber.set(backgroundWorker.id, subscriber); + // Start reading from the queue if we haven't already this.#enable(); } @@ -133,6 +149,16 @@ export class DevQueueConsumer { // We need to cancel all the in progress task run attempts and ack the messages so they will stop processing await this.#cancelInProgressAttempts(reason); + + // We need to unsubscribe from the background worker channels + for (const [id, subscriber] of this._backgroundWorkerSubscriber) { + logger.debug("Unsubscribing from background worker channel", { id }); + + await subscriber.stopListening(); + this._backgroundWorkerSubscriber.delete(id); + + logger.debug("Unsubscribed from background worker channel", { id }); + } } async #cancelInProgressAttempts(reason: string) { @@ -144,6 +170,10 @@ export class DevQueueConsumer { this._inProgressAttempts.clear(); + logger.debug("Cancelling in progress attempts", { + attempts: Array.from(inProgressAttempts.keys()), + }); + for (const [attemptId, messageId] of inProgressAttempts) { await this.#cancelInProgressAttempt(attemptId, messageId, service, cancelledAt, reason); } @@ -156,6 +186,8 @@ export class DevQueueConsumer { cancelledAt: Date, reason: string ) { + logger.debug("Cancelling in progress attempt", { attemptId, messageId }); + try { await cancelAttemptService.call(attemptId, messageId, cancelledAt, reason, this.env); } catch (e) { @@ -189,7 +221,7 @@ export class DevQueueConsumer { // Check if the trace has expired if ( this._perTraceCountdown === 0 || - Date.now() - this._lastNewTrace!.getTime() > this._options.traceTimeoutSeconds * 1000 || + Date.now() - this._lastNewTrace!.getTime() > this._traceTimeoutSeconds * 1000 || this._currentSpanContext === undefined || this._endSpanInNextIteration ) { @@ -309,6 +341,7 @@ export class DevQueueConsumer { data: { lockedAt: new Date(), lockedById: backgroundTask.id, + status: "EXECUTING", }, include: { attempts: { @@ -365,6 +398,7 @@ export class DevQueueConsumer { backgroundWorkerTaskId: backgroundTask.id, status: "EXECUTING" as const, queueId: queue.id, + runtimeEnvironmentId: this.env.id, }, }); @@ -441,6 +475,11 @@ export class DevQueueConsumer { }, }); + logger.debug("Saving the in progress attempt", { + taskRunAttempt: taskRunAttempt.id, + messageId: message.messageId, + }); + this._inProgressAttempts.set(taskRunAttempt.friendlyId, message.messageId); } catch (e) { if (e instanceof Error) { diff --git a/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts b/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts index 4413b604c..78ec796cd 100644 --- a/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts +++ b/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts @@ -31,9 +31,11 @@ const MessageBody = z.discriminatedUnion("type", [ z.object({ type: z.literal("RESUME"), completedAttemptIds: z.string().array(), + resumableAttemptId: z.string(), }), z.object({ type: z.literal("RESUME_AFTER_DURATION"), + resumableAttemptId: z.string(), }), ]); @@ -42,6 +44,8 @@ type BackgroundWorkerWithTasks = BackgroundWorker & { tasks: BackgroundWorkerTas export type SharedQueueConsumerOptions = { maximumItemsPerTrace?: number; traceTimeoutSeconds?: number; + nextTickInterval?: number; + interval?: number; }; export class SharedQueueConsumer { @@ -66,6 +70,8 @@ export class SharedQueueConsumer { this._options = { maximumItemsPerTrace: options.maximumItemsPerTrace ?? 1_000, // 1k items per trace traceTimeoutSeconds: options.traceTimeoutSeconds ?? 60, // 60 seconds + nextTickInterval: options.nextTickInterval ?? 1000, // 1 second + interval: options.interval ?? 100, // 100ms }; } @@ -233,7 +239,7 @@ export class SharedQueueConsumer { const message = await marqs?.dequeueMessageInSharedQueue(); if (!message) { - setTimeout(() => this.#doWork(), 1000); + setTimeout(() => this.#doWork(), this._options.nextTickInterval); return; } @@ -257,7 +263,7 @@ export class SharedQueueConsumer { envId, }); await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), 100); + setTimeout(() => this.#doWork(), this._options.interval); return; } @@ -272,7 +278,7 @@ export class SharedQueueConsumer { await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), 100); + setTimeout(() => this.#doWork(), this._options.interval); return; } @@ -290,7 +296,19 @@ export class SharedQueueConsumer { messageId: message.messageId, }); await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), 100); + setTimeout(() => this.#doWork(), this._options.interval); + return; + } + + if (existingTaskRun.status !== "PENDING") { + logger.debug("Task run is not pending, aborting", { + queueMessage: message.data, + messageId: message.messageId, + taskRun: existingTaskRun.id, + status: existingTaskRun.status, + }); + await marqs?.acknowledgeMessage(message.messageId); + setTimeout(() => this.#doWork(), this._options.interval); return; } @@ -321,7 +339,7 @@ export class SharedQueueConsumer { messageId: message.messageId, }); await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), 100); + setTimeout(() => this.#doWork(), this._options.interval); return; } @@ -332,7 +350,7 @@ export class SharedQueueConsumer { deployment: deployment.id, }); await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), 100); + setTimeout(() => this.#doWork(), this._options.interval); return; } @@ -351,7 +369,7 @@ export class SharedQueueConsumer { await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), 100); + setTimeout(() => this.#doWork(), this._options.interval); return; } @@ -383,7 +401,7 @@ export class SharedQueueConsumer { await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), 100); + setTimeout(() => this.#doWork(), this._options.interval); return; } @@ -398,7 +416,7 @@ export class SharedQueueConsumer { if (!queue) { await marqs?.nackMessage(message.messageId); - setTimeout(() => this.#doWork(), 1000); + setTimeout(() => this.#doWork(), this._options.nextTickInterval); return; } @@ -417,6 +435,7 @@ export class SharedQueueConsumer { backgroundWorkerTaskId: backgroundTask.id, status: "PENDING" as const, queueId: queue.id, + runtimeEnvironmentId: environment.id, }, }); @@ -462,7 +481,7 @@ export class SharedQueueConsumer { // Finally we need to nack the message so it can be retried await marqs?.nackMessage(message.messageId); } finally { - setTimeout(() => this.#doWork(), 100); + setTimeout(() => this.#doWork(), this._options.interval); } break; } @@ -474,41 +493,47 @@ export class SharedQueueConsumer { messageId: message.messageId, }); await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), 100); + setTimeout(() => this.#doWork(), this._options.interval); return; } - const resumableRun = await prisma.taskRun.findFirst({ + const resumableRun = await prisma.taskRun.findUnique({ where: { id: message.messageId, }, + }); + + if (!resumableRun) { + logger.error("Resumable run not found", { + queueMessage: message.data, + messageId: message.messageId, + }); + await marqs?.acknowledgeMessage(message.messageId); + setTimeout(() => this.#doWork(), this._options.interval); + return; + } + + const resumableAttempt = await prisma.taskRunAttempt.findUnique({ + where: { + id: messageBody.data.resumableAttemptId, + }, include: { - attempts: { + checkpoints: { + take: 1, orderBy: { createdAt: "desc", }, - take: 1, - include: { - checkpoints: { - take: 1, - orderBy: { - createdAt: "desc", - }, - }, - }, }, }, }); - const resumableAttempt = resumableRun?.attempts[0]; - if (!resumableAttempt) { logger.error("Resumable attempt not found", { queueMessage: message.data, messageId: message.messageId, }); await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), 100); + setTimeout(() => this.#doWork(), this._options.interval); return; } @@ -523,7 +548,7 @@ export class SharedQueueConsumer { if (!queue) { await marqs?.nackMessage(message.messageId); - setTimeout(() => this.#doWork(), 1000); + setTimeout(() => this.#doWork(), this._options.nextTickInterval); return; } @@ -543,7 +568,7 @@ export class SharedQueueConsumer { resumableAttemptId: resumableAttempt.id, }); await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), 100); + setTimeout(() => this.#doWork(), this._options.interval); return; } @@ -553,6 +578,13 @@ export class SharedQueueConsumer { }, data: { status: "EXECUTING", + taskRun: { + update: { + data: { + status: "EXECUTING", + }, + }, + }, }, }); @@ -565,7 +597,7 @@ export class SharedQueueConsumer { reason: latestCheckpoint.reason ?? undefined, }); - setTimeout(() => this.#doWork(), 100); + setTimeout(() => this.#doWork(), this._options.interval); return; } @@ -593,7 +625,7 @@ export class SharedQueueConsumer { messageId: message.messageId, }); await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), 100); + setTimeout(() => this.#doWork(), this._options.interval); return; } @@ -601,7 +633,7 @@ export class SharedQueueConsumer { if (!completion) { await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), 100); + setTimeout(() => this.#doWork(), this._options.interval); return; } @@ -614,7 +646,7 @@ export class SharedQueueConsumer { if (!executionPayload) { await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), 100); + setTimeout(() => this.#doWork(), this._options.interval); return; } @@ -641,43 +673,34 @@ export class SharedQueueConsumer { // Finally we need to nack the message so it can be retried await marqs?.nackMessage(message.messageId); } finally { - setTimeout(() => this.#doWork(), 100); + setTimeout(() => this.#doWork(), this._options.interval); } break; } // Resume after duration-based wait case "RESUME_AFTER_DURATION": { - const resumableRun = await prisma.taskRun.findFirst({ + const resumableAttempt = await prisma.taskRunAttempt.findUnique({ where: { - id: message.messageId, + id: messageBody.data.resumableAttemptId, }, include: { - attempts: { + checkpoints: { + take: 1, orderBy: { createdAt: "desc", }, - take: 1, - include: { - checkpoints: { - take: 1, - orderBy: { - createdAt: "desc", - }, - }, - }, }, + taskRun: true, }, }); - const resumableAttempt = resumableRun?.attempts[0]; - if (!resumableAttempt) { logger.error("Resumable attempt not found", { queueMessage: message.data, messageId: message.messageId, }); await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), 100); + setTimeout(() => this.#doWork(), this._options.interval); return; } @@ -687,7 +710,7 @@ export class SharedQueueConsumer { messageId: message.messageId, }); await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), 100); + setTimeout(() => this.#doWork(), this._options.interval); return; } @@ -702,7 +725,7 @@ export class SharedQueueConsumer { resumableAttemptId: resumableAttempt.id, }); await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), 100); + setTimeout(() => this.#doWork(), this._options.interval); return; } @@ -736,7 +759,7 @@ export class SharedQueueConsumer { // Finally we need to nack the message so it can be retried await marqs?.nackMessage(message.messageId); } finally { - setTimeout(() => this.#doWork(), 100); + setTimeout(() => this.#doWork(), this._options.interval); } break; } @@ -811,14 +834,14 @@ class SharedQueueTasks { include: { backgroundWorker: true, backgroundWorkerTask: true, + runtimeEnvironment: { + include: { + organization: true, + project: true, + }, + }, taskRun: { include: { - runtimeEnvironment: { - include: { - organization: true, - project: true, - }, - }, tags: true, batchItem: { include: { @@ -836,6 +859,30 @@ class SharedQueueTasks { return; } + if (attempt.status === "CANCELED") { + return; + } + + if (attempt.status === "FAILED") { + return; + } + + if (attempt.status === "COMPLETED") { + return; + } + + if (attempt.taskRun.status === "CANCELED") { + return; + } + + if (attempt.taskRun.status === "COMPLETED_SUCCESSFULLY") { + return; + } + + if (attempt.taskRun.status === "COMPLETED_WITH_ERRORS") { + return; + } + if (setToExecuting) { await prisma.taskRunAttempt.update({ where: { @@ -843,6 +890,13 @@ class SharedQueueTasks { }, data: { status: "EXECUTING", + taskRun: { + update: { + data: { + status: "EXECUTING", + }, + }, + }, }, }); } @@ -877,20 +931,20 @@ class SharedQueueTasks { name: queue.name, }, environment: { - id: taskRun.runtimeEnvironment.id, - slug: taskRun.runtimeEnvironment.slug, - type: taskRun.runtimeEnvironment.type, + id: attempt.runtimeEnvironment.id, + slug: attempt.runtimeEnvironment.slug, + type: attempt.runtimeEnvironment.type, }, organization: { - id: taskRun.runtimeEnvironment.organization.id, - slug: taskRun.runtimeEnvironment.organization.slug, - name: taskRun.runtimeEnvironment.organization.title, + id: attempt.runtimeEnvironment.organization.id, + slug: attempt.runtimeEnvironment.organization.slug, + name: attempt.runtimeEnvironment.organization.title, }, project: { - id: taskRun.runtimeEnvironment.project.id, - ref: taskRun.runtimeEnvironment.project.externalRef, - slug: taskRun.runtimeEnvironment.project.slug, - name: taskRun.runtimeEnvironment.project.name, + id: attempt.runtimeEnvironment.project.id, + ref: attempt.runtimeEnvironment.project.externalRef, + slug: attempt.runtimeEnvironment.project.slug, + name: attempt.runtimeEnvironment.project.name, }, batch: taskRun.batchItem?.batchTaskRun ? { id: taskRun.batchItem.batchTaskRun.friendlyId } @@ -904,8 +958,8 @@ class SharedQueueTasks { const environmentRepository = new EnvironmentVariablesRepository(); const variables = await environmentRepository.getEnvironmentVariables( - attempt.taskRun.runtimeEnvironment.projectId, - attempt.taskRun.runtimeEnvironmentId + attempt.runtimeEnvironment.projectId, + attempt.runtimeEnvironmentId ); const payload: ProdTaskRunExecutionPayload = { diff --git a/apps/webapp/app/v3/services/cancelAttempt.server.ts b/apps/webapp/app/v3/services/cancelAttempt.server.ts index 468f93de8..aea2abdda 100644 --- a/apps/webapp/app/v3/services/cancelAttempt.server.ts +++ b/apps/webapp/app/v3/services/cancelAttempt.server.ts @@ -5,6 +5,7 @@ import { BaseService } from "./baseService.server"; import { logger } from "~/services/logger.server"; import { PrismaClientOrTransaction, prisma } from "~/db.server"; +import { ResumeTaskRunDependenciesService } from "./resumeTaskRunDependencies.server"; export class CancelAttemptService extends BaseService { public async call( @@ -49,6 +50,14 @@ export class CancelAttemptService extends BaseService { }, data: { status: "CANCELED", + completedAt: cancelledAt, + taskRun: { + update: { + data: { + status: "INTERRUPTED", + }, + }, + }, }, }); @@ -65,6 +74,10 @@ export class CancelAttemptService extends BaseService { return eventRepository.cancelEvent(event, cancelledAt, reason); }) ); + + if (environment?.type !== "DEVELOPMENT") { + await ResumeTaskRunDependenciesService.enqueue(taskRunAttempt.id, this._prisma); + } }); } } @@ -78,14 +91,10 @@ async function getAuthenticatedEnvironmentFromAttempt( friendlyId, }, include: { - taskRun: { + runtimeEnvironment: { include: { - runtimeEnvironment: { - include: { - organization: true, - project: true, - }, - }, + organization: true, + project: true, }, }, }, @@ -95,5 +104,5 @@ async function getAuthenticatedEnvironmentFromAttempt( return; } - return taskRunAttempt?.taskRun.runtimeEnvironment; + return taskRunAttempt?.runtimeEnvironment; } diff --git a/apps/webapp/app/v3/services/cancelTaskRun.server.ts b/apps/webapp/app/v3/services/cancelTaskRun.server.ts new file mode 100644 index 000000000..72e5fcb6d --- /dev/null +++ b/apps/webapp/app/v3/services/cancelTaskRun.server.ts @@ -0,0 +1,129 @@ +import { TaskRun, TaskRunAttemptStatus, TaskRunStatus } from "@trigger.dev/database"; +import { eventRepository } from "../eventRepository.server"; +import { marqs } from "../marqs.server"; +import { devPubSub } from "../marqs/devPubSub.server"; +import { BaseService } from "./baseService.server"; +import { socketIo } from "../handleSocketIo.server"; +import { assertUnreachable } from "../utils/asserts.server"; +import { CancelAttemptService } from "./cancelAttempt.server"; +import { logger } from "~/services/logger.server"; + +const CANCELLABLE_STATUSES: Array = [ + "PENDING", + "EXECUTING", + "PAUSED", + "WAITING_TO_RESUME", + "PAUSED", + "RETRYING_AFTER_FAILURE", +]; + +const CANCELLABLE_ATTEMPT_STATUSES: Array = [ + "EXECUTING", + "PAUSED", + "PENDING", +]; + +export class CancelTaskRunService extends BaseService { + public async call(taskRun: TaskRun) { + // Make sure the task run is in a cancellable state + if (!CANCELLABLE_STATUSES.includes(taskRun.status)) { + 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", + }, + include: { + attempts: { + where: { + status: { + in: CANCELLABLE_ATTEMPT_STATUSES, + }, + }, + include: { + backgroundWorker: true, + runtimeEnvironment: true, + }, + }, + dependency: true, + runtimeEnvironment: true, + }, + }); + + const inProgressEvents = await eventRepository.queryIncompleteEvents({ + runId: taskRun.friendlyId, + }); + + logger.debug("Cancelling in-progress events", { + inProgressEvents: inProgressEvents.map((event) => event.id), + }); + + await Promise.all( + inProgressEvents.map((event) => { + return eventRepository.cancelEvent(event, new Date(), "Task run was cancelled by user"); + }) + ); + + // Cancel any in progress attempts + for (const attempt of cancelledTaskRun.attempts) { + if (attempt.runtimeEnvironment.type === "DEVELOPMENT") { + // Signal the task run attempt to stop + await devPubSub.publish( + `backgroundWorker:${attempt.backgroundWorkerId}:${attempt.id}`, + "CANCEL_ATTEMPT", + { + attemptId: attempt.friendlyId, + backgroundWorkerId: attempt.backgroundWorker.friendlyId, + taskRunId: cancelledTaskRun.friendlyId, + } + ); + } else { + switch (attempt.status) { + case "EXECUTING": { + // We need to send a cancel message to the coordinator + socketIo.coordinatorNamespace.emit("REQUEST_ATTEMPT_CANCELLATION", { + version: "v1", + attemptId: attempt.id, + }); + + break; + } + case "PENDING": + case "PAUSED": { + logger.debug("Cancelling pending or paused attempt", { + attempt, + }); + + const service = new CancelAttemptService(); + + await service.call( + attempt.friendlyId, + taskRun.id, + new Date(), + "Task run was cancelled by user" + ); + + break; + } + case "CANCELED": + case "COMPLETED": + case "FAILED": { + // Do nothing + break; + } + default: { + assertUnreachable(attempt.status); + } + } + } + } + } +} diff --git a/apps/webapp/app/v3/services/completeAttempt.server.ts b/apps/webapp/app/v3/services/completeAttempt.server.ts index 0c5b3c901..c36e5bd0d 100644 --- a/apps/webapp/app/v3/services/completeAttempt.server.ts +++ b/apps/webapp/app/v3/services/completeAttempt.server.ts @@ -1,72 +1,131 @@ +import { Attributes } from "@opentelemetry/api"; import { RetryOptions, TaskRunContext, TaskRunExecution, TaskRunExecutionResult, + TaskRunFailedExecutionResult, + TaskRunSuccessfulExecutionResult, defaultRetryOptions, flattenAttributes, } from "@trigger.dev/core/v3"; +import { PrismaClientOrTransaction } from "~/db.server"; import { AuthenticatedEnvironment } from "~/services/apiAuth.server"; +import { logger } from "~/services/logger.server"; +import { safeJsonParse } from "~/utils/json"; import { eventRepository } from "../eventRepository.server"; import { marqs } from "../marqs.server"; import { BaseService } from "./baseService.server"; -import { Attributes } from "@opentelemetry/api"; -import { logger } from "~/services/logger.server"; +import { ResumeTaskRunDependenciesService } from "./resumeTaskRunDependencies.server"; +import { CancelAttemptService } from "./cancelAttempt.server"; + +type FoundAttempt = Awaited>; export class CompleteAttemptService extends BaseService { public async call( completion: TaskRunExecutionResult, execution: TaskRunExecution, env?: AuthenticatedEnvironment - ): Promise<"ACKNOWLEDGED" | "RETRIED" | "FAILED"> { - const taskRunAttempt = completion.ok - ? await this._prisma.taskRunAttempt.update({ - where: { friendlyId: completion.id }, - data: { - status: "COMPLETED", - completedAt: new Date(), - output: completion.output, - outputType: completion.outputType, - }, - include: { - taskRun: { - include: { - batchItem: true, - dependency: { - include: { - dependentAttempt: true, - dependentBatchRun: true, - }, - }, - }, - }, - backgroundWorkerTask: true, - }, - }) - : await this._prisma.taskRunAttempt.update({ - where: { friendlyId: completion.id }, - data: { - status: "FAILED", - completedAt: new Date(), - error: completion.error, - }, - include: { - taskRun: { - include: { - batchItem: true, - dependency: { - include: { - dependentAttempt: true, - dependentBatchRun: true, - }, - }, - }, - }, - backgroundWorkerTask: true, - }, - }); + ) { + const taskRunAttempt = await findAttempt(this._prisma, completion.id); - if (!completion.ok && completion.retry !== undefined) { + if (!taskRunAttempt) { + logger.error("[CompleteAttemptService] Task run attempt not found", { id: completion.id }); + + // Update the task run to be failed + await this._prisma.taskRun.update({ + where: { + friendlyId: execution.run.id, + }, + data: { + status: "SYSTEM_FAILURE", + }, + }); + + return "FAILED"; + } + + if (completion.ok) { + return await this.#completeAttemptSuccessfully(completion, taskRunAttempt, env); + } else { + return await this.#completeAttemptFailed(completion, execution, taskRunAttempt, env); + } + } + + async #completeAttemptSuccessfully( + completion: TaskRunSuccessfulExecutionResult, + taskRunAttempt: NonNullable, + env?: AuthenticatedEnvironment + ) { + await this._prisma.taskRunAttempt.update({ + where: { friendlyId: completion.id }, + data: { + status: "COMPLETED", + completedAt: new Date(), + output: completion.output, + outputType: completion.outputType, + taskRun: { + update: { + data: { + status: "COMPLETED_SUCCESSFULLY", + }, + }, + }, + }, + }); + + logger.debug("Completed attempt successfully, 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(), + attributes: { + isError: false, + output: completion.output ? (safeJsonParse(completion.output) as Attributes) : undefined, + }, + }); + + if (!env || env.type !== "DEVELOPMENT") { + await ResumeTaskRunDependenciesService.enqueue(taskRunAttempt.id, this._prisma); + } + + return "ACKNOWLEDGED"; + } + + async #completeAttemptFailed( + completion: TaskRunFailedExecutionResult, + execution: TaskRunExecution, + taskRunAttempt: NonNullable, + env?: AuthenticatedEnvironment + ) { + if ( + completion.error.type === "INTERNAL_ERROR" && + completion.error.code === "TASK_RUN_CANCELLED" + ) { + // We need to cancel the task run instead of fail it + const cancelService = new CancelAttemptService(); + + return await cancelService.call( + taskRunAttempt.friendlyId, + taskRunAttempt.taskRunId, + new Date(), + "Cancelled by user", + env + ); + } + + await this._prisma.taskRunAttempt.update({ + where: { friendlyId: completion.id }, + data: { + status: "FAILED", + completedAt: new Date(), + error: completion.error, + }, + }); + + if (completion.retry !== undefined) { const retryConfig = taskRunAttempt.backgroundWorkerTask.retryConfig ? { ...defaultRetryOptions, @@ -111,6 +170,15 @@ export class CompleteAttemptService extends BaseService { logger.debug("Retrying", { taskRun: taskRunAttempt.taskRun.friendlyId }); + await this._prisma.taskRun.update({ + where: { + id: taskRunAttempt.taskRunId, + }, + data: { + status: "RETRYING_AFTER_FAILURE", + }, + }); + if (environment.type === "DEVELOPMENT") { // This is already an EXECUTE message so we can just NACK await marqs?.nackMessage(taskRunAttempt.taskRunId, completion.retry.timestamp); @@ -127,179 +195,31 @@ export class CompleteAttemptService extends BaseService { } return "RETRIED"; - } - // Attempt succeeded or this was the last retry - else { + } 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 - if (completion.ok) { - await eventRepository.completeEvent(taskRunAttempt.taskRun.spanId, { - endTime: new Date(), - attributes: { - isError: false, - output: completion.output ? (JSON.parse(completion.output) as Attributes) : undefined, - }, - }); - } else { - await eventRepository.completeEvent(taskRunAttempt.taskRun.spanId, { - endTime: new Date(), - attributes: { - isError: true, - }, - }); - } + await eventRepository.completeEvent(taskRunAttempt.taskRun.spanId, { + endTime: new Date(), + attributes: { + isError: true, + }, + }); - const { batchItem, dependency } = taskRunAttempt.taskRun; + await this._prisma.taskRun.update({ + where: { + id: taskRunAttempt.taskRunId, + }, + data: { + status: "COMPLETED_WITH_ERRORS", + }, + }); - // This run is part of a batch so we should update its status - if (batchItem) { - logger.debug("Completing attempt with batch item", { batchItem }); - - await this._prisma.batchTaskRunItem.update({ - where: { - id: batchItem.id, - }, - data: { - status: completion.ok ? "COMPLETED" : "FAILED", - }, - }); - - const finalizedBatchRun = await this._prisma.batchTaskRun.findFirst({ - where: { - id: batchItem.batchTaskRunId, - dependentTaskAttemptId: { - not: null, - }, - items: { - every: { - status: { - not: "PENDING", - }, - }, - }, - }, - include: { - dependentTaskAttempt: { - include: { - taskRun: true, - }, - }, - items: { - include: { - taskRun: { - include: { - attempts: { - orderBy: { - completedAt: "desc", - }, - take: 1, - select: { - id: true, - }, - }, - }, - }, - }, - }, - }, - }); - - // This batch has a dependent attempt and just finalized, we should resume that attempt - if (finalizedBatchRun && finalizedBatchRun.dependentTaskAttempt) { - const environment = - env ?? (await this.#getEnvironment(taskRunAttempt.taskRun.runtimeEnvironmentId)); - - if (!environment) { - logger.error("Environment not found", { - attemptId: taskRunAttempt.id, - envId: taskRunAttempt.taskRun.runtimeEnvironmentId, - }); - return "FAILED"; - } - - if (environment.type === "DEVELOPMENT") { - return "ACKNOWLEDGED"; - } - - const dependentRun = finalizedBatchRun.dependentTaskAttempt.taskRun; - - if (finalizedBatchRun.dependentTaskAttempt.status === "PAUSED") { - await marqs?.enqueueMessage( - environment, - dependentRun.queue, - dependentRun.id, - { - type: "RESUME", - completedAttemptIds: [taskRunAttempt.id], - }, - dependentRun.concurrencyKey ?? undefined - ); - } else { - await marqs?.replaceMessage(dependentRun.id, { - type: "RESUME", - completedAttemptIds: finalizedBatchRun.items.map( - (item) => item.taskRun.attempts[0]?.id - ), - }); - } - } - } - - if (dependency) { - logger.debug("Completing attempt with dependency", { dependency }); - - const environment = - env ?? (await this.#getEnvironment(taskRunAttempt.taskRun.runtimeEnvironmentId)); - - if (!environment) { - logger.error("Environment not found", { - attemptId: taskRunAttempt.id, - envId: taskRunAttempt.taskRun.runtimeEnvironmentId, - }); - return "FAILED"; - } - - if (environment.type === "DEVELOPMENT") { - return "ACKNOWLEDGED"; - } - - if (dependency.dependentAttempt) { - const dependentRun = await this._prisma.taskRun.findFirst({ - where: { - id: dependency.dependentAttempt.taskRunId, - }, - }); - - if (!dependentRun) { - logger.error("Dependent task run does not exist", { - attemptId: taskRunAttempt.id, - envId: taskRunAttempt.taskRun.runtimeEnvironmentId, - taskRunId: dependency.taskRunId, - }); - return "FAILED"; - } - - if (dependency.dependentAttempt.status === "PAUSED") { - await marqs?.enqueueMessage( - environment, - dependentRun.queue, - dependentRun.id, - { - type: "RESUME", - completedAttemptIds: [taskRunAttempt.id], - }, - dependentRun.concurrencyKey ?? undefined - ); - } else { - await marqs?.replaceMessage(dependentRun.id, { - type: "RESUME", - completedAttemptIds: [taskRunAttempt.id], - }); - } - } + if (!env || env.type !== "DEVELOPMENT") { + await ResumeTaskRunDependenciesService.enqueue(taskRunAttempt.id, this._prisma); } return "ACKNOWLEDGED"; @@ -329,3 +249,13 @@ export class CompleteAttemptService extends BaseService { }); } } + +async function findAttempt(prismaClient: PrismaClientOrTransaction, friendlyId: string) { + return prismaClient.taskRunAttempt.findUnique({ + where: { friendlyId }, + include: { + taskRun: true, + backgroundWorkerTask: true, + }, + }); +} diff --git a/apps/webapp/app/v3/services/createCheckpoint.server.ts b/apps/webapp/app/v3/services/createCheckpoint.server.ts index acebca7dc..70b6ac2d0 100644 --- a/apps/webapp/app/v3/services/createCheckpoint.server.ts +++ b/apps/webapp/app/v3/services/createCheckpoint.server.ts @@ -51,7 +51,7 @@ export class CreateCheckpointService { case "WAIT_FOR_DURATION": { await marqs?.replaceMessage( attempt.taskRunId, - { type: "RESUME_AFTER_DURATION" }, + { type: "RESUME_AFTER_DURATION", resumableAttemptId: attempt.id }, Date.now() + params.reason.ms ); break; diff --git a/apps/webapp/app/v3/services/resumeBatchRun.server.ts b/apps/webapp/app/v3/services/resumeBatchRun.server.ts new file mode 100644 index 000000000..deb19657c --- /dev/null +++ b/apps/webapp/app/v3/services/resumeBatchRun.server.ts @@ -0,0 +1,104 @@ +import { PrismaClientOrTransaction } from "~/db.server"; +import { workerQueue } from "~/services/worker.server"; +import { marqs } from "../marqs.server"; +import { BaseService } from "./baseService.server"; + +export class ResumeBatchRunService extends BaseService { + public async call(batchRunId: string, sourceTaskAttemptId: string) { + const batchRun = await this._prisma.batchTaskRun.findFirst({ + where: { + id: batchRunId, + dependentTaskAttemptId: { + not: null, + }, + status: "PENDING", + items: { + every: { + taskRunAttemptId: { + not: null, + }, + }, + }, + }, + include: { + dependentTaskAttempt: { + include: { + runtimeEnvironment: { + include: { + project: true, + organization: true, + }, + }, + taskRun: true, + }, + }, + items: true, + }, + }); + + if (!batchRun || !batchRun.dependentTaskAttempt) { + return; + } + + await this._prisma.batchTaskRun.update({ + where: { + id: batchRun.id, + }, + data: { + status: "COMPLETED", + }, + }); + + // We need to update the batchRun status so we don't resume it again + + // This batch has a dependent attempt and just finalized, we should resume that attempt + const environment = batchRun.dependentTaskAttempt.runtimeEnvironment; + + // If we are in development, we don't need to resume the dependent task (that will happen automatically) + if (environment.type === "DEVELOPMENT") { + return; + } + + const dependentRun = batchRun.dependentTaskAttempt.taskRun; + + if (batchRun.dependentTaskAttempt.status === "PAUSED") { + await marqs?.enqueueMessage( + environment, + dependentRun.queue, + dependentRun.id, + { + type: "RESUME", + completedAttemptIds: [sourceTaskAttemptId], + resumableAttemptId: batchRun.dependentTaskAttempt.id, + }, + dependentRun.concurrencyKey ?? undefined + ); + } else { + await marqs?.replaceMessage(dependentRun.id, { + type: "RESUME", + completedAttemptIds: batchRun.items.map((item) => item.taskRunAttemptId).filter(Boolean), + resumableAttemptId: batchRun.dependentTaskAttempt.id, + }); + } + } + + static async enqueue( + batchRunId: string, + sourceTaskAttemptId: string, + tx: PrismaClientOrTransaction, + runAt?: Date + ) { + return await workerQueue.enqueue( + "v3.resumeBatchRun", + { + batchRunId, + sourceTaskAttemptId, + }, + { + tx, + runAt, + queueName: `resumeBatchRun-${batchRunId}`, + } + ); + } +} diff --git a/apps/webapp/app/v3/services/resumeTaskDependency.server.ts b/apps/webapp/app/v3/services/resumeTaskDependency.server.ts new file mode 100644 index 000000000..cbfdfd3da --- /dev/null +++ b/apps/webapp/app/v3/services/resumeTaskDependency.server.ts @@ -0,0 +1,78 @@ +import { PrismaClientOrTransaction } from "~/db.server"; +import { workerQueue } from "~/services/worker.server"; +import { marqs } from "../marqs.server"; +import { BaseService } from "./baseService.server"; + +export class ResumeTaskDependencyService extends BaseService { + public async call(dependencyId: string, sourceTaskAttemptId: string) { + const dependency = await this._prisma.taskRunDependency.findUnique({ + where: { id: dependencyId }, + include: { + taskRun: { + include: { + runtimeEnvironment: { + include: { + project: true, + organization: true, + }, + }, + }, + }, + dependentAttempt: { + include: { + taskRun: true, + }, + }, + }, + }); + + // Dependencies with a dependentBatchRun are handled already by the ResumeBatchRunService + if (!dependency || !dependency.dependentAttempt) { + return; + } + + if (dependency.taskRun.runtimeEnvironment.type === "DEVELOPMENT") { + return; + } + const dependentRun = dependency.dependentAttempt.taskRun; + + if (dependency.dependentAttempt.status === "PAUSED") { + await marqs?.enqueueMessage( + dependency.taskRun.runtimeEnvironment, + dependentRun.queue, + dependentRun.id, + { + type: "RESUME", + completedAttemptIds: [sourceTaskAttemptId], + resumableAttemptId: dependency.dependentAttempt.id, + }, + dependentRun.concurrencyKey ?? undefined + ); + } else { + await marqs?.replaceMessage(dependentRun.id, { + type: "RESUME", + completedAttemptIds: [sourceTaskAttemptId], + resumableAttemptId: dependency.dependentAttempt.id, + }); + } + } + + static async enqueue( + dependencyId: string, + sourceTaskAttemptId: string, + tx: PrismaClientOrTransaction, + runAt?: Date + ) { + return await workerQueue.enqueue( + "v3.resumeTaskDependency", + { + dependencyId, + sourceTaskAttemptId, + }, + { + tx, + runAt, + } + ); + } +} diff --git a/apps/webapp/app/v3/services/resumeTaskRunDependencies.server.ts b/apps/webapp/app/v3/services/resumeTaskRunDependencies.server.ts new file mode 100644 index 000000000..d9849b4f1 --- /dev/null +++ b/apps/webapp/app/v3/services/resumeTaskRunDependencies.server.ts @@ -0,0 +1,86 @@ +import { BatchTaskRunItem, TaskRunAttempt, TaskRunDependency } from "@trigger.dev/database"; +import { $transaction, PrismaClientOrTransaction } from "~/db.server"; +import { workerQueue } from "~/services/worker.server"; +import { BaseService } from "./baseService.server"; +import { ResumeBatchRunService } from "./resumeBatchRun.server"; +import { ResumeTaskDependencyService } from "./resumeTaskDependency.server"; + +export class ResumeTaskRunDependenciesService extends BaseService { + public async call(attemptId: string) { + const taskAttempt = await this._prisma.taskRunAttempt.findUnique({ + where: { id: attemptId }, + include: { + taskRun: { + include: { + runtimeEnvironment: true, + batchItem: true, + dependency: { + include: { + dependentAttempt: true, + dependentBatchRun: true, + }, + }, + }, + }, + backgroundWorkerTask: true, + }, + }); + + if (!taskAttempt) { + return; + } + + if (taskAttempt.taskRun.runtimeEnvironment.type === "DEVELOPMENT") { + return; + } + + const { batchItem, dependency } = taskAttempt.taskRun; + + if (!batchItem && !dependency) { + return; + } + + if (batchItem) { + await this.#resumeBatchItem(batchItem, taskAttempt); + return; + } + + if (dependency && dependency.dependentAttempt) { + await this.#resumeDependency(dependency, taskAttempt); + } + } + + async #resumeBatchItem(batchItem: BatchTaskRunItem, taskAttempt: TaskRunAttempt) { + await $transaction(this._prisma, async (tx) => { + await tx.batchTaskRunItem.update({ + where: { + id: batchItem.id, + }, + data: { + status: "COMPLETED", + taskRunAttemptId: taskAttempt.id, + }, + }); + + await ResumeBatchRunService.enqueue(batchItem.batchTaskRunId, taskAttempt.id, tx); + }); + } + + async #resumeDependency(dependency: TaskRunDependency, taskAttempt: TaskRunAttempt) { + await ResumeTaskDependencyService.enqueue(dependency.id, taskAttempt.id, this._prisma); + } + + static async enqueue(attemptId: string, tx: PrismaClientOrTransaction, runAt?: Date) { + return await workerQueue.enqueue( + "v3.resumeTaskRunDependencies", + { + attemptId, + }, + { + tx, + runAt, + jobKey: `resumeTaskRunDependencies:${attemptId}`, + } + ); + } +} diff --git a/apps/webapp/app/v3/services/triggerTask.server.ts b/apps/webapp/app/v3/services/triggerTask.server.ts index 3525200cb..cc96e9caa 100644 --- a/apps/webapp/app/v3/services/triggerTask.server.ts +++ b/apps/webapp/app/v3/services/triggerTask.server.ts @@ -100,6 +100,7 @@ export class TriggerTaskService extends BaseService { const taskRun = await tx.taskRun.create({ data: { + status: "PENDING", number: counter.lastNumber, friendlyId: generateFriendlyId("run"), runtimeEnvironmentId: environment.id, diff --git a/apps/webapp/app/v3/sharedSocketConnection.ts b/apps/webapp/app/v3/sharedSocketConnection.ts index f47157c1a..1e7992dfc 100644 --- a/apps/webapp/app/v3/sharedSocketConnection.ts +++ b/apps/webapp/app/v3/sharedSocketConnection.ts @@ -47,7 +47,10 @@ export class SharedSocketConnection { }, }); - this._sharedConsumer = new SharedQueueConsumer(this._sender); + this._sharedConsumer = new SharedQueueConsumer(this._sender, { + interval: 100, + nextTickInterval: 1000, + }); socket.on("disconnect", this.#handleClose.bind(this)); socket.on("error", this.#handleError.bind(this)); diff --git a/apps/webapp/app/v3/utils/asserts.server.ts b/apps/webapp/app/v3/utils/asserts.server.ts new file mode 100644 index 000000000..966499f33 --- /dev/null +++ b/apps/webapp/app/v3/utils/asserts.server.ts @@ -0,0 +1,3 @@ +export function assertUnreachable(x: never): never { + throw new Error("Didn't expect to get here"); +} diff --git a/apps/webapp/app/v3/utils/zodPubSub.server.ts b/apps/webapp/app/v3/utils/zodPubSub.server.ts new file mode 100644 index 000000000..83f681db7 --- /dev/null +++ b/apps/webapp/app/v3/utils/zodPubSub.server.ts @@ -0,0 +1,117 @@ +import { Logger } from "@trigger.dev/core-backend"; +import { ZodMessageCatalogSchema, ZodMessageHandler, ZodMessageSender } from "@trigger.dev/core/v3"; +import Redis, { RedisOptions } from "ioredis"; +import { z } from "zod"; +import { logger } from "~/services/logger.server"; +import { safeJsonParse } from "~/utils/json"; + +export type ZodPubSubOptions = { + redis: RedisOptions; + schema: TMessageCatalog; +}; + +export interface ZodSubscriber { + on( + eventName: K, + listener: (payload: z.infer) => Promise + ): void; + + stopListening(): Promise; +} + +class RedisZodSubscriber + implements ZodSubscriber +{ + private _subscriber: Redis; + private _listeners: Map Promise> = new Map(); + private _messageHandler: ZodMessageHandler; + + constructor( + private readonly _pattern: string, + private readonly _options: ZodPubSubOptions, + private readonly _logger: Logger + ) { + this._subscriber = new Redis(_options.redis); + this._messageHandler = new ZodMessageHandler({ + schema: _options.schema, + }); + } + + async initialize() { + await this._subscriber.psubscribe(this._pattern); + this._subscriber.on("pmessage", this.#onMessage.bind(this)); + } + + public on( + eventName: K, + listener: (payload: z.infer) => Promise + ): void { + this._listeners.set(eventName as string, listener); + } + + public async stopListening(): Promise { + this._listeners.clear(); + await this._subscriber.unsubscribe(); + } + + async #onMessage(pattern: string, channel: string, serializedMessage: string) { + if (pattern !== this._pattern) { + return; + } + + const parsedMessage = safeJsonParse(serializedMessage); + + if (!parsedMessage) { + return; + } + + const message = this._messageHandler.parseMessage(parsedMessage); + + if (typeof message.type !== "string") { + return; + } + + const listener = this._listeners.get(message.type); + + if (!listener) { + this._logger.debug(`No listener for message type: ${message.type}`, { parsedMessage }); + + return; + } + + try { + await listener(message.payload); + } catch (error) { + this._logger.error("Error handling message", { error, message }); + } + } +} + +export class ZodPubSub { + private _publisher: Redis; + private _logger = logger.child({ module: "ZodPubSub" }); + + constructor(private _options: ZodPubSubOptions) { + this._publisher = new Redis(_options.redis); + } + + public async publish( + channel: string, + type: K, + payload: z.input + ): Promise { + try { + await this._publisher.publish(channel, JSON.stringify({ type, payload, version: "v1" })); + } catch (e) { + logger.error("Failed to publish message", { channel, type, payload, error: e }); + } + } + + public async subscribe(channel: string): Promise> { + const subscriber = new RedisZodSubscriber(channel, this._options, this._logger); + + await subscriber.initialize(); + + return subscriber; + } +} diff --git a/packages/cli-v3/src/workers/dev/backgroundWorker.ts b/packages/cli-v3/src/workers/dev/backgroundWorker.ts index 86ed0fabc..0778aeb38 100644 --- a/packages/cli-v3/src/workers/dev/backgroundWorker.ts +++ b/packages/cli-v3/src/workers/dev/backgroundWorker.ts @@ -108,9 +108,23 @@ export class BackgroundWorkerCoordinator { } async handleMessage(id: string, message: BackgroundWorkerServerMessages) { + logger.debug(`Received message from worker ${id}`, { workerMessage: message }); + switch (message.type) { case "EXECUTE_RUNS": { await Promise.all(message.payloads.map((payload) => this.#executeTaskRun(id, payload))); + break; + } + case "CANCEL_ATTEMPT": { + // Need to cancel the attempt somehow here + const worker = this._backgroundWorkers.get(id); + + if (!worker) { + logger.error(`Could not find worker ${id}`); + return; + } + + await worker.cancelRun(message.taskRunId); } } } @@ -158,7 +172,8 @@ export class BackgroundWorkerCoordinator { const resultText = !completion.ok ? completion.error.type === "INTERNAL_ERROR" && - completion.error.code === TaskRunErrorCodes.TASK_EXECUTION_ABORTED + (completion.error.code === TaskRunErrorCodes.TASK_EXECUTION_ABORTED || + completion.error.code === TaskRunErrorCodes.TASK_RUN_CANCELLED) ? chalk.yellow("cancelled") : chalk.red(`error${retryingText}`) : chalk.green("success"); @@ -214,6 +229,14 @@ class CleanupProcessError extends Error { } } +class CancelledProcessError extends Error { + constructor() { + super("Cancelled"); + + this.name = "CancelledProcessError"; + } +} + export type BackgroundWorkerParams = { env: Record; dependencies?: Record; @@ -367,6 +390,16 @@ export class BackgroundWorker { return this._taskRunProcesses.get(payload.execution.run.id) as TaskRunProcess; } + async cancelRun(taskRunId: string) { + const taskRunProcess = this._taskRunProcesses.get(taskRunId); + + if (!taskRunProcess) { + return; + } + + await taskRunProcess.cancel(); + } + // We need to fork the process before we can execute any tasks async executeTaskRun(payload: TaskRunExecutionPayload): Promise { try { @@ -393,6 +426,18 @@ export class BackgroundWorker { return result; } catch (e) { + if (e instanceof CancelledProcessError) { + return { + id: payload.execution.attempt.id, + ok: false, + retry: undefined, + error: { + type: "INTERNAL_ERROR", + code: TaskRunErrorCodes.TASK_RUN_CANCELLED, + }, + }; + } + if (e instanceof CleanupProcessError) { return { id: payload.execution.attempt.id, @@ -464,6 +509,7 @@ class TaskRunProcess { private _attemptStatuses: Map = new Map(); private _currentExecution: TaskRunExecution | undefined; private _isBeingKilled: boolean = false; + private _isBeingCancelled: boolean = false; public onTaskHeartbeat: Evt = new Evt(); public onExit: Evt = new Evt(); @@ -483,6 +529,12 @@ class TaskRunProcess { }); } + async cancel() { + this._isBeingCancelled = true; + + await this.cleanup(true); + } + async initialize() { this._child = fork(this.path, { stdio: [/*stdin*/ "ignore", /*stdout*/ "pipe", /*stderr*/ "pipe", "ipc"], @@ -617,7 +669,9 @@ class TaskRunProcess { const { rejecter } = attemptPromise; - if (this._isBeingKilled) { + if (this._isBeingCancelled) { + rejecter(new CancelledProcessError()); + } else if (this._isBeingKilled) { rejecter(new CleanupProcessError()); } else { rejecter(new UnexpectedExitError(code)); diff --git a/packages/cli-v3/src/workers/prod/backgroundWorker.ts b/packages/cli-v3/src/workers/prod/backgroundWorker.ts index ee47c76a0..8a4a1bf98 100644 --- a/packages/cli-v3/src/workers/prod/backgroundWorker.ts +++ b/packages/cli-v3/src/workers/prod/backgroundWorker.ts @@ -36,6 +36,14 @@ class CleanupProcessError extends Error { } } +class CancelledProcessError extends Error { + constructor() { + super("Cancelled"); + + this.name = "CancelledProcessError"; + } +} + type BackgroundWorkerParams = { env: Record; projectConfig: Config; @@ -58,7 +66,7 @@ export class ProdBackgroundWorker { public tasks: Array = []; - _taskRunProcesses: Map = new Map(); + _taskRunProcess: TaskRunProcess | undefined; private _closed: boolean = false; @@ -77,9 +85,7 @@ export class ProdBackgroundWorker { this.onTaskHeartbeat.detach(); // We need to close all the task run processes - for (const taskRunProcess of this._taskRunProcesses.values()) { - taskRunProcess.cleanup(true); - } + this._taskRunProcess?.cleanup(true); // Delete worker files this._onClose.post(); @@ -173,14 +179,11 @@ export class ProdBackgroundWorker { completion: TaskRunExecutionResult, execution: TaskRunExecution ) { - for (const taskRunProcess of this._taskRunProcesses.values()) { - taskRunProcess.taskRunCompletedNotification(completion, execution); - } + this._taskRunProcess?.taskRunCompletedNotification(completion, execution); } + async waitCompletedNotification() { - for (const taskRunProcess of this._taskRunProcesses.values()) { - taskRunProcess.waitCompletedNotification(); - } + this._taskRunProcess?.waitCompletedNotification(); } async #initializeTaskRunProcess(payload: ProdTaskRunExecutionPayload): Promise { @@ -189,47 +192,45 @@ export class ProdBackgroundWorker { payload.execution.worker.version ); - if (!this._taskRunProcesses.has(payload.execution.run.id)) { - const taskRunProcess = new TaskRunProcess( - this.path, - { - ...this.params.env, - ...(payload.environment ?? {}), - }, - metadata, - this.params - ); + const taskRunProcess = new TaskRunProcess( + this.path, + { + ...this.params.env, + ...(payload.environment ?? {}), + }, + metadata, + this.params + ); - taskRunProcess.onExit.attach(() => { - this._taskRunProcesses.delete(payload.execution.run.id); - }); + this._taskRunProcess = taskRunProcess; - taskRunProcess.onTaskHeartbeat.attach((id) => { - this.onTaskHeartbeat.post(id); - }); + taskRunProcess.onExit.attach(() => { + this._taskRunProcess = undefined; + }); - taskRunProcess.onWaitForBatch.attach((message) => { - this.onWaitForBatch.post(message); - }); + taskRunProcess.onTaskHeartbeat.attach((id) => { + this.onTaskHeartbeat.post(id); + }); - taskRunProcess.onWaitForDuration.attach((message) => { - this.onWaitForDuration.post(message); - }); + taskRunProcess.onWaitForBatch.attach((message) => { + this.onWaitForBatch.post(message); + }); - taskRunProcess.onWaitForTask.attach((message) => { - this.onWaitForTask.post(message); - }); + taskRunProcess.onWaitForDuration.attach((message) => { + this.onWaitForDuration.post(message); + }); - this.preCheckpointNotification.attach((message) => { - taskRunProcess.preCheckpointNotification.post(message); - }); + taskRunProcess.onWaitForTask.attach((message) => { + this.onWaitForTask.post(message); + }); - await taskRunProcess.initialize(); + this.preCheckpointNotification.attach((message) => { + taskRunProcess.preCheckpointNotification.post(message); + }); - this._taskRunProcesses.set(payload.execution.run.id, taskRunProcess); - } + await taskRunProcess.initialize(); - return this._taskRunProcesses.get(payload.execution.run.id) as TaskRunProcess; + return taskRunProcess; } // We need to fork the process before we can execute any tasks @@ -259,6 +260,18 @@ export class ProdBackgroundWorker { return result; } catch (e) { + if (e instanceof CancelledProcessError) { + return { + id: payload.execution.attempt.id, + ok: false, + retry: undefined, + error: { + type: "INTERNAL_ERROR", + code: TaskRunErrorCodes.TASK_RUN_CANCELLED, + }, + }; + } + if (e instanceof CleanupProcessError) { return { id: payload.execution.attempt.id, @@ -295,6 +308,10 @@ export class ProdBackgroundWorker { } } + async cancelAttempt(attemptId: string) { + await this._taskRunProcess?.cancel(); + } + async #correctError( error: TaskRunBuiltInError, execution: TaskRunExecution @@ -320,6 +337,7 @@ class TaskRunProcess { private _attemptStatuses: Map = new Map(); private _currentExecution: TaskRunExecution | undefined; private _isBeingKilled: boolean = false; + private _isBeingCancelled: boolean = false; public onTaskHeartbeat: Evt = new Evt(); public onExit: Evt = new Evt(); @@ -405,6 +423,12 @@ class TaskRunProcess { this._child.stderr?.on("data", this.#handleStdErr.bind(this)); } + async cancel() { + this._isBeingCancelled = true; + + await this.cleanup(true); + } + async cleanup(kill: boolean = false) { if (kill && this._isBeingKilled) { return; @@ -484,7 +508,9 @@ class TaskRunProcess { const { rejecter } = attemptPromise; - if (this._isBeingKilled) { + if (this._isBeingCancelled) { + rejecter(new CancelledProcessError()); + } else if (this._isBeingKilled) { rejecter(new CleanupProcessError()); } else { rejecter(new UnexpectedExitError(code)); diff --git a/packages/cli-v3/src/workers/prod/entry-point.ts b/packages/cli-v3/src/workers/prod/entry-point.ts index e262759e4..bee68e94e 100644 --- a/packages/cli-v3/src/workers/prod/entry-point.ts +++ b/packages/cli-v3/src/workers/prod/entry-point.ts @@ -200,6 +200,16 @@ class ProdWorker { process.exit(0); }, + REQUEST_ATTEMPT_CANCELLATION: async (message) => { + if (!this.executing) { + return; + } + + await this.#backgroundWorker.cancelAttempt(message.attemptId); + }, + REQUEST_EXIT: async () => { + process.exit(0); + }, }, onConnection: async (socket, handler, sender, logger) => { if (process.env.INDEX_TASKS === "true") { diff --git a/packages/core-backend/src/logger.ts b/packages/core-backend/src/logger.ts index 09dd5c0fe..d86c20e27 100644 --- a/packages/core-backend/src/logger.ts +++ b/packages/core-backend/src/logger.ts @@ -37,6 +37,16 @@ export class Logger { this.#additionalFields = additionalFields ?? (() => ({})); } + child(fields: Record) { + return new Logger( + this.#name, + logLevels[this.#level], + this.#filteredKeys, + this.#jsonReplacer, + () => ({ ...this.#additionalFields(), ...fields }) + ); + } + // Return a new Logger instance with the same name and a new log level // but filter out the keys from the log messages (at any level) filter(...keys: string[]) { diff --git a/packages/core/src/v3/schemas/common.ts b/packages/core/src/v3/schemas/common.ts index 123ea93dc..9a0427296 100644 --- a/packages/core/src/v3/schemas/common.ts +++ b/packages/core/src/v3/schemas/common.ts @@ -30,6 +30,7 @@ export const TaskRunErrorCodes = { TASK_EXECUTION_FAILED: "TASK_EXECUTION_FAILED", TASK_EXECUTION_ABORTED: "TASK_EXECUTION_ABORTED", TASK_PROCESS_EXITED_WITH_NON_ZERO_CODE: "TASK_PROCESS_EXITED_WITH_NON_ZERO_CODE", + TASK_RUN_CANCELLED: "TASK_RUN_CANCELLED", } as const; export const TaskRunInternalError = z.object({ @@ -41,6 +42,7 @@ export const TaskRunInternalError = z.object({ "TASK_EXECUTION_FAILED", "TASK_EXECUTION_ABORTED", "TASK_PROCESS_EXITED_WITH_NON_ZERO_CODE", + "TASK_RUN_CANCELLED", ]), }); diff --git a/packages/core/src/v3/schemas/messages.ts b/packages/core/src/v3/schemas/messages.ts index 020e4ee7f..456f4bb5a 100644 --- a/packages/core/src/v3/schemas/messages.ts +++ b/packages/core/src/v3/schemas/messages.ts @@ -32,6 +32,11 @@ export const BackgroundWorkerServerMessages = z.discriminatedUnion("type", [ type: z.literal("EXECUTE_RUNS"), payloads: z.array(TaskRunExecutionPayload), }), + z.object({ + type: z.literal("CANCEL_ATTEMPT"), + taskAttemptId: z.string(), + taskRunId: z.string(), + }), z.object({ type: z.literal("SCHEDULE_ATTEMPT"), id: z.string(), diff --git a/packages/core/src/v3/schemas/schemas.ts b/packages/core/src/v3/schemas/schemas.ts index 1353000c1..1c2d670b6 100644 --- a/packages/core/src/v3/schemas/schemas.ts +++ b/packages/core/src/v3/schemas/schemas.ts @@ -222,6 +222,12 @@ export const PlatformToCoordinatorMessages = { attemptId: z.string(), }), }, + REQUEST_ATTEMPT_CANCELLATION: { + message: z.object({ + version: z.literal("v1").default("v1"), + attemptId: z.string(), + }), + }, }; export const ClientToSharedQueueMessages = { @@ -375,6 +381,17 @@ export const CoordinatorToProdWorkerMessages = { executionPayload: ProdTaskRunExecutionPayload, }), }, + REQUEST_ATTEMPT_CANCELLATION: { + message: z.object({ + version: z.literal("v1").default("v1"), + attemptId: z.string(), + }), + }, + REQUEST_EXIT: { + message: z.object({ + version: z.literal("v1").default("v1"), + }), + }, }; export const ProdWorkerSocketData = z.object({ diff --git a/packages/core/src/v3/zodMessageHandler.ts b/packages/core/src/v3/zodMessageHandler.ts index 451949f1a..b4a64c1b6 100644 --- a/packages/core/src/v3/zodMessageHandler.ts +++ b/packages/core/src/v3/zodMessageHandler.ts @@ -17,7 +17,7 @@ export type ZodMessageHandlerOptions; }; -type MessageFromSchema< +export type MessageFromSchema< K extends keyof TMessageCatalog, TMessageCatalog extends ZodMessageCatalogSchema, > = { @@ -25,11 +25,11 @@ type MessageFromSchema< payload: z.input; }; -type MessageFromCatalog = { +export type MessageFromCatalog = { [K in keyof TMessageCatalog]: MessageFromSchema; }[keyof TMessageCatalog]; -const messageSchema = z.object({ +export const ZodMessageSchema = z.object({ version: z.literal("v1").default("v1"), type: z.string(), payload: z.unknown(), @@ -68,7 +68,7 @@ export class ZodMessageHandler } public parseMessage(message: unknown): MessageFromCatalog { - const parsedMessage = messageSchema.safeParse(message); + const parsedMessage = ZodMessageSchema.safeParse(message); if (!parsedMessage.success) { throw new Error(`Failed to parse message: ${JSON.stringify(parsedMessage.error)}`); @@ -161,6 +161,32 @@ export class ZodMessageSender { await this.#sender({ type, payload, version: "v1" }); } + + public async forwardMessage(message: unknown) { + const parsedMessage = ZodMessageSchema.safeParse(message); + + if (!parsedMessage.success) { + throw new Error(`Failed to parse message: ${JSON.stringify(parsedMessage.error)}`); + } + + const schema = this.#schema[parsedMessage.data.type]; + + if (!schema) { + throw new Error(`Unknown message type: ${parsedMessage.data.type}`); + } + + const parsedPayload = schema.safeParse(parsedMessage.data.payload); + + if (!parsedPayload.success) { + throw new Error(`Failed to parse message payload: ${JSON.stringify(parsedPayload.error)}`); + } + + await this.#sender({ + type: parsedMessage.data.type, + payload: parsedPayload.data, + version: "v1", + }); + } } export type MessageCatalogToSocketIoEvents = { diff --git a/packages/database/prisma/migrations/20240312095501_add_status_to_task_runs/migration.sql b/packages/database/prisma/migrations/20240312095501_add_status_to_task_runs/migration.sql new file mode 100644 index 000000000..e4b60577a --- /dev/null +++ b/packages/database/prisma/migrations/20240312095501_add_status_to_task_runs/migration.sql @@ -0,0 +1,5 @@ +-- CreateEnum +CREATE TYPE "TaskRunStatus" AS ENUM ('PENDING', 'EXECUTING', 'WAITING_TO_RESUME', 'RETRYING_AFTER_FAILURE', 'PAUSED', 'CANCELED', 'COMPLETED_SUCCESSFULLY', 'COMPLETED_WITH_ERRORS'); + +-- AlterTable +ALTER TABLE "TaskRun" ADD COLUMN "status" "TaskRunStatus" NOT NULL DEFAULT 'PENDING'; diff --git a/packages/database/prisma/migrations/20240312095826_add_interrupted_status_to_runs/migration.sql b/packages/database/prisma/migrations/20240312095826_add_interrupted_status_to_runs/migration.sql new file mode 100644 index 000000000..dd922e37e --- /dev/null +++ b/packages/database/prisma/migrations/20240312095826_add_interrupted_status_to_runs/migration.sql @@ -0,0 +1,2 @@ +-- AlterEnum +ALTER TYPE "TaskRunStatus" ADD VALUE 'INTERRUPTED'; diff --git a/packages/database/prisma/migrations/20240312105844_add_system_failure_status_to_task_run_status/migration.sql b/packages/database/prisma/migrations/20240312105844_add_system_failure_status_to_task_run_status/migration.sql new file mode 100644 index 000000000..9403da5ee --- /dev/null +++ b/packages/database/prisma/migrations/20240312105844_add_system_failure_status_to_task_run_status/migration.sql @@ -0,0 +1,2 @@ +-- AlterEnum +ALTER TYPE "TaskRunStatus" ADD VALUE 'SYSTEM_FAILURE'; diff --git a/packages/database/prisma/migrations/20240312125252_add_status_to_batch_run/migration.sql b/packages/database/prisma/migrations/20240312125252_add_status_to_batch_run/migration.sql new file mode 100644 index 000000000..3a296f89d --- /dev/null +++ b/packages/database/prisma/migrations/20240312125252_add_status_to_batch_run/migration.sql @@ -0,0 +1,5 @@ +-- CreateEnum +CREATE TYPE "BatchTaskRunStatus" AS ENUM ('PENDING', 'COMPLETED'); + +-- AlterTable +ALTER TABLE "BatchTaskRun" ADD COLUMN "status" "BatchTaskRunStatus" NOT NULL DEFAULT 'PENDING'; diff --git a/packages/database/prisma/migrations/20240312131122_add_task_run_attempt_to_batch_run_items/migration.sql b/packages/database/prisma/migrations/20240312131122_add_task_run_attempt_to_batch_run_items/migration.sql new file mode 100644 index 000000000..def157156 --- /dev/null +++ b/packages/database/prisma/migrations/20240312131122_add_task_run_attempt_to_batch_run_items/migration.sql @@ -0,0 +1,5 @@ +-- AlterTable +ALTER TABLE "BatchTaskRunItem" ADD COLUMN "taskRunAttemptId" TEXT; + +-- AddForeignKey +ALTER TABLE "BatchTaskRunItem" ADD CONSTRAINT "BatchTaskRunItem_taskRunAttemptId_fkey" FOREIGN KEY ("taskRunAttemptId") REFERENCES "TaskRunAttempt"("id") ON DELETE SET NULL ON UPDATE CASCADE; diff --git a/packages/database/prisma/migrations/20240313110150_add_environment_to_task_run_attempts/migration.sql b/packages/database/prisma/migrations/20240313110150_add_environment_to_task_run_attempts/migration.sql new file mode 100644 index 000000000..c0bf438dd --- /dev/null +++ b/packages/database/prisma/migrations/20240313110150_add_environment_to_task_run_attempts/migration.sql @@ -0,0 +1,11 @@ +/* + Warnings: + + - Added the required column `runtimeEnvironmentId` to the `TaskRunAttempt` table without a default value. This is not possible if the table is not empty. + +*/ +-- AlterTable +ALTER TABLE "TaskRunAttempt" ADD COLUMN "runtimeEnvironmentId" TEXT NOT NULL; + +-- AddForeignKey +ALTER TABLE "TaskRunAttempt" ADD CONSTRAINT "TaskRunAttempt_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 cd009cef0..390f6d872 100644 --- a/packages/database/prisma/schema.prisma +++ b/packages/database/prisma/schema.prisma @@ -393,6 +393,7 @@ model RuntimeEnvironment { checkpoints Checkpoint[] workerDeployments WorkerDeployment[] workerDeploymentPromotions WorkerDeploymentPromotion[] + taskRunAttempts TaskRunAttempt[] @@unique([projectId, slug, orgMemberId]) @@unique([projectId, shortcode]) @@ -1564,6 +1565,8 @@ model TaskRun { number Int @default(0) friendlyId String @unique + status TaskRunStatus @default(PENDING) + idempotencyKey String taskIdentifier String @@ -1606,6 +1609,38 @@ model TaskRun { @@unique([runtimeEnvironmentId, idempotencyKey]) } +enum TaskRunStatus { + /// Task is waiting to be executed by a worker + PENDING + + /// Task is currently being executed by a worker + EXECUTING + + /// Task has been paused by the system, and will be resumed by the system + WAITING_TO_RESUME + + /// Task has failed and is waiting to be retried + RETRYING_AFTER_FAILURE + + /// Task has been paused by the user, and can be resumed by the user + PAUSED + + /// Task has been canceled by the user + CANCELED + + /// Task was interrupted during execution, mostly this happens in development environments + INTERRUPTED + + /// Task has been completed successfully + COMPLETED_SUCCESSFULLY + + /// Task has been completed with errors + COMPLETED_WITH_ERRORS + + /// Task has failed to complete, due to an error in the system + SYSTEM_FAILURE +} + model TaskRunDependency { id String @id @default(cuid()) @@ -1660,6 +1695,9 @@ model TaskRunAttempt { backgroundWorkerTask BackgroundWorkerTask @relation(fields: [backgroundWorkerTaskId], references: [id], onDelete: Cascade, onUpdate: Cascade) backgroundWorkerTaskId String + runtimeEnvironment RuntimeEnvironment @relation(fields: [runtimeEnvironmentId], references: [id], onDelete: Cascade, onUpdate: Cascade) + runtimeEnvironmentId String + queue TaskQueue @relation(fields: [queueId], references: [id], onDelete: Cascade, onUpdate: Cascade) queueId String @@ -1678,7 +1716,8 @@ model TaskRunAttempt { taskRunDependency TaskRunDependency? @relation("dependentAttempt") batchTaskRunDependency BatchTaskRun? - checkpoints Checkpoint[] + checkpoints Checkpoint[] + batchTaskRunItems BatchTaskRunItem[] @@unique([taskRunId, number]) } @@ -1826,6 +1865,8 @@ model BatchTaskRun { friendlyId String @unique + status BatchTaskRunStatus @default(PENDING) + idempotencyKey String taskIdentifier String @@ -1844,6 +1885,11 @@ model BatchTaskRun { @@unique([runtimeEnvironmentId, idempotencyKey]) } +enum BatchTaskRunStatus { + PENDING + COMPLETED +} + model BatchTaskRunItem { id String @id @default(cuid()) @@ -1855,6 +1901,9 @@ model BatchTaskRunItem { taskRun TaskRun @relation(fields: [taskRunId], references: [id], onDelete: Cascade, onUpdate: Cascade) taskRunId String @unique + taskRunAttempt TaskRunAttempt? @relation(fields: [taskRunAttemptId], references: [id], onDelete: SetNull, onUpdate: Cascade) + taskRunAttemptId String? + createdAt DateTime @default(now()) updatedAt DateTime @updatedAt diff --git a/references/v3-catalog/src/trigger/longRunning.ts b/references/v3-catalog/src/trigger/longRunning.ts index b669e02b6..4d285aef6 100644 --- a/references/v3-catalog/src/trigger/longRunning.ts +++ b/references/v3-catalog/src/trigger/longRunning.ts @@ -13,3 +13,16 @@ export const longRunning = task({ }; }, }); + +export const longRunningParent = task({ + id: "long-running-parent", + run: async (payload: { message: string }) => { + logger.info("Long running parent", { payload }); + + await longRunning.triggerAndWait({ payload: { message: "child" } }); + + return { + finished: new Date().toISOString(), + }; + }, +}); diff --git a/references/v3-catalog/src/trigger/retries.ts b/references/v3-catalog/src/trigger/retries.ts index c9500f0d5..5053bcb86 100644 --- a/references/v3-catalog/src/trigger/retries.ts +++ b/references/v3-catalog/src/trigger/retries.ts @@ -4,6 +4,9 @@ import { interceptor } from "./utils/interceptor"; export const taskWithRetries = task({ id: "task-with-retries", + retry: { + maxAttempts: 4, + }, run: async (payload: any, { ctx }) => { const result = await retry.onThrow( async ({ attempt }) => { @@ -43,6 +46,13 @@ export const taskWithRetries = task({ }, }); +export const taskThatErrors = task({ + id: "task-that-errors", + run: async (payload: any, { ctx }) => { + throw new Error("failed"); + }, +}); + export const taskWithFetchRetries = task({ id: "task-with-fetch-retries", middleware: (payload: any, { next }) => {