Delayed run system

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

@@ -56,3 +56,27 @@ export function runStatusFromError(error: TaskRunError): TaskRunStatus {
       assertExhaustive(error.code);
   }
 }
+
+export class ServiceValidationError extends Error {
+  constructor(
+    message: string,
+    public status?: number
+  ) {
+    super(message);
+    this.name = "ServiceValidationError";
+  }
+}
+
+export class NotImplementedError extends Error {
+  constructor(message: string) {
+    console.error("This isn't implemented", { message });
+    super(message);
+  }
+}
+
+export class RunDuplicateIdempotencyKeyError extends Error {
+  constructor(message: string) {
+    super(message);
+    this.name = "RunDuplicateIdempotencyKeyError";
+  }
+}

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

@@ -52,6 +52,12 @@ import { SystemResources } from "./systems/systems.js";
 import { WaitpointSystem } from "./systems/waitpointSystem.js";
 import { EngineWorker, HeartbeatTimeouts, RunEngineOptions, TriggerParams } from "./types.js";
 import { workerCatalog } from "./workerCatalog.js";
+import {
+  NotImplementedError,
+  RunDuplicateIdempotencyKeyError,
+  ServiceValidationError,
+} from "./errors.js";
+import { DelayedRunSystem } from "./systems/delayedRunSystem.js";
 
 export class RunEngine {
   private runLockRedis: Redis;
@@ -75,6 +81,7 @@ export class RunEngine {
   batchSystem: BatchSystem;
   enqueueSystem: EnqueueSystem;
   checkpointSystem: CheckpointSystem;
+  delayedRunSystem: DelayedRunSystem;
 
   constructor(private readonly options: RunEngineOptions) {
     this.prisma = options.prisma;
@@ -159,7 +166,7 @@ export class RunEngine {
           });
         },
         enqueueDelayedRun: async ({ payload }) => {
-          await this.#enqueueDelayedRun({ runId: payload.runId });
+          await this.delayedRunSystem.enqueueDelayedRun({ runId: payload.runId });
         },
       },
     }).start();
@@ -248,6 +255,11 @@ export class RunEngine {
       executionSnapshotSystem: this.executionSnapshotSystem,
     });
 
+    this.delayedRunSystem = new DelayedRunSystem({
+      resources,
+      enqueueSystem: this.enqueueSystem,
+    });
+
     this.waitpointSystem = new WaitpointSystem({
       resources,
       executionSnapshotSystem: this.executionSnapshotSystem,
@@ -789,47 +801,11 @@ export class RunEngine {
     delayUntil: Date;
     tx?: PrismaClientOrTransaction;
   }): Promise<TaskRun> {
-    const prisma = tx ?? this.prisma;
-    return startSpan(
-      this.tracer,
-      "rescheduleRun",
-      async () => {
-        return await this.runLock.lock([runId], 5_000, async () => {
-          const snapshot = await getLatestExecutionSnapshot(prisma, runId);
-
-          //if the run isn't just created then we can't reschedule it
-          if (snapshot.executionStatus !== "RUN_CREATED") {
-            throw new ServiceValidationError("Cannot reschedule a run that is not delayed");
-          }
-
-          const updatedRun = await prisma.taskRun.update({
-            where: {
-              id: runId,
-            },
-            data: {
-              delayUntil: delayUntil,
-              executionSnapshots: {
-                create: {
-                  engine: "V2",
-                  executionStatus: "RUN_CREATED",
-                  description: "Delayed run was rescheduled to a future date",
-                  runStatus: "EXPIRED",
-                  environmentId: snapshot.environmentId,
-                  environmentType: snapshot.environmentType,
-                },
-              },
-            },
-          });
-
-          await this.worker.reschedule(`enqueueDelayedRun:${updatedRun.id}`, delayUntil);
-
-          return updatedRun;
-        });
-      },
-      {
-        attributes: { runId },
-      }
-    );
+    return this.delayedRunSystem.rescheduleDelayedRun({
+      runId,
+      delayUntil,
+      tx,
+    });
   }
 
   async lengthOfEnvQueue(environment: MinimalAuthenticatedEnvironment): Promise<number> {
@@ -1365,53 +1341,6 @@ export class RunEngine {
     });
   }
 
-  async #enqueueDelayedRun({ runId }: { runId: string }) {
-    const run = await this.prisma.taskRun.findFirst({
-      where: { id: runId },
-      include: {
-        runtimeEnvironment: {
-          include: {
-            project: true,
-            organization: true,
-          },
-        },
-      },
-    });
-
-    if (!run) {
-      throw new Error(`#enqueueDelayedRun: run not found: ${runId}`);
-    }
-
-    // Now we need to enqueue the run into the RunQueue
-    await this.enqueueSystem.enqueueRun({
-      run,
-      env: run.runtimeEnvironment,
-      timestamp: run.createdAt.getTime() - run.priorityMs,
-      batchId: run.batchId ?? undefined,
-    });
-
-    await this.prisma.taskRun.update({
-      where: { id: runId },
-      data: {
-        status: "PENDING",
-        queuedAt: new Date(),
-      },
-    });
-
-    if (run.ttl) {
-      const expireAt = parseNaturalLanguageDuration(run.ttl);
-
-      if (expireAt) {
-        await this.worker.enqueue({
-          id: `expireRun:${runId}`,
-          job: "expireRun",
-          payload: { runId },
-          availableAt: expireAt,
-        });
-      }
-    }
-  }
-
   async #queueRunsWaitingForWorker({ backgroundWorkerId }: { backgroundWorkerId: string }) {
     //It could be a lot of runs, so we will process them in a batch
     //if there are still more to process we will enqueue this function again
@@ -1636,27 +1565,3 @@ export class RunEngine {
     return `master-background-worker:${backgroundWorkerId}`;
   }
 }
-
-export class ServiceValidationError extends Error {
-  constructor(
-    message: string,
-    public status?: number
-  ) {
-    super(message);
-    this.name = "ServiceValidationError";
-  }
-}
-
-class NotImplementedError extends Error {
-  constructor(message: string) {
-    console.error("This isn't implemented", { message });
-    super(message);
-  }
-}
-
-export class RunDuplicateIdempotencyKeyError extends Error {
-  constructor(message: string) {
-    super(message);
-    this.name = "RunDuplicateIdempotencyKeyError";
-  }
-}

internal-packages/run-engine/src/engine/systems/checkpointSystem.ts

@@ -2,7 +2,6 @@ import { CheckpointInput, CreateCheckpointResult, ExecutionResult } from "@trigg
 import { CheckpointId } from "@trigger.dev/core/v3/isomorphic";
 import { PrismaClientOrTransaction } from "@trigger.dev/database";
 import { sendNotificationToWorker } from "../eventBus.js";
-import { ServiceValidationError } from "../index.js";
 import { isCheckpointable, isPendingExecuting } from "../statuses.js";
 import {
   getLatestExecutionSnapshot,
@@ -10,6 +9,7 @@ import {
   ExecutionSnapshotSystem,
 } from "./executionSnapshotSystem.js";
 import { SystemResources } from "./systems.js";
+import { ServiceValidationError } from "../errors.js";
 
 export type CheckpointSystemOptions = {
   resources: SystemResources;

internal-packages/run-engine/src/engine/systems/delayedRunSystem.ts

@@ -0,0 +1,124 @@
+import { startSpan } from "@internal/tracing";
+import { SystemResources } from "./systems.js";
+import { PrismaClientOrTransaction, TaskRun } from "@trigger.dev/database";
+import { getLatestExecutionSnapshot } from "./executionSnapshotSystem.js";
+import { parseNaturalLanguageDuration } from "@trigger.dev/core/v3/isomorphic";
+import { EnqueueSystem } from "./enqueueSystem.js";
+import { ServiceValidationError } from "../errors.js";
+
+export type DelayedRunSystemOptions = {
+  resources: SystemResources;
+  enqueueSystem: EnqueueSystem;
+};
+
+export class DelayedRunSystem {
+  private readonly $: SystemResources;
+  private readonly enqueueSystem: EnqueueSystem;
+
+  constructor(private readonly options: DelayedRunSystemOptions) {
+    this.$ = options.resources;
+    this.enqueueSystem = options.enqueueSystem;
+  }
+
+  /**
+   * Reschedules a delayed run where the run hasn't been queued yet
+   */
+  async rescheduleDelayedRun({
+    runId,
+    delayUntil,
+    tx,
+  }: {
+    runId: string;
+    delayUntil: Date;
+    tx?: PrismaClientOrTransaction;
+  }): Promise<TaskRun> {
+    const prisma = tx ?? this.$.prisma;
+    return startSpan(
+      this.$.tracer,
+      "rescheduleDelayedRun",
+      async () => {
+        return await this.$.runLock.lock([runId], 5_000, async () => {
+          const snapshot = await getLatestExecutionSnapshot(prisma, runId);
+
+          //if the run isn't just created then we can't reschedule it
+          if (snapshot.executionStatus !== "RUN_CREATED") {
+            throw new ServiceValidationError("Cannot reschedule a run that is not delayed");
+          }
+
+          const updatedRun = await prisma.taskRun.update({
+            where: {
+              id: runId,
+            },
+            data: {
+              delayUntil: delayUntil,
+              executionSnapshots: {
+                create: {
+                  engine: "V2",
+                  executionStatus: "RUN_CREATED",
+                  description: "Delayed run was rescheduled to a future date",
+                  runStatus: "EXPIRED",
+                  environmentId: snapshot.environmentId,
+                  environmentType: snapshot.environmentType,
+                },
+              },
+            },
+          });
+
+          await this.$.worker.reschedule(`enqueueDelayedRun:${updatedRun.id}`, delayUntil);
+
+          return updatedRun;
+        });
+      },
+      {
+        attributes: { runId },
+      }
+    );
+  }
+
+  async enqueueDelayedRun({ runId }: { runId: string }) {
+    const run = await this.$.prisma.taskRun.findFirst({
+      where: { id: runId },
+      include: {
+        runtimeEnvironment: {
+          include: {
+            project: true,
+            organization: true,
+          },
+        },
+      },
+    });
+
+    if (!run) {
+      throw new Error(`#enqueueDelayedRun: run not found: ${runId}`);
+    }
+
+    // Now we need to enqueue the run into the RunQueue
+    await this.enqueueSystem.enqueueRun({
+      run,
+      env: run.runtimeEnvironment,
+      timestamp: run.createdAt.getTime() - run.priorityMs,
+      batchId: run.batchId ?? undefined,
+    });
+
+    await this.$.prisma.taskRun.update({
+      where: { id: runId },
+      data: {
+        status: "PENDING",
+        queuedAt: new Date(),
+      },
+    });
+
+    if (run.ttl) {
+      const expireAt = parseNaturalLanguageDuration(run.ttl);
+
+      if (expireAt) {
+        await this.$.worker.enqueue({
+          id: `expireRun:${runId}`,
+          job: "expireRun",
+          payload: { runId },
+          availableAt: expireAt,
+        });
+      }
+    }
+  }
+}

internal-packages/run-engine/src/engine/systems/runAttemptSystem.ts

@@ -1,17 +1,4 @@
-import {
-  $transaction,
-  PrismaClient,
-  PrismaClientOrTransaction,
-  RuntimeEnvironmentType,
-  TaskRun,
-} from "@trigger.dev/database";
-import { Logger } from "@trigger.dev/core/logger";
-import { startSpan, Tracer } from "@internal/tracing";
-import {
-  executionResultFromSnapshot,
-  ExecutionSnapshotSystem,
-  getLatestExecutionSnapshot,
-} from "./executionSnapshotSystem.js";
+import { startSpan } from "@internal/tracing";
 import {
   CompleteRunAttemptResult,
   ExecutionResult,
@@ -23,20 +10,29 @@ import {
   TaskRunInternalError,
   TaskRunSuccessfulExecutionResult,
 } from "@trigger.dev/core/v3/schemas";
-import { RunLocker } from "../locking.js";
-import { EventBus, sendNotificationToWorker } from "../eventBus.js";
-import { ServiceValidationError } from "../index.js";
-import { retryOutcomeFromCompletion } from "../retrying.js";
-import { RunQueue } from "../../run-queue/index.js";
-import { isExecuting } from "../statuses.js";
-import { EngineWorker, RunEngineOptions } from "../types.js";
-import { runStatusFromError } from "../errors.js";
-import { BatchSystem } from "./batchSystem.js";
-import { WaitpointSystem } from "./waitpointSystem.js";
-import { MAX_TASK_RUN_ATTEMPTS } from "../consts.js";
-import { getMachinePreset } from "../machinePresets.js";
 import { parsePacket } from "@trigger.dev/core/v3/utils/ioSerialization";
+import {
+  $transaction,
+  PrismaClientOrTransaction,
+  RuntimeEnvironmentType,
+  TaskRun,
+} from "@trigger.dev/database";
+import { MAX_TASK_RUN_ATTEMPTS } from "../consts.js";
+import { runStatusFromError, ServiceValidationError } from "../errors.js";
+import { sendNotificationToWorker } from "../eventBus.js";
+import { getMachinePreset } from "../machinePresets.js";
+import { retryOutcomeFromCompletion } from "../retrying.js";
+import { isExecuting } from "../statuses.js";
+import { RunEngineOptions } from "../types.js";
+import { BatchSystem } from "./batchSystem.js";
+import {
+  executionResultFromSnapshot,
+  ExecutionSnapshotSystem,
+  getLatestExecutionSnapshot,
+} from "./executionSnapshotSystem.js";
 import { SystemResources } from "./systems.js";
+import { WaitpointSystem } from "./waitpointSystem.js";
+
 export type RunAttemptSystemOptions = {
   resources: SystemResources;
   executionSnapshotSystem: ExecutionSnapshotSystem;