Merge remote-tracking branch 'origin/main' into feat/supervisor

Author: nicktrn

apps/webapp/app/env.server.ts

@@ -571,6 +571,9 @@ const EnvironmentSchema = z.object({
   /** The max number of runs per API call that we'll dequeue in DEV */
   DEV_DEQUEUE_MAX_RUNS_PER_PULL: z.coerce.number().int().default(10),
 
+  /** The maximum concurrent local run processes executing at once in dev */
+  DEV_MAX_CONCURRENT_RUNS: z.coerce.number().int().default(25),
+
   LEGACY_RUN_ENGINE_WORKER_ENABLED: z.string().default(process.env.WORKER_ENABLED ?? "true"),
   LEGACY_RUN_ENGINE_WORKER_CONCURRENCY_WORKERS: z.coerce.number().int().default(2),
   LEGACY_RUN_ENGINE_WORKER_CONCURRENCY_TASKS_PER_WORKER: z.coerce.number().int().default(1),

apps/webapp/app/routes/engine.v1.dev.config.ts

@@ -20,6 +20,7 @@ export const loader = createLoaderApiRoute(
         environmentId: authentication.environment.id,
         dequeueIntervalWithRun: env.DEV_DEQUEUE_INTERVAL_WITH_RUN,
         dequeueIntervalWithoutRun: env.DEV_DEQUEUE_INTERVAL_WITHOUT_RUN,
+        maxConcurrentRuns: env.DEV_MAX_CONCURRENT_RUNS,
       });
     } catch (error) {
       logger.error("Failed to get dev settings", {

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

@@ -925,6 +925,7 @@ export class RunEngine {
 
             return {
               version: "1" as const,
+              dequeuedAt: new Date(),
               snapshot: {
                 id: newSnapshot.id,
                 friendlyId: newSnapshot.friendlyId,
@@ -2442,8 +2443,11 @@ export class RunEngine {
       },
     });
 
-    //extending is the same as creating a new heartbeat
-    await this.#setHeartbeatDeadline({ runId, snapshotId, status: latestSnapshot.executionStatus });
+    //extending the heartbeat
+    const intervalMs = this.#getHeartbeatIntervalMs(latestSnapshot.executionStatus);
+    if (intervalMs !== null) {
+      await this.worker.reschedule(`heartbeatSnapshot.${runId}`, new Date(Date.now() + intervalMs));
+    }
 
     return executionResultFromSnapshot(latestSnapshot);
   }
@@ -3647,11 +3651,15 @@ export class RunEngine {
 
     if (!error) {
       //set heartbeat (if relevant)
-      await this.#setHeartbeatDeadline({
-        status: newSnapshot.executionStatus,
-        runId: run.id,
-        snapshotId: newSnapshot.id,
-      });
+      const intervalMs = this.#getHeartbeatIntervalMs(newSnapshot.executionStatus);
+      if (intervalMs !== null) {
+        await this.worker.enqueue({
+          id: `heartbeatSnapshot.${run.id}`,
+          job: "heartbeatSnapshot",
+          payload: { snapshotId: newSnapshot.id, runId: run.id },
+          availableAt: new Date(Date.now() + intervalMs),
+        });
+      }
     }
 
     this.eventBus.emit("executionSnapshotCreated", {
@@ -3695,29 +3703,6 @@ export class RunEngine {
   //#endregion
 
   //#region Heartbeat
-  async #setHeartbeatDeadline({
-    runId,
-    snapshotId,
-    status,
-  }: {
-    runId: string;
-    snapshotId: string;
-    status: TaskRunExecutionStatus;
-  }) {
-    const intervalMs = this.#getHeartbeatIntervalMs(status);
-
-    if (intervalMs === null) {
-      return;
-    }
-
-    await this.worker.enqueue({
-      id: `heartbeatSnapshot.${runId}`,
-      job: "heartbeatSnapshot",
-      payload: { snapshotId, runId },
-      availableAt: new Date(Date.now() + intervalMs),
-    });
-  }
-
   async #handleStalledSnapshot({
     runId,
     snapshotId,

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

@@ -490,4 +490,123 @@ describe("RunEngine heartbeats", () => {
       await engine.quit();
     }
   });
+
+  containerTest(
+    "Heartbeat keeps run alive",
+    { timeout: 15_000 },
+    async ({ prisma, redisOptions }) => {
+      const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
+
+      const executingTimeout = 100;
+
+      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: {
+          EXECUTING: executingTimeout,
+        },
+        tracer: trace.getTracer("test", "0.0.0"),
+      });
+
+      try {
+        const taskIdentifier = "test-task";
+
+        //create background worker
+        const backgroundWorker = await setupBackgroundWorker(
+          prisma,
+          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",
+            queueName: "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
+        const attempt = await engine.startRunAttempt({
+          runId: dequeued[0].run.id,
+          snapshotId: dequeued[0].snapshot.id,
+        });
+
+        //should be executing
+        const executionData = await engine.getRunExecutionData({ runId: run.id });
+        assertNonNullable(executionData);
+        expect(executionData.snapshot.executionStatus).toBe("EXECUTING");
+        expect(executionData.run.status).toBe("EXECUTING");
+
+        // Send heartbeats every 50ms (half the timeout)
+        for (let i = 0; i < 6; i++) {
+          await setTimeout(50);
+          await engine.heartbeatRun({
+            runId: run.id,
+            snapshotId: attempt.snapshot.id,
+          });
+        }
+
+        // After 300ms (3x the timeout) the run should still be executing
+        // because we've been sending heartbeats
+        const executionData2 = await engine.getRunExecutionData({ runId: run.id });
+        assertNonNullable(executionData2);
+        expect(executionData2.snapshot.executionStatus).toBe("EXECUTING");
+        expect(executionData2.run.status).toBe("EXECUTING");
+
+        // Stop sending heartbeats and wait for timeout
+        await setTimeout(executingTimeout * 2);
+
+        // Now it should have timed out and be queued
+        const executionData3 = await engine.getRunExecutionData({ runId: run.id });
+        assertNonNullable(executionData3);
+        expect(executionData3.snapshot.executionStatus).toBe("QUEUED");
+      } finally {
+        await engine.quit();
+      }
+    }
+  );
 });

packages/cli-v3/package.json

@@ -110,6 +110,7 @@
     "nypm": "^0.3.9",
     "object-hash": "^3.0.0",
     "open": "^10.0.3",
+    "p-limit": "^6.2.0",
     "p-retry": "^6.1.0",
     "partysocket": "^1.0.2",
     "pkg-types": "^1.1.3",

packages/cli-v3/src/commands/dev.ts

@@ -19,6 +19,7 @@ const DevCommandOptions = CommonCommandOptions.extend({
   skipUpdateCheck: z.boolean().default(false),
   envFile: z.string().optional(),
   keepTmpFiles: z.boolean().default(false),
+  maxConcurrentRuns: z.coerce.number().optional(),
 });
 
 export type DevCommandOptions = z.infer<typeof DevCommandOptions>;
@@ -37,6 +38,10 @@ export function configureDevCommand(program: Command) {
         "--env-file <env file>",
         "Path to the .env file to use for the dev session. Defaults to .env in the project directory."
       )
+      .option(
+        "--max-concurrent-runs <max concurrent runs>",
+        "The maximum number of concurrent runs to allow in the dev session"
+      )
       .option("--debug-otel", "Enable OpenTelemetry debugging")
       .option("--skip-update-check", "Skip checking for @trigger.dev package updates")
       .option(

packages/cli-v3/src/dev/devSupervisor.ts

@@ -24,6 +24,7 @@ import {
   WorkerClientToServerEvents,
   WorkerServerToClientEvents,
 } from "@trigger.dev/core/v3/workers";
+import pLimit from "p-limit";
 
 export type WorkerRuntimeOptions = {
   name: string | undefined;
@@ -65,6 +66,8 @@ class DevSupervisor implements WorkerRuntime {
 
   private socketConnections = new Set<string>();
 
+  private runLimiter?: ReturnType<typeof pLimit>;
+
   constructor(public readonly options: WorkerRuntimeOptions) {}
 
   async init(): Promise<void> {
@@ -81,6 +84,15 @@ class DevSupervisor implements WorkerRuntime {
     logger.debug("[DevSupervisor] Got dev settings", { settings: settings.data });
     this.config = settings.data;
 
+    const maxConcurrentRuns = Math.min(
+      this.config.maxConcurrentRuns,
+      this.options.args.maxConcurrentRuns ?? this.config.maxConcurrentRuns
+    );
+
+    logger.debug("[DevSupervisor] Using maxConcurrentRuns", { maxConcurrentRuns });
+
+    this.runLimiter = pLimit(maxConcurrentRuns);
+
     this.#createSocket();
 
     //start an SSE connection for presence
@@ -178,6 +190,14 @@ class DevSupervisor implements WorkerRuntime {
       return;
     }
 
+    if (
+      this.runLimiter &&
+      this.runLimiter.activeCount + this.runLimiter.pendingCount > this.runLimiter.concurrency
+    ) {
+      logger.debug(`[DevSupervisor] dequeueRuns. Run limit reached, trying again later`);
+      setTimeout(() => this.#dequeueRuns(), this.config.dequeueIntervalWithoutRun);
+    }
+
     //get relevant versions
     //ignore deprecated and the latest worker
     const oldWorkerIds = this.#getActiveOldWorkers();
@@ -287,10 +307,16 @@ class DevSupervisor implements WorkerRuntime {
 
         this.runControllers.set(message.run.friendlyId, runController);
 
-        //don't await for run completion, we want to dequeue more runs
-        runController.start(message).then(() => {
-          logger.debug("[DevSupervisor] Run started", { runId: message.run.friendlyId });
-        });
+        if (this.runLimiter) {
+          this.runLimiter(() => runController.start(message)).then(() => {
+            logger.debug("[DevSupervisor] Run started", { runId: message.run.friendlyId });
+          });
+        } else {
+          //don't await for run completion, we want to dequeue more runs
+          runController.start(message).then(() => {
+            logger.debug("[DevSupervisor] Run started", { runId: message.run.friendlyId });
+          });
+        }
       }
 
       setTimeout(() => this.#dequeueRuns(), this.config.dequeueIntervalWithRun);