prevent infinite restores

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<string, unknown>,
-    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<void>
+    ): Result<void, Context>;
+
     acknowledgeMessage(
       parentQueue: string,
       messageKey: string,

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<typeof SharedQueueMessageBody>;
+
 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();

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);
+  }
 }