Add a heartbeat for SUSPENDED snapshots, where when stalled will attempt to continue the run if unblocked

Author: ericallam

apps/webapp/app/env.server.ts

@@ -452,6 +452,10 @@ const EnvironmentSchema = z.object({
   RUN_ENGINE_TIMEOUT_PENDING_CANCEL: z.coerce.number().int().default(60_000),
   RUN_ENGINE_TIMEOUT_EXECUTING: z.coerce.number().int().default(60_000),
   RUN_ENGINE_TIMEOUT_EXECUTING_WITH_WAITPOINTS: z.coerce.number().int().default(60_000),
+  RUN_ENGINE_TIMEOUT_SUSPENDED: z.coerce
+    .number()
+    .int()
+    .default(60_000 * 10),
   RUN_ENGINE_DEBUG_WORKER_NOTIFICATIONS: z.coerce.boolean().default(false),
   RUN_ENGINE_PARENT_QUEUE_LIMIT: z.coerce.number().int().default(1000),
   RUN_ENGINE_CONCURRENCY_LIMIT_BIAS: z.coerce.number().default(0.75),

apps/webapp/app/v3/runEngine.server.ts

@@ -76,6 +76,7 @@ function createRunEngine() {
       PENDING_CANCEL: env.RUN_ENGINE_TIMEOUT_PENDING_CANCEL,
       EXECUTING: env.RUN_ENGINE_TIMEOUT_EXECUTING,
       EXECUTING_WITH_WAITPOINTS: env.RUN_ENGINE_TIMEOUT_EXECUTING_WITH_WAITPOINTS,
+      SUSPENDED: env.RUN_ENGINE_TIMEOUT_SUSPENDED,
     },
     releaseConcurrency: {
       disabled: env.RUN_ENGINE_RELEASE_CONCURRENCY_ENABLED === "0",

internal-packages/run-engine/src/engine/index.ts

@@ -182,6 +182,7 @@ export class RunEngine {
       PENDING_CANCEL: 60_000,
       EXECUTING: 60_000,
       EXECUTING_WITH_WAITPOINTS: 60_000,
+      SUSPENDED: 60_000 * 10,
     };
     this.heartbeatTimeouts = {
       ...defaultHeartbeatTimeouts,
@@ -1274,9 +1275,29 @@ export class RunEngine {
           break;
         }
         case "SUSPENDED": {
-          //todo should we do a periodic check here for whether waitpoints are actually still blocking?
-          //we could at least log some things out if a run has been in this state for a long time
-          throw new NotImplementedError("Not implemented SUSPENDED");
+          const result = await this.waitpointSystem.continueRunIfUnblocked({ runId });
+
+          this.logger.info("handleStalledSnapshot SUSPENDED continueRunIfUnblocked", {
+            runId,
+            result,
+            snapshotId: latestSnapshot.id,
+          });
+
+          switch (result) {
+            case "blocked": {
+              // Reschedule the heartbeat
+              await this.executionSnapshotSystem.restartHeartbeatForRun({
+                runId,
+              });
+              break;
+            }
+            case "unblocked":
+            case "skipped": {
+              break;
+            }
+          }
+
+          break;
         }
         case "PENDING_CANCEL": {
           //if the run is waiting to cancel but the worker hasn't confirmed that,

internal-packages/run-engine/src/engine/systems/executionSnapshotSystem.ts

@@ -364,7 +364,7 @@ export class ExecutionSnapshotSystem {
       return executionResultFromSnapshot(latestSnapshot);
     }
 
-    if (latestSnapshot.workerId !== workerId) {
+    if (latestSnapshot.workerId && latestSnapshot.workerId !== workerId) {
       this.$.logger.debug("heartbeatRun: worker ID does not match the latest snapshot", {
         runId,
         snapshotId,
@@ -394,6 +394,38 @@ export class ExecutionSnapshotSystem {
     return executionResultFromSnapshot(latestSnapshot);
   }
 
+  public async restartHeartbeatForRun({
+    runId,
+    tx,
+  }: {
+    runId: string;
+    tx?: PrismaClientOrTransaction;
+  }): Promise<ExecutionResult> {
+    const prisma = tx ?? this.$.prisma;
+
+    const latestSnapshot = await getLatestExecutionSnapshot(prisma, runId);
+
+    //extending the heartbeat
+    const intervalMs = this.#getHeartbeatIntervalMs(latestSnapshot.executionStatus);
+
+    if (intervalMs !== null) {
+      this.$.logger.debug("restartHeartbeatForRun: enqueuing heartbeat", {
+        runId,
+        snapshotId: latestSnapshot.id,
+        intervalMs,
+      });
+
+      await this.$.worker.enqueue({
+        id: `heartbeatSnapshot.${runId}`,
+        job: "heartbeatSnapshot",
+        payload: { snapshotId: latestSnapshot.id, runId },
+        availableAt: new Date(Date.now() + intervalMs),
+      });
+    }
+
+    return executionResultFromSnapshot(latestSnapshot);
+  }
+
   #getHeartbeatIntervalMs(status: TaskRunExecutionStatus): number | null {
     switch (status) {
       case "PENDING_EXECUTING": {
@@ -408,6 +440,9 @@ export class ExecutionSnapshotSystem {
       case "EXECUTING_WITH_WAITPOINTS": {
         return this.heartbeatTimeouts.EXECUTING_WITH_WAITPOINTS;
       }
+      case "SUSPENDED": {
+        return this.heartbeatTimeouts.SUSPENDED;
+      }
       default: {
         return null;
       }

internal-packages/run-engine/src/engine/systems/waitpointSystem.ts

@@ -490,7 +490,11 @@ export class WaitpointSystem {
     });
   }
 
-  public async continueRunIfUnblocked({ runId }: { runId: string }) {
+  public async continueRunIfUnblocked({
+    runId,
+  }: {
+    runId: string;
+  }): Promise<"blocked" | "unblocked" | "skipped"> {
     this.$.logger.debug(`continueRunIfUnblocked: start`, {
       runId,
     });
@@ -516,7 +520,7 @@ export class WaitpointSystem {
         runId,
         blockingWaitpoints,
       });
-      return;
+      return "blocked";
     }
 
     // 3. Get the run with environment
@@ -553,7 +557,7 @@ export class WaitpointSystem {
           runId,
           snapshot,
         });
-        return;
+        return "skipped";
       }
 
       //run is still executing, send a message to the worker
@@ -677,6 +681,8 @@ export class WaitpointSystem {
         runId,
       });
     }
+
+    return "unblocked";
   }
 
   public async createRunAssociatedWaitpoint(

internal-packages/run-engine/src/engine/tests/heartbeats.test.ts

@@ -479,6 +479,149 @@ describe("RunEngine heartbeats", () => {
     }
   });
 
+  containerTest("Suspended", async ({ prisma, redisOptions }) => {
+    const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
+
+    const heartbeatTimeout = 1000;
+
+    const engine = new RunEngine({
+      prisma,
+      worker: {
+        redis: redisOptions,
+        workers: 1,
+        tasksPerWorker: 10,
+        pollIntervalMs: 100,
+      },
+      queue: {
+        redis: redisOptions,
+      },
+      runLock: {
+        redis: redisOptions,
+      },
+      machines: {
+        defaultMachine: "small-1x",
+        machines: {
+          "small-1x": {
+            name: "small-1x" as const,
+            cpu: 0.5,
+            memory: 0.5,
+            centsPerMs: 0.0001,
+          },
+        },
+        baseCostInCents: 0.0001,
+      },
+      heartbeatTimeoutsMs: {
+        SUSPENDED: heartbeatTimeout,
+      },
+      tracer: trace.getTracer("test", "0.0.0"),
+    });
+
+    try {
+      const taskIdentifier = "test-task";
+
+      //create background worker
+      const backgroundWorker = await setupBackgroundWorker(
+        engine,
+        authenticatedEnvironment,
+        taskIdentifier
+      );
+
+      //trigger the run
+      const run = await engine.trigger(
+        {
+          number: 1,
+          friendlyId: "run_1234",
+          environment: authenticatedEnvironment,
+          taskIdentifier,
+          payload: "{}",
+          payloadType: "application/json",
+          context: {},
+          traceContext: {},
+          traceId: "t12345",
+          spanId: "s12345",
+          masterQueue: "main",
+          queue: "task/test-task",
+          isTest: false,
+          tags: [],
+        },
+        prisma
+      );
+
+      //dequeue the run
+      const dequeued = await engine.dequeueFromMasterQueue({
+        consumerId: "test_12345",
+        masterQueue: run.masterQueue,
+        maxRunCount: 10,
+      });
+
+      //create an attempt
+      await engine.startRunAttempt({
+        runId: dequeued[0].run.id,
+        snapshotId: dequeued[0].snapshot.id,
+      });
+
+      //cancel run
+      //create a manual waitpoint
+      const waitpointResult = await engine.createManualWaitpoint({
+        environmentId: authenticatedEnvironment.id,
+        projectId: authenticatedEnvironment.projectId,
+      });
+      expect(waitpointResult.waitpoint.status).toBe("PENDING");
+
+      //block the run
+      const blockedResult = await engine.blockRunWithWaitpoint({
+        runId: run.id,
+        waitpoints: waitpointResult.waitpoint.id,
+        projectId: authenticatedEnvironment.projectId,
+        organizationId: authenticatedEnvironment.organizationId,
+      });
+
+      const blockedExecutionData = await engine.getRunExecutionData({ runId: run.id });
+      expect(blockedExecutionData?.snapshot.executionStatus).toBe("EXECUTING_WITH_WAITPOINTS");
+
+      // Create a checkpoint
+      const checkpointResult = await engine.createCheckpoint({
+        runId: run.id,
+        snapshotId: blockedResult.id,
+        checkpoint: {
+          type: "DOCKER",
+          reason: "TEST_CHECKPOINT",
+          location: "test-location",
+          imageRef: "test-image-ref",
+        },
+      });
+
+      expect(checkpointResult.ok).toBe(true);
+
+      const snapshot = checkpointResult.ok ? checkpointResult.snapshot : null;
+
+      assertNonNullable(snapshot);
+
+      // Verify checkpoint creation
+      expect(snapshot.executionStatus).toBe("SUSPENDED");
+
+      // Now wait for the heartbeat to timeout, but it should retry later
+      await setTimeout(heartbeatTimeout * 1.5);
+
+      // Simulate a suspended run without any blocking waitpoints by deleting any blocking task run waitpoints
+      await prisma.taskRunWaitpoint.deleteMany({
+        where: {
+          taskRunId: run.id,
+        },
+      });
+
+      // Now wait for the heartbeat to timeout again
+      await setTimeout(heartbeatTimeout * 2);
+
+      // Expect the run to be queued
+      const executionData2 = await engine.getRunExecutionData({ runId: run.id });
+      assertNonNullable(executionData2);
+      expect(executionData2.snapshot.executionStatus).toBe("QUEUED");
+    } finally {
+      await engine.quit();
+    }
+  });
+
   containerTest("Heartbeat keeps run alive", async ({ prisma, redisOptions }) => {
     const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
 

internal-packages/run-engine/src/engine/types.ts

@@ -68,6 +68,7 @@ export type HeartbeatTimeouts = {
   PENDING_CANCEL: number;
   EXECUTING: number;
   EXECUTING_WITH_WAITPOINTS: number;
+  SUSPENDED: number;
 };
 
 export type TriggerParams = {