Reliabily resolve task queue for a run and ack runs where we can't find the queue

ericallam · ericallam · commit 7a1b5b1d24cc · 2025-01-21T12:13:29.000Z
diff --git a/apps/webapp/app/models/taskQueue.server.ts b/apps/webapp/app/models/taskQueue.server.ts
@@ -0,0 +1,63 @@
+import { QueueOptions } from "@trigger.dev/core/v3/schemas";
+import { TaskQueue } from "@trigger.dev/database";
+import { prisma } from "~/db.server";
+
+export async function findQueueInEnvironment(
+  queueName: string,
+  environmentId: string,
+  backgroundWorkerTaskId?: string,
+  backgroundTask?: { queueConfig?: unknown }
+): Promise<TaskQueue | undefined> {
+  const sanitizedQueueName = sanitizeQueueName(queueName);
+
+  const queue = await prisma.taskQueue.findFirst({
+    where: {
+      runtimeEnvironmentId: environmentId,
+      name: sanitizedQueueName,
+    },
+  });
+
+  if (queue) {
+    return queue;
+  }
+
+  const task = backgroundTask
+    ? backgroundTask
+    : backgroundWorkerTaskId
+    ? await prisma.backgroundWorkerTask.findFirst({
+        where: {
+          id: backgroundWorkerTaskId,
+        },
+      })
+    : undefined;
+
+  if (!task) {
+    return;
+  }
+
+  const queueConfig = QueueOptions.safeParse(task.queueConfig);
+
+  if (queueConfig.success) {
+    const taskQueueName = queueConfig.data.name
+      ? sanitizeQueueName(queueConfig.data.name)
+      : undefined;
+
+    if (taskQueueName && taskQueueName !== sanitizedQueueName) {
+      const queue = await prisma.taskQueue.findFirst({
+        where: {
+          runtimeEnvironmentId: environmentId,
+          name: taskQueueName,
+        },
+      });
+
+      if (queue) {
+        return queue;
+      }
+    }
+  }
+}
+
+// Only allow alphanumeric characters, underscores, hyphens, and slashes (and only the first 128 characters)
+export function sanitizeQueueName(queueName: string) {
+  return queueName.replace(/[^a-zA-Z0-9_\-\/]/g, "").substring(0, 128);
+}
diff --git a/apps/webapp/app/v3/marqs/devQueueConsumer.server.ts b/apps/webapp/app/v3/marqs/devQueueConsumer.server.ts
@@ -13,14 +13,15 @@ import { prisma } from "~/db.server";
 import { createNewSession, disconnectSession } from "~/models/runtimeEnvironment.server";
 import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { logger } from "~/services/logger.server";
-import { marqs, sanitizeQueueName } from "~/v3/marqs/index.server";
+import { marqs } from "~/v3/marqs/index.server";
 import { resolveVariablesForEnvironment } from "../environmentVariables/environmentVariablesRepository.server";
 import { FailedTaskRunService } from "../failedTaskRun.server";
 import { CancelDevSessionRunsService } from "../services/cancelDevSessionRuns.server";
 import { CompleteAttemptService } from "../services/completeAttempt.server";
 import { attributesFromAuthenticatedEnv, tracer } from "../tracer.server";
 import { getMaxDuration } from "../utils/maxDuration";
 import { DevSubscriber, devPubSub } from "./devPubSub.server";
+import { findQueueInEnvironment, sanitizeQueueName } from "~/models/taskQueue.server";
 
 const MessageBody = z.discriminatedUnion("type", [
   z.object({
@@ -432,14 +433,12 @@ export class DevQueueConsumer {
       return;
     }
 
-    const queue = await prisma.taskQueue.findUnique({
-      where: {
-        runtimeEnvironmentId_name: {
-          runtimeEnvironmentId: this.env.id,
-          name: sanitizeQueueName(lockedTaskRun.queue),
-        },
-      },
-    });
+    const queue = await findQueueInEnvironment(
+      lockedTaskRun.queue,
+      this.env.id,
+      backgroundTask.id,
+      backgroundTask
+    );
 
     if (!queue) {
       logger.debug("[DevQueueConsumer] Failed to find queue", {
diff --git a/apps/webapp/app/v3/marqs/index.server.ts b/apps/webapp/app/v3/marqs/index.server.ts
@@ -1962,8 +1962,3 @@ function getMarQSClient() {
     }
   }
 }
-
-// Only allow alphanumeric characters, underscores, hyphens, and slashes (and only the first 128 characters)
-export function sanitizeQueueName(queueName: string) {
-  return queueName.replace(/[^a-zA-Z0-9_\-\/]/g, "").substring(0, 128);
-}
diff --git a/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts b/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts
@@ -13,6 +13,7 @@ import {
   MachinePreset,
   ProdTaskRunExecution,
   ProdTaskRunExecutionPayload,
+  QueueOptions,
   TaskRunError,
   TaskRunErrorCodes,
   TaskRunExecution,
@@ -28,6 +29,7 @@ import {
   BackgroundWorker,
   BackgroundWorkerTask,
   Prisma,
+  TaskQueue,
   TaskRunStatus,
 } from "@trigger.dev/database";
 import { z } from "zod";
@@ -37,7 +39,7 @@ import { findEnvironmentById } from "~/models/runtimeEnvironment.server";
 import { generateJWTTokenForEnvironment } from "~/services/apiAuth.server";
 import { logger } from "~/services/logger.server";
 import { singleton } from "~/utils/singleton";
-import { marqs, sanitizeQueueName } from "~/v3/marqs/index.server";
+import { marqs } from "~/v3/marqs/index.server";
 import {
   RuntimeEnvironmentForEnvRepo,
   RuntimeEnvironmentForEnvRepoPayload,
@@ -65,6 +67,7 @@ import {
 import { tracer } from "../tracer.server";
 import { getMaxDuration } from "../utils/maxDuration";
 import { MessagePayload } from "./types";
+import { findQueueInEnvironment, sanitizeQueueName } from "~/models/taskQueue.server";
 
 const WithTraceContext = z.object({
   traceparent: z.string().optional(),
@@ -744,12 +747,12 @@ export class SharedQueueConsumer {
       };
     }
 
-    const queue = await prisma.taskQueue.findFirst({
-      where: {
-        runtimeEnvironmentId: lockedTaskRun.runtimeEnvironmentId,
-        name: sanitizeQueueName(lockedTaskRun.queue),
-      },
-    });
+    const queue = await findQueueInEnvironment(
+      lockedTaskRun.queue,
+      lockedTaskRun.runtimeEnvironmentId,
+      lockedTaskRun.lockedById ?? undefined,
+      backgroundTask
+    );
 
     if (!queue) {
       logger.debug("SharedQueueConsumer queue not found, so nacking message", {
@@ -759,7 +762,7 @@ export class SharedQueueConsumer {
       });
 
       return {
-        action: "nack_and_do_more_work",
+        action: "ack_and_do_more_work",
         reason: "queue_not_found",
         attrs: {
           queue_name: sanitizeQueueName(lockedTaskRun.queue),
@@ -1031,12 +1034,11 @@ export class SharedQueueConsumer {
       };
     }
 
-    const queue = await prisma.taskQueue.findFirst({
-      where: {
-        runtimeEnvironmentId: resumableAttempt.runtimeEnvironmentId,
-        name: sanitizeQueueName(resumableRun.queue),
-      },
-    });
+    const queue = await findQueueInEnvironment(
+      resumableRun.queue,
+      resumableRun.runtimeEnvironmentId,
+      resumableRun.lockedById ?? undefined
+    );
 
     if (!queue) {
       logger.debug("SharedQueueConsumer queue not found, so nacking message", {
@@ -1045,7 +1047,7 @@ export class SharedQueueConsumer {
       });
 
       return {
-        action: "nack_and_do_more_work",
+        action: "ack_and_do_more_work",
         reason: "queue_not_found",
         attrs: {
           queue_name: sanitizeQueueName(resumableRun.queue),
diff --git a/apps/webapp/app/v3/models/workerDeployment.server.ts b/apps/webapp/app/v3/models/workerDeployment.server.ts
@@ -44,6 +44,7 @@ type WorkerDeploymentWithWorkerTasks = Prisma.WorkerDeploymentGetPayload<{
             triggerSource: true;
             machineConfig: true;
             maxDurationInSeconds: true;
+            queueConfig: true;
           };
         };
       };
diff --git a/apps/webapp/app/v3/services/createBackgroundWorker.server.ts b/apps/webapp/app/v3/services/createBackgroundWorker.server.ts
@@ -7,7 +7,7 @@ import type { BackgroundWorker } from "@trigger.dev/database";
 import { Prisma, PrismaClientOrTransaction } from "~/db.server";
 import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { logger } from "~/services/logger.server";
-import { marqs, sanitizeQueueName } from "~/v3/marqs/index.server";
+import { marqs } from "~/v3/marqs/index.server";
 import { generateFriendlyId } from "../friendlyIdentifiers";
 import { calculateNextBuildVersion } from "../utils/calculateNextBuildVersion";
 import { BaseService } from "./baseService.server";
@@ -16,6 +16,7 @@ import { RegisterNextTaskScheduleInstanceService } from "./registerNextTaskSched
 import cronstrue from "cronstrue";
 import { CheckScheduleService } from "./checkSchedule.server";
 import { clampMaxDuration } from "../utils/maxDuration";
+import { sanitizeQueueName } from "~/models/taskQueue.server";
 
 export class CreateBackgroundWorkerService extends BaseService {
   public async call(
diff --git a/apps/webapp/app/v3/services/createTaskRunAttempt.server.ts b/apps/webapp/app/v3/services/createTaskRunAttempt.server.ts
@@ -10,6 +10,7 @@ import { machinePresetFromConfig, machinePresetFromRun } from "../machinePresets
 import { BaseService, ServiceValidationError } from "./baseService.server";
 import { CrashTaskRunService } from "./crashTaskRun.server";
 import { ExpireEnqueuedRunService } from "./expireEnqueuedRun.server";
+import { findQueueInEnvironment } from "~/models/taskQueue.server";
 
 export class CreateTaskRunAttemptService extends BaseService {
   public async call({
@@ -95,18 +96,18 @@ export class CreateTaskRunAttemptService extends BaseService {
         throw new ServiceValidationError("Task run is cancelled", 400);
       }
 
-      if (!taskRun.lockedBy) {
+      const lockedBy = taskRun.lockedBy;
+
+      if (!lockedBy) {
         throw new ServiceValidationError("Task run is not locked", 400);
       }
 
-      const queue = await this._prisma.taskQueue.findUnique({
-        where: {
-          runtimeEnvironmentId_name: {
-            runtimeEnvironmentId: environment.id,
-            name: taskRun.queue,
-          },
-        },
-      });
+      const queue = await findQueueInEnvironment(
+        taskRun.queue,
+        environment.id,
+        lockedBy.id,
+        lockedBy
+      );
 
       if (!queue) {
         throw new ServiceValidationError("Queue not found", 404);
@@ -121,7 +122,7 @@ export class CreateTaskRunAttemptService extends BaseService {
       if (nextAttemptNumber > MAX_TASK_RUN_ATTEMPTS) {
         const service = new CrashTaskRunService(this._prisma);
         await service.call(taskRun.id, {
-          reason: taskRun.lockedBy.worker.supportsLazyAttempts
+          reason: lockedBy.worker.supportsLazyAttempts
             ? "Max attempts reached."
             : "Max attempts reached. Please upgrade your CLI and SDK.",
         });
@@ -136,8 +137,8 @@ export class CreateTaskRunAttemptService extends BaseService {
             friendlyId: generateFriendlyId("attempt"),
             taskRunId: taskRun.id,
             startedAt: new Date(),
-            backgroundWorkerId: taskRun.lockedBy!.worker.id,
-            backgroundWorkerTaskId: taskRun.lockedBy!.id,
+            backgroundWorkerId: lockedBy.worker.id,
+            backgroundWorkerTaskId: lockedBy.id,
             status: setToExecuting ? "EXECUTING" : "PENDING",
             queueId: queue.id,
             runtimeEnvironmentId: environment.id,
@@ -174,8 +175,7 @@ export class CreateTaskRunAttemptService extends BaseService {
       }
 
       const machinePreset =
-        machinePresetFromRun(taskRun) ??
-        machinePresetFromConfig(taskRun.lockedBy.machineConfig ?? {});
+        machinePresetFromRun(taskRun) ?? machinePresetFromConfig(lockedBy.machineConfig ?? {});
 
       const metadata = await parsePacket({
         data: taskRun.metadata ?? undefined,
@@ -184,16 +184,16 @@ export class CreateTaskRunAttemptService extends BaseService {
 
       const execution: TaskRunExecution = {
         task: {
-          id: taskRun.lockedBy.slug,
-          filePath: taskRun.lockedBy.filePath,
-          exportName: taskRun.lockedBy.exportName,
+          id: lockedBy.slug,
+          filePath: lockedBy.filePath,
+          exportName: lockedBy.exportName,
         },
         attempt: {
           id: taskRunAttempt.friendlyId,
           number: taskRunAttempt.number,
           startedAt: taskRunAttempt.startedAt ?? taskRunAttempt.createdAt,
-          backgroundWorkerId: taskRun.lockedBy.worker.id,
-          backgroundWorkerTaskId: taskRun.lockedBy.id,
+          backgroundWorkerId: lockedBy.worker.id,
+          backgroundWorkerTaskId: lockedBy.id,
           status: "EXECUTING" as const,
         },
         run: {
@@ -210,7 +210,7 @@ export class CreateTaskRunAttemptService extends BaseService {
           costInCents: taskRun.costInCents,
           baseCostInCents: taskRun.baseCostInCents,
           maxAttempts: taskRun.maxAttempts ?? undefined,
-          version: taskRun.lockedBy.worker.version,
+          version: lockedBy.worker.version,
           metadata,
           maxDuration: taskRun.maxDurationInSeconds ?? undefined,
         },
diff --git a/apps/webapp/app/v3/services/finalizeTaskRun.server.ts b/apps/webapp/app/v3/services/finalizeTaskRun.server.ts
@@ -1,7 +1,7 @@
 import { FlushedRunMetadata, sanitizeError, TaskRunError } from "@trigger.dev/core/v3";
 import { type Prisma, type TaskRun } from "@trigger.dev/database";
 import { logger } from "~/services/logger.server";
-import { marqs, sanitizeQueueName } from "~/v3/marqs/index.server";
+import { marqs } from "~/v3/marqs/index.server";
 import { generateFriendlyId } from "../friendlyIdentifiers";
 import {
   FINAL_ATTEMPT_STATUSES,
@@ -17,6 +17,7 @@ import { socketIo } from "../handleSocketIo.server";
 import { ResumeBatchRunService } from "./resumeBatchRun.server";
 import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { updateMetadataService } from "~/services/metadata/updateMetadata.server";
+import { findQueueInEnvironment, sanitizeQueueName } from "~/models/taskQueue.server";
 
 type BaseInput = {
   id: string;
@@ -291,6 +292,7 @@ export class FinalizeTaskRunService extends BaseService {
         id: true,
         workerId: true,
         runtimeEnvironmentId: true,
+        queueConfig: true,
       },
       where: {
         id: run.lockedById,
@@ -302,14 +304,12 @@ export class FinalizeTaskRunService extends BaseService {
       return;
     }
 
-    const queue = await this._prisma.taskQueue.findUnique({
-      where: {
-        runtimeEnvironmentId_name: {
-          runtimeEnvironmentId: workerTask.runtimeEnvironmentId,
-          name: sanitizeQueueName(run.queue),
-        },
-      },
-    });
+    const queue = await findQueueInEnvironment(
+      run.queue,
+      workerTask.runtimeEnvironmentId,
+      workerTask.id,
+      workerTask
+    );
 
     if (!queue) {
       logger.error("FinalizeTaskRunService: No queue found", { runId: run.id });
diff --git a/apps/webapp/app/v3/services/triggerTask.server.ts b/apps/webapp/app/v3/services/triggerTask.server.ts
@@ -9,7 +9,7 @@ import { env } from "~/env.server";
 import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { autoIncrementCounter } from "~/services/autoIncrementCounter.server";
 import { workerQueue } from "~/services/worker.server";
-import { marqs, sanitizeQueueName } from "~/v3/marqs/index.server";
+import { marqs } from "~/v3/marqs/index.server";
 import { eventRepository } from "../eventRepository.server";
 import { generateFriendlyId } from "../friendlyIdentifiers";
 import { uploadPacketToObjectStore } from "../r2.server";
@@ -27,6 +27,7 @@ import { guardQueueSizeLimitsForEnv } from "../queueSizeLimits.server";
 import { clampMaxDuration } from "../utils/maxDuration";
 import { resolveIdempotencyKeyTTL } from "~/utils/idempotencyKeys.server";
 import { Prisma } from "@trigger.dev/database";
+import { sanitizeQueueName } from "~/models/taskQueue.server";
 
 export type TriggerTaskServiceOptions = {
   idempotencyKey?: string;
diff --git a/references/v3-catalog/src/trigger/queues.ts b/references/v3-catalog/src/trigger/queues.ts
@@ -30,3 +30,13 @@ export const queuesTest = task({
     await wait.for({ seconds: payload.waitSeconds ?? 1 });
   },
 });
+
+export const namedQueueTask = task({
+  id: "queues/named-queue",
+  queue: {
+    name: "named-queue",
+  },
+  run: async () => {
+    logger.info("named-queue");
+  },
+});

Original file line number	Diff line number	Diff line change
`@@ -1962,8 +1962,3 @@ function getMarQSClient() {`
`1962`	`1962`	`}`
`1963`	`1963`	`}`
`1964`	`1964`	`}`
`1965`		`-`
`1966`		`-// Only allow alphanumeric characters, underscores, hyphens, and slashes (and only the first 128 characters)`
`1967`		`-export function sanitizeQueueName(queueName: string) {`
`1968`		`- return queueName.replace(/[^a-zA-Z0-9_\-\/]/g, "").substring(0, 128);`
`1969`		`-}`