Added RESUME_AFTER_DEPENDENCY_WITH_ACK

Author: matt-aitken

apps/coordinator/src/index.ts

@@ -162,6 +162,33 @@ class TaskCoordinator {
 
           taskSocket.emit("RESUME_AFTER_DEPENDENCY", message);
         },
+        RESUME_AFTER_DEPENDENCY_WITH_ACK: async (message) => {
+          const taskSocket = await this.#getAttemptSocket(message.attemptFriendlyId);
+
+          if (!taskSocket) {
+            logger.log("Socket for attempt not found", {
+              attemptFriendlyId: message.attemptFriendlyId,
+            });
+            return {
+              success: false,
+              error: {
+                name: "SocketNotFoundError",
+                message: "Socket for attempt not found",
+              },
+            };
+          }
+
+          await chaosMonkey.call();
+
+          // In case the task resumed faster than we could checkpoint
+          this.#cancelCheckpoint(message.runId);
+
+          taskSocket.emit("RESUME_AFTER_DEPENDENCY", message);
+
+          return {
+            success: true,
+          };
+        },
         RESUME_AFTER_DURATION: async (message) => {
           const taskSocket = await this.#getAttemptSocket(message.attemptFriendlyId);
 

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

@@ -725,20 +725,47 @@ export class SharedQueueConsumer {
         }
 
         try {
-          logger.debug("Broadcasting RESUME_AFTER_DEPENDENCY", {
-            runId: resumableAttempt.taskRunId,
-            attemptId: resumableAttempt.id,
-          });
-
-          // The attempt should still be running so we can broadcast to all coordinators to resume immediately
-          socketIo.coordinatorNamespace.emit("RESUME_AFTER_DEPENDENCY", {
-            version: "v1",
+          const resumeMessage = {
+            version: "v1" as const,
             runId: resumableAttempt.taskRunId,
             attemptId: resumableAttempt.id,
             attemptFriendlyId: resumableAttempt.friendlyId,
             completions,
             executions,
+          };
+
+          logger.debug("Broadcasting RESUME_AFTER_DEPENDENCY_WITH_ACK", { resumeMessage, message });
+
+          // The attempt should still be running so we can broadcast to all coordinators to resume immediately
+          const responses = await socketIo.coordinatorNamespace
+            .timeout(10_000)
+            .emitWithAck("RESUME_AFTER_DEPENDENCY_WITH_ACK", resumeMessage);
+
+          logger.debug("RESUME_AFTER_DEPENDENCY_WITH_ACK received", {
+            resumeMessage,
+            responses,
+            message,
           });
+
+          if (responses.length === 0) {
+            logger.error("RESUME_AFTER_DEPENDENCY_WITH_ACK no response", {
+              resumeMessage,
+              message,
+            });
+            await this.#nackAndDoMoreWork(message.messageId, this._options.nextTickInterval, 1_000);
+            return;
+          }
+
+          const failed = responses.filter((response) => !response.success);
+          if (failed.length > 0) {
+            logger.error("RESUME_AFTER_DEPENDENCY_WITH_ACK failed", {
+              resumeMessage,
+              failed,
+              message,
+            });
+            await this.#nackAndDoMoreWork(message.messageId, this._options.nextTickInterval, 1_000);
+            return;
+          }
         } catch (e) {
           if (e instanceof Error) {
             this._currentSpan?.recordException(e);

packages/core/src/v3/schemas/messages.ts

@@ -15,6 +15,21 @@ import {
   WaitReason,
 } from "./schemas.js";
 
+const ackCallbackResult = z.discriminatedUnion("success", [
+  z.object({
+    success: z.literal(false),
+    error: z.object({
+      name: z.string(),
+      message: z.string(),
+      stack: z.string().optional(),
+      stderr: z.string().optional(),
+    }),
+  }),
+  z.object({
+    success: z.literal(true),
+  }),
+]);
+
 export const BackgroundWorkerServerMessages = z.discriminatedUnion("type", [
   z.object({
     type: z.literal("CANCEL_ATTEMPT"),
@@ -269,20 +284,7 @@ export const PlatformToProviderMessages = {
       projectId: z.string(),
       deploymentId: z.string(),
     }),
-    callback: z.discriminatedUnion("success", [
-      z.object({
-        success: z.literal(false),
-        error: z.object({
-          name: z.string(),
-          message: z.string(),
-          stack: z.string().optional(),
-          stderr: z.string().optional(),
-        }),
-      }),
-      z.object({
-        success: z.literal(true),
-      }),
-    ]),
+    callback: ackCallbackResult,
   },
   RESTORE: {
     message: z.object({
@@ -504,6 +506,7 @@ export const CoordinatorToPlatformMessages = {
 };
 
 export const PlatformToCoordinatorMessages = {
+  /** @deprecated use RESUME_AFTER_DEPENDENCY_WITH_ACK instead  */
   RESUME_AFTER_DEPENDENCY: {
     message: z.object({
       version: z.literal("v1").default("v1"),
@@ -514,6 +517,17 @@ export const PlatformToCoordinatorMessages = {
       executions: TaskRunExecution.array(),
     }),
   },
+  RESUME_AFTER_DEPENDENCY_WITH_ACK: {
+    message: z.object({
+      version: z.literal("v1").default("v1"),
+      runId: z.string(),
+      attemptId: z.string(),
+      attemptFriendlyId: z.string(),
+      completions: TaskRunExecutionResult.array(),
+      executions: TaskRunExecution.array(),
+    }),
+    callback: ackCallbackResult,
+  },
   RESUME_AFTER_DURATION: {
     message: z.object({
       version: z.literal("v1").default("v1"),