diff --git a/apps/webapp/app/presenters/v3/QueueListPresenter.server.ts b/apps/webapp/app/presenters/v3/QueueListPresenter.server.ts index 0dc3daa9856..7db2a6d2e39 100644 --- a/apps/webapp/app/presenters/v3/QueueListPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/QueueListPresenter.server.ts @@ -70,6 +70,7 @@ function buildQueueListWhere( return { runtimeEnvironmentId: environmentId, + role: "QUEUE" as const, version: "V2", name: trimmedQuery ? { diff --git a/apps/webapp/app/presenters/v3/QueueRetrievePresenter.server.ts b/apps/webapp/app/presenters/v3/QueueRetrievePresenter.server.ts index f6918394e5c..6385b388d1f 100644 --- a/apps/webapp/app/presenters/v3/QueueRetrievePresenter.server.ts +++ b/apps/webapp/app/presenters/v3/QueueRetrievePresenter.server.ts @@ -31,6 +31,7 @@ export async function getQueue( where: { friendlyId: queue, runtimeEnvironmentId: environment.id, + role: "QUEUE", }, }) ); @@ -44,6 +45,7 @@ export async function getQueue( where: { name: queueName, runtimeEnvironmentId: environment.id, + role: "QUEUE", }, }) ); diff --git a/apps/webapp/app/routes/api.v1.queues.$queueParam.concurrency.combined.override.ts b/apps/webapp/app/routes/api.v1.queues.$queueParam.concurrency.combined.override.ts new file mode 100644 index 00000000000..c643b77965a --- /dev/null +++ b/apps/webapp/app/routes/api.v1.queues.$queueParam.concurrency.combined.override.ts @@ -0,0 +1,98 @@ +import { json } from "@remix-run/server-runtime"; +import { type RetrieveQueueParam, RetrieveQueueType } from "@trigger.dev/core/v3"; +import { z } from "zod"; +import { toQueueItem } from "~/presenters/v3/QueueRetrievePresenter.server"; +import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server"; +import { concurrencySystem } from "~/v3/services/concurrencySystemInstance.server"; + +const BodySchema = z.object({ + type: RetrieveQueueType.default("id"), + concurrencyLimit: z.number().int().min(0).max(100000), +}); + +const route = createActionApiRoute( + { + body: BodySchema, + params: z.object({ + queueParam: z.string().transform((val) => val.replace(/%2F/g, "/")), + }), + authorization: { + action: "write", + resource: () => ({ type: "queues" }), + }, + }, + async ({ params, body, authentication }) => { + const input: RetrieveQueueParam = + body.type === "id" + ? params.queueParam + : { + type: body.type, + name: decodeURIComponent(params.queueParam).replace(/%2F/g, "/"), + }; + + return concurrencySystem.queues + .overrideTotalConcurrencyLimit(authentication.environment, input, body.concurrencyLimit) + .match( + (queue) => { + return json( + toQueueItem({ + friendlyId: queue.friendlyId, + name: queue.name, + type: queue.type, + running: queue.running, + queued: queue.queued, + concurrencyLimit: queue.concurrencyLimit, + concurrencyLimitBase: queue.concurrencyLimitBase, + concurrencyLimitOverriddenAt: queue.concurrencyLimitOverriddenAt, + concurrencyLimitOverriddenBy: null, + paused: queue.paused, + }), + { status: 200 } + ); + }, + (error) => { + switch (error.type) { + case "queue_not_found": { + return json({ error: "Queue not found" }, { status: 404 }); + } + case "invalid_override": + case "concurrency_limit_exceeds_maximum": { + return json({ error: error.message }, { status: 400 }); + } + case "queue_update_failed": { + return json( + { error: "Failed to update queue total concurrency limit" }, + { status: 500 } + ); + } + case "sync_queue_concurrency_to_engine_failed": { + return json({ error: "Failed to sync the total concurrency limit" }, { status: 500 }); + } + case "get_queue_stats_failed": { + return json({ error: "Failed to read queue stats" }, { status: 500 }); + } + case "other": { + return json( + { error: "Failed to update queue total concurrency limit" }, + { + status: 500, + } + ); + } + default: { + return json( + { error: "Failed to update queue total concurrency limit" }, + { + status: 500, + } + ); + } + } + } + ); + } +); + +export const action = route.action; +/** The builder's loader answers non-POST methods with a 405. */ +export const loader = route.loader; diff --git a/apps/webapp/app/routes/api.v1.queues.$queueParam.concurrency.combined.reset.ts b/apps/webapp/app/routes/api.v1.queues.$queueParam.concurrency.combined.reset.ts new file mode 100644 index 00000000000..b2841f1efe6 --- /dev/null +++ b/apps/webapp/app/routes/api.v1.queues.$queueParam.concurrency.combined.reset.ts @@ -0,0 +1,99 @@ +import { json } from "@remix-run/server-runtime"; +import { type RetrieveQueueParam, RetrieveQueueType } from "@trigger.dev/core/v3"; +import { z } from "zod"; +import { toQueueItem } from "~/presenters/v3/QueueRetrievePresenter.server"; +import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server"; +import { concurrencySystem } from "~/v3/services/concurrencySystemInstance.server"; + +const BodySchema = z.object({ + type: RetrieveQueueType.default("id"), +}); + +const route = createActionApiRoute( + { + body: BodySchema, + params: z.object({ + queueParam: z.string().transform((val) => val.replace(/%2F/g, "/")), + }), + authorization: { + action: "write", + resource: () => ({ type: "queues" }), + }, + }, + async ({ params, body, authentication }) => { + const input: RetrieveQueueParam = + body.type === "id" + ? params.queueParam + : { + type: body.type, + name: decodeURIComponent(params.queueParam).replace(/%2F/g, "/"), + }; + + return concurrencySystem.queues + .resetTotalConcurrencyLimit(authentication.environment, input) + .match( + (queue) => { + return json( + toQueueItem({ + friendlyId: queue.friendlyId, + name: queue.name, + type: queue.type, + running: queue.running, + queued: queue.queued, + concurrencyLimit: queue.concurrencyLimit, + concurrencyLimitBase: queue.concurrencyLimitBase, + concurrencyLimitOverriddenAt: queue.concurrencyLimitOverriddenAt, + concurrencyLimitOverriddenBy: null, + paused: queue.paused, + }), + { status: 200 } + ); + }, + (error) => { + switch (error.type) { + case "queue_not_found": { + return json({ error: "Queue not found" }, { status: 404 }); + } + case "queue_not_overridden": { + return json( + { error: "The queue total concurrency limit is not overridden" }, + { status: 400 } + ); + } + case "queue_update_failed": { + return json( + { error: "Failed to reset the queue total concurrency limit" }, + { status: 500 } + ); + } + case "sync_queue_concurrency_to_engine_failed": { + return json({ error: "Failed to sync the total concurrency limit" }, { status: 500 }); + } + case "get_queue_stats_failed": { + return json({ error: "Failed to read queue stats" }, { status: 500 }); + } + case "other": { + return json( + { error: "Failed to reset the queue total concurrency limit" }, + { + status: 500, + } + ); + } + default: { + return json( + { error: "Failed to reset the queue total concurrency limit" }, + { + status: 500, + } + ); + } + } + } + ); + } +); + +export const action = route.action; +/** The builder's loader answers non-POST methods with a 405. */ +export const loader = route.loader; diff --git a/apps/webapp/app/runEngine/concerns/queues.server.ts b/apps/webapp/app/runEngine/concerns/queues.server.ts index 9214e91b9b3..b42f6c83d4e 100644 --- a/apps/webapp/app/runEngine/concerns/queues.server.ts +++ b/apps/webapp/app/runEngine/concerns/queues.server.ts @@ -220,13 +220,65 @@ export class DefaultQueueManager implements QueueManager { queueName = sanitizedQueueName; } - const requestedGates = request.body.options?.gates ?? taskGates ?? undefined; - const gates = requestedGates - ?.flatMap((gate) => { - const sanitized = sanitizeQueueName(gate.queue); - return sanitized ? [{ queue: sanitized, concurrencyKey: gate.concurrencyKey }] : []; - }) - .slice(0, 2); + const triggerLimits = request.body.options?.concurrency; + + for (const name of triggerLimits ?? []) { + if (!/^[a-zA-Z0-9_-]{1,122}$/.test(name)) { + throw new ServiceValidationError( + `Invalid concurrency limit name "${name}": names are 1-122 characters using only letters, numbers, underscores and hyphens.` + ); + } + } + + /** + * Trigger-time names replace the task's declared NAMED limits only. The task's + * inline limit rides in its stored gates as an anonymous "limit/task/" gate and + * always applies, so it is carried over into the replacement (an empty array + * clears the named limits but keeps the inline one). + */ + const inlineTaskGates = (taskGates ?? []).filter((gate) => + gate.queue.startsWith("limit/task/") + ); + const concurrencyGates = triggerLimits + ? [ + ...inlineTaskGates, + ...triggerLimits.map((name): { queue: string; concurrencyKey?: string } => ({ + queue: `limit/${name}`, + })), + ] + : undefined; + + /** + * The raw gates option replaces stored gates the same way concurrency does, so + * it also carries the inline gate over; a replay resending the stored gates + * collapses back to the original set through the dedupe below. + */ + const rawGates = request.body.options?.gates; + const requestedGates = + concurrencyGates ?? + (rawGates ? [...inlineTaskGates, ...rawGates] : undefined) ?? + taskGates ?? + undefined; + + const seenGates = new Set(); + const gates = requestedGates?.flatMap((gate) => { + const sanitized = sanitizeQueueName(gate.queue); + if (!sanitized) { + return []; + } + const dedupeKey = `${sanitized}${gate.concurrencyKey ?? ""}`; + if (seenGates.has(dedupeKey)) { + return []; + } + seenGates.add(dedupeKey); + return [{ queue: sanitized, concurrencyKey: gate.concurrencyKey }]; + }); + + if (gates && gates.length > 3) { + throw new ServiceValidationError( + `A run can hold at most three gates (the task's inline limit plus two named limits); this request resolves to ${gates.length}.` + ); + } return { queueName, diff --git a/apps/webapp/app/v3/services/concurrencySystem.server.ts b/apps/webapp/app/v3/services/concurrencySystem.server.ts index 51c51674234..099f22bdd69 100644 --- a/apps/webapp/app/v3/services/concurrencySystem.server.ts +++ b/apps/webapp/app/v3/services/concurrencySystem.server.ts @@ -3,7 +3,12 @@ import { errAsync, fromPromise, okAsync } from "neverthrow"; import type { PrismaClientOrTransaction } from "~/db.server"; import type { AuthenticatedEnvironment } from "~/services/apiAuth.server"; import { logger } from "~/services/logger.server"; -import { removeQueueConcurrencyLimits, updateQueueConcurrencyLimits } from "../runQueue.server"; +import { + removeQueueConcurrencyLimits, + removeQueueTotalConcurrencyLimits, + updateQueueConcurrencyLimits, + updateQueueTotalConcurrencyLimits, +} from "../runQueue.server"; import { engine } from "../runEngine.server"; export type ConcurrencySystemOptions = { @@ -77,6 +82,32 @@ export class ConcurrencySystem { .andThen((queue) => syncQueueConcurrencyToEngine(environment, queue)) .andThen((queue) => getQueueStats(environment, queue)); }, + overrideTotalConcurrencyLimit: ( + environment: AuthenticatedEnvironment, + queue: QueueInput, + totalConcurrencyLimit: number, + overriddenBy?: User + ) => { + return findQueueFromInput(this.db, environment, queue) + .andThen((queue) => + overrideQueueTotalConcurrencyLimit( + this.db, + environment, + queue, + totalConcurrencyLimit, + overriddenBy + ) + ) + .andThen((queue) => syncQueueTotalConcurrencyToEngine(environment, queue)) + .andThen((queue) => getQueueStats(environment, queue)); + }, + resetTotalConcurrencyLimit: (environment: AuthenticatedEnvironment, queue: QueueInput) => { + return findQueueFromInput(this.db, environment, queue) + .andThen((queue) => syncQueueTotalConcurrencyResetToEngine(environment, queue)) + .andThen((queue) => resetQueueTotalConcurrencyLimit(this.db, queue)) + .andThen((queue) => syncQueueTotalConcurrencyToEngine(environment, queue)) + .andThen((queue) => getQueueStats(environment, queue)); + }, /** * Recalculates the materialized limit of every percent-based override in the environment * against its CURRENT maximumConcurrencyLimit and syncs changed queues to the run engine. @@ -159,6 +190,7 @@ function findQueueByFriendlyId( where: { runtimeEnvironmentId: environment.id, friendlyId, + role: "QUEUE", }, }), (error) => ({ @@ -183,6 +215,7 @@ function findQueueByName( where: { runtimeEnvironmentId: environment.id, name: queue, + role: "QUEUE", }, }), (error) => ({ @@ -316,6 +349,125 @@ function syncQueueConcurrencyToEngine(environment: AuthenticatedEnvironment, que } } +function overrideQueueTotalConcurrencyLimit( + db: PrismaClientOrTransaction, + environment: AuthenticatedEnvironment, + queue: TaskQueue, + totalConcurrencyLimit: number, + overriddenBy?: User +) { + const maximum = environment.maximumConcurrencyLimit; + + if (!Number.isFinite(totalConcurrencyLimit) || totalConcurrencyLimit < 0) { + return errAsync({ + type: "invalid_override" as const, + message: "Combined concurrency limit must be a non-negative number", + }); + } + + if (totalConcurrencyLimit > maximum) { + return errAsync({ + type: "concurrency_limit_exceeds_maximum" as const, + message: `Combined concurrency limit (${totalConcurrencyLimit}) cannot exceed the environment limit (${maximum})`, + }); + } + + const totalConcurrencyLimitBase = queue.totalConcurrencyLimitOverriddenAt + ? queue.totalConcurrencyLimitBase + : queue.totalConcurrencyLimit; + + return fromPromise( + db.taskQueue.update({ + where: { id: queue.id }, + data: { + totalConcurrencyLimit, + totalConcurrencyLimitBase: totalConcurrencyLimitBase ?? null, + totalConcurrencyLimitOverriddenAt: new Date(), + totalConcurrencyLimitOverriddenBy: overriddenBy?.id ?? null, + }, + }), + (error) => ({ + type: "queue_update_failed" as const, + cause: error, + }) + ); +} + +/** + * Enforce first, then persist: syncs the engine to the declared base BEFORE clearing + * the override marker, so an engine failure leaves the marker set and a retry + * converges instead of being rejected while the overridden limit stays enforced. + */ +function syncQueueTotalConcurrencyResetToEngine( + environment: AuthenticatedEnvironment, + queue: TaskQueue +) { + if (queue.totalConcurrencyLimitOverriddenAt === null) { + return errAsync({ type: "queue_not_overridden" as const }); + } + + if (typeof queue.totalConcurrencyLimitBase === "number") { + return fromPromise( + updateQueueTotalConcurrencyLimits(environment, queue.name, queue.totalConcurrencyLimitBase), + (error) => ({ + type: "sync_queue_concurrency_to_engine_failed" as const, + cause: error, + }) + ).andThen(() => okAsync(queue)); + } + + return fromPromise(removeQueueTotalConcurrencyLimits(environment, queue.name), (error) => ({ + type: "sync_queue_concurrency_to_engine_failed" as const, + cause: error, + })).andThen(() => okAsync(queue)); +} + +function resetQueueTotalConcurrencyLimit(db: PrismaClientOrTransaction, queue: TaskQueue) { + if (queue.totalConcurrencyLimitOverriddenAt === null) { + return errAsync({ type: "queue_not_overridden" as const }); + } + + return fromPromise( + db.taskQueue.update({ + where: { id: queue.id }, + data: { + totalConcurrencyLimit: queue.totalConcurrencyLimitBase, + totalConcurrencyLimitBase: null, + totalConcurrencyLimitOverriddenAt: null, + totalConcurrencyLimitOverriddenBy: null, + }, + }), + (error) => ({ + type: "queue_update_failed" as const, + cause: error, + }) + ); +} + +/** + * The total limit key is separate from the per-queue limit key that pause zeroes, + * so it syncs regardless of the paused state. + */ +function syncQueueTotalConcurrencyToEngine( + environment: AuthenticatedEnvironment, + queue: TaskQueue +) { + if (typeof queue.totalConcurrencyLimit === "number") { + return fromPromise( + updateQueueTotalConcurrencyLimits(environment, queue.name, queue.totalConcurrencyLimit), + (error) => ({ + type: "sync_queue_concurrency_to_engine_failed" as const, + cause: error, + }) + ).andThen(() => okAsync(queue)); + } + + return fromPromise(removeQueueTotalConcurrencyLimits(environment, queue.name), (error) => ({ + type: "sync_queue_concurrency_to_engine_failed" as const, + cause: error, + })).andThen(() => okAsync(queue)); +} + function getQueueStats(environment: AuthenticatedEnvironment, queue: TaskQueue) { return fromPromise( Promise.all([ diff --git a/apps/webapp/app/v3/services/createBackgroundWorker.server.ts b/apps/webapp/app/v3/services/createBackgroundWorker.server.ts index 386f8f038d6..1ea03d35e40 100644 --- a/apps/webapp/app/v3/services/createBackgroundWorker.server.ts +++ b/apps/webapp/app/v3/services/createBackgroundWorker.server.ts @@ -16,7 +16,13 @@ import { stringifyDuration, } from "@trigger.dev/core/v3/isomorphic"; import { randomBytes } from "node:crypto"; -import type { BackgroundWorker, TaskQueue, TaskQueueType } from "@trigger.dev/database"; +import type { + BackgroundWorker, + TaskQueue, + TaskQueueConcurrencyVersion, + TaskQueueRole, + TaskQueueType, +} from "@trigger.dev/database"; import cronstrue from "cronstrue"; import type { PrismaClientOrTransaction, WebhookDatabase } from "~/db.server"; import { $transaction, Prisma, boundedIn, webhookPrisma } from "~/db.server"; @@ -91,6 +97,8 @@ export class CreateBackgroundWorkerService extends BaseService { }, }); + validateWorkerConcurrencyDeclarations(body.metadata); + const latestBackgroundWorker = project.backgroundWorkers[0]; if (latestBackgroundWorker?.contentHash === body.metadata.contentHash) { @@ -337,6 +345,7 @@ export async function createWorkerResources( // Create the queues const queues = await createWorkerQueues(metadata, worker, environment, prisma); + await createWorkerConcurrencyLimits(metadata, worker, environment, prisma); // Create the tasks const taskEntries = await createWorkerTasks( @@ -391,10 +400,59 @@ async function createWorkerTask( ): Promise { // Hoisted so the P2002 catch branch can return the same entry shape. let queue: TaskQueue | undefined; + let compiledGates: Array<{ queue: string; concurrencyKey?: string }> = []; let resolvedTriggerSource: "SCHEDULED" | "AGENT" | "WEBHOOK" | "STANDARD" | undefined; let resolvedTtl: string | null | undefined; try { + const concurrency = task.concurrency; + + if (task.queue?.name) { + assertNotReservedQueueName(task.queue.name, `Task "${task.id}"`); + } + + if (concurrency && typeof task.queue?.concurrencyLimit === "number") { + throw new ServiceValidationError( + `Task "${task.id}" declares both a queue concurrencyLimit and the concurrency option; use concurrency.` + ); + } + + compiledGates = (concurrency?.limits ?? []).map((name) => ({ + queue: concurrencyLimitQueueName(name), + })); + + let queueConcurrencyLimit = task.queue?.concurrencyLimit; + let queueTotalConcurrencyLimit = task.queue?.combinedConcurrencyLimit; + + if (concurrency?.inline) { + if (!task.queue?.name) { + queueConcurrencyLimit = concurrency.inline.perKey ?? concurrency.inline.total; + queueTotalConcurrencyLimit = concurrency.inline.total; + } else { + if (compiledGates.length > 1) { + throw new ServiceValidationError( + `Task "${task.id}": an inline limit on a shared queue uses a gate slot, so at most one named limit can be combined with it.` + ); + } + const anonymousQueueName = anonymousConcurrencyLimitQueueName(task.id); + await createWorkerQueue( + { + name: anonymousQueueName, + concurrencyLimit: concurrency.inline.perKey ?? concurrency.inline.total ?? null, + combinedConcurrencyLimit: concurrency.inline.total ?? null, + }, + `task/${task.id}`, + "NAMED", + worker, + environment, + prisma, + "LIMIT", + "V2" + ); + compiledGates = [{ queue: anonymousQueueName }, ...compiledGates]; + } + } + queue = queues.find((queue) => queue.name === task.queue?.name); if (!queue) { @@ -402,14 +460,22 @@ async function createWorkerTask( queue = await createWorkerQueue( { name: task.queue?.name ?? `task/${task.id}`, - concurrencyLimit: task.queue?.concurrencyLimit, - combinedConcurrencyLimit: task.queue?.combinedConcurrencyLimit, + concurrencyLimit: queueConcurrencyLimit, + combinedConcurrencyLimit: queueTotalConcurrencyLimit, }, task.queue?.name ?? task.id, task.queue?.name ? "NAMED" : "VIRTUAL", worker, environment, - prisma + prisma, + "QUEUE", + /** + * V2 marks rows whose limit fields hold the new perKey/total vocabulary. That + * only happens when an inline limit compiles into the task's own default queue; + * a shared queue keeps V1 regardless of task concurrency (which lives in gates + * and LIMIT rows), matching rows materialized from queue() declarations. + */ + concurrency?.inline && !task.queue?.name ? "V2" : "V1" ); } @@ -437,7 +503,7 @@ async function createWorkerTask( exportName: task.exportName, retryConfig: task.retry, queueConfig: task.queue, - gates: task.gates, + gates: compiledGates.length > 0 ? compiledGates : task.gates, machineConfig: task.machine, triggerSource: resolvedTriggerSource, config: task.agentConfig ? (task.agentConfig as any) : undefined, @@ -455,9 +521,12 @@ async function createWorkerTask( triggerSource: resolvedTriggerSource, queueId: queue.id, queueName: queue.name, - gates: task.gates ?? null, + gates: compiledGates.length > 0 ? compiledGates : (task.gates ?? null), }; } catch (error) { + if (error instanceof ServiceValidationError) { + throw error; + } if (error instanceof Prisma.PrismaClientKnownRequestError) { // The error code for unique constraint violation in Prisma is P2002 if (error.code === "P2002") { @@ -479,7 +548,7 @@ async function createWorkerTask( triggerSource: resolvedTriggerSource, queueId: queue.id, queueName: queue.name, - gates: task.gates ?? null, + gates: compiledGates.length > 0 ? compiledGates : (task.gates ?? null), }; } } else { @@ -531,6 +600,7 @@ async function createWorkerQueues( const chunk = metadata.queues.slice(i, i + CHUNK_SIZE); const queueChunk = await Promise.all( chunk.map(async (queue) => { + assertNotReservedQueueName(queue.name, `Queue "${queue.name}"`); return createWorkerQueue(queue, queue.name, "NAMED", worker, environment, prisma); }) ); @@ -540,13 +610,146 @@ async function createWorkerQueues( return allQueues; } +/** + * Rejects invalid concurrency declarations before any worker rows are written, so a + * failed deploy leaves nothing behind for a same-content retry to return. + */ +export function validateWorkerConcurrencyDeclarations(metadata: BackgroundWorkerMetadata): void { + for (const queue of metadata.queues ?? []) { + assertNotReservedQueueName(queue.name, `Queue "${queue.name}"`); + } + + for (const limit of metadata.concurrencyLimits ?? []) { + assertValidConcurrencyLimitName(limit.name); + } + + for (const task of metadata.tasks) { + if (task.queue?.name) { + assertNotReservedQueueName(task.queue.name, `Task "${task.id}"`); + } + + const concurrency = task.concurrency; + if (!concurrency) { + continue; + } + + if (typeof task.queue?.concurrencyLimit === "number") { + throw new ServiceValidationError( + `Task "${task.id}" declares both a queue concurrencyLimit and the concurrency option; use concurrency.` + ); + } + + for (const name of concurrency.limits ?? []) { + assertValidConcurrencyLimitName(name); + } + + if (concurrency.inline && task.queue?.name) { + if ((concurrency.limits ?? []).length > 1) { + throw new ServiceValidationError( + `Task "${task.id}": an inline limit on a shared queue uses a gate slot, so at most one named limit can be combined with it.` + ); + } + } + } +} + +/** Queue rows that back named concurrency limits live under this reserved prefix so + * they can never collide with a user's queue names. */ +const CONCURRENCY_LIMIT_QUEUE_PREFIX = "limit/"; + +const CONCURRENCY_LIMIT_NAME_MAX_LENGTH = 128 - CONCURRENCY_LIMIT_QUEUE_PREFIX.length; + +/** + * Named limits require a strict charset so distinct declared names can never merge + * onto one row after queue-name sanitization (which strips disallowed characters). + */ +function assertValidConcurrencyLimitName(name: string): void { + if (!new RegExp(`^[a-zA-Z0-9_-]{1,${CONCURRENCY_LIMIT_NAME_MAX_LENGTH}}$`).test(name)) { + throw new ServiceValidationError( + `Concurrency limit name "${name}" must be 1-${CONCURRENCY_LIMIT_NAME_MAX_LENGTH} characters using only letters, numbers, underscores and hyphens.` + ); + } +} + +function concurrencyLimitQueueName(limitName: string): string { + assertValidConcurrencyLimitName(limitName); + return `${CONCURRENCY_LIMIT_QUEUE_PREFIX}${limitName}`; +} + +/** + * Row name for a task's anonymous inline limit. Task ids are not charset-restricted, + * so when sanitization would be lossy (or the name would overflow the 128-char queue + * name limit) a hash of the raw id keeps distinct task ids on distinct rows. + */ +function anonymousConcurrencyLimitQueueName(taskId: string): string { + const sanitized = sanitizeQueueName(taskId); + const name = `${CONCURRENCY_LIMIT_QUEUE_PREFIX}task/${sanitized}`; + if (sanitized === taskId && name.length <= 128) { + return name; + } + const hash = createHash("sha256").update(taskId).digest("hex").slice(0, 8); + const budget = 128 - `${CONCURRENCY_LIMIT_QUEUE_PREFIX}task/`.length - hash.length - 1; + return `${CONCURRENCY_LIMIT_QUEUE_PREFIX}task/${sanitized.slice(0, budget)}-${hash}`; +} + +/** User queue names may not claim the reserved limit/ namespace. */ +function assertNotReservedQueueName(name: string, context: string): void { + if (sanitizeQueueName(name).startsWith(CONCURRENCY_LIMIT_QUEUE_PREFIX)) { + throw new ServiceValidationError( + `${context}: queue names starting with "${CONCURRENCY_LIMIT_QUEUE_PREFIX}" are reserved for concurrency limits.` + ); + } +} + +/** + * Materializes the worker's declared named concurrency limits (plus any names tasks + * reference without declaring, created uncapped) as LIMIT-role TaskQueue rows. A + * total-only limit stores the total as its per-key limit too, so no single key (or + * the keyless pool) can exceed it even before the group check applies. + */ +async function createWorkerConcurrencyLimits( + metadata: BackgroundWorkerMetadata, + worker: BackgroundWorker, + environment: AuthenticatedEnvironment, + prisma: PrismaClientOrTransaction +) { + const declared = new Map((metadata.concurrencyLimits ?? []).map((l) => [l.name, l])); + + for (const task of metadata.tasks) { + for (const name of task.concurrency?.limits ?? []) { + if (!declared.has(name)) { + declared.set(name, { name }); + } + } + } + + for (const limit of declared.values()) { + await createWorkerQueue( + { + name: concurrencyLimitQueueName(limit.name), + concurrencyLimit: limit.perKey ?? limit.total ?? null, + combinedConcurrencyLimit: limit.total ?? null, + }, + limit.name, + "NAMED", + worker, + environment, + prisma, + "LIMIT", + "V2" + ); + } +} + async function createWorkerQueue( queue: QueueManifest, orderableName: string, queueType: TaskQueueType, worker: BackgroundWorker, environment: AuthenticatedEnvironment, - prisma: PrismaClientOrTransaction + prisma: PrismaClientOrTransaction, + role: TaskQueueRole = "QUEUE", + concurrencyVersion: TaskQueueConcurrencyVersion = "V1" ) { let queueName = sanitizeQueueName(queue.name); @@ -562,56 +765,96 @@ async function createWorkerQueue( orderableName, queueType, worker, - prisma + prisma, + 0, + role, + concurrencyVersion ); - const newConcurrencyLimit = taskQueue.concurrencyLimit; - - /** - * The total limit key is separate from the per-queue limit key that pause zeroes, - * so it is safe to sync it regardless of the paused state. The engine clamps it - * to the environment limit at read time, so the raw declared value is stored. - */ - if (typeof taskQueue.totalConcurrencyLimit === "number") { - await updateQueueTotalConcurrencyLimits( - environment, - taskQueue.name, - taskQueue.totalConcurrencyLimit - ); - } else { - await removeQueueTotalConcurrencyLimits(environment, taskQueue.name); - } + const syncQueueLimitsToEngine = async (row: { + name: string; + paused: boolean; + concurrencyLimit: number | null; + totalConcurrencyLimit: number | null; + }) => { + /** + * The total limit key is separate from the per-queue limit key that pause zeroes, + * so it is safe to sync it regardless of the paused state. The engine clamps it + * to the environment limit at read time, so the raw declared value is stored. + */ + if (typeof row.totalConcurrencyLimit === "number") { + await updateQueueTotalConcurrencyLimits(environment, row.name, row.totalConcurrencyLimit); + } else { + await removeQueueTotalConcurrencyLimits(environment, row.name); + } - if (!taskQueue.paused) { - if (typeof newConcurrencyLimit === "number") { - logger.debug("createWorkerQueue: updating concurrency limit", { + if (!row.paused) { + logger.debug("createWorkerQueue: syncing concurrency limit", { workerId: worker.id, - taskQueue, + taskQueue: row, orgId: environment.organizationId, projectId: environment.projectId, environmentId: environment.id, - concurrencyLimit: newConcurrencyLimit, + concurrencyLimit: row.concurrencyLimit, }); - await updateQueueConcurrencyLimits(environment, taskQueue.name, newConcurrencyLimit); + if (typeof row.concurrencyLimit === "number") { + await updateQueueConcurrencyLimits(environment, row.name, row.concurrencyLimit); + } else { + await removeQueueConcurrencyLimits(environment, row.name); + } } else { - logger.debug("createWorkerQueue: removing concurrency limit", { + logger.debug("createWorkerQueue: queue is paused, not updating concurrency limit", { workerId: worker.id, - taskQueue, + taskQueue: row, orgId: environment.organizationId, projectId: environment.projectId, environmentId: environment.id, - concurrencyLimit: newConcurrencyLimit, }); - await removeQueueConcurrencyLimits(environment, taskQueue.name); } - } else { - logger.debug("createWorkerQueue: queue is paused, not updating concurrency limit", { - workerId: worker.id, - taskQueue, - orgId: environment.organizationId, - projectId: environment.projectId, - environmentId: environment.id, + }; + + await syncQueueLimitsToEngine(taskQueue); + + /** + * The optimistic markers only guard the Postgres write; an override or reset can + * still land between that write and the engine sync above, which would leave the + * engine holding this deploy's stale values. Re-read the markers and re-sync from + * the fresh row until they stop moving (bounded): every actor writes Postgres + * before its own engine sync, so re-syncing whatever is freshest converges. A + * marker moving after the final read is healed by that actor's own engine sync + * or the next deploy. + */ + let syncedMarkers = { + concurrency: taskQueue.concurrencyLimitOverriddenAt?.getTime(), + total: taskQueue.totalConcurrencyLimitOverriddenAt?.getTime(), + }; + + for (let i = 0; i < 3; i++) { + const freshQueue = await prisma.taskQueue.findFirst({ + where: { id: taskQueue.id }, + select: { + name: true, + paused: true, + concurrencyLimit: true, + totalConcurrencyLimit: true, + concurrencyLimitOverriddenAt: true, + totalConcurrencyLimitOverriddenAt: true, + }, }); + + if ( + !freshQueue || + (freshQueue.concurrencyLimitOverriddenAt?.getTime() === syncedMarkers.concurrency && + freshQueue.totalConcurrencyLimitOverriddenAt?.getTime() === syncedMarkers.total) + ) { + break; + } + + await syncQueueLimitsToEngine(freshQueue); + syncedMarkers = { + concurrency: freshQueue.concurrencyLimitOverriddenAt?.getTime(), + total: freshQueue.totalConcurrencyLimitOverriddenAt?.getTime(), + }; } return taskQueue; @@ -625,7 +868,9 @@ async function upsertWorkerQueueRecord( queueType: TaskQueueType, worker: BackgroundWorker, prisma: PrismaClientOrTransaction, - attempt: number = 0 + attempt: number = 0, + role: TaskQueueRole = "QUEUE", + concurrencyVersion: TaskQueueConcurrencyVersion = "V1" ): Promise { if (attempt > 3) { throw new Error("Failed to insert queue record"); @@ -644,6 +889,8 @@ async function upsertWorkerQueueRecord( data: { friendlyId: generateFriendlyId("queue"), version: "V2", + role, + concurrencyVersion, name: queueName, orderableName, concurrencyLimit, @@ -660,27 +907,44 @@ async function upsertWorkerQueueRecord( }); } else { const hasOverride = taskQueue.concurrencyLimitOverriddenAt !== null; - + const hasTotalOverride = taskQueue.totalConcurrencyLimitOverriddenAt !== null; + + /** + * The override markers in the where clause make this an optimistic-concurrency + * update: a concurrent override/reset between the read above and this write + * changes a marker, the update misses (P2025) and the catch below retries with + * a fresh read, so a deploy can never clobber an operator's override. + */ taskQueue = await prisma.taskQueue.update({ where: { id: taskQueue.id, + concurrencyLimitOverriddenAt: taskQueue.concurrencyLimitOverriddenAt, + totalConcurrencyLimitOverriddenAt: taskQueue.totalConcurrencyLimitOverriddenAt, }, data: { workers: { connect: { id: worker.id } }, version: "V2", + concurrencyVersion, orderableName, // If overridden, keep current limit and update base; otherwise update limit normally concurrencyLimit: hasOverride ? undefined : concurrencyLimit, concurrencyLimitBase: hasOverride ? concurrencyLimit : undefined, - totalConcurrencyLimit, + totalConcurrencyLimit: hasTotalOverride ? undefined : totalConcurrencyLimit, + totalConcurrencyLimitBase: hasTotalOverride ? totalConcurrencyLimit : undefined, }, }); } return taskQueue; } catch (error) { - // If the queue already exists, let's try again - if (error instanceof Prisma.PrismaClientKnownRequestError && error.code === "P2002") { + /** + * P2002: the queue was created concurrently. P2025: an override/reset moved a + * marker under the optimistic update. Both re-read and retry. + */ + if ( + error instanceof Prisma.PrismaClientKnownRequestError && + (error.code === "P2002" || error.code === "P2025") + ) { return await upsertWorkerQueueRecord( queueName, concurrencyLimit, @@ -689,7 +953,9 @@ async function upsertWorkerQueueRecord( queueType, worker, prisma, - attempt + 1 + attempt + 1, + role, + concurrencyVersion ); } throw error; diff --git a/apps/webapp/app/v3/services/createDeploymentBackgroundWorkerV4.server.ts b/apps/webapp/app/v3/services/createDeploymentBackgroundWorkerV4.server.ts index bf2c95745f3..f6100a00b7c 100644 --- a/apps/webapp/app/v3/services/createDeploymentBackgroundWorkerV4.server.ts +++ b/apps/webapp/app/v3/services/createDeploymentBackgroundWorkerV4.server.ts @@ -15,6 +15,7 @@ import { createWorkerResources, syncDeclarativeSchedules, syncDeclarativeWebhooks, + validateWorkerConcurrencyDeclarations, } from "./createBackgroundWorker.server"; import { findOrCreateBackgroundWorker } from "./createDeploymentBackgroundWorkerV4/findOrCreateBackgroundWorker.server"; import { TimeoutDeploymentService } from "./timeoutDeployment.server"; @@ -67,6 +68,23 @@ export class CreateDeploymentBackgroundWorkerServiceV4 extends BaseService { return; } + /** + * Reject invalid concurrency declarations before any worker rows exist. Queue + * rows and their engine limit keys are per-environment, so a mid-creation + * failure would leave the running version's limits already mutated. + */ + try { + validateWorkerConcurrencyDeclarations(body.metadata); + } catch (concurrencyError) { + if (concurrencyError instanceof ServiceValidationError) { + logger.warn("Invalid worker concurrency declarations", { + error: concurrencyError.message, + }); + await this.#failBackgroundWorkerDeployment(deployment, concurrencyError, environment); + } + throw concurrencyError; + } + // Handle multi-platform builds const deploymentPlatforms = deployment.imagePlatform?.split(",") ?? []; if (deploymentPlatforms.length > 1) { diff --git a/internal-packages/database/prisma/migrations/20260829150000_add_concurrency_overrides/migration.sql b/internal-packages/database/prisma/migrations/20260829150000_add_concurrency_overrides/migration.sql new file mode 100644 index 00000000000..4b8b6ec6077 --- /dev/null +++ b/internal-packages/database/prisma/migrations/20260829150000_add_concurrency_overrides/migration.sql @@ -0,0 +1,4 @@ +-- AlterTable +ALTER TABLE "TaskQueue" ADD COLUMN "totalConcurrencyLimitOverriddenAt" TIMESTAMP(3); +ALTER TABLE "TaskQueue" ADD COLUMN "totalConcurrencyLimitOverriddenBy" TEXT; +ALTER TABLE "TaskQueue" ADD COLUMN "totalConcurrencyLimitBase" INTEGER; diff --git a/internal-packages/database/prisma/migrations/20260906140000_add_task_queue_concurrency_version_and_role/migration.sql b/internal-packages/database/prisma/migrations/20260906140000_add_task_queue_concurrency_version_and_role/migration.sql new file mode 100644 index 00000000000..faff8327eca --- /dev/null +++ b/internal-packages/database/prisma/migrations/20260906140000_add_task_queue_concurrency_version_and_role/migration.sql @@ -0,0 +1,9 @@ +-- CreateEnum +CREATE TYPE "TaskQueueConcurrencyVersion" AS ENUM ('V1', 'V2'); + +-- CreateEnum +CREATE TYPE "TaskQueueRole" AS ENUM ('QUEUE', 'LIMIT'); + +-- AlterTable +ALTER TABLE "TaskQueue" ADD COLUMN "concurrencyVersion" "TaskQueueConcurrencyVersion" NOT NULL DEFAULT 'V1', +ADD COLUMN "role" "TaskQueueRole" NOT NULL DEFAULT 'QUEUE'; diff --git a/internal-packages/database/prisma/schema.prisma b/internal-packages/database/prisma/schema.prisma index c88ba5b5887..b0621525231 100644 --- a/internal-packages/database/prisma/schema.prisma +++ b/internal-packages/database/prisma/schema.prisma @@ -1963,6 +1963,13 @@ model TaskQueue { version TaskQueueVersion @default(V1) orderableName String? + /// How the queue's concurrency was declared: V1 = legacy concurrencyLimit semantics + /// (per key when keyed, whole queue otherwise), V2 = the explicit perKey/total shape. + concurrencyVersion TaskQueueConcurrencyVersion @default(V1) + /// QUEUE rows are real queues; LIMIT rows back named concurrency limits and are + /// excluded from every queue-facing surface. + role TaskQueueRole @default(QUEUE) + project Project @relation(fields: [projectId], references: [id], onDelete: Cascade, onUpdate: Cascade) projectId String @@ -1983,7 +1990,13 @@ model TaskQueue { concurrencyLimitOverridePercent Decimal? @db.Decimal(5, 2) /// Caps total concurrent runs across ALL concurrencyKey values of this queue /// (concurrencyLimit applies per key value). Null = no total cap. - totalConcurrencyLimit Int? + totalConcurrencyLimit Int? + /// When the total concurrency limit was overridden + totalConcurrencyLimitOverriddenAt DateTime? + /// Who overrode the total concurrency limit (null when overridden via the API) + totalConcurrencyLimitOverriddenBy String? + /// If totalConcurrencyLimit is overridden, the declared value it reverts to on reset + totalConcurrencyLimitBase Int? rateLimit Json? paused Boolean @default(false) @@ -1995,9 +2008,11 @@ model TaskQueue { tasks BackgroundWorkerTask[] workers BackgroundWorker[] + @@unique([runtimeEnvironmentId, name]) } + enum TaskQueueType { VIRTUAL NAMED @@ -2008,6 +2023,16 @@ enum TaskQueueVersion { V2 } +enum TaskQueueConcurrencyVersion { + V1 + V2 +} + +enum TaskQueueRole { + QUEUE + LIMIT +} + model BatchTaskRun { id String @id @default(cuid()) friendlyId String @unique diff --git a/internal-packages/run-engine/src/run-queue/index.ts b/internal-packages/run-engine/src/run-queue/index.ts index f8e28d55731..1dcbcb33a88 100644 --- a/internal-packages/run-engine/src/run-queue/index.ts +++ b/internal-packages/run-engine/src/run-queue/index.ts @@ -116,12 +116,18 @@ local function __gateReconcile(setKey, msgKeyPrefix, reconcileKeyPrefix) end end -local function __gatesHaveCapacity(gatesKeyPrefix, msg, messageId, envLimit, msgKeyPrefix) +local function __gatesHaveCapacity(gatesKeyPrefix, msg, messageId, envLimit, msgKeyPrefix, ckOverridesEnabled) if not msg.gates then return true end for _, gate in ipairs(msg.gates) do local base, variant, gateKey = __gateKeys(gatesKeyPrefix, msg, gate) local occupancy = tonumber(redis.call('SCARD', variant .. ':currentConcurrency') or '0') local perKeyLimit = math.min(tonumber(redis.call('GET', base .. ':concurrency') or '1000000'), envLimit) + if ckOverridesEnabled and gateKey and gateKey ~= '' then + local gateOverride = redis.call('HGET', base .. ':ckLimits', string.sub(variant, #gatesKeyPrefix + 1)) + if gateOverride then + perKeyLimit = math.min(tonumber(gateOverride), envLimit) + end + end if occupancy >= perKeyLimit and redis.call('SISMEMBER', variant .. ':currentConcurrency', messageId) == 0 then __gateReconcile(variant .. ':currentConcurrency', msgKeyPrefix, gatesKeyPrefix) return false @@ -255,6 +261,13 @@ export interface RunQueueMetricsEmitter { emitGauge(shardKey: string, fields: Record): void; } +export class RunQueueConcurrencyKeyLimitExceededError extends Error { + constructor(message: string) { + super(message); + this.name = "RunQueueConcurrencyKeyLimitExceededError"; + } +} + export type RunQueueOptions = { name: string; tracer: Tracer; @@ -303,6 +316,10 @@ export type RunQueueOptions = { * that dead-lettered or suspended through a mirror-less path. Enabling only after * every instance runs this build avoids the noise but is no longer load-bearing * for correctness. + * + * Per-concurrency-key limit overrides are part of the same concurrency-limits + * feature and are deliberately enforced behind this flag too: writes are always + * accepted and durable, and enforcement of both arrives together. */ totalConcurrencyEnabled?: boolean; /** @@ -314,6 +331,8 @@ export type RunQueueOptions = { * the total cap covering releases from builds without the mirror. */ gatesEnabled?: boolean; + /** Cap on per-concurrency-key limit overrides stored per queue. Default 1000. */ + maxConcurrencyKeyOverridesPerQueue?: number; workerOptions?: { pollIntervalMs?: number; immediatePollIntervalMs?: number; @@ -427,6 +446,7 @@ export class RunQueue { private queueSelectionStrategy: RunQueueSelectionStrategy; private shardCount: number; private counterTtlSeconds: number; + private maxConcurrencyKeyOverridesPerQueue: number; private abortController: AbortController; private worker: Worker; private workerQueueResolver: WorkerQueueResolver; @@ -437,6 +457,7 @@ export class RunQueue { constructor(public readonly options: RunQueueOptions) { this.shardCount = options.shardCount ?? 2; this.counterTtlSeconds = options.counterTtlSeconds ?? 86400; + this.maxConcurrencyKeyOverridesPerQueue = options.maxConcurrencyKeyOverridesPerQueue ?? 1000; this.retryOptions = options.retryOptions ?? defaultRetrySettings; this.redis = createRedisClient(options.redis, { onError: (error) => { @@ -631,6 +652,62 @@ export class RunQueue { return this.redis.scard(this.keys.queueGroupConcurrencyKey(env, queue)); } + /** + * Sets a per-concurrency-key limit override for a queue. The stored value is the + * raw requested limit; admit paths clamp to the environment limit at read time. + * Throws RunQueueConcurrencyKeyLimitExceededError when a NEW key would push the + * queue past maxConcurrencyKeyOverridesPerQueue (updates to existing keys always + * succeed). + */ + public async updateQueueConcurrencyKeyLimit( + env: MinimalAuthenticatedEnvironment, + queue: string, + concurrencyKey: string, + limit: number + ) { + const result = await this.redis.setQueueConcurrencyKeyLimit( + this.keys.queueCkLimitsKey(env, queue), + this.keys.queueKey(env, queue, concurrencyKey), + String(limit), + String(this.maxConcurrencyKeyOverridesPerQueue) + ); + + if (result === 0) { + throw new RunQueueConcurrencyKeyLimitExceededError( + `Cannot add a concurrency key override to queue ${queue}: the queue already has ${this.maxConcurrencyKeyOverridesPerQueue} overrides` + ); + } + } + + public async removeQueueConcurrencyKeyLimit( + env: MinimalAuthenticatedEnvironment, + queue: string, + concurrencyKey: string + ) { + return this.redis.hdel( + this.keys.queueCkLimitsKey(env, queue), + this.keys.queueKey(env, queue, concurrencyKey) + ); + } + + /** Returns the raw per-concurrency-key limit overrides for a queue, keyed by concurrency key value. */ + public async getQueueConcurrencyKeyLimits( + env: MinimalAuthenticatedEnvironment, + queue: string + ): Promise> { + const raw = await this.redis.hgetall(this.keys.queueCkLimitsKey(env, queue)); + + const limits: Record = {}; + for (const [variantName, value] of Object.entries(raw)) { + const ckIndex = variantName.indexOf(":ck:"); + if (ckIndex === -1) { + continue; + } + limits[variantName.slice(ckIndex + 4)] = Number(value); + } + return limits; + } + public async updateEnvConcurrencyLimits(env: MinimalAuthenticatedEnvironment) { await this.#callUpdateEnvironmentConcurrencyLimits({ envConcurrencyLimitKey: this.keys.envConcurrencyLimitKey(env), @@ -2380,6 +2457,7 @@ export class RunQueue { const totalConcurrencyLimitKey = this.keys.queueTotalConcurrencyLimitKeyFromQueue( message.queue ); + const ckLimitsKey = this.keys.queueCkLimitsKeyFromQueue(message.queue); const totalConcurrencyEnabledArg = this.options.totalConcurrencyEnabled ? "1" : "0"; if (ttlInfo) { @@ -2403,6 +2481,7 @@ export class RunQueue { baseQueueKey, groupConcurrencyKey, totalConcurrencyLimitKey, + ckLimitsKey, // args queueName, messageId, @@ -2442,6 +2521,7 @@ export class RunQueue { baseQueueKey, groupConcurrencyKey, totalConcurrencyLimitKey, + ckLimitsKey, // args queueName, messageId, @@ -2491,6 +2571,7 @@ export class RunQueue { enableFastPathArg, this.options.redis.keyPrefix ?? "", this.options.gatesEnabled ? "1" : "0", + this.options.totalConcurrencyEnabled ? "1" : "0", metricsGaugeArg ); } else { @@ -2520,6 +2601,7 @@ export class RunQueue { enableFastPathArg, this.options.redis.keyPrefix ?? "", this.options.gatesEnabled ? "1" : "0", + this.options.totalConcurrencyEnabled ? "1" : "0", metricsGaugeArg ); } @@ -2601,6 +2683,7 @@ export class RunQueue { this.options.redis.keyPrefix ?? "", String(maxCount), this.options.gatesEnabled ? "1" : "0", + this.options.totalConcurrencyEnabled ? "1" : "0", metricsGaugeArg ); @@ -2729,6 +2812,7 @@ export class RunQueue { runningCounterKey, this.keys.queueGroupConcurrencyKeyFromQueue(ckWildcardQueue), this.keys.queueTotalConcurrencyLimitKeyFromQueue(ckWildcardQueue), + this.keys.queueCkLimitsKeyFromQueue(ckWildcardQueue), //args ckWildcardQueue, String(Date.now()), @@ -3597,6 +3681,7 @@ local currentTime = ARGV[8] local enableFastPath = ARGV[9] local keyPrefix = ARGV[10] local gatesEnabled = ARGV[11] == '1' +local totalConcurrencyEnabled = ARGV[12] == '1' ${QUEUE_METRICS_GAUGE_PRELUDE} ${QUEUE_GATES_LUA_HELPERS} @@ -3624,7 +3709,7 @@ if enableFastPath == '1' then local okDecode, decoded = pcall(cjson.decode, messageData) if okDecode and type(decoded) == 'table' and decoded.gates then gateMsg = decoded - gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, messageId, envLimit, nil) + gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, messageId, envLimit, nil, totalConcurrencyEnabled) end end @@ -3711,6 +3796,7 @@ local currentTime = ARGV[10] local enableFastPath = ARGV[11] local keyPrefix = ARGV[12] local gatesEnabled = ARGV[13] == '1' +local totalConcurrencyEnabled = ARGV[14] == '1' ${QUEUE_METRICS_GAUGE_PRELUDE} ${QUEUE_GATES_LUA_HELPERS} @@ -3738,7 +3824,7 @@ if enableFastPath == '1' then local okDecode, decoded = pcall(cjson.decode, messageData) if okDecode and type(decoded) == 'table' and decoded.gates then gateMsg = decoded - gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, messageId, envLimit, nil) + gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, messageId, envLimit, nil, totalConcurrencyEnabled) end end @@ -4017,7 +4103,7 @@ return __qmret(0) // *Tracked variants of dequeueMessageFromKey and the ack/nack/dlq/release/clear // scripts. this.redis.defineCommand("enqueueMessageCkTracked", { - numberOfKeys: 17, + numberOfKeys: 18, lua: ` local masterQueueKey = KEYS[1] local queueKey = KEYS[2] @@ -4039,6 +4125,7 @@ local baseQueueKey = KEYS[15] -- Total-cap keys (KEYS 16-17) local groupConcurrencyKey = KEYS[16] local totalConcurrencyLimitKey = KEYS[17] +local ckLimitsKey = KEYS[18] local queueName = ARGV[1] local messageId = ARGV[2] @@ -4076,6 +4163,12 @@ if enableFastPath == '1' then tonumber(redis.call('GET', queueConcurrencyLimitKey) or '1000000'), envLimit ) + if totalConcurrencyEnabled then + local perKeyOverride = redis.call('HGET', ckLimitsKey, queueName) + if perKeyOverride then + queueLimit = math.min(tonumber(perKeyOverride), envLimit) + end + end if queueCurrent < queueLimit then -- Total-cap gate: a fast-path admit consumes a group slot, so it must @@ -4098,7 +4191,7 @@ if enableFastPath == '1' then local okDecode, decoded = pcall(cjson.decode, messageData) if okDecode and type(decoded) == 'table' and decoded.gates then gateMsg = decoded - gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, messageId, envLimit, nil) + gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, messageId, envLimit, nil, totalConcurrencyEnabled) end end @@ -4187,7 +4280,7 @@ return __qmret(0) }); this.redis.defineCommand("enqueueMessageWithTtlCkTracked", { - numberOfKeys: 18, + numberOfKeys: 19, lua: ` local masterQueueKey = KEYS[1] local queueKey = KEYS[2] @@ -4210,6 +4303,7 @@ local baseQueueKey = KEYS[16] -- Total-cap keys (KEYS 17-18) local groupConcurrencyKey = KEYS[17] local totalConcurrencyLimitKey = KEYS[18] +local ckLimitsKey = KEYS[19] local queueName = ARGV[1] local messageId = ARGV[2] @@ -4249,6 +4343,12 @@ if enableFastPath == '1' then tonumber(redis.call('GET', queueConcurrencyLimitKey) or '1000000'), envLimit ) + if totalConcurrencyEnabled then + local perKeyOverride = redis.call('HGET', ckLimitsKey, queueName) + if perKeyOverride then + queueLimit = math.min(tonumber(perKeyOverride), envLimit) + end + end if queueCurrent < queueLimit then -- Total-cap gate: see enqueueMessageCkTracked. @@ -4269,7 +4369,7 @@ if enableFastPath == '1' then local okDecode, decoded = pcall(cjson.decode, messageData) if okDecode and type(decoded) == 'table' and decoded.gates then gateMsg = decoded - gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, messageId, envLimit, nil) + gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, messageId, envLimit, nil, totalConcurrencyEnabled) end end @@ -4599,6 +4699,7 @@ local defaultEnvConcurrencyBurstFactor = ARGV[4] local keyPrefix = ARGV[5] local maxCount = tonumber(ARGV[6] or '1') local gatesEnabled = ARGV[7] == '1' +local totalConcurrencyEnabled = ARGV[8] == '1' ${QUEUE_METRICS_GAUGE_PRELUDE} ${QUEUE_GATES_LUA_HELPERS} ${QUEUE_METRICS_GAUGE_LUA} @@ -4671,7 +4772,7 @@ for i = 1, #messages, 2 do else local gatesAllow = true if gatesEnabled then - gatesAllow = __gatesHaveCapacity(keyPrefix, messageData, messageId, envConcurrencyLimit, messageKeyPrefix) + gatesAllow = __gatesHaveCapacity(keyPrefix, messageData, messageId, envConcurrencyLimit, messageKeyPrefix, totalConcurrencyEnabled) end if gatesAllow then @@ -4873,7 +4974,7 @@ return results // (normal dequeue, TTL-expired, or stale-orphan path — all of which were // counted at enqueue time). this.redis.defineCommand("dequeueMessagesFromCkQueueTracked", { - numberOfKeys: 13, + numberOfKeys: 14, lua: ` local ckIndexKey = KEYS[1] local queueConcurrencyLimitKey = KEYS[2] @@ -4888,6 +4989,7 @@ local lengthCounterKey = KEYS[10] local runningCounterKey = KEYS[11] local groupConcurrencyKey = KEYS[12] local totalConcurrencyLimitKey = KEYS[13] +local ckLimitsKey = KEYS[14] local ckWildcardName = ARGV[1] local currentTime = tonumber(ARGV[2]) @@ -4973,11 +5075,28 @@ for _, ckQueueName in ipairs(ckQueues) do end local fullQueueKey = keyPrefix .. ckQueueName + local blockedByGates = false local ckConcurrencyKey = fullQueueKey .. ':currentConcurrency' local ckCurrentConcurrency = tonumber(redis.call('SCARD', ckConcurrencyKey) or '0') - if ckCurrentConcurrency < queueConcurrencyLimit then + local perKeyLimit = queueConcurrencyLimit + if totalConcurrencyEnabled then + local perKeyOverride = redis.call('HGET', ckLimitsKey, ckQueueName) + if perKeyOverride then + perKeyLimit = math.min(tonumber(perKeyOverride), envConcurrencyLimit) + end + end + + if ckCurrentConcurrency >= perKeyLimit then + -- Back a blocked variant off so it cannot pin the bounded candidate window + -- and starve later keys (acute with a zero per-key override, which never + -- self-clears). Acks and nacks rebalance the score back to the oldest + -- message, so the key is eligible again the moment capacity frees. + redis.call('ZADD', ckIndexKey, currentTime + 1000, ckQueueName) + end + + if ckCurrentConcurrency < perKeyLimit then local messages = redis.call('ZRANGEBYSCORE', fullQueueKey, '-inf', tostring(currentTime), 'WITHSCORES', 'LIMIT', 0, 1) if #messages >= 2 then @@ -5002,7 +5121,10 @@ for _, ckQueueName in ipairs(ckQueues) do else local gatesAllow = true if gatesEnabled then - gatesAllow = __gatesHaveCapacity(keyPrefix, messageData, messageId, envConcurrencyLimit, messageKeyPrefix) + gatesAllow = __gatesHaveCapacity(keyPrefix, messageData, messageId, envConcurrencyLimit, messageKeyPrefix, totalConcurrencyEnabled) + end + if not gatesAllow then + blockedByGates = true end local alreadyInGroup = false @@ -5046,11 +5168,15 @@ for _, ckQueueName in ipairs(ckQueues) do decrLengthCounter() end - local earliest = redis.call('ZRANGE', fullQueueKey, 0, 0, 'WITHSCORES') - if #earliest == 0 then - redis.call('ZREM', ckIndexKey, ckQueueName) + if blockedByGates then + redis.call('ZADD', ckIndexKey, currentTime + 1000, ckQueueName) else - redis.call('ZADD', ckIndexKey, earliest[2], ckQueueName) + local earliest = redis.call('ZRANGE', fullQueueKey, 0, 0, 'WITHSCORES') + if #earliest == 0 then + redis.call('ZREM', ckIndexKey, ckQueueName) + else + redis.call('ZADD', ckIndexKey, earliest[2], ckQueueName) + end end else local any = redis.call('ZRANGE', fullQueueKey, 0, 0, 'WITHSCORES') @@ -5893,6 +6019,26 @@ __gatesRelease(keyPrefix, redis.call('GET', messageKey), messageId) `, }); + this.redis.defineCommand("setQueueConcurrencyKeyLimit", { + numberOfKeys: 1, + lua: ` +local ckLimitsKey = KEYS[1] + +local fieldName = ARGV[1] +local limit = ARGV[2] +local maxFields = tonumber(ARGV[3]) + +if redis.call('HEXISTS', ckLimitsKey, fieldName) == 0 then + if redis.call('HLEN', ckLimitsKey) >= maxFields then + return 0 + end +end + +redis.call('HSET', ckLimitsKey, fieldName, limit) +return 1 +`, + }); + this.redis.defineCommand("updateEnvironmentConcurrencyLimits", { numberOfKeys: 2, lua: ` @@ -6083,6 +6229,7 @@ declare module "@internal/redis" { enableFastPath: string, keyPrefix: string, gatesEnabled: string, + totalConcurrencyEnabled: string, metricsEnabled: string, callback?: Callback<[number, number[] | null]> ): Result<[number, number[] | null], Context>; @@ -6116,6 +6263,7 @@ declare module "@internal/redis" { enableFastPath: string, keyPrefix: string, gatesEnabled: string, + totalConcurrencyEnabled: string, metricsEnabled: string, callback?: Callback<[number, number[] | null]> ): Result<[number, number[] | null], Context>; @@ -6154,6 +6302,7 @@ declare module "@internal/redis" { keyPrefix: string, maxCount: string, gatesEnabled: string, + totalConcurrencyEnabled: string, metricsEnabled: string, callback?: Callback<[string[] | null, number[] | null]> ): Result<[string[] | null, number[] | null], Context>; @@ -6254,6 +6403,14 @@ declare module "@internal/redis" { callback?: Callback ): Result; + setQueueConcurrencyKeyLimit( + ckLimitsKey: string, + fieldName: string, + limit: string, + maxFields: string, + callback?: Callback + ): Result; + updateEnvironmentConcurrencyLimits( // keys envConcurrencyLimitKey: string, @@ -6440,6 +6597,7 @@ declare module "@internal/redis" { baseQueueKey: string, groupConcurrencyKey: string, totalConcurrencyLimitKey: string, + ckLimitsKey: string, queueName: string, messageId: string, messageData: string, @@ -6477,6 +6635,7 @@ declare module "@internal/redis" { baseQueueKey: string, groupConcurrencyKey: string, totalConcurrencyLimitKey: string, + ckLimitsKey: string, queueName: string, messageId: string, messageData: string, @@ -6511,6 +6670,7 @@ declare module "@internal/redis" { runningCounterKey: string, groupConcurrencyKey: string, totalConcurrencyLimitKey: string, + ckLimitsKey: string, ckWildcardName: string, currentTime: string, defaultEnvConcurrencyLimit: string, diff --git a/internal-packages/run-engine/src/run-queue/keyProducer.ts b/internal-packages/run-engine/src/run-queue/keyProducer.ts index 98028f5af7b..120e04f8c38 100644 --- a/internal-packages/run-engine/src/run-queue/keyProducer.ts +++ b/internal-packages/run-engine/src/run-queue/keyProducer.ts @@ -366,6 +366,14 @@ export class RunQueueFullKeyProducer implements RunQueueKeyProducer { return `${this.baseQueueKeyFromQueue(queue)}:${constants.TOTAL_CONCURRENCY_LIMIT_PART}`; } + queueCkLimitsKey(env: RunQueueKeyProducerEnvironment, queue: string): string { + return `${this.queueKey(env, queue)}:ckLimits`; + } + + queueCkLimitsKeyFromQueue(queue: string): string { + return `${this.baseQueueKeyFromQueue(queue)}:ckLimits`; + } + isCkWildcard(queue: string): boolean { return queue.endsWith(":ck:*"); } diff --git a/internal-packages/run-engine/src/run-queue/types.ts b/internal-packages/run-engine/src/run-queue/types.ts index df7cc50b221..f8da1248368 100644 --- a/internal-packages/run-engine/src/run-queue/types.ts +++ b/internal-packages/run-engine/src/run-queue/types.ts @@ -112,6 +112,9 @@ export interface RunQueueKeyProducer { queueTotalConcurrencyLimitKey(env: RunQueueKeyProducerEnvironment, queue: string): string; queueTotalConcurrencyLimitKeyFromQueue(queue: string): string; + queueCkLimitsKey(env: RunQueueKeyProducerEnvironment, queue: string): string; + queueCkLimitsKeyFromQueue(queue: string): string; + //env oncurrency envCurrentConcurrencyKey(env: EnvDescriptor): string; envCurrentConcurrencyKey(env: RunQueueKeyProducerEnvironment): string;