diff --git a/.changeset/queue-total-concurrency-limit.md b/.changeset/queue-total-concurrency-limit.md new file mode 100644 index 00000000000..742f4925744 --- /dev/null +++ b/.changeset/queue-total-concurrency-limit.md @@ -0,0 +1,18 @@ +--- +"@trigger.dev/sdk": patch +"@trigger.dev/core": patch +--- + +Cap a queue's total concurrency across all of its `concurrencyKey` values with the new `totalConcurrencyLimit` queue option. On a keyed queue, `concurrencyLimit` applies to each key value independently, so ten active keys with a limit of 5 can run 50 at once. `totalConcurrencyLimit` bounds the whole queue while each key still gets at most `concurrencyLimit`. + +```ts +import { queue } from "@trigger.dev/sdk"; + +export const perUserQueue = queue({ + name: "per-user-queue", + concurrencyLimit: 1, + totalConcurrencyLimit: 10, +}); +``` + +Enforcement happens server-side and only applies to runs triggered with a `concurrencyKey`. Servers that have not enabled total concurrency limits accept the option but do not enforce it yet. diff --git a/apps/webapp/app/env.server.ts b/apps/webapp/app/env.server.ts index 2b1fba86980..8b660127648 100644 --- a/apps/webapp/app/env.server.ts +++ b/apps/webapp/app/env.server.ts @@ -1377,6 +1377,7 @@ const EnvironmentSchema = z RUN_ENGINE_RUN_QUEUE_LOG_LEVEL: z .enum(["log", "error", "warn", "info", "debug"]) .default("info"), + RUN_ENGINE_TOTAL_CONCURRENCY_LIMITS_ENABLED: z.string().default("0"), RUN_ENGINE_TREAT_PRODUCTION_EXECUTION_STALLS_AS_OOM: z.string().default("0"), RUN_ENGINE_READ_REPLICA_SNAPSHOTS_SINCE_ENABLED: z.string().default("0"), RUN_ENGINE_SNAPSHOTS_SINCE_REPLICA_RETRY_MIN_MS: z.coerce.number().int().default(50), diff --git a/apps/webapp/app/v3/runEngine.server.ts b/apps/webapp/app/v3/runEngine.server.ts index 76ac35f349b..7a6cb0f8d3d 100644 --- a/apps/webapp/app/v3/runEngine.server.ts +++ b/apps/webapp/app/v3/runEngine.server.ts @@ -62,6 +62,7 @@ function createRunEngine() { queue: { defaultEnvConcurrency: env.DEFAULT_ENV_EXECUTION_CONCURRENCY_LIMIT, defaultEnvConcurrencyBurstFactor: env.DEFAULT_ENV_EXECUTION_CONCURRENCY_BURST_FACTOR, + totalConcurrencyEnabled: env.RUN_ENGINE_TOTAL_CONCURRENCY_LIMITS_ENABLED === "1", logLevel: env.RUN_ENGINE_RUN_QUEUE_LOG_LEVEL, redis: { keyPrefix: "engine:", diff --git a/apps/webapp/app/v3/runQueue.server.ts b/apps/webapp/app/v3/runQueue.server.ts index 304819f9f36..2005b059be4 100644 --- a/apps/webapp/app/v3/runQueue.server.ts +++ b/apps/webapp/app/v3/runQueue.server.ts @@ -40,6 +40,23 @@ export async function updateQueueConcurrencyLimits( await engine.runQueue.updateQueueConcurrencyLimits(environment, queueName, concurrency); } +/** Updates the RunQueue total concurrency limit for a queue (the cap across all concurrency-key values) */ +export async function updateQueueTotalConcurrencyLimits( + environment: AuthenticatedEnvironment, + queueName: string, + totalConcurrency: number +) { + await engine.runQueue.updateQueueTotalConcurrencyLimits(environment, queueName, totalConcurrency); +} + +/** Removes the RunQueue total concurrency limit for a queue */ +export async function removeQueueTotalConcurrencyLimits( + environment: AuthenticatedEnvironment, + queueName: string +) { + await engine.runQueue.removeQueueTotalConcurrencyLimits(environment, queueName); +} + /** Removes the RunQueue limits for a queue */ export async function removeQueueConcurrencyLimits( environment: AuthenticatedEnvironment, diff --git a/apps/webapp/app/v3/services/createBackgroundWorker.server.ts b/apps/webapp/app/v3/services/createBackgroundWorker.server.ts index f8a69e83d2b..f1cebe640d5 100644 --- a/apps/webapp/app/v3/services/createBackgroundWorker.server.ts +++ b/apps/webapp/app/v3/services/createBackgroundWorker.server.ts @@ -33,8 +33,10 @@ import { generateFriendlyId } from "../friendlyIdentifiers"; import { engine } from "../runEngine.server"; import { removeQueueConcurrencyLimits, + removeQueueTotalConcurrencyLimits, updateEnvConcurrencyLimits, updateQueueConcurrencyLimits, + updateQueueTotalConcurrencyLimits, } from "../runQueue.server"; import { scheduleEngine } from "../scheduleEngine.server"; import { normalizeScheduleWindow } from "../scheduleWindow.server"; @@ -401,6 +403,7 @@ async function createWorkerTask( { name: task.queue?.name ?? `task/${task.id}`, concurrencyLimit: task.queue?.concurrencyLimit, + totalConcurrencyLimit: task.queue?.totalConcurrencyLimit, }, task.id, task.queue?.name ? "NAMED" : "VIRTUAL", @@ -552,6 +555,7 @@ async function createWorkerQueue( const taskQueue = await upsertWorkerQueueRecord( queueName, baseConcurrencyLimit ?? null, + queue.totalConcurrencyLimit ?? null, orderableName, queueType, worker, @@ -560,6 +564,21 @@ async function createWorkerQueue( 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); + } + if (!taskQueue.paused) { if (typeof newConcurrencyLimit === "number") { logger.debug("createWorkerQueue: updating concurrency limit", { @@ -598,6 +617,7 @@ async function createWorkerQueue( async function upsertWorkerQueueRecord( queueName: string, concurrencyLimit: number | null, + totalConcurrencyLimit: number | null, orderableName: string, queueType: TaskQueueType, worker: BackgroundWorker, @@ -624,6 +644,7 @@ async function upsertWorkerQueueRecord( name: queueName, orderableName, concurrencyLimit, + totalConcurrencyLimit, runtimeEnvironmentId: worker.runtimeEnvironmentId, projectId: worker.projectId, type: queueType, @@ -648,6 +669,7 @@ async function upsertWorkerQueueRecord( // If overridden, keep current limit and update base; otherwise update limit normally concurrencyLimit: hasOverride ? undefined : concurrencyLimit, concurrencyLimitBase: hasOverride ? concurrencyLimit : undefined, + totalConcurrencyLimit, }, }); } @@ -659,6 +681,7 @@ async function upsertWorkerQueueRecord( return await upsertWorkerQueueRecord( queueName, concurrencyLimit, + totalConcurrencyLimit, orderableName, queueType, worker, diff --git a/internal-packages/database/prisma/migrations/20260827120000_add_task_queue_total_concurrency_limit/migration.sql b/internal-packages/database/prisma/migrations/20260827120000_add_task_queue_total_concurrency_limit/migration.sql new file mode 100644 index 00000000000..21cdb8664d9 --- /dev/null +++ b/internal-packages/database/prisma/migrations/20260827120000_add_task_queue_total_concurrency_limit/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "TaskQueue" ADD COLUMN "totalConcurrencyLimit" INTEGER; diff --git a/internal-packages/database/prisma/schema.prisma b/internal-packages/database/prisma/schema.prisma index a1423f340c0..ac0ac09b2f0 100644 --- a/internal-packages/database/prisma/schema.prisma +++ b/internal-packages/database/prisma/schema.prisma @@ -1974,6 +1974,9 @@ model TaskQueue { /// percentage (the source of truth). The absolute concurrencyLimit is materialized from it. /// Decimal(5,2) allows fractional percentages like 12.50% (0.01–100.00). 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? rateLimit Json? paused Boolean @default(false) diff --git a/internal-packages/run-engine/src/engine/index.ts b/internal-packages/run-engine/src/engine/index.ts index 7917ba68303..fc4d76de6db 100644 --- a/internal-packages/run-engine/src/engine/index.ts +++ b/internal-packages/run-engine/src/engine/index.ts @@ -209,6 +209,7 @@ export class RunEngine { queueSelectionStrategy: new FairQueueSelectionStrategy(queueSelectionStrategyOptions), defaultEnvConcurrency: options.queue?.defaultEnvConcurrency ?? 10, defaultEnvConcurrencyBurstFactor: options.queue?.defaultEnvConcurrencyBurstFactor, + totalConcurrencyEnabled: options.queue?.totalConcurrencyEnabled, logger: new Logger("RunQueue", options.queue?.logLevel ?? "info"), redis: { ...options.queue.redis, keyPrefix: `${options.queue.redis.keyPrefix}runqueue:` }, retryOptions: options.queue?.retryOptions, diff --git a/internal-packages/run-engine/src/engine/types.ts b/internal-packages/run-engine/src/engine/types.ts index 2516776373d..ed1b3a0c2c9 100644 --- a/internal-packages/run-engine/src/engine/types.ts +++ b/internal-packages/run-engine/src/engine/types.ts @@ -91,6 +91,8 @@ export type RunEngineOptions = { defaultEnvConcurrency?: number; defaultEnvConcurrencyBurstFactor?: number; logLevel?: LogLevel; + /** Enforce per-queue total concurrency limits across concurrency-key variants. See RunQueueOptions.totalConcurrencyEnabled. */ + totalConcurrencyEnabled?: boolean; /** Optional queue-metrics emitter; enables gauge + counter emission from the RunQueue. */ queueMetrics?: RunQueueMetricsEmitter; queueSelectionStrategyOptions?: Pick< diff --git a/internal-packages/run-engine/src/run-queue/index.ts b/internal-packages/run-engine/src/run-queue/index.ts index 4afa5b2bab9..d3575ab273f 100644 --- a/internal-packages/run-engine/src/run-queue/index.ts +++ b/internal-packages/run-engine/src/run-queue/index.ts @@ -179,6 +179,26 @@ export type RunQueueOptions = { * CK operation re-anchors from ckIndex. Default: 86400 (24h). */ counterTtlSeconds?: number; + /** + * When true, concurrency-keyed queues maintain a per-base-queue groupConcurrency SET + * (total in-flight across all key variants) and enforce the queue's total concurrency + * limit at admit time. Default false: admit paths are byte-identical to before, and + * only the release-side SREM mirror runs (a no-op on an absent set), so the flag can + * be flipped on a fleet that has fully rolled onto this build without draining queues. + * + * Runs already in flight when the flag turns on are not in the set, so a queue can + * transiently exceed its total limit by the number of those runs. The excess is + * one-time and self-corrects as each pre-flag run completes (its release mirror + * no-ops), after which the limit is enforced exactly. + * + * A release path that misses the group mirror (an instance on an older build during + * rollout) leaves the member behind, briefly under-admitting. The dequeue gate + * reconciles: when a queue sits at its total, members whose message key no longer + * exists are pruned, so such leaks clear within seconds instead of blocking the + * queue. Enabling only after every instance runs this build avoids the noise but is + * no longer load-bearing for correctness. + */ + totalConcurrencyEnabled?: boolean; workerOptions?: { pollIntervalMs?: number; immediatePollIntervalMs?: number; @@ -464,6 +484,38 @@ export class RunQueue { return result ? Number(result) : undefined; } + public async updateQueueTotalConcurrencyLimits( + env: MinimalAuthenticatedEnvironment, + queue: string, + totalConcurrency: number + ) { + return this.redis.set(this.keys.queueTotalConcurrencyLimitKey(env, queue), totalConcurrency); + } + + public async removeQueueTotalConcurrencyLimits( + env: MinimalAuthenticatedEnvironment, + queue: string + ) { + return this.redis.del(this.keys.queueTotalConcurrencyLimitKey(env, queue)); + } + + public async getQueueTotalConcurrencyLimit(env: MinimalAuthenticatedEnvironment, queue: string) { + const result = await this.redis.get(this.keys.queueTotalConcurrencyLimitKey(env, queue)); + + return result ? Number(result) : undefined; + } + + /** + * Total in-flight runs across all concurrency-key variants of a queue (the + * groupConcurrency SET cardinality). Admits only populate the set while + * totalConcurrencyEnabled is on. After the flag is turned off the set drains + * to zero through the release-side mirrors, so a nonzero read reflects real + * runs admitted while it was on, never stale state. + */ + public async totalConcurrencyOfQueue(env: MinimalAuthenticatedEnvironment, queue: string) { + return this.redis.scard(this.keys.queueGroupConcurrencyKey(env, queue)); + } + public async updateEnvConcurrencyLimits(env: MinimalAuthenticatedEnvironment) { await this.#callUpdateEnvironmentConcurrencyLimits({ envConcurrencyLimitKey: this.keys.envConcurrencyLimitKey(env), @@ -1240,6 +1292,7 @@ export class RunQueue { this.keys.envCurrentDequeuedKeyFromQueue(message.queue), this.keys.queueRunningCounterKeyFromQueue(message.queue), this.keys.ckIndexKeyFromQueue(message.queue), + this.keys.queueGroupConcurrencyKeyFromQueue(message.queue), messageId, this.options.redis.keyPrefix ?? "", String(this.counterTtlSeconds) @@ -2189,6 +2242,11 @@ export class RunQueue { const lengthCounterKey = this.keys.queueLengthCounterKeyFromQueue(message.queue); const baseQueueKey = this.keys.baseQueueKeyFromQueue(message.queue); const ckKeyPrefix = this.options.redis.keyPrefix ?? ""; + const groupConcurrencyKey = this.keys.queueGroupConcurrencyKeyFromQueue(message.queue); + const totalConcurrencyLimitKey = this.keys.queueTotalConcurrencyLimitKeyFromQueue( + message.queue + ); + const totalConcurrencyEnabledArg = this.options.totalConcurrencyEnabled ? "1" : "0"; if (ttlInfo) { result = await this.redis.enqueueMessageWithTtlCkTracked( @@ -2209,6 +2267,8 @@ export class RunQueue { envConcurrencyLimitBurstFactorKey, lengthCounterKey, baseQueueKey, + groupConcurrencyKey, + totalConcurrencyLimitKey, // args queueName, messageId, @@ -2224,6 +2284,7 @@ export class RunQueue { enableFastPathArg, ckKeyPrefix, String(this.counterTtlSeconds), + totalConcurrencyEnabledArg, metricsGaugeArg ); } else { @@ -2244,6 +2305,8 @@ export class RunQueue { envConcurrencyLimitBurstFactorKey, lengthCounterKey, baseQueueKey, + groupConcurrencyKey, + totalConcurrencyLimitKey, // args queueName, messageId, @@ -2257,6 +2320,7 @@ export class RunQueue { enableFastPathArg, ckKeyPrefix, String(this.counterTtlSeconds), + totalConcurrencyEnabledArg, metricsGaugeArg ); } @@ -2522,6 +2586,8 @@ export class RunQueue { ttlQueueKey, lengthCounterKey, runningCounterKey, + this.keys.queueGroupConcurrencyKeyFromQueue(ckWildcardQueue), + this.keys.queueTotalConcurrencyLimitKeyFromQueue(ckWildcardQueue), //args ckWildcardQueue, String(Date.now()), @@ -2529,6 +2595,7 @@ export class RunQueue { String(this.options.defaultEnvConcurrencyBurstFactor ?? 1), this.options.redis.keyPrefix ?? "", String(maxCount), + this.options.totalConcurrencyEnabled ? "1" : "0", metricsGaugeArg ); @@ -2778,6 +2845,7 @@ export class RunQueue { ckIndexKey, lengthCounterKey, runningCounterKey, + this.keys.queueGroupConcurrencyKeyFromQueue(messageQueue), messageId, messageQueue, messageKeyValue, @@ -2841,6 +2909,7 @@ export class RunQueue { envCurrentDequeuedKey, this.keys.queueRunningCounterKeyFromQueue(queue), this.keys.ckIndexKeyFromQueue(queue), + this.keys.queueGroupConcurrencyKeyFromQueue(queue), messageId, this.options.redis.keyPrefix ?? "", String(this.counterTtlSeconds) @@ -2907,6 +2976,7 @@ export class RunQueue { ckIndexKey, lengthCounterKey, runningCounterKey, + this.keys.queueGroupConcurrencyKeyFromQueue(messageQueue), //args messageId, messageQueue, @@ -2970,6 +3040,7 @@ export class RunQueue { ckIndexKey, lengthCounterKey, runningCounterKey, + this.keys.queueGroupConcurrencyKeyFromQueue(messageQueue), messageId, messageQueue, ckWildcardName @@ -3758,7 +3829,7 @@ return __qmret(0) // *Tracked variants of dequeueMessageFromKey and the ack/nack/dlq/release/clear // scripts. this.redis.defineCommand("enqueueMessageCkTracked", { - numberOfKeys: 15, + numberOfKeys: 17, lua: ` local masterQueueKey = KEYS[1] local queueKey = KEYS[2] @@ -3777,6 +3848,9 @@ local envConcurrencyLimitBurstFactorKey = KEYS[13] -- Counter keys (KEYS 14-15) local lengthCounterKey = KEYS[14] local baseQueueKey = KEYS[15] +-- Total-cap keys (KEYS 16-17) +local groupConcurrencyKey = KEYS[16] +local totalConcurrencyLimitKey = KEYS[17] local queueName = ARGV[1] local messageId = ARGV[2] @@ -3793,6 +3867,7 @@ local enableFastPath = ARGV[10] local keyPrefix = ARGV[11] -- TTL (seconds) applied to counter lazy-init SETs local counterTtl = ARGV[12] +local totalConcurrencyEnabled = ARGV[13] == '1' ${QUEUE_METRICS_GAUGE_PRELUDE} @@ -3813,15 +3888,34 @@ if enableFastPath == '1' then ) if queueCurrent < queueLimit then - redis.call('SET', messageKey, messageData) - redis.call('SADD', queueCurrentConcurrencyKey, messageId) - redis.call('SADD', envCurrentConcurrencyKey, messageId) - redis.call('RPUSH', workerQueueKey, messageKeyValue) + -- Total-cap gate: a fast-path admit consumes a group slot, so it must + -- respect the env-clamped total limit. At the cap we fall through to the + -- slow path (the message queues; the dequeue gate holds it). + local totalAllowsFastPath = true + if totalConcurrencyEnabled then + local rawTotalLimit = redis.call('GET', totalConcurrencyLimitKey) + if rawTotalLimit then + local totalLimit = math.min(tonumber(rawTotalLimit), envLimit) + if tonumber(redis.call('SCARD', groupConcurrencyKey) or '0') >= totalLimit then + totalAllowsFastPath = false + end + end + end + + if totalAllowsFastPath then + redis.call('SET', messageKey, messageData) + redis.call('SADD', queueCurrentConcurrencyKey, messageId) + redis.call('SADD', envCurrentConcurrencyKey, messageId) + if totalConcurrencyEnabled then + redis.call('SADD', groupConcurrencyKey, messageId) + end + redis.call('RPUSH', workerQueueKey, messageKeyValue) ${QUEUE_METRICS_CK_ENQUEUE_FASTPATH_GAUGE_LUA} - -- Fast-path skips the CK variant zset entirely; lengthCounter is unchanged. - -- runningCounter is bumped later by dequeueMessageFromKeyTracked when the - -- worker pulls the message from the worker queue. - return __qmret(1) + -- Fast-path skips the CK variant zset entirely; lengthCounter is unchanged. + -- runningCounter is bumped later by dequeueMessageFromKeyTracked when the + -- worker pulls the message from the worker queue. + return __qmret(1) + end end end end @@ -3873,8 +3967,12 @@ if queueName ~= ckWildcardName then redis.call('ZREM', masterQueueKey, queueName) end --- Update the concurrency keys -redis.call('SREM', queueCurrentConcurrencyKey, messageId) +-- Update the concurrency keys. The groupConcurrency SREM mirrors the per-CK SREM +-- unconditionally (no flag check) so a disabled flag still drains the group set. +local removedFromCurrentConcurrency = redis.call('SREM', queueCurrentConcurrencyKey, messageId) +if removedFromCurrentConcurrency == 1 then + redis.call('SREM', groupConcurrencyKey, messageId) +end redis.call('SREM', envCurrentConcurrencyKey, messageId) redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) @@ -3885,7 +3983,7 @@ return __qmret(0) }); this.redis.defineCommand("enqueueMessageWithTtlCkTracked", { - numberOfKeys: 16, + numberOfKeys: 18, lua: ` local masterQueueKey = KEYS[1] local queueKey = KEYS[2] @@ -3905,6 +4003,9 @@ local envConcurrencyLimitBurstFactorKey = KEYS[14] -- Counter keys (KEYS 15-16) local lengthCounterKey = KEYS[15] local baseQueueKey = KEYS[16] +-- Total-cap keys (KEYS 17-18) +local groupConcurrencyKey = KEYS[17] +local totalConcurrencyLimitKey = KEYS[18] local queueName = ARGV[1] local messageId = ARGV[2] @@ -3923,6 +4024,7 @@ local enableFastPath = ARGV[12] local keyPrefix = ARGV[13] -- TTL (seconds) applied to counter lazy-init SETs local counterTtl = ARGV[14] +local totalConcurrencyEnabled = ARGV[15] == '1' ${QUEUE_METRICS_GAUGE_PRELUDE} @@ -3943,12 +4045,29 @@ if enableFastPath == '1' then ) if queueCurrent < queueLimit then - redis.call('SET', messageKey, messageData) - redis.call('SADD', queueCurrentConcurrencyKey, messageId) - redis.call('SADD', envCurrentConcurrencyKey, messageId) - redis.call('RPUSH', workerQueueKey, messageKeyValue) + -- Total-cap gate: see enqueueMessageCkTracked. + local totalAllowsFastPath = true + if totalConcurrencyEnabled then + local rawTotalLimit = redis.call('GET', totalConcurrencyLimitKey) + if rawTotalLimit then + local totalLimit = math.min(tonumber(rawTotalLimit), envLimit) + if tonumber(redis.call('SCARD', groupConcurrencyKey) or '0') >= totalLimit then + totalAllowsFastPath = false + end + end + end + + if totalAllowsFastPath then + redis.call('SET', messageKey, messageData) + redis.call('SADD', queueCurrentConcurrencyKey, messageId) + redis.call('SADD', envCurrentConcurrencyKey, messageId) + if totalConcurrencyEnabled then + redis.call('SADD', groupConcurrencyKey, messageId) + end + redis.call('RPUSH', workerQueueKey, messageKeyValue) ${QUEUE_METRICS_CK_ENQUEUE_FASTPATH_GAUGE_LUA} - return __qmret(1) + return __qmret(1) + end end end end @@ -3996,8 +4115,12 @@ if queueName ~= ckWildcardName then redis.call('ZREM', masterQueueKey, queueName) end --- Update the concurrency keys -redis.call('SREM', queueCurrentConcurrencyKey, messageId) +-- Update the concurrency keys. The groupConcurrency SREM mirrors the per-CK SREM +-- unconditionally (no flag check) so a disabled flag still drains the group set. +local removedFromCurrentConcurrency = redis.call('SREM', queueCurrentConcurrencyKey, messageId) +if removedFromCurrentConcurrency == 1 then + redis.call('SREM', groupConcurrencyKey, messageId) +end redis.call('SREM', envCurrentConcurrencyKey, messageId) redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) @@ -4180,7 +4303,7 @@ for i, member in ipairs(expiredMembers) do local concurrencyKey = queueKey .. ":currentConcurrency" local dequeuedKey = queueKey .. ":currentDequeued" - redis.call('SREM', concurrencyKey, runId) + local removedFromCurrent = redis.call('SREM', concurrencyKey, runId) local removedFromDequeued = redis.call('SREM', dequeuedKey, runId) local projMatch = string.match(rawQueueKey, ":proj:([^:]+):env:") @@ -4200,6 +4323,10 @@ for i, member in ipairs(expiredMembers) do if removedFromDequeued == 1 then decrFloored(runningCounterKey) end + -- Mirror the per-CK currentConcurrency SREM into the base groupConcurrency set + if removedFromCurrent == 1 then + redis.call('SREM', keyPrefix .. ckMatch .. ":groupConcurrency", runId) + end local ckIndexKey = keyPrefix .. ckMatch .. ":ckIndex" local earliest = redis.call('ZRANGE', queueKey, 0, 0, 'WITHSCORES') @@ -4514,7 +4641,7 @@ return results // (normal dequeue, TTL-expired, or stale-orphan path — all of which were // counted at enqueue time). this.redis.defineCommand("dequeueMessagesFromCkQueueTracked", { - numberOfKeys: 11, + numberOfKeys: 13, lua: ` local ckIndexKey = KEYS[1] local queueConcurrencyLimitKey = KEYS[2] @@ -4527,6 +4654,8 @@ local masterQueueKey = KEYS[8] local ttlQueueKey = KEYS[9] local lengthCounterKey = KEYS[10] local runningCounterKey = KEYS[11] +local groupConcurrencyKey = KEYS[12] +local totalConcurrencyLimitKey = KEYS[13] local ckWildcardName = ARGV[1] local currentTime = tonumber(ARGV[2]) @@ -4534,6 +4663,7 @@ local defaultEnvConcurrencyLimit = ARGV[3] local defaultEnvConcurrencyBurstFactor = ARGV[4] local keyPrefix = ARGV[5] local maxCount = tonumber(ARGV[6] or '1') +local totalConcurrencyEnabled = ARGV[7] == '1' ${QUEUE_METRICS_GAUGE_PRELUDE} ${QUEUE_METRICS_CK_DEQUEUE_GAUGE_LUA} @@ -4558,6 +4688,46 @@ local queueConcurrencyLimit = math.min(tonumber(redis.call('GET', queueConcurren local envAvailableCapacity = envConcurrencyLimitWithBurstFactor - envCurrentConcurrency local actualMaxCount = math.min(maxCount, envAvailableCapacity) +-- Total-cap gate: bound this batch by the remaining headroom across ALL ck +-- variants (groupConcurrency SCARD vs the env-clamped total limit). Each admit +-- below SADDs into the group set and bumps dequeuedCount, and dequeuedCount is +-- bounded by actualMaxCount, so tightening here is sufficient to prevent +-- over-admitting past the cap within a single batch. +if totalConcurrencyEnabled then + local rawTotalLimit = redis.call('GET', totalConcurrencyLimitKey) + if rawTotalLimit then + local totalConcurrencyLimit = math.min(tonumber(rawTotalLimit), envConcurrencyLimit) + local groupCurrentConcurrency = tonumber(redis.call('SCARD', groupConcurrencyKey) or '0') + + -- Self-heal before holding the queue at its limit. A terminal release path + -- that misses the group mirror (an older build, or a future script) leaves + -- a member behind, but every terminal path deletes the run's message key, + -- so a member with no message key is provably dead. Members of re-queued + -- runs keep their message key and clear through the mirrored ack when the + -- run completes. The short lock bounds a saturated queue to one pass per + -- interval, and SSCAN with a persisted cursor bounds each pass to one + -- batch so a large set never blocks Redis for a full traversal; successive + -- passes cover the whole set. + if groupCurrentConcurrency >= totalConcurrencyLimit then + local reconcileLockKey = groupConcurrencyKey .. ':reconcileLock' + if redis.call('SET', reconcileLockKey, '1', 'NX', 'EX', '10') then + local reconcileCursorKey = groupConcurrencyKey .. ':reconcileCursor' + local reconcileCursor = redis.call('GET', reconcileCursorKey) or '0' + local scanResult = redis.call('SSCAN', groupConcurrencyKey, reconcileCursor, 'COUNT', '500') + redis.call('SET', reconcileCursorKey, scanResult[1], 'EX', '3600') + for _, groupMemberId in ipairs(scanResult[2]) do + if redis.call('EXISTS', messageKeyPrefix .. groupMemberId) == 0 then + redis.call('SREM', groupConcurrencyKey, groupMemberId) + end + end + groupCurrentConcurrency = tonumber(redis.call('SCARD', groupConcurrencyKey) or '0') + end + end + + actualMaxCount = math.min(actualMaxCount, totalConcurrencyLimit - groupCurrentConcurrency) + end +end + if actualMaxCount <= 0 then return __qmret(nil) end @@ -4615,6 +4785,9 @@ for _, ckQueueName in ipairs(ckQueues) do decrLengthCounter() redis.call('SADD', ckConcurrencyKey, messageId) redis.call('SADD', envCurrentConcurrencyKey, messageId) + if totalConcurrencyEnabled then + redis.call('SADD', groupConcurrencyKey, messageId) + end if ttlQueueKey and ttlQueueKey ~= '' and ttlExpiresAt then local ttlMember = ckQueueName .. '|' .. messageId .. '|' .. (messageData.orgId or '') @@ -5104,7 +5277,7 @@ redis.call('SREM', envCurrentDequeuedKey, messageId) // removed something) and runningCounter (when SREM currentDequeued actually // removed something). this.redis.defineCommand("acknowledgeMessageCkTracked", { - numberOfKeys: 12, + numberOfKeys: 13, lua: ` -- Keys: local masterQueueKey = KEYS[1] @@ -5119,6 +5292,7 @@ local workerQueueKey = KEYS[9] local ckIndexKey = KEYS[10] local lengthCounterKey = KEYS[11] local runningCounterKey = KEYS[12] +local groupConcurrencyKey = KEYS[13] -- Args: local messageId = ARGV[1] @@ -5171,7 +5345,12 @@ end -- Update the concurrency keys. DECR runningCounter only when SREM -- currentDequeued actually removed an entry (the message was in flight). -redis.call('SREM', queueCurrentConcurrencyKey, messageId) +-- The groupConcurrency SREM mirrors the per-CK SREM so the group set drains +-- on every release path. +local removedFromCurrentConcurrency = redis.call('SREM', queueCurrentConcurrencyKey, messageId) +if removedFromCurrentConcurrency == 1 then + redis.call('SREM', groupConcurrencyKey, messageId) +end redis.call('SREM', envCurrentConcurrencyKey, messageId) local removedFromDequeued = redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) @@ -5190,7 +5369,7 @@ end // runningCounter (floored); ZADD back to the variant zset INCRs // lengthCounter only when ZADD reported a new entry. this.redis.defineCommand("nackMessageCkTracked", { - numberOfKeys: 11, + numberOfKeys: 12, lua: ` -- Keys: local masterQueueKey = KEYS[1] @@ -5204,6 +5383,7 @@ local envQueueKey = KEYS[8] local ckIndexKey = KEYS[9] local lengthCounterKey = KEYS[10] local runningCounterKey = KEYS[11] +local groupConcurrencyKey = KEYS[12] -- Args: local messageId = ARGV[1] @@ -5229,7 +5409,11 @@ redis.call('SET', messageKey, messageData) -- so we skip the eager lazy-init here (unlike releaseConcurrencyTracked, which -- mirrors the same DECR pattern with init). A post-TTL nack's floored DECR -- no-ops; the next dequeueMessageFromKeyTracked reseeds from current state. -redis.call('SREM', queueCurrentConcurrencyKey, messageId) +-- The groupConcurrency SREM mirrors the per-CK SREM. +local removedFromCurrentConcurrency = redis.call('SREM', queueCurrentConcurrencyKey, messageId) +if removedFromCurrentConcurrency == 1 then + redis.call('SREM', groupConcurrencyKey, messageId) +end redis.call('SREM', envCurrentConcurrencyKey, messageId) local removedFromDequeued = redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) @@ -5285,7 +5469,7 @@ end // Tracked variant: same as moveToDeadLetterQueueCk. ZREM may DECR // lengthCounter (defensive); SREM currentDequeued may DECR runningCounter. this.redis.defineCommand("moveToDeadLetterQueueCkTracked", { - numberOfKeys: 12, + numberOfKeys: 13, lua: ` -- Keys: local masterQueueKey = KEYS[1] @@ -5300,6 +5484,7 @@ local deadLetterQueueKey = KEYS[9] local ckIndexKey = KEYS[10] local lengthCounterKey = KEYS[11] local runningCounterKey = KEYS[12] +local groupConcurrencyKey = KEYS[13] -- Args: local messageId = ARGV[1] @@ -5349,8 +5534,12 @@ end redis.call('ZADD', deadLetterQueueKey, tonumber(redis.call('TIME')[1]), messageId) -- Update the concurrency keys. DECR runningCounter only when SREM --- currentDequeued actually removed an entry. -redis.call('SREM', queueCurrentConcurrencyKey, messageId) +-- currentDequeued actually removed an entry. The groupConcurrency SREM mirrors +-- the per-CK SREM. +local removedFromCurrentConcurrency = redis.call('SREM', queueCurrentConcurrencyKey, messageId) +if removedFromCurrentConcurrency == 1 then + redis.call('SREM', groupConcurrencyKey, messageId) +end redis.call('SREM', envCurrentConcurrencyKey, messageId) local removedFromDequeued = redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) @@ -5385,7 +5574,7 @@ redis.call('SREM', envCurrentDequeuedKey, messageId) // something. Caller should only invoke this variant for CK queues — non-CK // queues should keep calling releaseConcurrency. this.redis.defineCommand("releaseConcurrencyTracked", { - numberOfKeys: 6, + numberOfKeys: 7, lua: ` -- Keys: local queueCurrentConcurrencyKey = KEYS[1] @@ -5394,6 +5583,7 @@ local queueCurrentDequeuedKey = KEYS[3] local envCurrentDequeuedKey = KEYS[4] local runningCounterKey = KEYS[5] local ckIndexKey = KEYS[6] +local groupConcurrencyKey = KEYS[7] -- Args: local messageId = ARGV[1] @@ -5415,7 +5605,10 @@ if redis.call('EXISTS', runningCounterKey) == 0 then redis.call('SET', runningCounterKey, total, 'EX', counterTtl) end -redis.call('SREM', queueCurrentConcurrencyKey, messageId) +local removedFromCurrentConcurrency = redis.call('SREM', queueCurrentConcurrencyKey, messageId) +if removedFromCurrentConcurrency == 1 then + redis.call('SREM', groupConcurrencyKey, messageId) +end redis.call('SREM', envCurrentConcurrencyKey, messageId) local removedFromDequeued = redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) @@ -5530,7 +5723,7 @@ redis.call('SREM', envCurrentDequeuedKey, messageId) // Tracked variant of clearMessageFromConcurrencySets — see releaseConcurrencyTracked // for the contract. Only invoke for CK queues. this.redis.defineCommand("clearMessageFromConcurrencySetsTracked", { - numberOfKeys: 6, + numberOfKeys: 7, lua: ` -- Keys: local queueCurrentConcurrencyKey = KEYS[1] @@ -5539,6 +5732,7 @@ local queueCurrentDequeuedKey = KEYS[3] local envCurrentDequeuedKey = KEYS[4] local runningCounterKey = KEYS[5] local ckIndexKey = KEYS[6] +local groupConcurrencyKey = KEYS[7] -- Args: local messageId = ARGV[1] @@ -5556,7 +5750,10 @@ if redis.call('EXISTS', runningCounterKey) == 0 then redis.call('SET', runningCounterKey, total, 'EX', counterTtl) end -redis.call('SREM', queueCurrentConcurrencyKey, messageId) +local removedFromCurrentConcurrency = redis.call('SREM', queueCurrentConcurrencyKey, messageId) +if removedFromCurrentConcurrency == 1 then + redis.call('SREM', groupConcurrencyKey, messageId) +end redis.call('SREM', envCurrentConcurrencyKey, messageId) local removedFromDequeued = redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) @@ -5950,6 +6147,8 @@ declare module "@internal/redis" { envConcurrencyLimitBurstFactorKey: string, lengthCounterKey: string, baseQueueKey: string, + groupConcurrencyKey: string, + totalConcurrencyLimitKey: string, queueName: string, messageId: string, messageData: string, @@ -5962,6 +6161,7 @@ declare module "@internal/redis" { enableFastPath: string, keyPrefix: string, counterTtl: string, + totalConcurrencyEnabled: string, metricsEnabled: string, callback?: Callback<[number, number[] | null]> ): Result<[number, number[] | null], Context>; @@ -5983,6 +6183,8 @@ declare module "@internal/redis" { envConcurrencyLimitBurstFactorKey: string, lengthCounterKey: string, baseQueueKey: string, + groupConcurrencyKey: string, + totalConcurrencyLimitKey: string, queueName: string, messageId: string, messageData: string, @@ -5997,6 +6199,7 @@ declare module "@internal/redis" { enableFastPath: string, keyPrefix: string, counterTtl: string, + totalConcurrencyEnabled: string, metricsEnabled: string, callback?: Callback<[number, number[] | null]> ): Result<[number, number[] | null], Context>; @@ -6013,12 +6216,15 @@ declare module "@internal/redis" { ttlQueueKey: string, lengthCounterKey: string, runningCounterKey: string, + groupConcurrencyKey: string, + totalConcurrencyLimitKey: string, ckWildcardName: string, currentTime: string, defaultEnvConcurrencyLimit: string, defaultEnvConcurrencyBurstFactor: string, keyPrefix: string, maxCount: string, + totalConcurrencyEnabled: string, metricsEnabled: string, callback?: Callback<[string[] | null, number[] | null]> ): Result<[string[] | null, number[] | null], Context>; @@ -6043,6 +6249,7 @@ declare module "@internal/redis" { ckIndexKey: string, lengthCounterKey: string, runningCounterKey: string, + groupConcurrencyKey: string, messageId: string, messageQueueName: string, messageKeyValue: string, @@ -6063,6 +6270,7 @@ declare module "@internal/redis" { ckIndexKey: string, lengthCounterKey: string, runningCounterKey: string, + groupConcurrencyKey: string, messageId: string, messageQueueName: string, messageData: string, @@ -6086,6 +6294,7 @@ declare module "@internal/redis" { ckIndexKey: string, lengthCounterKey: string, runningCounterKey: string, + groupConcurrencyKey: string, messageId: string, messageQueueName: string, ckWildcardName: string, @@ -6111,6 +6320,7 @@ declare module "@internal/redis" { envCurrentDequeuedKey: string, runningCounterKey: string, ckIndexKey: string, + groupConcurrencyKey: string, messageId: string, keyPrefix: string, counterTtl: string, @@ -6124,6 +6334,7 @@ declare module "@internal/redis" { envCurrentDequeuedKey: string, runningCounterKey: string, ckIndexKey: string, + groupConcurrencyKey: string, messageId: string, keyPrefix: string, counterTtl: string, diff --git a/internal-packages/run-engine/src/run-queue/keyProducer.ts b/internal-packages/run-engine/src/run-queue/keyProducer.ts index 0609b3d719b..98028f5af7b 100644 --- a/internal-packages/run-engine/src/run-queue/keyProducer.ts +++ b/internal-packages/run-engine/src/run-queue/keyProducer.ts @@ -24,6 +24,8 @@ const constants = { CK_INDEX_PART: "ckIndex", LENGTH_COUNTER_PART: "lengthCounter", RUNNING_COUNTER_PART: "runningCounter", + GROUP_CONCURRENCY_PART: "groupConcurrency", + TOTAL_CONCURRENCY_LIMIT_PART: "totalConcurrency", } as const; export class RunQueueFullKeyProducer implements RunQueueKeyProducer { @@ -338,6 +340,32 @@ export class RunQueueFullKeyProducer implements RunQueueKeyProducer { return `${this.baseQueueKeyFromQueue(queue)}:${constants.RUNNING_COUNTER_PART}`; } + /** + * SET of in-flight messageIds across ALL concurrency-key variants of a base queue. + * SCARD of this set is the queue's total running count, gated by the total + * concurrency limit. Lives at the base queue so every ck variant shares it. + */ + queueGroupConcurrencyKey(env: RunQueueKeyProducerEnvironment, queue: string): string { + return `${this.queueKey(env, queue)}:${constants.GROUP_CONCURRENCY_PART}`; + } + + queueGroupConcurrencyKeyFromQueue(queue: string): string { + return `${this.baseQueueKeyFromQueue(queue)}:${constants.GROUP_CONCURRENCY_PART}`; + } + + /** + * String key holding the queue's total concurrency limit (the cap across all + * concurrency-key variants). Absent = no total cap. Readers clamp to the + * environment limit; the raw requested value is what's stored. + */ + queueTotalConcurrencyLimitKey(env: RunQueueKeyProducerEnvironment, queue: string): string { + return `${this.queueKey(env, queue)}:${constants.TOTAL_CONCURRENCY_LIMIT_PART}`; + } + + queueTotalConcurrencyLimitKeyFromQueue(queue: string): string { + return `${this.baseQueueKeyFromQueue(queue)}:${constants.TOTAL_CONCURRENCY_LIMIT_PART}`; + } + isCkWildcard(queue: string): boolean { return queue.endsWith(":ck:*"); } diff --git a/internal-packages/run-engine/src/run-queue/tests/totalConcurrency.test.ts b/internal-packages/run-engine/src/run-queue/tests/totalConcurrency.test.ts new file mode 100644 index 00000000000..cd6d299f151 --- /dev/null +++ b/internal-packages/run-engine/src/run-queue/tests/totalConcurrency.test.ts @@ -0,0 +1,412 @@ +import { assertNonNullable, redisTest } from "@internal/testcontainers"; +import { trace } from "@internal/tracing"; +import { setTimeout } from "node:timers/promises"; +import { describe } from "vitest"; +import { FairQueueSelectionStrategy } from "../fairQueueSelectionStrategy.js"; +import { RunQueue } from "../index.js"; +import { RunQueueFullKeyProducer } from "../keyProducer.js"; +import type { InputPayload } from "../types.js"; +import { Decimal } from "@trigger.dev/database"; + +const testOptions = { + name: "rq", + tracer: trace.getTracer("rq"), + workers: 1, + defaultEnvConcurrency: 25, + retryOptions: { + maxAttempts: 5, + factor: 1.1, + minTimeoutInMs: 100, + maxTimeoutInMs: 1_000, + randomize: true, + }, + keys: new RunQueueFullKeyProducer(), +}; + +const authenticatedEnvDev = { + id: "e1234", + type: "DEVELOPMENT" as const, + maximumConcurrencyLimit: 10, + concurrencyLimitBurstFactor: new Decimal(2.0), + project: { id: "p1234" }, + organization: { id: "o1234" }, +}; + +function createQueue(redisContainer: any, totalConcurrencyEnabled: boolean) { + return new RunQueue({ + ...testOptions, + totalConcurrencyEnabled, + queueSelectionStrategy: new FairQueueSelectionStrategy({ + redis: { + keyPrefix: "runqueue:test:", + host: redisContainer.getHost(), + port: redisContainer.getPort(), + }, + keys: testOptions.keys, + }), + redis: { + keyPrefix: "runqueue:test:", + host: redisContainer.getHost(), + port: redisContainer.getPort(), + }, + }); +} + +function makeMessage(overrides: Partial = {}): InputPayload { + return { + runId: "r1", + taskIdentifier: "task/my-task", + orgId: "o1234", + projectId: "p1234", + environmentId: "e1234", + environmentType: "DEVELOPMENT", + queue: "task/my-task", + timestamp: Date.now(), + attempt: 0, + ...overrides, + }; +} + +async function waitFor(condition: () => Promise, timeoutMs = 20_000): Promise { + const deadline = Date.now() + timeoutMs; + while (Date.now() < deadline) { + if (await condition()) { + return true; + } + await setTimeout(250); + } + return condition(); +} + +vi.setConfig({ testTimeout: 60_000 }); + +describe("RunQueue total concurrency limit", () => { + redisTest( + "caps in-flight runs across concurrency keys at the total limit", + async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueTotalConcurrencyLimits(authenticatedEnvDev, "task/my-task", 2); + + const now = Date.now(); + for (const [i, ck] of ["ck-a", "ck-a", "ck-b", "ck-b"].entries()) { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: `r${i}`, + concurrencyKey: ck, + timestamp: now - 1000 + i, + }), + workerQueue: "main", + }); + } + + const admittedTwo = await waitFor( + async () => + (await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")) === 2 + ); + expect(admittedTwo).toBe(true); + + /** + * The remaining two messages must stay queued: give the master consumers a + * couple of extra polling cycles to prove the gate holds, not just that it + * hadn't caught up yet. + */ + await setTimeout(2000); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(2); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(2); + + const dequeued1 = await queue.dequeueMessageFromWorkerQueue("consumer-1", "main"); + assertNonNullable(dequeued1); + const dequeued2 = await queue.dequeueMessageFromWorkerQueue("consumer-1", "main"); + assertNonNullable(dequeued2); + + await queue.acknowledgeMessage(authenticatedEnvDev.organization.id, dequeued1.messageId); + + const thirdAdmitted = await waitFor(async () => { + const total = await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task"); + const queued = await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task"); + return total === 2 && queued === 1; + }); + expect(thirdAdmitted).toBe(true); + } finally { + await queue.quit(); + } + } + ); + + redisTest( + "still enforces the per-key limit under the total limit", + async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 1); + await queue.updateQueueTotalConcurrencyLimits(authenticatedEnvDev, "task/my-task", 10); + + const now = Date.now(); + for (const i of [0, 1]) { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: `r${i}`, + concurrencyKey: "ck-a", + timestamp: now - 1000 + i, + }), + workerQueue: "main", + }); + } + + const oneAdmitted = await waitFor( + async () => + (await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")) === 1 + ); + expect(oneAdmitted).toBe(true); + + await setTimeout(2000); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + expect( + await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "task/my-task", "ck-a") + ).toBe(1); + } finally { + await queue.quit(); + } + } + ); + + redisTest( + "ignores the stored total limit and maintains no group set when disabled", + async ({ redisContainer }) => { + const queue = createQueue(redisContainer, false); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueTotalConcurrencyLimits(authenticatedEnvDev, "task/my-task", 1); + + const now = Date.now(); + for (const [i, ck] of ["ck-a", "ck-b", "ck-c"].entries()) { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: `r${i}`, + concurrencyKey: ck, + timestamp: now - 1000 + i, + }), + workerQueue: "main", + }); + } + + const allAdmitted = await waitFor( + async () => (await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")) === 0 + ); + expect(allAdmitted).toBe(true); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(0); + } finally { + await queue.quit(); + } + } + ); + + redisTest("enqueue fast path respects the total limit", async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueTotalConcurrencyLimits(authenticatedEnvDev, "task/my-task", 1); + + const now = Date.now(); + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ runId: "r0", concurrencyKey: "ck-a", timestamp: now - 1000 }), + workerQueue: "main", + enableFastPath: true, + skipDequeueProcessing: true, + }); + + /** The fast path admits synchronously, so the group slot is taken immediately. */ + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(0); + + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ runId: "r1", concurrencyKey: "ck-b", timestamp: now - 999 }), + workerQueue: "main", + enableFastPath: true, + skipDequeueProcessing: true, + }); + + /** At the total limit the fast path must fall through to a normal enqueue. */ + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + } finally { + await queue.quit(); + } + }); + + redisTest("nacking releases the total slot", async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueTotalConcurrencyLimits(authenticatedEnvDev, "task/my-task", 1); + + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ runId: "r0", concurrencyKey: "ck-a", timestamp: Date.now() - 1000 }), + workerQueue: "main", + }); + + const admitted = await waitFor( + async () => (await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")) === 1 + ); + expect(admitted).toBe(true); + + /** A second run on another key waits behind the total limit of 1. */ + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ runId: "r1", concurrencyKey: "ck-b", timestamp: Date.now() - 500 }), + workerQueue: "main", + }); + + await setTimeout(2000); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + + const dequeued = await queue.dequeueMessageFromWorkerQueue("consumer-1", "main"); + assertNonNullable(dequeued); + expect(dequeued.messageId).toBe("r0"); + + /** + * Nack r0 with a far-future retryAt so it cannot immediately reclaim the + * slot. If the nack released r0's group slot, r1 is the only eligible run + * and must be admitted; if the slot leaked, the queue stays blocked and r1 + * never surfaces. + */ + await queue.nackMessage({ + orgId: authenticatedEnvDev.organization.id, + messageId: "r0", + retryAt: Date.now() + 120_000, + }); + + const r1Admitted = await waitFor(async () => { + const next = await queue.dequeueMessageFromWorkerQueue("consumer-1", "main", { + blockingPop: false, + }); + return next?.messageId === "r1"; + }); + expect(r1Admitted).toBe(true); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + } finally { + await queue.quit(); + } + }); + + redisTest( + "reconciles a leaked group member instead of blocking the queue", + async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + const keys = testOptions.keys; + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueTotalConcurrencyLimits(authenticatedEnvDev, "task/my-task", 1); + + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: "r0", + concurrencyKey: "ck-a", + timestamp: Date.now() - 1000, + }), + workerQueue: "main", + }); + + const admitted = await waitFor( + async () => + (await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")) === 1 + ); + expect(admitted).toBe(true); + + const dequeued = await queue.dequeueMessageFromWorkerQueue("consumer-1", "main"); + assertNonNullable(dequeued); + expect(dequeued.messageId).toBe("r0"); + + /** + * Simulate a terminal release from a build without the group mirror: the + * message key is deleted and the per-key and env sets are cleared, but the + * group member is left behind. + */ + await queue.redis.del(keys.messageKey(authenticatedEnvDev.organization.id, "r0")); + await queue.redis.srem( + keys.queueCurrentConcurrencyKey(authenticatedEnvDev, "task/my-task", "ck-a"), + "r0" + ); + await queue.redis.srem(keys.envCurrentConcurrencyKey(authenticatedEnvDev), "r0"); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + + /** The next run must still be admitted: the gate prunes the dead member. */ + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: "r1", + concurrencyKey: "ck-b", + timestamp: Date.now() - 500, + }), + workerQueue: "main", + }); + + const r1Admitted = await waitFor(async () => { + const next = await queue.dequeueMessageFromWorkerQueue("consumer-1", "main", { + blockingPop: false, + }); + return next?.messageId === "r1"; + }); + expect(r1Admitted).toBe(true); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + } finally { + await queue.quit(); + } + } + ); + + redisTest( + "reconciles a large leaked backlog across bounded passes", + async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + const keys = testOptions.keys; + const groupKey = keys.queueGroupConcurrencyKey(authenticatedEnvDev, "task/my-task"); + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueTotalConcurrencyLimits(authenticatedEnvDev, "task/my-task", 1); + + /** 1,200 dead members: more than one SSCAN batch, none with a message key. */ + const dead = Array.from({ length: 1200 }, (_, i) => `dead-${i}`); + await queue.redis.sadd(groupKey, ...dead); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1200); + + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: "r0", + concurrencyKey: "ck-a", + timestamp: Date.now() - 1000, + }), + workerQueue: "main", + }); + + /** + * Each dequeue attempt reconciles at most one SSCAN batch behind a 10s + * lock; dropping the lock between polls lets the passes run back to + * back instead of waiting out the interval. + */ + const r0Admitted = await waitFor(async () => { + await queue.redis.del(`${groupKey}:reconcileLock`); + const next = await queue.dequeueMessageFromWorkerQueue("consumer-1", "main", { + blockingPop: false, + }); + return next?.messageId === "r0"; + }, 30_000); + expect(r0Admitted).toBe(true); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + } finally { + await queue.quit(); + } + } + ); +}); diff --git a/internal-packages/run-engine/src/run-queue/types.ts b/internal-packages/run-engine/src/run-queue/types.ts index 8a7d3c93ec5..10575d8b5d3 100644 --- a/internal-packages/run-engine/src/run-queue/types.ts +++ b/internal-packages/run-engine/src/run-queue/types.ts @@ -91,6 +91,10 @@ export interface RunQueueKeyProducer { queue: string, concurrencyKey?: string ): string; + queueGroupConcurrencyKey(env: RunQueueKeyProducerEnvironment, queue: string): string; + queueGroupConcurrencyKeyFromQueue(queue: string): string; + queueTotalConcurrencyLimitKey(env: RunQueueKeyProducerEnvironment, queue: string): string; + queueTotalConcurrencyLimitKeyFromQueue(queue: string): string; //env oncurrency envCurrentConcurrencyKey(env: EnvDescriptor): string; diff --git a/packages/core/src/v3/schemas/schemas.ts b/packages/core/src/v3/schemas/schemas.ts index 505622d1ef0..ff669193c42 100644 --- a/packages/core/src/v3/schemas/schemas.ts +++ b/packages/core/src/v3/schemas/schemas.ts @@ -178,6 +178,12 @@ export const QueueManifest = z.object({ * * If this property is omitted, the task can potentially use up the full concurrency of an environment */ concurrencyLimit: z.number().int().min(0).max(100000).optional().nullable(), + /** An optional property that caps the total number of concurrent run executions across ALL + * `concurrencyKey` values of this queue. On a queue with a `concurrencyKey`, `concurrencyLimit` + * applies per key value; this is the ceiling for the whole queue. + * + * Only enforced for runs triggered with a `concurrencyKey`, and requires server-side support. */ + totalConcurrencyLimit: z.number().int().min(0).max(100000).optional().nullable(), }); export type QueueManifest = z.infer; diff --git a/packages/core/src/v3/types/queues.ts b/packages/core/src/v3/types/queues.ts index 9e87f136e2e..1fba786c868 100644 --- a/packages/core/src/v3/types/queues.ts +++ b/packages/core/src/v3/types/queues.ts @@ -35,4 +35,28 @@ export type QueueOptions = { * * If this property is omitted, the task can potentially use up the full concurrency of an environment */ concurrencyLimit?: number; + /** An optional property that caps the total number of concurrent run executions across ALL + * `concurrencyKey` values of this queue. + * + * On a queue used with a `concurrencyKey`, `concurrencyLimit` applies to each key value + * independently — ten active keys with `concurrencyLimit: 5` can run 50 at once. Setting + * `totalConcurrencyLimit: 20` bounds the whole queue to 20 while each key still gets at + * most `concurrencyLimit`. + * + * @example + * + * ```ts + * const perUserQueue = queue({ + name: "per-user-queue", + concurrencyLimit: 1, + totalConcurrencyLimit: 10, + }); + * ``` + * + * Only enforced for runs triggered with a `concurrencyKey`, and requires server-side support. + * + * Omit for no total cap. Like `concurrencyLimit`, a value of `0` holds every keyed run in + * the queue rather than removing the cap. + */ + totalConcurrencyLimit?: number; }; diff --git a/packages/core/src/v3/types/tasks.ts b/packages/core/src/v3/types/tasks.ts index 224b7604db3..070cc96f3df 100644 --- a/packages/core/src/v3/types/tasks.ts +++ b/packages/core/src/v3/types/tasks.ts @@ -225,6 +225,7 @@ type CommonTaskOptions< queue?: { name?: string; concurrencyLimit?: number; + totalConcurrencyLimit?: number; }; /** Configure the spec of the [machine](https://trigger.dev/docs/machines) you want your task to run on. * diff --git a/packages/trigger-sdk/src/v3/shared.ts b/packages/trigger-sdk/src/v3/shared.ts index 69b51f0ed3c..b962697d559 100644 --- a/packages/trigger-sdk/src/v3/shared.ts +++ b/packages/trigger-sdk/src/v3/shared.ts @@ -277,6 +277,7 @@ export function createTask< resourceCatalog.registerQueueMetadata({ name: queue.name, concurrencyLimit: queue.concurrencyLimit, + totalConcurrencyLimit: queue.totalConcurrencyLimit, }); } @@ -432,6 +433,7 @@ export function createSchemaTask< resourceCatalog.registerQueueMetadata({ name: queue.name, concurrencyLimit: queue.concurrencyLimit, + totalConcurrencyLimit: queue.totalConcurrencyLimit, }); }