Support custom queue when triggering a task (#1138)

* Created a v3-catalog test script for queues

* SDK: Fix for calling trigger and passing a custom queue

* Support custom queue in TriggerTaskService

* Improved the script in the catalog so it’s clearer what’s going on

* Remove the concurrencyLimit from a queue if the limit is null

* Fix for the test code… stupid

Author: matt-aitken

.changeset/itchy-chairs-itch.md

@@ -0,0 +1,5 @@
+---
+"@trigger.dev/sdk": patch
+---
+
+Fix for calling trigger and passing a custom queue

apps/webapp/app/v3/marqs/index.server.ts

@@ -82,6 +82,10 @@ export class MarQS {
     return this.redis.set(this.keys.queueConcurrencyLimitKey(env, queue), concurrency);
   }
 
+  public async removeQueueConcurrencyLimits(env: AuthenticatedEnvironment, queue: string) {
+    return this.redis.del(this.keys.queueConcurrencyLimitKey(env, queue));
+  }
+
   public async updateEnvConcurrencyLimits(env: AuthenticatedEnvironment) {
     await this.#callUpdateGlobalConcurrencyLimits({
       envConcurrencyLimitKey: this.keys.envConcurrencyLimitKey(env),

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

@@ -179,6 +179,8 @@ export async function createBackgroundTasks(
           taskQueue.name,
           taskQueue.concurrencyLimit
         );
+      } else {
+        await marqs?.removeQueueConcurrencyLimits(environment, taskQueue.name);
       }
     } catch (error) {
       if (error instanceof Prisma.PrismaClientKnownRequestError) {

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

@@ -5,11 +5,11 @@ import {
   packetRequiresOffloading,
 } from "@trigger.dev/core/v3";
 import { createHash } from "node:crypto";
-import { $transaction } from "~/db.server";
+import { $transaction, prisma } from "~/db.server";
 import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
+import { marqs, sanitizeQueueName } from "~/v3/marqs/index.server";
 import { eventRepository } from "../eventRepository.server";
 import { generateFriendlyId } from "../friendlyIdentifiers";
-import { marqs, sanitizeQueueName } from "~/v3/marqs/index.server";
 import { uploadToObjectStore } from "../r2.server";
 import { BaseService } from "./baseService.server";
 
@@ -111,7 +111,12 @@ export class TriggerTaskService extends BaseService {
               select: { lastNumber: true },
             });
 
-            const queueName = sanitizeQueueName(body.options?.queue?.name ?? `task/${taskId}`);
+            let queueName = sanitizeQueueName(body.options?.queue?.name ?? `task/${taskId}`);
+
+            // Check that the queuename is not an empty string
+            if (!queueName) {
+              queueName = sanitizeQueueName(`task/${taskId}`);
+            }
 
             event.setAttribute("queueName", queueName);
             span.setAttribute("queueName", queueName);
@@ -182,6 +187,43 @@ export class TriggerTaskService extends BaseService {
               }
             }
 
+            if (body.options?.queue) {
+              const concurrencyLimit = body.options.queue.concurrencyLimit
+                ? Math.max(0, body.options.queue.concurrencyLimit)
+                : null;
+              const taskQueue = await prisma.taskQueue.upsert({
+                where: {
+                  runtimeEnvironmentId_name: {
+                    runtimeEnvironmentId: environment.id,
+                    name: queueName,
+                  },
+                },
+                update: {
+                  concurrencyLimit,
+                  rateLimit: body.options.queue.rateLimit,
+                },
+                create: {
+                  friendlyId: generateFriendlyId("queue"),
+                  name: queueName,
+                  concurrencyLimit,
+                  runtimeEnvironmentId: environment.id,
+                  projectId: environment.projectId,
+                  rateLimit: body.options.queue.rateLimit,
+                  type: "NAMED",
+                },
+              });
+
+              if (typeof taskQueue.concurrencyLimit === "number") {
+                await marqs?.updateQueueConcurrencyLimits(
+                  environment,
+                  taskQueue.name,
+                  taskQueue.concurrencyLimit
+                );
+              } else {
+                await marqs?.removeQueueConcurrencyLimits(environment, taskQueue.name);
+              }
+            }
+
             return taskRun;
           });
 

packages/trigger-sdk/src/v3/shared.ts

@@ -353,7 +353,7 @@ export function createTask<TInput = void, TOutput = unknown, TInitOutput extends
             {
               payload: payloadPacket.data,
               options: {
-                queue: params.queue,
+                queue: options?.queue ?? params.queue,
                 concurrencyKey: options?.concurrencyKey,
                 test: taskContext.ctx?.run.isTest,
                 payloadType: payloadPacket.dataType,
@@ -482,7 +482,7 @@ export function createTask<TInput = void, TOutput = unknown, TInitOutput extends
             options: {
               dependentAttempt: ctx.attempt.id,
               lockToVersion: taskContext.worker?.version, // Lock to current version because we're waiting for it to finish
-              queue: params.queue,
+              queue: options?.queue ?? params.queue,
               concurrencyKey: options?.concurrencyKey,
               test: taskContext.ctx?.run.isTest,
               payloadType: payloadPacket.dataType,

references/v3-catalog/package.json

@@ -4,7 +4,8 @@
   "private": true,
   "scripts": {
     "dev:trigger": "triggerdev dev",
-    "management": "ts-node -r tsconfig-paths/register ./src/management.ts"
+    "management": "ts-node -r tsconfig-paths/register ./src/management.ts",
+    "queues": "ts-node -r tsconfig-paths/register ./src/queues.ts"
   },
   "dependencies": {
     "@ffmpeg-installer/ffmpeg": "^1.1.0",
@@ -54,4 +55,4 @@
     "tsconfig-paths": "^4.2.0",
     "typescript": "^5.3.0"
   }
-}
+}

references/v3-catalog/src/queues.ts

@@ -0,0 +1,73 @@
+import dotenv from "dotenv";
+import { simpleChildTask } from "./trigger/subtasks";
+import { wait } from "@trigger.dev/sdk/v3";
+import { setTimeout } from "timers/promises";
+
+dotenv.config();
+
+export async function run() {
+  await simpleChildTask.trigger({ message: "Regular queue" });
+  await simpleChildTask.trigger(
+    { message: "Simple alt queue 1" },
+    {
+      queue: {
+        name: "queue-concurrency-3",
+        concurrencyLimit: 1,
+      },
+    }
+  );
+  await simpleChildTask.trigger(
+    { message: "Simple alt queue 2" },
+    {
+      queue: {
+        name: "queue-concurrency-3",
+        concurrencyLimit: 1,
+      },
+    }
+  );
+
+  await simpleChildTask.batchTrigger([{ payload: { message: "Regular queue" } }]);
+  await simpleChildTask.batchTrigger([
+    {
+      payload: { message: "Batched alt queue 1" },
+      options: {
+        queue: {
+          name: "queue-concurrency-3",
+          concurrencyLimit: 1,
+        },
+      },
+    },
+    {
+      payload: { message: "Batched alt queue 2" },
+      options: {
+        queue: {
+          name: "queue-concurrency-3",
+          concurrencyLimit: 1,
+        },
+      },
+    },
+    {
+      payload: { message: "Batched alt queue 3" },
+      options: {
+        queue: {
+          name: "queue-concurrency-3",
+          concurrencyLimit: 1,
+        },
+      },
+    },
+  ]);
+
+  await setTimeout(10_000);
+
+  //this should set the concurrencyLimit back to none
+  await simpleChildTask.trigger(
+    { message: "Simple alt queue 2" },
+    {
+      queue: {
+        name: "queue-concurrency-3",
+      },
+    }
+  );
+}
+
+run().catch(console.error);