Run engine: Using root queue timestamp to prioritize completing runs (#1818)

* Improve the new run engine Trigger/Batch trigger service class names

* centralize queue timestamp logic in EnqueueSystem, adding queueTimestamp support and propagation

* Fixing the create tags in the new run engine trigger task service

Author: ericallam

apps/webapp/app/models/taskRunTag.server.ts

@@ -1,11 +1,15 @@
 import { prisma } from "~/db.server";
 import { generateFriendlyId } from "~/v3/friendlyIdentifiers";
+import { PrismaClientOrTransaction } from "@trigger.dev/database";
 
 export const MAX_TAGS_PER_RUN = 10;
 
-export async function createTag({ tag, projectId }: { tag: string; projectId: string }) {
+export async function createTag(
+  { tag, projectId }: { tag: string; projectId: string },
+  prismaClient: PrismaClientOrTransaction = prisma
+) {
   if (tag.trim().length === 0) return;
-  return prisma.taskRunTag.upsert({
+  return prismaClient.taskRunTag.upsert({
     where: {
       projectId_name: {
         projectId: projectId,
@@ -21,6 +25,48 @@ export async function createTag({ tag, projectId }: { tag: string; projectId: st
   });
 }
 
+export type TagRecord = {
+  id: string;
+  name: string;
+};
+
+export async function createTags(
+  {
+    tags,
+    projectId,
+  }: {
+    tags: string | string[] | undefined;
+    projectId: string;
+  },
+  prismaClient: PrismaClientOrTransaction = prisma
+): Promise<TagRecord[]> {
+  if (!tags) {
+    return [];
+  }
+
+  const tagsArray = typeof tags === "string" ? [tags] : tags;
+
+  if (tagsArray.length === 0) {
+    return [];
+  }
+
+  const tagRecords: TagRecord[] = [];
+  for (const tag of tagsArray) {
+    const tagRecord = await createTag(
+      {
+        tag,
+        projectId,
+      },
+      prismaClient
+    );
+    if (tagRecord) {
+      tagRecords.push({ id: tagRecord.id, name: tagRecord.name });
+    }
+  }
+
+  return tagRecords;
+}
+
 export async function getTagsForRunId({
   friendlyId,
   environmentId,

apps/webapp/app/routes/api.v2.tasks.batch.ts

@@ -10,9 +10,9 @@ import { logger } from "~/services/logger.server";
 import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server";
 import { ServiceValidationError } from "~/v3/services/baseService.server";
 import { BatchProcessingStrategy } from "~/v3/services/batchTriggerV3.server";
-import { BatchTriggerV4Service } from "~/v3/services/batchTriggerV4.server";
 import { OutOfEntitlementError } from "~/v3/services/triggerTask.server";
 import { HeadersSchema } from "./api.v1.tasks.$taskId.trigger";
+import { RunEngineBatchTriggerService } from "~/runEngine/services/batchTrigger.server";
 
 const { action, loader } = createActionApiRoute(
   {
@@ -74,7 +74,7 @@ const { action, loader } = createActionApiRoute(
         ? { traceparent, tracestate }
         : undefined;
 
-    const service = new BatchTriggerV4Service(batchProcessingStrategy ?? undefined);
+    const service = new RunEngineBatchTriggerService(batchProcessingStrategy ?? undefined);
 
     try {
       const batch = await service.call(authentication.environment, body, {

…app/v3/services/batchTriggerV4.server.ts …unEngine/services/batchTrigger.server.tsapps/webapp/app/v3/services/batchTriggerV4.server.ts renamed to apps/webapp/app/runEngine/services/batchTrigger.server.ts apps/webapp/app/v3/services/batchTriggerV4.server.ts renamed to apps/webapp/app/runEngine/services/batchTrigger.server.ts

@@ -15,10 +15,10 @@ import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { logger } from "~/services/logger.server";
 import { getEntitlement } from "~/services/platform.v3.server";
 import { workerQueue } from "~/services/worker.server";
-import { downloadPacketFromObjectStore, uploadPacketToObjectStore } from "../r2.server";
-import { startActiveSpan } from "../tracer.server";
-import { ServiceValidationError, WithRunEngine } from "./baseService.server";
-import { OutOfEntitlementError, TriggerTaskService } from "./triggerTask.server";
+import { downloadPacketFromObjectStore, uploadPacketToObjectStore } from "../../v3/r2.server";
+import { startActiveSpan } from "../../v3/tracer.server";
+import { ServiceValidationError, WithRunEngine } from "../../v3/services/baseService.server";
+import { OutOfEntitlementError, TriggerTaskService } from "../../v3/services/triggerTask.server";
 
 const PROCESSING_BATCH_SIZE = 50;
 const ASYNC_BATCH_PROCESS_SIZE_THRESHOLD = 20;
@@ -49,7 +49,7 @@ export type BatchTriggerTaskServiceOptions = {
 /**
  * Larger batches, used in Run Engine v2
  */
-export class BatchTriggerV4Service extends WithRunEngine {
+export class RunEngineBatchTriggerService extends WithRunEngine {
   private _batchProcessingStrategy: BatchProcessingStrategy;
 
   constructor(
@@ -643,7 +643,7 @@ export class BatchTriggerV4Service extends WithRunEngine {
   }
 
   async #enqueueBatchTaskRun(options: BatchProcessingOptions, tx?: PrismaClientOrTransaction) {
-    await workerQueue.enqueue("v3.processBatchTaskRunV3", options, {
+    await workerQueue.enqueue("runengine.processBatchTaskRun", options, {
       tx,
       jobKey: `BatchTriggerV3Service.process:${options.batchId}:${options.processingId}`,
     });

…/app/v3/services/triggerTaskV2.server.ts …runEngine/services/triggerTask.server.tsapps/webapp/app/v3/services/triggerTaskV2.server.ts renamed to apps/webapp/app/runEngine/services/triggerTask.server.ts apps/webapp/app/v3/services/triggerTaskV2.server.ts renamed to apps/webapp/app/runEngine/services/triggerTask.server.ts

@@ -1,4 +1,4 @@
-import { RunEngine, RunDuplicateIdempotencyKeyError } from "@internal/run-engine";
+import { RunDuplicateIdempotencyKeyError, RunEngine } from "@internal/run-engine";
 import {
   IOPacket,
   packetRequiresOffloading,
@@ -14,34 +14,33 @@ import {
   sanitizeQueueName,
   stringifyDuration,
 } from "@trigger.dev/core/v3/isomorphic";
-import { Prisma, TaskRun } from "@trigger.dev/database";
+import { Prisma } from "@trigger.dev/database";
 import { env } from "~/env.server";
-import { createTag, MAX_TAGS_PER_RUN } from "~/models/taskRunTag.server";
+import { createTags, MAX_TAGS_PER_RUN } from "~/models/taskRunTag.server";
 import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { autoIncrementCounter } from "~/services/autoIncrementCounter.server";
 import { logger } from "~/services/logger.server";
 import { getEntitlement } from "~/services/platform.v3.server";
 import { parseDelay } from "~/utils/delays";
 import { resolveIdempotencyKeyTTL } from "~/utils/idempotencyKeys.server";
 import { handleMetadataPacket } from "~/utils/packets";
-import { eventRepository } from "../eventRepository.server";
-import { findCurrentWorkerFromEnvironment } from "../models/workerDeployment.server";
-import { uploadPacketToObjectStore } from "../r2.server";
-import { isFinalRunStatus } from "../taskStatus";
-import { startActiveSpan } from "../tracer.server";
-import { clampMaxDuration } from "../utils/maxDuration";
-import { ServiceValidationError, WithRunEngine } from "./baseService.server";
+import { eventRepository } from "../../v3/eventRepository.server";
+import { findCurrentWorkerFromEnvironment } from "../../v3/models/workerDeployment.server";
+import { uploadPacketToObjectStore } from "../../v3/r2.server";
+import { getTaskEventStore } from "../../v3/taskEventStore.server";
+import { isFinalRunStatus } from "../../v3/taskStatus";
+import { startActiveSpan } from "../../v3/tracer.server";
+import { clampMaxDuration } from "../../v3/utils/maxDuration";
+import { ServiceValidationError, WithRunEngine } from "../../v3/services/baseService.server";
 import {
   MAX_ATTEMPTS,
   OutOfEntitlementError,
   TriggerTaskServiceOptions,
   TriggerTaskServiceResult,
-} from "./triggerTask.server";
-import { WorkerGroupService } from "./worker/workerGroupService.server";
-import { getTaskEventStore } from "../taskEventStore.server";
+} from "../../v3/services/triggerTask.server";
+import { WorkerGroupService } from "../../v3/services/worker/workerGroupService.server";
 
-/** @deprecated Use TriggerTaskService in `triggerTask.server.ts` instead. */
-export class TriggerTaskServiceV2 extends WithRunEngine {
+export class RunEngineTriggerTaskService extends WithRunEngine {
   public async call({
     taskId,
     environment,
@@ -299,20 +298,13 @@ export class TriggerTaskServiceV2 extends WithRunEngine {
                 span.setAttribute("queueName", queueName);
 
                 //upsert tags
-                let tags: { id: string; name: string }[] = [];
-                const bodyTags =
-                  typeof body.options?.tags === "string" ? [body.options.tags] : body.options?.tags;
-                if (bodyTags && bodyTags.length > 0) {
-                  for (const tag of bodyTags) {
-                    const tagRecord = await createTag({
-                      tag,
-                      projectId: environment.projectId,
-                    });
-                    if (tagRecord) {
-                      tags.push(tagRecord);
-                    }
-                  }
-                }
+                const tags = await createTags(
+                  {
+                    tags: body.options?.tags,
+                    projectId: environment.projectId,
+                  },
+                  this._prisma
+                );
 
                 const depth = parentRun ? parentRun.depth + 1 : 0;
 
@@ -372,6 +364,10 @@ export class TriggerTaskServiceV2 extends WithRunEngine {
                     machine: body.options?.machine,
                     priorityMs: body.options?.priority ? body.options.priority * 1_000 : undefined,
                     releaseConcurrency: body.options?.releaseConcurrency,
+                    queueTimestamp:
+                      parentRun && body.options?.resumeParentOnCompletion
+                        ? parentRun.queueTimestamp ?? undefined
+                        : undefined,
                   },
                   this._prisma
                 );

apps/webapp/app/services/worker.server.ts

@@ -29,9 +29,9 @@ import { reportInvocationUsage } from "./platform.v3.server";
 import { logger } from "./logger.server";
 import { BatchProcessingOptions, BatchTriggerV3Service } from "~/v3/services/batchTriggerV3.server";
 import {
-  BatchProcessingOptions as BatchProcessingOptionsV4,
-  BatchTriggerV4Service,
-} from "~/v3/services/batchTriggerV4.server";
+  BatchProcessingOptions as RunEngineBatchProcessingOptions,
+  RunEngineBatchTriggerService,
+} from "~/runEngine/services/batchTrigger.server";
 
 const workerCatalog = {
   scheduleEmail: DeliverEmailSchema,
@@ -99,7 +99,7 @@ const workerCatalog = {
   }),
   "v3.cancelDevSessionRuns": CancelDevSessionRunsServiceOptions,
   "v3.processBatchTaskRun": BatchProcessingOptions,
-  "v3.processBatchTaskRunV3": BatchProcessingOptionsV4,
+  "runengine.processBatchTaskRun": RunEngineBatchProcessingOptions,
 };
 
 let workerQueue: ZodWorker<typeof workerCatalog>;
@@ -341,11 +341,11 @@ function getWorkerQueue() {
           await service.processBatchTaskRun(payload);
         },
       },
-      "v3.processBatchTaskRunV3": {
+      "runengine.processBatchTaskRun": {
         priority: 0,
         maxAttempts: 5,
         handler: async (payload, job) => {
-          const service = new BatchTriggerV4Service(payload.strategy);
+          const service = new RunEngineBatchTriggerService(payload.strategy);
 
           await service.processBatchTaskRun(payload);
         },

apps/webapp/app/v3/services/triggerTask.server.ts

@@ -4,7 +4,7 @@ import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { determineEngineVersion } from "../engineVersion.server";
 import { WithRunEngine } from "./baseService.server";
 import { TriggerTaskServiceV1 } from "./triggerTaskV1.server";
-import { TriggerTaskServiceV2 } from "./triggerTaskV2.server";
+import { RunEngineTriggerTaskService } from "~/runEngine/services/triggerTask.server";
 
 export type TriggerTaskServiceOptions = {
   idempotencyKey?: string;
@@ -78,7 +78,7 @@ export class TriggerTaskService extends WithRunEngine {
     body: TriggerTaskRequestBody,
     options: TriggerTaskServiceOptions = {}
   ): Promise<TriggerTaskServiceResult | undefined> {
-    const service = new TriggerTaskServiceV2({
+    const service = new RunEngineTriggerTaskService({
       prisma: this._prisma,
       engine: this._engine,
     });

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

@@ -334,6 +334,7 @@ export class RunEngine {
       maxAttempts,
       taskEventStore,
       priorityMs,
+      queueTimestamp,
       ttl,
       tags,
       parentTaskRunId,
@@ -414,6 +415,7 @@ export class RunEngine {
               maxAttempts,
               taskEventStore,
               priorityMs,
+              queueTimestamp: queueTimestamp ?? delayUntil ?? new Date(),
               ttl,
               tags:
                 tags.length === 0
@@ -520,7 +522,6 @@ export class RunEngine {
             await this.enqueueSystem.enqueueRun({
               run: taskRun,
               env: environment,
-              timestamp: Date.now() - taskRun.priorityMs,
               workerId,
               runnerId,
               tx: prisma,

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

@@ -161,7 +161,6 @@ export class CheckpointSystem {
         const newSnapshot = await this.enqueueSystem.enqueueRun({
           run,
           env: run.runtimeEnvironment,
-          timestamp: run.createdAt.getTime() - run.priorityMs,
           snapshot: {
             status: "QUEUED",
             description:

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

@@ -98,7 +98,6 @@ export class DelayedRunSystem {
     await this.enqueueSystem.enqueueRun({
       run,
       env: run.runtimeEnvironment,
-      timestamp: run.createdAt.getTime() - run.priorityMs,
       batchId: run.batchId ?? undefined,
     });
 

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

@@ -25,7 +25,6 @@ export class EnqueueSystem {
   public async enqueueRun({
     run,
     env,
-    timestamp,
     tx,
     snapshot,
     previousSnapshotId,
@@ -37,7 +36,6 @@ export class EnqueueSystem {
   }: {
     run: TaskRun;
     env: MinimalAuthenticatedEnvironment;
-    timestamp: number;
     tx?: PrismaClientOrTransaction;
     snapshot?: {
       status?: Extract<TaskRunExecutionStatus, "QUEUED" | "QUEUED_EXECUTING">;
@@ -81,6 +79,8 @@ export class EnqueueSystem {
         masterQueues.push(run.secondaryMasterQueue);
       }
 
+      const timestamp = (run.queueTimestamp ?? run.createdAt).getTime() - run.priorityMs;
+
       await this.$.runQueue.enqueueMessage({
         env,
         masterQueues,