diff --git a/apps/webapp/app/v3/marqs/index.server.ts b/apps/webapp/app/v3/marqs/index.server.ts index 554976879..b46628d07 100644 --- a/apps/webapp/app/v3/marqs/index.server.ts +++ b/apps/webapp/app/v3/marqs/index.server.ts @@ -231,7 +231,7 @@ export class MarQS { return; } - const message = await this.#readMessage(messageData.messageId); + const message = await this.readMessage(messageData.messageId); if (message) { span.setAttributes({ @@ -308,7 +308,7 @@ export class MarQS { return; } - const message = await this.#readMessage(messageData.messageId); + const message = await this.readMessage(messageData.messageId); if (message) { span.setAttributes({ @@ -336,7 +336,7 @@ export class MarQS { return this.#trace( "acknowledgeMessage", async (span) => { - const message = await this.#readMessage(messageId); + const message = await this.readMessage(messageId); if (!message) { return; @@ -374,12 +374,13 @@ export class MarQS { public async replaceMessage( messageId: string, messageData: Record, - timestamp?: number + timestamp?: number, + inplace?: boolean ) { return this.#trace( "replaceMessage", async (span) => { - const oldMessage = await this.#readMessage(messageId); + const oldMessage = await this.readMessage(messageId); if (!oldMessage) { return; @@ -392,6 +393,27 @@ export class MarQS { [SemanticAttributes.PARENT_QUEUE]: oldMessage.parentQueue, }); + const traceContext = { + traceparent: oldMessage.data.traceparent, + tracestate: oldMessage.data.tracestate, + }; + + const newMessage: MessagePayload = { + version: "1", + // preserve original trace context + data: { ...messageData, ...traceContext }, + queue: oldMessage.queue, + concurrencyKey: oldMessage.concurrencyKey, + timestamp: timestamp ?? Date.now(), + messageId, + parentQueue: oldMessage.parentQueue, + }; + + if (inplace) { + await this.#callReplaceMessage(newMessage); + return; + } + await this.#callAcknowledgeMessage({ parentQueue: oldMessage.parentQueue, messageKey: this.keys.messageKey(messageId), @@ -403,16 +425,6 @@ export class MarQS { messageId, }); - const newMessage: MessagePayload = { - version: "1", - data: messageData, - queue: oldMessage.queue, - concurrencyKey: oldMessage.concurrencyKey, - timestamp: timestamp ?? Date.now(), - messageId, - parentQueue: oldMessage.parentQueue, - }; - await this.#callEnqueueMessage(newMessage); }, { @@ -455,7 +467,7 @@ export class MarQS { return this.#trace( "nackMessage", async (span) => { - const message = await this.#readMessage(messageId); + const message = await this.readMessage(messageId); if (!message) { return; @@ -505,7 +517,7 @@ export class MarQS { return this.options.visibilityTimeoutInMs ?? 300000; } - async #readMessage(messageId: string) { + async readMessage(messageId: string) { return this.#trace( "readMessage", async (span) => { @@ -881,6 +893,17 @@ export class MarQS { }; } + async #callReplaceMessage(message: MessagePayload) { + logger.debug("Calling replaceMessage", { + messagePayload: message, + }); + + return this.redis.replaceMessage( + this.keys.messageKey(message.messageId), + JSON.stringify(message) + ); + } + async #callAcknowledgeMessage({ parentQueue, messageKey, @@ -1185,6 +1208,25 @@ return {messageId, messageScore} -- Return message details `, }); + this.redis.defineCommand("replaceMessage", { + numberOfKeys: 1, + lua: ` +local messageKey = KEYS[1] +local messageData = ARGV[1] + +-- Check if message exists +local existingMessage = redis.call('GET', messageKey) + +-- Do nothing if it doesn't +if #existingMessage == nil then + return nil +end + +-- Replace the message +redis.call('SET', messageKey, messageData, 'GET') + `, + }); + this.redis.defineCommand("acknowledgeMessage", { numberOfKeys: 7, lua: ` @@ -1406,6 +1448,12 @@ declare module "ioredis" { callback?: Callback<[string, string]> ): Result<[string, string] | null, Context>; + replaceMessage( + messageKey: string, + messageData: string, + callback?: Callback + ): Result; + acknowledgeMessage( parentQueue: string, messageKey: string, diff --git a/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts b/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts index 3816c6b34..dbd908e7f 100644 --- a/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts +++ b/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts @@ -28,13 +28,14 @@ import { socketIo } from "../handleSocketIo.server"; import { findCurrentWorkerDeployment } from "../models/workerDeployment.server"; import { RestoreCheckpointService } from "../services/restoreCheckpoint.server"; import { tracer } from "../tracer.server"; +import { CrashTaskRunService } from "../services/crashTaskRun.server"; const WithTraceContext = z.object({ traceparent: z.string().optional(), tracestate: z.string().optional(), }); -const MessageBody = z.discriminatedUnion("type", [ +export const SharedQueueMessageBody = z.discriminatedUnion("type", [ WithTraceContext.extend({ type: z.literal("EXECUTE"), taskIdentifier: z.string(), @@ -51,8 +52,14 @@ const MessageBody = z.discriminatedUnion("type", [ resumableAttemptId: z.string(), checkpointEventId: z.string(), }), + WithTraceContext.extend({ + type: z.literal("FAIL"), + reason: z.string(), + }), ]); +export type SharedQueueMessageBody = z.infer; + type BackgroundWorkerWithTasks = BackgroundWorker & { tasks: BackgroundWorkerTask[] }; export type SharedQueueConsumerOptions = { @@ -233,7 +240,7 @@ export class SharedQueueConsumer { logger.log("dequeueMessageInSharedQueue()", { queueMessage: message }); - const messageBody = MessageBody.safeParse(message.data); + const messageBody = SharedQueueMessageBody.safeParse(message.data); if (!messageBody.success) { logger.error("Failed to parse message", { @@ -739,6 +746,34 @@ export class SharedQueueConsumer { break; } + // Fail for whatever reason, usually runs that have been resumed but stopped heartbeating + case "FAIL": { + const existingTaskRun = await prisma.taskRun.findUnique({ + where: { + id: message.messageId, + }, + }); + + if (!existingTaskRun) { + logger.error("No existing task run to fail", { + queueMessage: messageBody, + messageId: message.messageId, + }); + + await this.#ackAndDoMoreWork(message.messageId); + return; + } + + // TODO: Consider failing the attempt and retrying instead. This may not be a good idea, as dequeued FAIL messages tend to point towards critical, persistent errors. + const service = new CrashTaskRunService(); + await service.call(existingTaskRun.id, { + crashAttempts: true, + reason: messageBody.data.reason, + }); + + await this.#ackAndDoMoreWork(message.messageId); + return; + } } this.#doMoreWork(); diff --git a/apps/webapp/app/v3/services/resumeAttempt.server.ts b/apps/webapp/app/v3/services/resumeAttempt.server.ts index 2e8b62735..42c14ba23 100644 --- a/apps/webapp/app/v3/services/resumeAttempt.server.ts +++ b/apps/webapp/app/v3/services/resumeAttempt.server.ts @@ -2,13 +2,14 @@ import { CoordinatorToPlatformMessages, TaskRunExecution, TaskRunExecutionResult, + WaitReason, } from "@trigger.dev/core/v3"; import type { InferSocketMessageSchema } from "@trigger.dev/core/v3/zodSocket"; import { $transaction, PrismaClientOrTransaction } from "~/db.server"; import { logger } from "~/services/logger.server"; import { marqs } from "~/v3/marqs/index.server"; import { socketIo } from "../handleSocketIo.server"; -import { sharedQueueTasks } from "../marqs/sharedQueueConsumer.server"; +import { SharedQueueMessageBody, sharedQueueTasks } from "../marqs/sharedQueueConsumer.server"; import { BaseService } from "./baseService.server"; import { TaskRunAttempt } from "@trigger.dev/database"; @@ -149,6 +150,9 @@ export class ResumeAttemptService extends BaseService { break; } } + + // Prevent infinite restores by failing runs that don't heartbeat after post-restore resume requests + await this.#replaceResumeWithFailMessage(attempt.taskRunId, params.type); }); } @@ -249,4 +253,50 @@ export class ResumeAttemptService extends BaseService { }, }); } + + async #replaceResumeWithFailMessage(messageId: string, waitReason: WaitReason) { + const currentMessage = await marqs?.readMessage(messageId); + + if (!currentMessage) { + logger.debug("No message to replace", { messageId, waitReason }); + return; + } + + const currentBody = SharedQueueMessageBody.safeParse(currentMessage.data); + + if (!currentBody.success) { + logger.debug("Invalid message body", { messageId, waitReason, currentBody }); + return; + } + + const currentType = currentBody.data.type; + + if (currentType !== "RESUME" && currentType !== "RESUME_AFTER_DURATION") { + logger.debug("Not a resume message", { messageId, waitReason, currentBody }); + return; + } + + let reason = "Worker unresponsive after restore"; + + switch (waitReason) { + case "WAIT_FOR_DURATION": + reason = "Worker unresponsive after waiting for duration"; + break; + case "WAIT_FOR_TASK": + reason = "Worker unresponsive after waiting for task"; + break; + case "WAIT_FOR_BATCH": + reason = "Worker unresponsive after waiting for batch task"; + break; + default: + break; + } + + const failMessage: SharedQueueMessageBody = { + type: "FAIL", + reason, + }; + + return await marqs?.replaceMessage(messageId, failMessage, undefined, true); + } }