prep for snapshots since

Author: nicktrn

packages/cli-v3/src/entryPoints/managed/controller.ts

@@ -18,7 +18,7 @@ type ManagedRunControllerOptions = {
   env: EnvObject;
 };
 
-type SupervisorSocket = Socket<WorkloadServerToClientEvents, WorkloadClientToServerEvents>;
+export type SupervisorSocket = Socket<WorkloadServerToClientEvents, WorkloadClientToServerEvents>;
 
 export class ManagedRunController {
   private readonly env: RunnerEnv;
@@ -31,6 +31,9 @@ export class ManagedRunController {
   private warmStartCount = 0;
   private restoreCount = 0;
 
+  private notificationCount = 0;
+  private lastNotificationAt: Date | null = null;
+
   private currentExecution: RunExecution | null = null;
 
   constructor(opts: ManagedRunControllerOptions) {
@@ -91,6 +94,8 @@ export class ManagedRunController {
     return {
       warmStartCount: this.warmStartCount,
       restoreCount: this.restoreCount,
+      notificationCount: this.notificationCount,
+      lastNotificationAt: this.lastNotificationAt,
     };
   }
 
@@ -188,12 +193,16 @@ export class ManagedRunController {
         this.currentExecution = null;
       }
 
+      // Remove all run notification listeners just to be safe
+      this.socket.removeAllListeners("run:notify");
+
       if (!this.currentExecution || !this.currentExecution.canExecute) {
         this.currentExecution = new RunExecution({
           workerManifest: this.workerManifest,
           env: this.env,
           httpClient: this.httpClient,
           logger: this.logger,
+          supervisorSocket: this.socket,
         });
       }
 
@@ -224,8 +233,8 @@ export class ManagedRunController {
 
     const metrics = this.currentExecution?.metrics;
 
-    if (metrics?.restoreCount) {
-      this.restoreCount += metrics.restoreCount;
+    if (metrics?.execution?.restoreCount) {
+      this.restoreCount += metrics.execution.restoreCount;
     }
 
     this.lockedRunExecution = null;
@@ -288,6 +297,7 @@ export class ManagedRunController {
           env: this.env,
           httpClient: this.httpClient,
           logger: this.logger,
+          supervisorSocket: this.socket,
         }).prepareForExecution({
           taskRunEnv: previousTaskRunEnv,
         });
@@ -392,116 +402,12 @@ export class ManagedRunController {
   createSupervisorSocket(): SupervisorSocket {
     const wsUrl = new URL("/workload", this.workerApiUrl);
 
-    const socket = io(wsUrl.href, {
+    const socket: SupervisorSocket = io(wsUrl.href, {
       transports: ["websocket"],
       extraHeaders: {
         [WORKLOAD_HEADERS.DEPLOYMENT_ID]: this.env.TRIGGER_DEPLOYMENT_ID,
         [WORKLOAD_HEADERS.RUNNER_ID]: this.env.TRIGGER_RUNNER_ID,
       },
-    }) satisfies SupervisorSocket;
-
-    socket.on("run:notify", async ({ version, run }) => {
-      // Generate a unique ID for the notification
-      const notificationId = Math.random().toString(36).substring(2, 15);
-
-      // Use this to track the notification incl. any processing
-      const notification = {
-        id: notificationId,
-        runId: run.friendlyId,
-        version,
-      };
-
-      // Lock this to the current run and snapshot IDs
-      const controller = {
-        runFriendlyId: this.runFriendlyId,
-        snapshotFriendlyId: this.snapshotFriendlyId,
-      };
-
-      this.sendDebugLog({
-        runId: run.friendlyId,
-        message: "run:notify received by runner",
-        properties: {
-          notification,
-          controller,
-        },
-      });
-
-      if (!controller.runFriendlyId) {
-        this.sendDebugLog({
-          runId: run.friendlyId,
-          message: "run:notify: ignoring notification, no run ID",
-          properties: {
-            notification,
-            controller,
-          },
-        });
-        return;
-      }
-
-      if (!controller.snapshotFriendlyId) {
-        this.sendDebugLog({
-          runId: run.friendlyId,
-          message: "run:notify: ignoring notification, no snapshot ID",
-          properties: { notification, controller },
-        });
-      }
-
-      if (run.friendlyId !== controller.runFriendlyId) {
-        this.sendDebugLog({
-          runId: run.friendlyId,
-          message: "run:notify: ignoring notification for different run",
-          properties: {
-            notification,
-            controller,
-          },
-        });
-        return;
-      }
-
-      const latestSnapshot = await this.httpClient.getRunExecutionData(controller.runFriendlyId);
-
-      if (!latestSnapshot.success) {
-        this.sendDebugLog({
-          runId: run.friendlyId,
-          message: "run:notify: failed to get latest snapshot data",
-          properties: {
-            notification,
-            controller,
-            error: latestSnapshot.error,
-          },
-        });
-        return;
-      }
-
-      const runExecutionData = latestSnapshot.data.execution;
-
-      if (!this.currentExecution) {
-        this.sendDebugLog({
-          runId: run.friendlyId,
-          message: "run:notify: no current execution",
-          properties: {
-            notification,
-            controller,
-          },
-        });
-        return;
-      }
-
-      const [error] = await tryCatch(
-        this.currentExecution.enqueueSnapshotChangeAndWait(runExecutionData)
-      );
-
-      if (error) {
-        this.sendDebugLog({
-          runId: run.friendlyId,
-          message: "run:notify: unexpected error",
-          properties: {
-            notification,
-            controller,
-            error: error.message,
-          },
-        });
-      }
     });
 
     socket.on("connect", () => {

packages/cli-v3/src/entryPoints/managed/execution.ts

@@ -20,6 +20,8 @@ import { assertExhaustive, tryCatch } from "@trigger.dev/core/utils";
 import { MetadataClient } from "./overrides.js";
 import { randomBytes } from "node:crypto";
 import { SnapshotManager, SnapshotState } from "./snapshot.js";
+import type { SupervisorSocket } from "./controller.js";
+import { RunNotifier } from "./notifier.js";
 
 class ExecutionAbortError extends Error {
   constructor(message: string) {
@@ -33,6 +35,7 @@ type RunExecutionOptions = {
   env: RunnerEnv;
   httpClient: WorkloadHttpClient;
   logger: RunLogger;
+  supervisorSocket: SupervisorSocket;
 };
 
 type RunExecutionPrepareOptions = {
@@ -71,12 +74,16 @@ export class RunExecution {
   private isShuttingDown = false;
   private shutdownReason?: string;
 
+  private supervisorSocket: SupervisorSocket;
+  private notifier?: RunNotifier;
+
   constructor(opts: RunExecutionOptions) {
     this.id = randomBytes(4).toString("hex");
     this.workerManifest = opts.workerManifest;
     this.env = opts.env;
     this.httpClient = opts.httpClient;
     this.logger = opts.logger;
+    this.supervisorSocket = opts.supervisorSocket;
 
     this.restoreCount = 0;
     this.executionAbortController = new AbortController();
@@ -439,10 +446,16 @@ export class RunExecution {
     this.snapshotPoller = new RunExecutionSnapshotPoller({
       runFriendlyId: this.runFriendlyId,
       snapshotFriendlyId: this.snapshotManager.snapshotId,
-      httpClient: this.httpClient,
       logger: this.logger,
       snapshotPollIntervalSeconds: this.env.TRIGGER_SNAPSHOT_POLL_INTERVAL_SECONDS,
-      handleSnapshotChange: this.enqueueSnapshotChangeAndWait.bind(this),
+      onPoll: this.fetchAndEnqueueSnapshotChange.bind(this),
+    }).start();
+
+    this.notifier = new RunNotifier({
+      runFriendlyId: this.runFriendlyId,
+      supervisorSocket: this.supervisorSocket,
+      onNotify: this.fetchAndEnqueueSnapshotChange.bind(this),
+      logger: this.logger,
     }).start();
 
     const [startError, start] = await tryCatch(
@@ -947,7 +960,12 @@ export class RunExecution {
 
   public get metrics() {
     return {
-      restoreCount: this.restoreCount,
+      execution: {
+        restoreCount: this.restoreCount,
+        lastHeartbeat: this.lastHeartbeat,
+      },
+      poller: this.snapshotPoller?.metrics,
+      notifier: this.notifier?.metrics,
     };
   }
 
@@ -981,6 +999,8 @@ export class RunExecution {
     this.snapshotPoller?.stop();
     this.snapshotManager?.dispose();
 
+    this.notifier?.stop();
+
     this.taskRunProcess?.unsafeDetachEvtHandlers();
   }
 
@@ -1056,4 +1076,40 @@ export class RunExecution {
 
     this.sendDebugLog("suspending, any day now 🚬", { suspendableSnapshot });
   }
+
+  /**
+   * Fetches the latest execution data and enqueues snapshot changes. Used by both poller and notification handlers.
+   * @param source string - where this call originated (e.g. 'poller', 'notification')
+   */
+  public async fetchAndEnqueueSnapshotChange(source: string): Promise<void> {
+    if (!this.runFriendlyId) {
+      this.sendDebugLog(`fetchAndEnqueueSnapshotChange: missing runFriendlyId`, { source });
+      return;
+    }
+
+    const latestSnapshot = await this.httpClient.getRunExecutionData(this.runFriendlyId);
+
+    if (!latestSnapshot.success) {
+      this.sendDebugLog(`fetchAndEnqueueSnapshotChange: failed to get latest snapshot data`, {
+        source,
+        error: latestSnapshot.error,
+      });
+      return;
+    }
+
+    const [error] = await tryCatch(
+      this.enqueueSnapshotChangeAndWait(latestSnapshot.data.execution)
+    );
+
+    if (error) {
+      this.sendDebugLog(
+        `fetchAndEnqueueSnapshotChange: failed to enqueue and process snapshot change`,
+        {
+          source,
+          error: error.message,
+        }
+      );
+      return;
+    }
+  }
 }

packages/cli-v3/src/entryPoints/managed/notifier.ts

@@ -0,0 +1,89 @@
+import type { SupervisorSocket } from "./controller.js";
+import type { RunLogger, SendDebugLogOptions } from "./logger.js";
+
+type OnNotify = (source: string) => Promise<void>;
+
+type RunNotifierOptions = {
+  runFriendlyId: string;
+  supervisorSocket: SupervisorSocket;
+  onNotify: OnNotify;
+  logger: RunLogger;
+};
+
+export class RunNotifier {
+  private runFriendlyId: string;
+  private socket: SupervisorSocket;
+  private onNotify: OnNotify;
+  private logger: RunLogger;
+
+  private lastNotificationAt: Date | null = null;
+  private notificationCount = 0;
+
+  private lastInvalidNotificationAt: Date | null = null;
+  private invalidNotificationCount = 0;
+
+  constructor(opts: RunNotifierOptions) {
+    this.runFriendlyId = opts.runFriendlyId;
+    this.socket = opts.supervisorSocket;
+    this.onNotify = opts.onNotify;
+    this.logger = opts.logger;
+  }
+
+  start(): RunNotifier {
+    this.sendDebugLog("start");
+
+    this.socket.on("run:notify", async ({ version, run }) => {
+      // Generate a unique ID for the notification
+      const notificationId = Math.random().toString(36).substring(2, 15);
+
+      // Use this to track the notification incl. any processing
+      const notification = {
+        id: notificationId,
+        runId: run.friendlyId,
+        version,
+      };
+
+      if (run.friendlyId !== this.runFriendlyId) {
+        this.sendDebugLog("run:notify received invalid notification", { notification });
+
+        this.invalidNotificationCount++;
+        this.lastInvalidNotificationAt = new Date();
+
+        return;
+      }
+
+      this.sendDebugLog("run:notify received by runner", { notification });
+
+      this.notificationCount++;
+      this.lastNotificationAt = new Date();
+
+      await this.onNotify(`notifier:${notificationId}`);
+    });
+
+    return this;
+  }
+
+  stop() {
+    this.sendDebugLog("stop");
+    this.socket.removeAllListeners("run:notify");
+  }
+
+  get metrics() {
+    return {
+      lastNotificationAt: this.lastNotificationAt,
+      notificationCount: this.notificationCount,
+      lastInvalidNotificationAt: this.lastInvalidNotificationAt,
+      invalidNotificationCount: this.invalidNotificationCount,
+    };
+  }
+
+  private sendDebugLog(message: string, properties?: SendDebugLogOptions["properties"]) {
+    this.logger?.sendDebugLog({
+      message: `[notifier] ${message}`,
+      properties: {
+        ...properties,
+        ...this.metrics,
+      },
+    });
+  }
+}