Fix several restore and resume bugs (#1418)

* try to correct resume messages with missing checkpoint

* prevent creating checkpoints for outdated task waits

* prevent creating checkpoints for outdated batch waits

* use heartbeats to check for and clean up any leftover containers

* lint

* improve exec logging

* improve resume attempt logs

* fix for resuming parents of canceled child runs

* separate SIGTERM from maybe OOM errors

* pretty errors can have magic dashboard links

* prevent uncancellable checkpoints

* simplify task run error code enum export

* grab the last, not the first child run

* Revert "prevent creating checkpoints for outdated batch waits"

This reverts commit f2b5c2a.

* Revert "grab the last, not the first child run"

This reverts commit 89ec5c8.

* Revert "prevent creating checkpoints for outdated task waits"

This reverts commit 11066b4.

* more logs for resume message handling

* add magic error link comment

* add changeset

Author: nicktrn

.changeset/many-plants-destroy.md

@@ -0,0 +1,5 @@
+---
+"@trigger.dev/core": patch
+---
+
+SIGTERM detection and prettier errors

apps/coordinator/src/checkpointer.ts

@@ -436,7 +436,10 @@ export class Checkpointer {
         this.#logger.error("Error during cleanup", { ...metadata, error });
       }
 
-      this.#abortControllers.delete(runId);
+      // Ensure only the current controller is removed
+      if (this.#abortControllers.get(runId) === controller) {
+        this.#abortControllers.delete(runId);
+      }
       controller.signal.removeEventListener("abort", onAbort);
     };
 

apps/coordinator/src/exec.ts

@@ -64,7 +64,18 @@ export class Exec {
       command,
       argsRaw: args,
       argsTrimmed,
-      ...output,
+      globalOpts: {
+        trimArgs: this.trimArgs,
+        neverThrow: this.neverThrow,
+        hasAbortSignal: !!this.abortSignal,
+      },
+      localOpts: opts,
+      stdout: output.stdout,
+      stderr: output.stderr,
+      pid: result.pid,
+      exitCode: result.exitCode,
+      aborted: result.aborted,
+      killed: result.killed,
     };
 
     if (this.logOutput) {

apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.v3.$projectParam.runs.$runParam.spans.$spanParam/route.tsx

@@ -1,4 +1,10 @@
-import { CheckIcon, ClockIcon, CloudArrowDownIcon, QueueListIcon } from "@heroicons/react/20/solid";
+import {
+  CheckIcon,
+  ClockIcon,
+  CloudArrowDownIcon,
+  EnvelopeIcon,
+  QueueListIcon,
+} from "@heroicons/react/20/solid";
 import { Link } from "@remix-run/react";
 import { LoaderFunctionArgs } from "@remix-run/server-runtime";
 import {
@@ -13,6 +19,7 @@ import { typedjson, useTypedFetcher } from "remix-typedjson";
 import { ExitIcon } from "~/assets/icons/ExitIcon";
 import { CodeBlock } from "~/components/code/CodeBlock";
 import { EnvironmentLabel } from "~/components/environments/EnvironmentLabel";
+import { Feedback } from "~/components/Feedback";
 import { Button, LinkButton } from "~/components/primitives/Buttons";
 import { Callout } from "~/components/primitives/Callout";
 import { DateTime, DateTimeAccurate } from "~/components/primitives/DateTime";
@@ -963,11 +970,26 @@ function RunError({ error }: { error: TaskRunError }) {
         <div className="flex flex-col gap-2 rounded-sm border border-rose-500/50 px-3 pb-3 pt-2">
           <Header3 className="text-rose-500">{name}</Header3>
           {enhancedError.message && <Callout variant="error">{enhancedError.message}</Callout>}
-          {enhancedError.link && (
-            <Callout variant="docs" to={enhancedError.link.href}>
-              {enhancedError.link.name}
-            </Callout>
-          )}
+          {enhancedError.link &&
+            (enhancedError.link.magic === "CONTACT_FORM" ? (
+              <Feedback
+                button={
+                  <Button
+                    variant="tertiary/medium"
+                    LeadingIcon={EnvelopeIcon}
+                    leadingIconClassName="text-blue-400"
+                    fullWidth
+                    textAlignLeft
+                  >
+                    {enhancedError.link.name}
+                  </Button>
+                }
+              />
+            ) : (
+              <Callout variant="docs" to={enhancedError.link.href}>
+                {enhancedError.link.name}
+              </Callout>
+            ))}
           {enhancedError.stackTrace && (
             <CodeBlock
               showCopyButton={false}

apps/webapp/app/v3/environmentVariables/environmentVariablesRepository.server.ts

@@ -136,7 +136,7 @@ export class EnvironmentVariablesRepository implements Repository {
 
     try {
       for (const variable of values) {
-      const result = await $transaction(this.prismaClient, async (tx) => {
+        const result = await $transaction(this.prismaClient, async (tx) => {
           const environmentVariable = await tx.environmentVariable.upsert({
             where: {
               projectId_key: {

apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts

@@ -21,7 +21,7 @@ import {
   TaskRunStatus,
 } from "@trigger.dev/database";
 import { z } from "zod";
-import { prisma } from "~/db.server";
+import { $replica, prisma } from "~/db.server";
 import { findEnvironmentById } from "~/models/runtimeEnvironment.server";
 import { logger } from "~/services/logger.server";
 import { singleton } from "~/utils/singleton";
@@ -43,7 +43,12 @@ import { generateJWTTokenForEnvironment } from "~/services/apiAuth.server";
 import { EnvironmentVariable } from "../environmentVariables/repository";
 import { machinePresetFromConfig } from "../machinePresets.server";
 import { env } from "~/env.server";
-import { isFinalAttemptStatus, isFinalRunStatus } from "../taskStatus";
+import {
+  FINAL_ATTEMPT_STATUSES,
+  FINAL_RUN_STATUSES,
+  isFinalAttemptStatus,
+  isFinalRunStatus,
+} from "../taskStatus";
 import { getMaxDuration } from "../utils/maxDuration";
 
 const WithTraceContext = z.object({
@@ -620,6 +625,9 @@ export class SharedQueueConsumer {
         const resumableRun = await prisma.taskRun.findUnique({
           where: {
             id: message.messageId,
+            status: {
+              notIn: FINAL_RUN_STATUSES,
+            },
           },
         });
 
@@ -633,6 +641,14 @@ export class SharedQueueConsumer {
           return;
         }
 
+        if (resumableRun.status !== "EXECUTING") {
+          logger.warn("Run is not executing, will try to resume anyway", {
+            queueMessage: message.data,
+            messageId: message.messageId,
+            runStatus: resumableRun.status,
+          });
+        }
+
         const resumableAttempt = await prisma.taskRunAttempt.findUnique({
           where: {
             id: messageBody.data.resumableAttemptId,
@@ -740,7 +756,11 @@ export class SharedQueueConsumer {
             executions,
           };
 
-          logger.debug("Broadcasting RESUME_AFTER_DEPENDENCY_WITH_ACK", { resumeMessage, message });
+          logger.debug("Broadcasting RESUME_AFTER_DEPENDENCY_WITH_ACK", {
+            resumeMessage,
+            message,
+            resumableRun,
+          });
 
           // The attempt should still be running so we can broadcast to all coordinators to resume immediately
           const responses = await socketIo.coordinatorNamespace
@@ -763,15 +783,91 @@ export class SharedQueueConsumer {
           }
 
           const hasSuccess = responses.some((response) => response.success);
-          if (!hasSuccess) {
-            logger.warn("RESUME_AFTER_DEPENDENCY_WITH_ACK failed", {
-              resumeMessage,
-              responses,
-              message,
-            });
-            await this.#nackAndDoMoreWork(message.messageId, this._options.nextTickInterval, 5_000);
+
+          if (hasSuccess) {
+            this.#doMoreWork();
             return;
           }
+
+          // No coordinator was able to resume the run
+          logger.warn("RESUME_AFTER_DEPENDENCY_WITH_ACK failed", {
+            resumeMessage,
+            responses,
+            message,
+          });
+
+          // Let's check if the run is frozen
+          if (resumableRun.status === "WAITING_TO_RESUME") {
+            logger.debug("RESUME_AFTER_DEPENDENCY_WITH_ACK run is waiting to be restored", {
+              queueMessage: message.data,
+              messageId: message.messageId,
+            });
+
+            try {
+              const restoreService = new RestoreCheckpointService();
+
+              const checkpointEvent = await restoreService.getLastCheckpointEventIfUnrestored(
+                resumableRun.id
+              );
+
+              if (checkpointEvent) {
+                // The last checkpoint hasn't been restored yet, so restore it
+                const checkpoint = await restoreService.call({
+                  eventId: checkpointEvent.id,
+                });
+
+                if (!checkpoint) {
+                  logger.debug("RESUME_AFTER_DEPENDENCY_WITH_ACK failed to restore checkpoint", {
+                    queueMessage: message.data,
+                    messageId: message.messageId,
+                  });
+
+                  await this.#ackAndDoMoreWork(message.messageId);
+                  return;
+                }
+
+                logger.debug("RESUME_AFTER_DEPENDENCY_WITH_ACK restored checkpoint", {
+                  queueMessage: message.data,
+                  messageId: message.messageId,
+                  checkpoint,
+                });
+
+                this.#doMoreWork();
+                return;
+              } else {
+                logger.debug(
+                  "RESUME_AFTER_DEPENDENCY_WITH_ACK run is frozen without last checkpoint event",
+                  {
+                    queueMessage: message.data,
+                    messageId: message.messageId,
+                  }
+                );
+              }
+            } catch (e) {
+              if (e instanceof Error) {
+                this._currentSpan?.recordException(e);
+              } else {
+                this._currentSpan?.recordException(new Error(String(e)));
+              }
+
+              this._endSpanInNextIteration = true;
+
+              await this.#nackAndDoMoreWork(
+                message.messageId,
+                this._options.nextTickInterval,
+                5_000
+              );
+              return;
+            }
+          }
+
+          logger.debug("RESUME_AFTER_DEPENDENCY_WITH_ACK retrying", {
+            queueMessage: message.data,
+            messageId: message.messageId,
+          });
+
+          await this.#nackAndDoMoreWork(message.messageId, this._options.nextTickInterval, 5_000);
+          return;
         } catch (e) {
           if (e instanceof Error) {
             this._currentSpan?.recordException(e);
@@ -896,7 +992,7 @@ class SharedQueueTasks {
       where: {
         id,
         status: {
-          in: ["COMPLETED", "FAILED"],
+          in: FINAL_ATTEMPT_STATUSES,
         },
       },
       include: {
@@ -1237,13 +1333,13 @@ class SharedQueueTasks {
       return;
     }
 
-    await marqs?.heartbeatMessage(taskRunAttempt.taskRunId);
+    await this.#heartbeat(taskRunAttempt.taskRunId);
   }
 
   async taskRunHeartbeat(runId: string) {
     logger.debug("[SharedQueueConsumer] taskRunHeartbeat()", { runId });
 
-    await marqs?.heartbeatMessage(runId);
+    await this.#heartbeat(runId);
   }
 
   public async taskRunFailed(completion: TaskRunFailedExecutionResult) {
@@ -1254,6 +1350,66 @@ class SharedQueueTasks {
     await service.call(completion.id, completion);
   }
 
+  async #heartbeat(runId: string) {
+    await marqs?.heartbeatMessage(runId);
+
+    try {
+      // There can be a lot of calls per minute and the data doesn't have to be accurate, so use the read replica
+      const taskRun = await $replica.taskRun.findFirst({
+        where: {
+          id: runId,
+        },
+        select: {
+          id: true,
+          status: true,
+          runtimeEnvironment: {
+            select: {
+              type: true,
+            },
+          },
+          lockedToVersion: {
+            select: {
+              supportsLazyAttempts: true,
+            },
+          },
+        },
+      });
+
+      if (!taskRun) {
+        logger.error("SharedQueueTasks.#heartbeat: Task run not found", {
+          runId,
+        });
+
+        return;
+      }
+
+      if (taskRun.runtimeEnvironment.type === "DEVELOPMENT") {
+        return;
+      }
+
+      if (isFinalRunStatus(taskRun.status)) {
+        logger.debug("SharedQueueTasks.#heartbeat: Task run is in final status", {
+          runId,
+          status: taskRun.status,
+        });
+
+        // Signal to exit any leftover containers
+        socketIo.coordinatorNamespace.emit("REQUEST_RUN_CANCELLATION", {
+          version: "v1",
+          runId: taskRun.id,
+          // Give the run a few seconds to exit to complete any flushing etc
+          delayInMs: taskRun.lockedToVersion?.supportsLazyAttempts ? 5_000 : undefined,
+        });
+        return;
+      }
+    } catch (error) {
+      logger.error("SharedQueueTasks.#heartbeat: Error signaling run cancellation", {
+        runId,
+        error: error instanceof Error ? error.message : error,
+      });
+    }
+  }
+
   async #buildEnvironmentVariables(
     environment: RuntimeEnvironment,
     runId: string,