Cancel runs and attempts, in prod and dev

Author: ericallam

apps/coordinator/package.json

@@ -4,7 +4,6 @@
   "version": "0.0.1",
   "description": "",
   "main": "dist/index.cjs",
-  "type": "module",
   "scripts": {
     "build": "npm run build:bundle",
     "build:bundle": "esbuild src/index.ts --bundle --outfile=dist/index.mjs --platform=node --format=esm --target=esnext --banner:js=\"const require = createRequire(import.meta.url);\"",
@@ -31,4 +30,4 @@
     "tsx": "^4.7.0",
     "typescript": "^5.3.3"
   }
-}
+}

apps/coordinator/src/index.ts

@@ -301,6 +301,16 @@ class TaskCoordinator {
 
           taskSocket.emit("RESUME_AFTER_DURATION", message);
         },
+        REQUEST_ATTEMPT_CANCELLATION: async (message) => {
+          const taskSocket = await this.#getAttemptSocket(message.attemptId);
+
+          if (!taskSocket) {
+            logger.log("Socket for attempt not found", { attemptId: message.attemptId });
+            return;
+          }
+
+          taskSocket.emit("REQUEST_ATTEMPT_CANCELLATION", message);
+        },
       },
     });
 
@@ -385,11 +395,21 @@ class TaskCoordinator {
 
           if (!executionAck) {
             logger.error("no execution ack", { attemptId: socket.data.attemptId });
+
+            socket.emit("REQUEST_EXIT", {
+              version: "v1",
+            });
+
             return;
           }
 
           if (!executionAck.success) {
             logger.error("execution unsuccessful", { attemptId: socket.data.attemptId });
+
+            socket.emit("REQUEST_EXIT", {
+              version: "v1",
+            });
+
             return;
           }
 

apps/coordinator/tsconfig.json

@@ -9,10 +9,10 @@
     "strict": true,
     "skipLibCheck": true,
     "paths": {
-      "@trigger.dev/core/v3": ["../core/src/v3"],
-      "@trigger.dev/core/v3/*": ["../core/src/v3/*"],
-      "@trigger.dev/core-apps": ["../core-apps/src"],
-      "@trigger.dev/core-apps/*": ["../core-apps/src/*"]
+      "@trigger.dev/core/v3": ["../../packages/core/src/v3"],
+      "@trigger.dev/core/v3/*": ["../../packages/core/src/v3/*"],
+      "@trigger.dev/core-apps": ["../../packages/core-apps/src"],
+      "@trigger.dev/core-apps/*": ["../../packages/core-apps/src/*"]
     }
   }
 }

apps/docker-provider/package.json

@@ -4,7 +4,6 @@
   "version": "0.0.1",
   "description": "",
   "main": "dist/index.cjs",
-  "type": "module",
   "scripts": {
     "build": "npm run build:bundle",
     "build:bundle": "esbuild src/index.ts --bundle --outfile=dist/index.mjs --platform=node --format=esm --target=esnext --banner:js=\"const require = createRequire(import.meta.url);\"",
@@ -29,4 +28,4 @@
     "tsx": "^4.7.0",
     "typescript": "^5.3.3"
   }
-}
+}

apps/docker-provider/tsconfig.json

@@ -7,10 +7,10 @@
     "strict": true,
     "skipLibCheck": true,
     "paths": {
-      "@trigger.dev/core/v3": ["../core/src/v3"],
-      "@trigger.dev/core/v3/*": ["../core/src/v3/*"],
-      "@trigger.dev/core-apps": ["../core-apps/src"],
-      "@trigger.dev/core-apps/*": ["../core-apps/src/*"]
+      "@trigger.dev/core/v3": ["../../packages/core/src/v3"],
+      "@trigger.dev/core/v3/*": ["../../packages/core/src/v3/*"],
+      "@trigger.dev/core-apps": ["../../packages/core-apps/src"],
+      "@trigger.dev/core-apps/*": ["../../packages/core-apps/src/*"]
     }
   }
 }

apps/webapp/app/routes/api.v2.runs.$runParam.cancel.ts

@@ -0,0 +1,52 @@
+import type { ActionFunctionArgs } from "@remix-run/server-runtime";
+import { json } from "@remix-run/server-runtime";
+import { z } from "zod";
+import { prisma } from "~/db.server";
+import { authenticateApiRequest } from "~/services/apiAuth.server";
+import { CancelTaskRunService } from "~/v3/services/cancelTaskRun.server";
+
+const ParamsSchema = z.object({
+  runParam: z.string(),
+});
+
+export async function action({ request, params }: ActionFunctionArgs) {
+  // Ensure this is a POST request
+  if (request.method.toUpperCase() !== "POST") {
+    return { status: 405, body: "Method Not Allowed" };
+  }
+
+  // Authenticate the request
+  const authenticationResult = await authenticateApiRequest(request);
+
+  if (!authenticationResult) {
+    return json({ error: "Invalid or Missing API Key" }, { status: 401 });
+  }
+
+  const parsed = ParamsSchema.safeParse(params);
+
+  if (!parsed.success) {
+    return json({ error: "Invalid or Missing runId" }, { status: 400 });
+  }
+
+  const { runParam } = parsed.data;
+
+  const taskRun = await prisma.taskRun.findUnique({
+    where: {
+      friendlyId: runParam,
+    },
+  });
+
+  if (!taskRun) {
+    return json({ error: "Run not found" }, { status: 404 });
+  }
+
+  const service = new CancelTaskRunService();
+
+  try {
+    await service.call(taskRun);
+  } catch (error) {
+    return json({ error: "Internal Server Error" }, { status: 500 });
+  }
+
+  return json({ message: "Run cancelled" }, { status: 200 });
+}

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

@@ -10,6 +10,7 @@ import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { logger } from "~/services/logger.server";
 import { DevQueueConsumer } from "./marqs/devQueueConsumer.server";
 import type { WebSocket, MessageEvent, CloseEvent, ErrorEvent } from "ws";
+import { env } from "~/env.server";
 
 export class AuthenticatedSocketConnection {
   public id: string;
@@ -26,6 +27,10 @@ export class AuthenticatedSocketConnection {
       schema: serverWebsocketMessages,
       sender: async (message) => {
         return new Promise((resolve, reject) => {
+          if (!ws.OPEN) {
+            return reject(new Error("Websocket is not open"));
+          }
+
           ws.send(JSON.stringify(message), {}, (err) => {
             if (err) {
               reject(err);
@@ -84,6 +89,8 @@ export class AuthenticatedSocketConnection {
   }
 
   async #handleClose(ev: CloseEvent) {
+    logger.debug("[AuthenticatedSocketConnection] Websocket closed", { ev });
+
     await this._consumer.stop();
 
     this.onClose.post(ev);

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

@@ -44,6 +44,7 @@ export type TraceAttributes = Partial<
     CreatableEvent,
     | "attemptId"
     | "isError"
+    | "isCancelled"
     | "runId"
     | "runIsTest"
     | "output"

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

@@ -0,0 +1,38 @@
+import { z } from "zod";
+import { singleton } from "~/utils/singleton";
+import { ZodPubSub, ZodSubscriber } from "../utils/zodPubSub.server";
+import { env } from "~/env.server";
+
+const messageCatalog = {
+  CANCEL_ATTEMPT: z.object({
+    version: z.literal("v1").default("v1"),
+    backgroundWorkerId: z.string(),
+    attemptId: z.string(),
+    taskRunId: z.string(),
+  }),
+};
+
+export type DevSubscriber = ZodSubscriber<typeof messageCatalog>;
+
+export const devPubSub = singleton("devPubSub", initializeDevPubSub);
+
+function initializeDevPubSub() {
+  return new ZodPubSub({
+    redis: {
+      port: env.REDIS_PORT,
+      host: env.REDIS_HOST,
+      username: env.REDIS_USERNAME,
+      password: env.REDIS_PASSWORD,
+      enableAutoPipelining: true,
+      ...(env.REDIS_TLS_DISABLED === "true" ? {} : { tls: {} }),
+    },
+    schema: {
+      CANCEL_ATTEMPT: z.object({
+        version: z.literal("v1").default("v1"),
+        backgroundWorkerId: z.string(),
+        attemptId: z.string(),
+        taskRunId: z.string(),
+      }),
+    },
+  });
+}

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

@@ -17,6 +17,7 @@ import { marqs } from "../marqs.server";
 import { CancelAttemptService } from "../services/cancelAttempt.server";
 import { CompleteAttemptService } from "../services/completeAttempt.server";
 import { attributesFromAuthenticatedEnv } from "../tracer.server";
+import { DevSubscriber, devPubSub } from "./devPubSub.server";
 
 const tracer = trace.getTracer("devQueueConsumer");
 
@@ -36,9 +37,11 @@ export type DevQueueConsumerOptions = {
 
 export class DevQueueConsumer {
   private _backgroundWorkers: Map<string, BackgroundWorkerWithTasks> = new Map();
+  private _backgroundWorkerSubscriber: Map<string, DevSubscriber> = new Map();
   private _deprecatedWorkers: Map<string, BackgroundWorkerWithTasks> = new Map();
   private _enabled = false;
-  private _options: Required<DevQueueConsumerOptions>;
+  private _maximumItemsPerTrace: number;
+  private _traceTimeoutSeconds: number;
   private _perTraceCountdown: number | undefined;
   private _lastNewTrace: Date | undefined;
   private _currentSpanContext: Context | undefined;
@@ -51,12 +54,10 @@ export class DevQueueConsumer {
   constructor(
     public env: AuthenticatedEnvironment,
     private _sender: ZodMessageSender<typeof serverWebsocketMessages>,
-    options: DevQueueConsumerOptions = {}
+    private _options: DevQueueConsumerOptions = {}
   ) {
-    this._options = {
-      maximumItemsPerTrace: options.maximumItemsPerTrace ?? 1_000, // 1k items per trace
-      traceTimeoutSeconds: options.traceTimeoutSeconds ?? 60, // 60 seconds
-    };
+    this._traceTimeoutSeconds = _options.traceTimeoutSeconds ?? 60;
+    this._maximumItemsPerTrace = _options.maximumItemsPerTrace ?? 1_000;
   }
 
   // This method is called when a background worker is deprecated and will no longer be used unless a run is locked to it
@@ -87,6 +88,21 @@ export class DevQueueConsumer {
 
     logger.debug("Registered background worker", { backgroundWorker: backgroundWorker.id });
 
+    const subscriber = await devPubSub.subscribe(`backgroundWorker:${backgroundWorker.id}:*`);
+
+    subscriber.on("CANCEL_ATTEMPT", async (message) => {
+      await this._sender.send("BACKGROUND_WORKER_MESSAGE", {
+        backgroundWorkerId: backgroundWorker.friendlyId,
+        data: {
+          type: "CANCEL_ATTEMPT",
+          taskAttemptId: message.attemptId,
+          taskRunId: message.taskRunId,
+        },
+      });
+    });
+
+    this._backgroundWorkerSubscriber.set(backgroundWorker.id, subscriber);
+
     // Start reading from the queue if we haven't already
     this.#enable();
   }
@@ -133,6 +149,16 @@ export class DevQueueConsumer {
 
     // We need to cancel all the in progress task run attempts and ack the messages so they will stop processing
     await this.#cancelInProgressAttempts(reason);
+
+    // We need to unsubscribe from the background worker channels
+    for (const [id, subscriber] of this._backgroundWorkerSubscriber) {
+      logger.debug("Unsubscribing from background worker channel", { id });
+
+      await subscriber.stopListening();
+      this._backgroundWorkerSubscriber.delete(id);
+
+      logger.debug("Unsubscribed from background worker channel", { id });
+    }
   }
 
   async #cancelInProgressAttempts(reason: string) {
@@ -144,6 +170,10 @@ export class DevQueueConsumer {
 
     this._inProgressAttempts.clear();
 
+    logger.debug("Cancelling in progress attempts", {
+      attempts: Array.from(inProgressAttempts.keys()),
+    });
+
     for (const [attemptId, messageId] of inProgressAttempts) {
       await this.#cancelInProgressAttempt(attemptId, messageId, service, cancelledAt, reason);
     }
@@ -156,6 +186,8 @@ export class DevQueueConsumer {
     cancelledAt: Date,
     reason: string
   ) {
+    logger.debug("Cancelling in progress attempt", { attemptId, messageId });
+
     try {
       await cancelAttemptService.call(attemptId, messageId, cancelledAt, reason, this.env);
     } catch (e) {
@@ -189,7 +221,7 @@ export class DevQueueConsumer {
     // Check if the trace has expired
     if (
       this._perTraceCountdown === 0 ||
-      Date.now() - this._lastNewTrace!.getTime() > this._options.traceTimeoutSeconds * 1000 ||
+      Date.now() - this._lastNewTrace!.getTime() > this._traceTimeoutSeconds * 1000 ||
       this._currentSpanContext === undefined ||
       this._endSpanInNextIteration
     ) {
@@ -366,6 +398,7 @@ export class DevQueueConsumer {
         backgroundWorkerTaskId: backgroundTask.id,
         status: "EXECUTING" as const,
         queueId: queue.id,
+        runtimeEnvironmentId: this.env.id,
       },
     });
 
@@ -442,6 +475,11 @@ export class DevQueueConsumer {
         },
       });
 
+      logger.debug("Saving the in progress attempt", {
+        taskRunAttempt: taskRunAttempt.id,
+        messageId: message.messageId,
+      });
+
       this._inProgressAttempts.set(taskRunAttempt.friendlyId, message.messageId);
     } catch (e) {
       if (e instanceof Error) {