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
This commit is contained in:
Matt Aitken
2024-05-31 11:24:13 +01:00
committed by GitHub
parent ba71f959e2
commit 51bb4c887a
7 changed files with 134 additions and 7 deletions
+5
View File
@@ -0,0 +1,5 @@
---
"@trigger.dev/sdk": patch
---
Fix for calling trigger and passing a custom queue
+4
View File
@@ -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),
@@ -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) {
@@ -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;
});
+2 -2
View File
@@ -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,
+3 -2
View File
@@ -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"
}
}
}
+73
View File
@@ -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);