Some fixes from badge conflict resolution

Author: ericallam

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

@@ -1,34 +1,36 @@
 import {
   IOPacket,
+  packetRequiresOffloading,
   QueueOptions,
   SemanticInternalAttributes,
+  taskRunErrorEnhancer,
+  taskRunErrorToString,
   TriggerTaskRequestBody,
-  packetRequiresOffloading,
 } from "@trigger.dev/core/v3";
+import { parseNaturalLanguageDuration } from "@trigger.dev/core/v3/apps";
+import { Prisma, TaskRun } from "@trigger.dev/database";
 import { env } from "~/env.server";
+import { sanitizeQueueName } from "~/models/taskQueue.server";
+import { createTag, MAX_TAGS_PER_RUN } from "~/models/taskRunTag.server";
 import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { autoIncrementCounter } from "~/services/autoIncrementCounter.server";
-import { workerQueue } from "~/services/worker.server";
+import { logger } from "~/services/logger.server";
+import { getEntitlement } from "~/services/platform.v3.server";
+import { resolveIdempotencyKeyTTL } from "~/utils/idempotencyKeys.server";
+import { handleMetadataPacket } from "~/utils/packets";
 import { marqs } from "~/v3/marqs/index.server";
 import { eventRepository } from "../eventRepository.server";
 import { generateFriendlyId } from "../friendlyIdentifiers";
-import { uploadPacketToObjectStore } from "../r2.server";
-import { startActiveSpan } from "../tracer.server";
-import { getEntitlement } from "~/services/platform.v3.server";
-import { BaseService, ServiceValidationError } from "./baseService.server";
-import { logger } from "~/services/logger.server";
-import { isFinalAttemptStatus, isFinalRunStatus } from "../taskStatus";
-import { createTag, MAX_TAGS_PER_RUN } from "~/models/taskRunTag.server";
 import { findCurrentWorkerFromEnvironment } from "../models/workerDeployment.server";
-import { handleMetadataPacket } from "~/utils/packets";
-import { parseNaturalLanguageDuration } from "@trigger.dev/core/v3/apps";
-import { ExpireEnqueuedRunService } from "./expireEnqueuedRun.server";
 import { guardQueueSizeLimitsForEnv } from "../queueSizeLimits.server";
+import { uploadPacketToObjectStore } from "../r2.server";
+import { isFinalAttemptStatus, isFinalRunStatus } from "../taskStatus";
+import { startActiveSpan } from "../tracer.server";
 import { clampMaxDuration } from "../utils/maxDuration";
-import { resolveIdempotencyKeyTTL } from "~/utils/idempotencyKeys.server";
-import { Prisma, TaskRun } from "@trigger.dev/database";
-import { sanitizeQueueName } from "~/models/taskQueue.server";
+import { BaseService, ServiceValidationError } from "./baseService.server";
 import { EnqueueDelayedRunService } from "./enqueueDelayedRun.server";
+import { enqueueRun } from "./enqueueRun.server";
+import { ExpireEnqueuedRunService } from "./expireEnqueuedRun.server";
 import { getTaskEventStore } from "../taskEventStore.server";
 
 export type TriggerTaskServiceOptions = {
@@ -187,6 +189,7 @@ export class TriggerTaskService extends BaseService {
                   rootTaskRunId: true,
                   depth: true,
                   queueTimestamp: true,
+                  queue: true,
                 },
               },
             },
@@ -244,6 +247,7 @@ export class TriggerTaskService extends BaseService {
                       rootTaskRunId: true,
                       depth: true,
                       queueTimestamp: true,
+                      queue: true,
                     },
                   },
                 },
@@ -296,7 +300,7 @@ export class TriggerTaskService extends BaseService {
         : undefined;
 
       try {
-        return await eventRepository.traceEvent(
+        const result = await eventRepository.traceEvent(
           taskId,
           {
             context: options.traceContext,
@@ -556,40 +560,61 @@ export class TriggerTaskService extends BaseService {
               this._prisma
             );
 
-            // If this is a triggerAndWait or batchTriggerAndWait,
-            // we need to add the parent run to the reserve concurrency set
-            // to free up concurrency for the children to run
-            if (dependentAttempt) {
-              await marqs?.reserveConcurrency(dependentAttempt.taskRun.id);
-            } else if (dependentBatchRun?.dependentTaskAttempt) {
-              await marqs?.reserveConcurrency(dependentBatchRun.dependentTaskAttempt.taskRun.id);
-            }
-
             if (!run) {
               return;
             }
 
-            // We need to enqueue the task run into the appropriate queue. This is done after the tx completes to prevent a race condition where the task run hasn't been created yet by the time we dequeue.
+            // Now enqueue the run if it's not delayed
             if (run.status === "PENDING") {
-              await marqs?.enqueueMessage(
-                environment,
-                run.queue,
-                run.id,
-                {
-                  type: "EXECUTE",
-                  taskIdentifier: taskId,
-                  projectId: environment.projectId,
-                  environmentId: environment.id,
-                  environmentType: environment.type,
-                },
-                body.options?.concurrencyKey,
-                run.queueTimestamp ?? undefined
-              );
+              const enqueueResult = await enqueueRun({
+                env: environment,
+                run,
+                dependentRun:
+                  dependentAttempt?.taskRun ?? dependentBatchRun?.dependentTaskAttempt?.taskRun,
+              });
+
+              if (!enqueueResult.ok) {
+                // Now we need to fail the run with enqueueResult.error and make sure and
+                // set the traced event to failed as well
+                await this._prisma.taskRun.update({
+                  where: { id: run.id },
+                  data: {
+                    status: "SYSTEM_FAILURE",
+                    completedAt: new Date(),
+                    error: enqueueResult.error,
+                  },
+                });
+
+                event.failWithError(enqueueResult.error);
+
+                return {
+                  run,
+                  isCached: false,
+                  error: enqueueResult.error,
+                };
+              }
             }
 
             return { run, isCached: false };
           }
         );
+
+        if (result?.error) {
+          throw new ServiceValidationError(
+            taskRunErrorToString(taskRunErrorEnhancer(result.error))
+          );
+        }
+
+        const run = result?.run;
+
+        if (!run) {
+          return;
+        }
+
+        return {
+          run,
+          isCached: result?.isCached,
+        };
       } catch (error) {
         // Detect a prisma transaction Unique constraint violation
         if (error instanceof Prisma.PrismaClientKnownRequestError) {

internal-packages/database/prisma/schema.prisma

@@ -1731,6 +1731,8 @@ model TaskRun {
 
   taskEventStore String @default("taskEvent")
 
+  queueTimestamp DateTime?
+
   batchItems             BatchTaskRunItem[]
   dependency             TaskRunDependency?
   CheckpointRestoreEvent CheckpointRestoreEvent[]