Stop heartbeats when creating andWait checkpoints, fix for concurrency tracker (#1268)

* RESUME messages were missing some data that is needed for the concurrency tracker

* Cancel heartbeats when checkpoints are created for triggerAndWait and batchTriggerAndWait

* Added some logs when canceling heartbeats

* Improved the nested dependencies test task

Author: matt-aitken

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

@@ -521,6 +521,27 @@ export class MarQS {
     );
   }
 
+  public async cancelHeartbeat(messageId: string) {
+    return this.#trace(
+      "cancelHeartbeat",
+      async (span) => {
+        span.setAttributes({
+          [SemanticAttributes.MESSAGE_ID]: messageId,
+        });
+
+        await this.options.visibilityTimeoutStrategy.cancelHeartbeat(messageId);
+      },
+      {
+        kind: SpanKind.CONSUMER,
+        attributes: {
+          [SEMATTRS_MESSAGING_OPERATION]: "cancelHeartbeat",
+          [SEMATTRS_MESSAGE_ID]: messageId,
+          [SEMATTRS_MESSAGING_SYSTEM]: "marqs",
+        },
+      }
+    );
+  }
+
   async #trace<T>(
     name: string,
     fn: (span: Span) => Promise<T>,

apps/webapp/app/v3/services/createCheckpoint.server.ts

@@ -5,7 +5,6 @@ import { logger } from "~/services/logger.server";
 import { marqs } from "~/v3/marqs/index.server";
 import { generateFriendlyId } from "../friendlyIdentifiers";
 import {
-  FINAL_ATTEMPT_STATUSES,
   isFinalAttemptStatus,
   isFinalRunStatus,
   isFreezableAttemptStatus,
@@ -135,14 +134,32 @@ export class CreateCheckpointService extends BaseService {
     const { reason } = params;
 
     let checkpointEvent: CheckpointRestoreEvent | undefined;
-    let keepRunAlive = false;
 
     switch (reason.type) {
       case "WAIT_FOR_DURATION": {
         checkpointEvent = await eventService.checkpoint({
           checkpointId: checkpoint.id,
         });
 
+        if (checkpointEvent) {
+          await marqs?.replaceMessage(
+            attempt.taskRunId,
+            {
+              type: "RESUME_AFTER_DURATION",
+              resumableAttemptId: attempt.id,
+              checkpointEventId: checkpointEvent.id,
+            },
+            reason.now + reason.ms
+          );
+
+          return {
+            success: true,
+            checkpoint,
+            event: checkpointEvent,
+            keepRunAlive: false,
+          };
+        }
+
         break;
       }
       case "WAIT_FOR_TASK": {
@@ -152,6 +169,14 @@ export class CreateCheckpointService extends BaseService {
         });
 
         if (checkpointEvent) {
+          //heartbeats will start again when the run resumes
+          logger.log("CreateCheckpointService: Canceling heartbeat", {
+            attemptId: attempt.id,
+            taskRunId: attempt.taskRunId,
+            type: "WAIT_FOR_TASK",
+          });
+          await marqs?.cancelHeartbeat(attempt.taskRunId);
+
           const dependency = await this._prisma.taskRunDependency.findFirst({
             select: {
               id: true,
@@ -270,6 +295,7 @@ export class CreateCheckpointService extends BaseService {
             };
           }
 
+          //resume the dependent task
           await ResumeTaskDependencyService.enqueue(dependency.id, lastAttempt.id, this._prisma);
 
           return {
@@ -289,6 +315,14 @@ export class CreateCheckpointService extends BaseService {
         });
 
         if (checkpointEvent) {
+          //heartbeats will start again when the run resumes
+          logger.log("CreateCheckpointService: Canceling heartbeat", {
+            attemptId: attempt.id,
+            taskRunId: attempt.taskRunId,
+            type: "WAIT_FOR_BATCH",
+          });
+          await marqs?.cancelHeartbeat(attempt.taskRunId);
+
           const batchRun = await this._prisma.batchTaskRun.findFirst({
             select: {
               id: true,
@@ -348,23 +382,11 @@ export class CreateCheckpointService extends BaseService {
       };
     }
 
-    if (reason.type === "WAIT_FOR_DURATION") {
-      await marqs?.replaceMessage(
-        attempt.taskRunId,
-        {
-          type: "RESUME_AFTER_DURATION",
-          resumableAttemptId: attempt.id,
-          checkpointEventId: checkpointEvent.id,
-        },
-        reason.now + reason.ms
-      );
-    }
-
     return {
       success: true,
       checkpoint,
       event: checkpointEvent,
-      keepRunAlive,
+      keepRunAlive: false,
     };
   }
 }

apps/webapp/app/v3/services/resumeBatchRun.server.ts

@@ -80,6 +80,7 @@ export class ResumeBatchRunService extends BaseService {
           completedAttemptIds: [],
           resumableAttemptId: batchRun.dependentTaskAttempt.id,
           checkpointEventId: batchRun.checkpointEventId,
+          taskIdentifier: batchRun.dependentTaskAttempt.taskRun.taskIdentifier,
           projectId: batchRun.dependentTaskAttempt.runtimeEnvironment.projectId,
           environmentId: batchRun.dependentTaskAttempt.runtimeEnvironment.id,
           environmentType: batchRun.dependentTaskAttempt.runtimeEnvironment.type,
@@ -109,6 +110,10 @@ export class ResumeBatchRunService extends BaseService {
         completedAttemptIds: batchRun.items.map((item) => item.taskRunAttemptId).filter(Boolean),
         resumableAttemptId: batchRun.dependentTaskAttempt.id,
         checkpointEventId: batchRun.checkpointEventId ?? undefined,
+        taskIdentifier: batchRun.dependentTaskAttempt.taskRun.taskIdentifier,
+        projectId: batchRun.dependentTaskAttempt.runtimeEnvironment.projectId,
+        environmentId: batchRun.dependentTaskAttempt.runtimeEnvironment.id,
+        environmentType: batchRun.dependentTaskAttempt.runtimeEnvironment.type,
       });
     }
   }

apps/webapp/app/v3/services/resumeTaskDependency.server.ts

@@ -48,6 +48,7 @@ export class ResumeTaskDependencyService extends BaseService {
           completedAttemptIds: [sourceTaskAttemptId],
           resumableAttemptId: dependency.dependentAttempt.id,
           checkpointEventId: dependency.checkpointEventId,
+          taskIdentifier: dependency.taskRun.taskIdentifier,
           projectId: dependency.taskRun.runtimeEnvironment.projectId,
           environmentId: dependency.taskRun.runtimeEnvironment.id,
           environmentType: dependency.taskRun.runtimeEnvironment.type,
@@ -77,6 +78,10 @@ export class ResumeTaskDependencyService extends BaseService {
         completedAttemptIds: [sourceTaskAttemptId],
         resumableAttemptId: dependency.dependentAttempt.id,
         checkpointEventId: dependency.checkpointEventId ?? undefined,
+        taskIdentifier: dependency.taskRun.taskIdentifier,
+        projectId: dependency.taskRun.runtimeEnvironment.projectId,
+        environmentId: dependency.taskRun.runtimeEnvironment.id,
+        environmentType: dependency.taskRun.runtimeEnvironment.type,
       });
     }
   }

references/v3-catalog/src/trigger/checkpoints.ts

@@ -61,16 +61,18 @@ export const nestedDependencies = task({
     const triggerOrBatch = depth % 2 === 0;
 
     if (triggerOrBatch) {
-      const result = await nestedDependencies.triggerAndWait({
-        depth: depth + 1,
-        maxDepth,
-        waitSeconds,
-        failAttemptChance,
-      });
-      logger.log(`Triggered complete`);
-
-      if (!result.ok && failParents) {
-        throw new Error(`Failed at ${depth}/${maxDepth} depth`);
+      for (let i = 0; i < batchSize; i++) {
+        const result = await nestedDependencies.triggerAndWait({
+          depth: depth + 1,
+          maxDepth,
+          waitSeconds,
+          failAttemptChance,
+        });
+        logger.log(`Triggered complete ${i + 1}/${batchSize}`);
+
+        if (!result.ok && failParents) {
+          throw new Error(`Failed at ${depth}/${maxDepth} depth`);
+        }
       }
     } else {
       const results = await nestedDependencies.batchTriggerAndWait(