Fix issue with heavy contention on TaskQueue updating concurrency limit

ericallam · ericallam · commit 750c868fa59c · 2025-01-30T11:13:27.000Z
diff --git a/.changeset/funny-emus-pay.md b/.changeset/funny-emus-pay.md
@@ -0,0 +1,5 @@
+---
+"@trigger.dev/core": patch
+---
+
+Allow setting concurrencyLimit to null to signal removing the concurrency limit on the queue
diff --git a/apps/webapp/app/v3/services/createBackgroundWorker.server.ts b/apps/webapp/app/v3/services/createBackgroundWorker.server.ts
@@ -180,29 +180,29 @@ export async function createBackgroundTasks(
               ),
               0
             )
-          : null;
+          : task.queue?.concurrencyLimit;
 
-      const taskQueue = await prisma.taskQueue.upsert({
+      let taskQueue = await prisma.taskQueue.findFirst({
         where: {
-          runtimeEnvironmentId_name: {
-            runtimeEnvironmentId: worker.runtimeEnvironmentId,
-            name: queueName,
-          },
-        },
-        update: {
-          concurrencyLimit,
-        },
-        create: {
-          friendlyId: generateFriendlyId("queue"),
-          name: queueName,
-          concurrencyLimit,
           runtimeEnvironmentId: worker.runtimeEnvironmentId,
-          projectId: worker.projectId,
-          type: task.queue?.name ? "NAMED" : "VIRTUAL",
+          name: queueName,
         },
       });
 
-      if (typeof taskQueue.concurrencyLimit === "number") {
+      if (!taskQueue) {
+        taskQueue = await prisma.taskQueue.create({
+          data: {
+            friendlyId: generateFriendlyId("queue"),
+            name: queueName,
+            concurrencyLimit,
+            runtimeEnvironmentId: worker.runtimeEnvironmentId,
+            projectId: worker.projectId,
+            type: task.queue?.name ? "NAMED" : "VIRTUAL",
+          },
+        });
+      }
+
+      if (typeof concurrencyLimit === "number") {
         logger.debug("CreateBackgroundWorkerService: updating concurrency limit", {
           workerId: worker.id,
           taskQueue,
@@ -212,12 +212,8 @@ export async function createBackgroundTasks(
           concurrencyLimit,
           taskidentifier: task.id,
         });
-        await marqs?.updateQueueConcurrencyLimits(
-          environment,
-          taskQueue.name,
-          taskQueue.concurrencyLimit
-        );
-      } else {
+        await marqs?.updateQueueConcurrencyLimits(environment, taskQueue.name, concurrencyLimit);
+      } else if (concurrencyLimit === null) {
         logger.debug("CreateBackgroundWorkerService: removing concurrency limit", {
           workerId: worker.id,
           taskQueue,
@@ -227,6 +223,7 @@ export async function createBackgroundTasks(
           concurrencyLimit,
           taskidentifier: task.id,
         });
+
         await marqs?.removeQueueConcurrencyLimits(environment, taskQueue.name);
       }
     } catch (error) {
diff --git a/apps/webapp/app/v3/services/replayTaskRun.server.ts b/apps/webapp/app/v3/services/replayTaskRun.server.ts
@@ -96,7 +96,6 @@ export class ReplayTaskRunService extends BaseService {
             queue: taskQueue
               ? {
                   name: taskQueue.name,
-                  concurrencyLimit: taskQueue.concurrencyLimit ?? undefined,
                 }
               : undefined,
             concurrencyKey: existingTaskRun.concurrencyKey ?? undefined,
diff --git a/apps/webapp/app/v3/services/triggerTask.server.ts b/apps/webapp/app/v3/services/triggerTask.server.ts
@@ -450,7 +450,7 @@ export class TriggerTaskService extends BaseService {
                           ),
                           0
                         )
-                      : null;
+                      : body.options.queue?.concurrencyLimit;
 
                   let taskQueue = await tx.taskQueue.findFirst({
                     where: {
@@ -459,74 +459,47 @@ export class TriggerTaskService extends BaseService {
                     },
                   });
 
-                  const existingConcurrencyLimit =
-                    typeof taskQueue?.concurrencyLimit === "number"
-                      ? taskQueue.concurrencyLimit
-                      : undefined;
-
-                  if (taskQueue) {
-                    if (existingConcurrencyLimit !== concurrencyLimit) {
-                      taskQueue = await tx.taskQueue.update({
-                        where: {
-                          id: taskQueue.id,
-                        },
-                        data: {
-                          concurrencyLimit:
-                            typeof concurrencyLimit === "number" ? concurrencyLimit : null,
-                        },
-                      });
-
-                      if (typeof taskQueue.concurrencyLimit === "number") {
-                        logger.debug("TriggerTaskService: updating concurrency limit", {
-                          runId: taskRun.id,
-                          friendlyId: taskRun.friendlyId,
-                          taskQueue,
-                          orgId: environment.organizationId,
-                          projectId: environment.projectId,
-                          existingConcurrencyLimit,
-                          concurrencyLimit,
-                          queueOptions: body.options?.queue,
-                        });
-                        await marqs?.updateQueueConcurrencyLimits(
-                          environment,
-                          taskQueue.name,
-                          taskQueue.concurrencyLimit
-                        );
-                      } else {
-                        logger.debug("TriggerTaskService: removing concurrency limit", {
-                          runId: taskRun.id,
-                          friendlyId: taskRun.friendlyId,
-                          taskQueue,
-                          orgId: environment.organizationId,
-                          projectId: environment.projectId,
-                          existingConcurrencyLimit,
-                          concurrencyLimit,
-                          queueOptions: body.options?.queue,
-                        });
-                        await marqs?.removeQueueConcurrencyLimits(environment, taskQueue.name);
-                      }
-                    }
-                  } else {
-                    const queueId = generateFriendlyId("queue");
-
+                  if (!taskQueue) {
+                    // handle conflicts with existing queues
                     taskQueue = await tx.taskQueue.create({
                       data: {
-                        friendlyId: queueId,
+                        friendlyId: generateFriendlyId("queue"),
                         name: queueName,
                         concurrencyLimit,
                         runtimeEnvironmentId: environment.id,
                         projectId: environment.projectId,
                         type: "NAMED",
                       },
                     });
+                  }
 
-                    if (typeof taskQueue.concurrencyLimit === "number") {
-                      await marqs?.updateQueueConcurrencyLimits(
-                        environment,
-                        taskQueue.name,
-                        taskQueue.concurrencyLimit
-                      );
-                    }
+                  if (typeof concurrencyLimit === "number") {
+                    logger.debug("TriggerTaskService: updating concurrency limit", {
+                      runId: taskRun.id,
+                      friendlyId: taskRun.friendlyId,
+                      taskQueue,
+                      orgId: environment.organizationId,
+                      projectId: environment.projectId,
+                      concurrencyLimit,
+                      queueOptions: body.options?.queue,
+                    });
+
+                    await marqs?.updateQueueConcurrencyLimits(
+                      environment,
+                      taskQueue.name,
+                      concurrencyLimit
+                    );
+                  } else if (concurrencyLimit === null) {
+                    logger.debug("TriggerTaskService: removing concurrency limit", {
+                      runId: taskRun.id,
+                      friendlyId: taskRun.friendlyId,
+                      taskQueue,
+                      orgId: environment.organizationId,
+                      projectId: environment.projectId,
+                      queueOptions: body.options?.queue,
+                    });
+
+                    await marqs?.removeQueueConcurrencyLimits(environment, taskQueue.name);
                   }
                 }
 
@@ -623,6 +596,18 @@ export class TriggerTaskService extends BaseService {
               throw new ServiceValidationError(
                 `Cannot trigger ${taskId} with a one-time use token as it has already been used.`
               );
+            } else if (
+              Array.isArray(target) &&
+              target.length == 2 &&
+              typeof target[0] === "string" &&
+              typeof target[1] === "string" &&
+              target[0] == "runtimeEnvironmentId" &&
+              target[1] == "name" &&
+              error.message.includes("prisma.taskQueue.create")
+            ) {
+              throw new Error(
+                `Failed to trigger ${taskId} as the queue could not be created do to a unique constraint error, please try again.`
+              );
             } else {
               throw new ServiceValidationError(
                 `Cannot trigger ${taskId} as it has already been triggered with the same idempotency key.`
diff --git a/packages/core/src/v3/schemas/schemas.ts b/packages/core/src/v3/schemas/schemas.ts
@@ -134,8 +134,8 @@ export const QueueOptions = z.object({
   name: z.string().optional(),
   /** An optional property that specifies the maximum number of concurrent run executions.
    *
-   * If this property is omitted, the task can potentially use up the full concurrency of an environment. */
-  concurrencyLimit: z.number().int().min(0).max(1000).optional(),
+   * If this property is omitted, the task can potentially use up the full concurrency of an environment */
+  concurrencyLimit: z.number().int().min(0).max(1000).optional().nullable(),
 });
 
 export type QueueOptions = z.infer<typeof QueueOptions>;
diff --git a/references/v3-catalog/src/trigger/queues.ts b/references/v3-catalog/src/trigger/queues.ts
@@ -11,16 +11,53 @@ export const queuesController = task({
     length?: number;
     waitSeconds?: number;
   }) => {
-    await queuesTest.batchTriggerAndWait(
-      Array.from({ length }, (_, i) => ({
-        payload: { waitSeconds },
-        options: {
+    await Promise.all([
+      queuesTest.trigger(
+        { waitSeconds },
+        {
           queue: {
-            name: `queue-${i % numberOfQueues}`,
+            name: "controller-3",
+            concurrencyLimit: 9,
           },
-        },
-      }))
-    );
+        }
+      ),
+      queuesTest.trigger(
+        { waitSeconds },
+        {
+          queue: {
+            name: "controller-3",
+            concurrencyLimit: 9,
+          },
+        }
+      ),
+      queuesTest.trigger(
+        { waitSeconds },
+        {
+          queue: {
+            name: "controller-3",
+            concurrencyLimit: 9,
+          },
+        }
+      ),
+      queuesTest.trigger(
+        { waitSeconds },
+        {
+          queue: {
+            name: "controller-3",
+            concurrencyLimit: 9,
+          },
+        }
+      ),
+      queuesTest.trigger(
+        { waitSeconds },
+        {
+          queue: {
+            name: "controller-3",
+            concurrencyLimit: 9,
+          },
+        }
+      ),
+    ]);
   },
 });
 
@@ -34,9 +71,10 @@ export const queuesTest = task({
 export const namedQueueTask = task({
   id: "queues/named-queue",
   queue: {
-    name: "named-queue",
+    name: "controller",
+    concurrencyLimit: 9,
   },
   run: async () => {
-    logger.info("named-queue");
+    logger.info("named-queue 2");
   },
 });

-Original file line number
+Diff line change
@@ @@ -0,0 +1,5 @@ @@
 +---
 +"@trigger.dev/core": patch
 +---
++
 +Allow setting concurrencyLimit to null to signal removing the concurrency limit on the queue
Original file line number	Diff line number	Diff line change
`@@ -96,7 +96,6 @@ export class ReplayTaskRunService extends BaseService {`
`96`	`96`	`queue: taskQueue`
`97`	`97`	`? {`
`98`	`98`	`name: taskQueue.name,`
`99`		`- concurrencyLimit: taskQueue.concurrencyLimit ?? undefined,`
`100`	`99`	`}`
`101`	`100`	`: undefined,`
`102`	`101`	`concurrencyKey: existingTaskRun.concurrencyKey ?? undefined,`