diff --git a/apps/webapp/app/env.server.ts b/apps/webapp/app/env.server.ts index 8b660127648..a9dcfebc784 100644 --- a/apps/webapp/app/env.server.ts +++ b/apps/webapp/app/env.server.ts @@ -1378,6 +1378,7 @@ const EnvironmentSchema = z .enum(["log", "error", "warn", "info", "debug"]) .default("info"), RUN_ENGINE_TOTAL_CONCURRENCY_LIMITS_ENABLED: z.string().default("0"), + RUN_ENGINE_QUEUE_GATES_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 7a6cb0f8d3d..caeb3662933 100644 --- a/apps/webapp/app/v3/runEngine.server.ts +++ b/apps/webapp/app/v3/runEngine.server.ts @@ -63,6 +63,7 @@ function createRunEngine() { defaultEnvConcurrency: env.DEFAULT_ENV_EXECUTION_CONCURRENCY_LIMIT, defaultEnvConcurrencyBurstFactor: env.DEFAULT_ENV_EXECUTION_CONCURRENCY_BURST_FACTOR, totalConcurrencyEnabled: env.RUN_ENGINE_TOTAL_CONCURRENCY_LIMITS_ENABLED === "1", + gatesEnabled: env.RUN_ENGINE_QUEUE_GATES_ENABLED === "1", logLevel: env.RUN_ENGINE_RUN_QUEUE_LOG_LEVEL, redis: { keyPrefix: "engine:", diff --git a/internal-packages/run-engine/src/engine/index.ts b/internal-packages/run-engine/src/engine/index.ts index fc4d76de6db..478786dd230 100644 --- a/internal-packages/run-engine/src/engine/index.ts +++ b/internal-packages/run-engine/src/engine/index.ts @@ -210,6 +210,7 @@ export class RunEngine { defaultEnvConcurrency: options.queue?.defaultEnvConcurrency ?? 10, defaultEnvConcurrencyBurstFactor: options.queue?.defaultEnvConcurrencyBurstFactor, totalConcurrencyEnabled: options.queue?.totalConcurrencyEnabled, + gatesEnabled: options.queue?.gatesEnabled, 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 ed1b3a0c2c9..d8df88c4d2a 100644 --- a/internal-packages/run-engine/src/engine/types.ts +++ b/internal-packages/run-engine/src/engine/types.ts @@ -93,6 +93,8 @@ export type RunEngineOptions = { logLevel?: LogLevel; /** Enforce per-queue total concurrency limits across concurrency-key variants. See RunQueueOptions.totalConcurrencyEnabled. */ totalConcurrencyEnabled?: boolean; + /** Enforce the gates carried in message payloads. See RunQueueOptions.gatesEnabled. */ + gatesEnabled?: 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 d3575ab273f..b6ef93a9ea3 100644 --- a/internal-packages/run-engine/src/run-queue/index.ts +++ b/internal-packages/run-engine/src/run-queue/index.ts @@ -58,6 +58,106 @@ const SemanticAttributes = { ORG_ID: "runqueue.orgId", }; +/** + * Gate helpers, spliced into scripts that admit or release runs. A gate is another + * queue named in the message payload that the run must also hold a slot in while it + * executes (payload shape: gates = [{queue, concurrencyKey?}]). Keys are built from + * the payload's own org/proj/env, so gates always live in the message's hash slot. + * + * __gatesHaveCapacity / __gatesAcquire run on admit paths behind the gatesEnabled + * flag. __gatesRelease runs on every release path UNCONDITIONALLY and is payload- + * driven (a cheap substring probe before decoding), so slots acquired while the flag + * was on always drain, and gateless messages pay near zero. __gateReconcile is the + * same bounded self-heal as the total-cap gate: a member whose message key is gone + * was terminally released by a path that missed the mirror and is provably dead. + */ +const QUEUE_GATES_LUA_HELPERS = ` +local function __gateKeys(gatesKeyPrefix, msg, gate) + local base = gatesKeyPrefix .. '{org:' .. msg.orgId .. '}:proj:' .. msg.projectId .. ':env:' .. msg.environmentId .. ':queue:' .. gate.queue + local gateKey = gate.concurrencyKey + if (not gateKey or gateKey == '') and msg.concurrencyKey and msg.concurrencyKey ~= '' then + gateKey = msg.concurrencyKey + end + local variant = base + if gateKey and gateKey ~= '' then + variant = base .. ':ck:' .. gateKey + end + return base, variant, gateKey +end + +local function __gateReconcile(setKey, msgKeyPrefix, reconcileKeyPrefix) + if not msgKeyPrefix then return end + if redis.call('SET', setKey .. ':reconcileLock', '1', 'NX', 'EX', '10') then + local cursorKey = setKey .. ':reconcileCursor' + local cursor = redis.call('GET', cursorKey) or '0' + local scanResult = redis.call('SSCAN', setKey, cursor, 'COUNT', '100') + redis.call('SET', cursorKey, scanResult[1], 'EX', '3600') + for _, memberId in ipairs(scanResult[2]) do + local rawMemberPayload = redis.call('GET', msgKeyPrefix .. memberId) + if not rawMemberPayload then + redis.call('SREM', setKey, memberId) + elseif reconcileKeyPrefix then + local okMember, member = pcall(cjson.decode, rawMemberPayload) + if okMember and type(member) == 'table' and type(member.queue) == 'string' then + if redis.call('ZSCORE', reconcileKeyPrefix .. member.queue, memberId) then + redis.call('SREM', setKey, memberId) + end + end + end + end + end +end + +local function __gatesHaveCapacity(gatesKeyPrefix, msg, messageId, envLimit, msgKeyPrefix) + 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 occupancy >= perKeyLimit and redis.call('SISMEMBER', variant .. ':currentConcurrency', messageId) == 0 then + __gateReconcile(variant .. ':currentConcurrency', msgKeyPrefix, gatesKeyPrefix) + return false + end + if gateKey and gateKey ~= '' then + local rawTotal = redis.call('GET', base .. ':totalConcurrency') + if rawTotal then + local totalLimit = math.min(tonumber(rawTotal), envLimit) + local groupKey = base .. ':groupConcurrency' + if tonumber(redis.call('SCARD', groupKey) or '0') >= totalLimit and redis.call('SISMEMBER', groupKey, messageId) == 0 then + __gateReconcile(groupKey, msgKeyPrefix, gatesKeyPrefix) + return false + end + end + end + end + return true +end + +local function __gatesAcquire(gatesKeyPrefix, msg, messageId) + if not msg.gates then return end + for _, gate in ipairs(msg.gates) do + local base, variant, gateKey = __gateKeys(gatesKeyPrefix, msg, gate) + redis.call('SADD', variant .. ':currentConcurrency', messageId) + if gateKey and gateKey ~= '' then + redis.call('SADD', base .. ':groupConcurrency', messageId) + end + end +end + +local function __gatesRelease(gatesKeyPrefix, rawPayload, messageId) + if not rawPayload or rawPayload == false then return end + if not string.find(rawPayload, '"gates"', 1, true) then return end + local ok, msg = pcall(cjson.decode, rawPayload) + if not ok or type(msg) ~= 'table' or not msg.gates then return end + for _, gate in ipairs(msg.gates) do + local base, variant, gateKey = __gateKeys(gatesKeyPrefix, msg, gate) + local removed = redis.call('SREM', variant .. ':currentConcurrency', messageId) + if removed == 1 and gateKey and gateKey ~= '' then + redis.call('SREM', base .. ':groupConcurrency', messageId) + end + end +end`; + // Prelude spliced at the top of every gauge-carrying script: declares the gauge slot and // the return wrapper. A splice fills __qm_g; every return goes through __qmret so the reply // is always {original, gauge}. A nil original becomes false, else Lua drops it from the @@ -199,6 +299,15 @@ export type RunQueueOptions = { * no longer load-bearing for correctness. */ totalConcurrencyEnabled?: boolean; + /** + * When true, admit paths enforce the gates carried in a message's payload: the run + * must also have capacity in, and holds a slot in, each gate queue until release. + * Default false: gated messages admit as if they had no gates. Release-side gate + * removal is payload-driven and always on, so slots acquired while the flag was on + * drain correctly after it is turned off, with the same message-key self-heal as + * the total cap covering releases from builds without the mirror. + */ + gatesEnabled?: boolean; workerOptions?: { pollIntervalMs?: number; immediatePollIntervalMs?: number; @@ -1293,6 +1402,7 @@ export class RunQueue { this.keys.queueRunningCounterKeyFromQueue(message.queue), this.keys.ckIndexKeyFromQueue(message.queue), this.keys.queueGroupConcurrencyKeyFromQueue(message.queue), + this.keys.messageKey(message.orgId, messageId), messageId, this.options.redis.keyPrefix ?? "", String(this.counterTtlSeconds) @@ -1304,7 +1414,9 @@ export class RunQueue { this.keys.envCurrentConcurrencyKeyFromQueue(message.queue), this.keys.queueCurrentDequeuedKeyFromQueue(message.queue), this.keys.envCurrentDequeuedKeyFromQueue(message.queue), - messageId + this.keys.messageKey(message.orgId, messageId), + messageId, + this.options.redis.keyPrefix ?? "" ); }, { @@ -2285,6 +2397,7 @@ export class RunQueue { ckKeyPrefix, String(this.counterTtlSeconds), totalConcurrencyEnabledArg, + this.options.gatesEnabled ? "1" : "0", metricsGaugeArg ); } else { @@ -2321,6 +2434,7 @@ export class RunQueue { ckKeyPrefix, String(this.counterTtlSeconds), totalConcurrencyEnabledArg, + this.options.gatesEnabled ? "1" : "0", metricsGaugeArg ); } @@ -2353,6 +2467,8 @@ export class RunQueue { defaultEnvConcurrencyBurstFactor, currentTime, enableFastPathArg, + this.options.redis.keyPrefix ?? "", + this.options.gatesEnabled ? "1" : "0", metricsGaugeArg ); } else { @@ -2380,6 +2496,8 @@ export class RunQueue { defaultEnvConcurrencyBurstFactor, currentTime, enableFastPathArg, + this.options.redis.keyPrefix ?? "", + this.options.gatesEnabled ? "1" : "0", metricsGaugeArg ); } @@ -2460,6 +2578,7 @@ export class RunQueue { String(this.options.defaultEnvConcurrencyBurstFactor ?? 1), this.options.redis.keyPrefix ?? "", String(maxCount), + this.options.gatesEnabled ? "1" : "0", metricsGaugeArg ); @@ -2596,6 +2715,7 @@ export class RunQueue { this.options.redis.keyPrefix ?? "", String(maxCount), this.options.totalConcurrencyEnabled ? "1" : "0", + this.options.gatesEnabled ? "1" : "0", metricsGaugeArg ); @@ -2850,7 +2970,8 @@ export class RunQueue { messageQueue, messageKeyValue, removeFromWorkerQueue ? "1" : "0", - ckWildcardName + ckWildcardName, + this.options.redis.keyPrefix ?? "" ); } @@ -2867,7 +2988,8 @@ export class RunQueue { messageId, messageQueue, messageKeyValue, - removeFromWorkerQueue ? "1" : "0" + removeFromWorkerQueue ? "1" : "0", + this.options.redis.keyPrefix ?? "" ); } @@ -2910,6 +3032,7 @@ export class RunQueue { this.keys.queueRunningCounterKeyFromQueue(queue), this.keys.ckIndexKeyFromQueue(queue), this.keys.queueGroupConcurrencyKeyFromQueue(queue), + messageKey, messageId, this.options.redis.keyPrefix ?? "", String(this.counterTtlSeconds) @@ -2921,7 +3044,9 @@ export class RunQueue { envCurrentConcurrencyKey, queueCurrentDequeuedKey, envCurrentDequeuedKey, - messageId + messageKey, + messageId, + this.options.redis.keyPrefix ?? "" ); } @@ -3001,7 +3126,8 @@ export class RunQueue { messageId, messageQueue, JSON.stringify(message), - String(messageScore) + String(messageScore), + this.options.redis.keyPrefix ?? "" ); } } @@ -3043,7 +3169,8 @@ export class RunQueue { this.keys.queueGroupConcurrencyKeyFromQueue(messageQueue), messageId, messageQueue, - ckWildcardName + ckWildcardName, + this.options.redis.keyPrefix ?? "" ); } else { await this.redis.moveToDeadLetterQueue( @@ -3057,7 +3184,8 @@ export class RunQueue { envQueueKey, deadLetterQueueKey, messageId, - messageQueue + messageQueue, + this.options.redis.keyPrefix ?? "" ); } } @@ -3445,8 +3573,11 @@ local defaultEnvConcurrencyLimit = ARGV[6] local defaultEnvConcurrencyBurstFactor = ARGV[7] local currentTime = ARGV[8] local enableFastPath = ARGV[9] +local keyPrefix = ARGV[10] +local gatesEnabled = ARGV[11] == '1' ${QUEUE_METRICS_GAUGE_PRELUDE} +${QUEUE_GATES_LUA_HELPERS} -- Fast path: check if we can skip the queue and go directly to worker queue if enableFastPath == '1' then @@ -3465,12 +3596,27 @@ 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) + local gateMsg = nil + local gatesAllowFastPath = true + if gatesEnabled and string.find(messageData, '"gates"', 1, true) 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) + end + end + + if gatesAllowFastPath then + redis.call('SET', messageKey, messageData) + redis.call('SADD', queueCurrentConcurrencyKey, messageId) + redis.call('SADD', envCurrentConcurrencyKey, messageId) + if gateMsg then + __gatesAcquire(keyPrefix, gateMsg, messageId) + end + redis.call('RPUSH', workerQueueKey, messageKeyValue) ${QUEUE_METRICS_ENQUEUE_FASTPATH_GAUGE_LUA} - return __qmret(1) + return __qmret(1) + end end end end @@ -3500,6 +3646,7 @@ redis.call('SREM', queueCurrentConcurrencyKey, messageId) redis.call('SREM', envCurrentConcurrencyKey, messageId) redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) +__gatesRelease(keyPrefix, messageData, messageId) ${QUEUE_METRICS_GAUGE_LUA} return __qmret(0) @@ -3540,8 +3687,11 @@ local defaultEnvConcurrencyLimit = ARGV[8] local defaultEnvConcurrencyBurstFactor = ARGV[9] local currentTime = ARGV[10] local enableFastPath = ARGV[11] +local keyPrefix = ARGV[12] +local gatesEnabled = ARGV[13] == '1' ${QUEUE_METRICS_GAUGE_PRELUDE} +${QUEUE_GATES_LUA_HELPERS} -- Fast path: check if we can skip the queue and go directly to worker queue if enableFastPath == '1' then @@ -3560,13 +3710,28 @@ 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) + local gateMsg = nil + local gatesAllowFastPath = true + if gatesEnabled and string.find(messageData, '"gates"', 1, true) 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) + end + end + + if gatesAllowFastPath then + redis.call('SET', messageKey, messageData) + redis.call('SADD', queueCurrentConcurrencyKey, messageId) + redis.call('SADD', envCurrentConcurrencyKey, messageId) + if gateMsg then + __gatesAcquire(keyPrefix, gateMsg, messageId) + end + redis.call('RPUSH', workerQueueKey, messageKeyValue) ${QUEUE_METRICS_ENQUEUE_FASTPATH_GAUGE_LUA} -- Skip TTL sorted set: the expireRun worker job handles TTL expiry independently - return __qmret(1) + return __qmret(1) + end end end end @@ -3599,6 +3764,7 @@ redis.call('SREM', queueCurrentConcurrencyKey, messageId) redis.call('SREM', envCurrentConcurrencyKey, messageId) redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) +__gatesRelease(keyPrefix, messageData, messageId) ${QUEUE_METRICS_GAUGE_LUA} return __qmret(0) @@ -3868,8 +4034,10 @@ local keyPrefix = ARGV[11] -- TTL (seconds) applied to counter lazy-init SETs local counterTtl = ARGV[12] local totalConcurrencyEnabled = ARGV[13] == '1' +local gatesEnabled = ARGV[14] == '1' ${QUEUE_METRICS_GAUGE_PRELUDE} +${QUEUE_GATES_LUA_HELPERS} -- Fast path: check if we can skip the queue and go directly to worker queue if enableFastPath == '1' then @@ -3902,13 +4070,26 @@ if enableFastPath == '1' then end end - if totalAllowsFastPath then + local gateMsg = nil + local gatesAllowFastPath = true + if gatesEnabled and string.find(messageData, '"gates"', 1, true) 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) + end + end + + if totalAllowsFastPath and gatesAllowFastPath 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 + if gateMsg then + __gatesAcquire(keyPrefix, gateMsg, 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. @@ -3976,6 +4157,7 @@ end redis.call('SREM', envCurrentConcurrencyKey, messageId) redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) +__gatesRelease(keyPrefix, messageData, messageId) ${QUEUE_METRICS_CK_ENQUEUE_GAUGE_LUA} return __qmret(0) @@ -4025,8 +4207,10 @@ local keyPrefix = ARGV[13] -- TTL (seconds) applied to counter lazy-init SETs local counterTtl = ARGV[14] local totalConcurrencyEnabled = ARGV[15] == '1' +local gatesEnabled = ARGV[16] == '1' ${QUEUE_METRICS_GAUGE_PRELUDE} +${QUEUE_GATES_LUA_HELPERS} -- Fast path: check if we can skip the queue and go directly to worker queue if enableFastPath == '1' then @@ -4057,13 +4241,26 @@ if enableFastPath == '1' then end end - if totalAllowsFastPath then + local gateMsg = nil + local gatesAllowFastPath = true + if gatesEnabled and string.find(messageData, '"gates"', 1, true) 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) + end + end + + if totalAllowsFastPath and gatesAllowFastPath 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 + if gateMsg then + __gatesAcquire(keyPrefix, gateMsg, messageId) + end redis.call('RPUSH', workerQueueKey, messageKeyValue) ${QUEUE_METRICS_CK_ENQUEUE_FASTPATH_GAUGE_LUA} return __qmret(1) @@ -4124,6 +4321,7 @@ end redis.call('SREM', envCurrentConcurrencyKey, messageId) redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) +__gatesRelease(keyPrefix, messageData, messageId) ${QUEUE_METRICS_CK_ENQUEUE_GAUGE_LUA} return __qmret(0) @@ -4259,6 +4457,7 @@ local function decrFloored(key) redis.call('DECR', key) end end +${QUEUE_GATES_LUA_HELPERS} local expiredMembers = redis.call('ZRANGEBYSCORE', ttlQueueKey, '-inf', currentTime, 'LIMIT', 0, batchSize) @@ -4290,6 +4489,7 @@ for i, member in ipairs(expiredMembers) do local messageKey = keyPrefix .. "{org:" .. orgFromQueue .. "}:message:" .. runId + local rawPayload = redis.call('GET', messageKey) redis.call('DEL', messageKey) -- ZREM from queue; if successful AND this is a CK variant, DECR lengthCounter. @@ -4305,6 +4505,7 @@ for i, member in ipairs(expiredMembers) do local dequeuedKey = queueKey .. ":currentDequeued" local removedFromCurrent = redis.call('SREM', concurrencyKey, runId) local removedFromDequeued = redis.call('SREM', dequeuedKey, runId) + __gatesRelease(keyPrefix, rawPayload, runId) local projMatch = string.match(rawQueueKey, ":proj:([^:]+):env:") local envConcurrencyKey = keyPrefix .. "{org:" .. orgFromQueue .. "}:proj:" .. (projMatch or "") .. ":env:" .. (envMatch or "") .. ":currentConcurrency" @@ -4375,7 +4576,9 @@ local defaultEnvConcurrencyLimit = ARGV[3] local defaultEnvConcurrencyBurstFactor = ARGV[4] local keyPrefix = ARGV[5] local maxCount = tonumber(ARGV[6] or '1') +local gatesEnabled = ARGV[7] == '1' ${QUEUE_METRICS_GAUGE_PRELUDE} +${QUEUE_GATES_LUA_HELPERS} ${QUEUE_METRICS_GAUGE_LUA} -- Check current env concurrency against the limit @@ -4444,24 +4647,31 @@ for i = 1, #messages, 2 do redis.call('ZADD', ttlQueueKey, ttlExpiresAt, ttlMember) end else - -- Not expired - process normally - redis.call('ZREM', queueKey, messageId) - redis.call('ZREM', envQueueKey, messageId) - redis.call('SADD', queueCurrentConcurrencyKey, messageId) - redis.call('SADD', envCurrentConcurrencyKey, messageId) - - -- Remove from TTL set if provided (run is being executed, not expired) - if ttlQueueKey and ttlQueueKey ~= '' and ttlExpiresAt then - local ttlMember = queueName .. '|' .. messageId .. '|' .. (messageData.orgId or '') - redis.call('ZREM', ttlQueueKey, ttlMember) + local gatesAllow = true + if gatesEnabled then + gatesAllow = __gatesHaveCapacity(keyPrefix, messageData, messageId, envConcurrencyLimit, messageKeyPrefix) end - -- Add to results - table.insert(results, messageId) - table.insert(results, messageScore) - table.insert(results, messagePayload) - - dequeuedCount = dequeuedCount + 1 + if gatesAllow then + redis.call('ZREM', queueKey, messageId) + redis.call('ZREM', envQueueKey, messageId) + redis.call('SADD', queueCurrentConcurrencyKey, messageId) + redis.call('SADD', envCurrentConcurrencyKey, messageId) + if gatesEnabled then + __gatesAcquire(keyPrefix, messageData, messageId) + end + + if ttlQueueKey and ttlQueueKey ~= '' and ttlExpiresAt then + local ttlMember = queueName .. '|' .. messageId .. '|' .. (messageData.orgId or '') + redis.call('ZREM', ttlQueueKey, ttlMember) + end + + table.insert(results, messageId) + table.insert(results, messageScore) + table.insert(results, messagePayload) + + dequeuedCount = dequeuedCount + 1 + end end else -- Stale entry: message key was already deleted (e.g. acknowledged), @@ -4664,7 +4874,9 @@ local defaultEnvConcurrencyBurstFactor = ARGV[4] local keyPrefix = ARGV[5] local maxCount = tonumber(ARGV[6] or '1') local totalConcurrencyEnabled = ARGV[7] == '1' +local gatesEnabled = ARGV[8] == '1' ${QUEUE_METRICS_GAUGE_PRELUDE} +${QUEUE_GATES_LUA_HELPERS} ${QUEUE_METRICS_CK_DEQUEUE_GAUGE_LUA} local function decrLengthCounter() @@ -4688,43 +4900,30 @@ 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. +-- Total-cap gate: track the remaining headroom across ALL ck variants +-- (groupConcurrency SCARD vs the env-clamped total limit). Admission below +-- decrements the headroom per new member, and a message that is ALREADY a +-- group member passes even at zero headroom: it holds its own slot (an +-- unmirrored release from an older build can leave a queued run's membership +-- behind, and blocking on it would deadlock the run against itself). +local totalHeadroom = nil 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. + -- Self-heal before holding the queue at its limit: a member with no message + -- key was terminally released without the mirror and is dead, and a member + -- whose message is QUEUED (in its variant zset) holds no legitimate slot, + -- since queued and in-flight are mutually exclusive on every path. Both are + -- pruned by the shared bounded reconcile. 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 + __gateReconcile(groupConcurrencyKey, messageKeyPrefix, keyPrefix) + groupCurrentConcurrency = tonumber(redis.call('SCARD', groupConcurrencyKey) or '0') end - actualMaxCount = math.min(actualMaxCount, totalConcurrencyLimit - groupCurrentConcurrency) + totalHeadroom = totalConcurrencyLimit - groupCurrentConcurrency end end @@ -4780,25 +4979,45 @@ for _, ckQueueName in ipairs(ckQueues) do redis.call('ZADD', ttlQueueKey, ttlExpiresAt, ttlMember) end else - redis.call('ZREM', fullQueueKey, messageId) - redis.call('ZREM', envQueueKey, messageId) - decrLengthCounter() - redis.call('SADD', ckConcurrencyKey, messageId) - redis.call('SADD', envCurrentConcurrencyKey, messageId) - if totalConcurrencyEnabled then - redis.call('SADD', groupConcurrencyKey, messageId) + local gatesAllow = true + if gatesEnabled then + gatesAllow = __gatesHaveCapacity(keyPrefix, messageData, messageId, envConcurrencyLimit, messageKeyPrefix) end - if ttlQueueKey and ttlQueueKey ~= '' and ttlExpiresAt then - local ttlMember = ckQueueName .. '|' .. messageId .. '|' .. (messageData.orgId or '') - redis.call('ZREM', ttlQueueKey, ttlMember) + local alreadyInGroup = false + local totalAllows = true + if totalHeadroom ~= nil then + alreadyInGroup = redis.call('SISMEMBER', groupConcurrencyKey, messageId) == 1 + totalAllows = alreadyInGroup or totalHeadroom > 0 end - table.insert(results, messageId) - table.insert(results, messageScore) - table.insert(results, messagePayload) + if gatesAllow and totalAllows then + redis.call('ZREM', fullQueueKey, messageId) + redis.call('ZREM', envQueueKey, messageId) + decrLengthCounter() + redis.call('SADD', ckConcurrencyKey, messageId) + redis.call('SADD', envCurrentConcurrencyKey, messageId) + if totalConcurrencyEnabled then + redis.call('SADD', groupConcurrencyKey, messageId) + if totalHeadroom ~= nil and not alreadyInGroup then + totalHeadroom = totalHeadroom - 1 + end + end + if gatesEnabled then + __gatesAcquire(keyPrefix, messageData, messageId) + end - dequeuedCount = dequeuedCount + 1 + if ttlQueueKey and ttlQueueKey ~= '' and ttlExpiresAt then + local ttlMember = ckQueueName .. '|' .. messageId .. '|' .. (messageData.orgId or '') + redis.call('ZREM', ttlQueueKey, ttlMember) + end + + table.insert(results, messageId) + table.insert(results, messageScore) + table.insert(results, messagePayload) + + dequeuedCount = dequeuedCount + 1 + end end else redis.call('ZREM', fullQueueKey, messageId) @@ -4974,6 +5193,10 @@ local messageId = ARGV[1] local messageQueueName = ARGV[2] local messageKeyValue = ARGV[3] local removeFromWorkerQueue = ARGV[4] +local keyPrefix = ARGV[5] +${QUEUE_GATES_LUA_HELPERS} + +local rawPayload = redis.call('GET', messageKey) -- Remove the message from the message key redis.call('DEL', messageKey) @@ -4995,6 +5218,7 @@ redis.call('SREM', queueCurrentConcurrencyKey, messageId) redis.call('SREM', envCurrentConcurrencyKey, messageId) redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) +__gatesRelease(keyPrefix, rawPayload, messageId) -- Remove the message from the worker queue if removeFromWorkerQueue == '1' then @@ -5021,6 +5245,8 @@ local messageId = ARGV[1] local messageQueueName = ARGV[2] local messageData = ARGV[3] local messageScore = tonumber(ARGV[4]) +local keyPrefix = ARGV[5] +${QUEUE_GATES_LUA_HELPERS} -- Update the message data redis.call('SET', messageKey, messageData) @@ -5030,6 +5256,7 @@ redis.call('SREM', queueCurrentConcurrencyKey, messageId) redis.call('SREM', envCurrentConcurrencyKey, messageId) redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) +__gatesRelease(keyPrefix, messageData, messageId) -- Enqueue the message into the queue redis.call('ZADD', messageQueueKey, messageScore, messageId) @@ -5062,6 +5289,10 @@ local deadLetterQueueKey = KEYS[9] -- Args: local messageId = ARGV[1] local messageQueueName = ARGV[2] +local keyPrefix = ARGV[3] +${QUEUE_GATES_LUA_HELPERS} + +local rawPayload = redis.call('GET', messageKey) -- Remove the message from the queue redis.call('ZREM', messageQueue, messageId) @@ -5083,6 +5314,7 @@ redis.call('SREM', queueCurrentConcurrencyKey, messageId) redis.call('SREM', envCurrentConcurrencyKey, messageId) redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) +__gatesRelease(keyPrefix, rawPayload, messageId) `, }); @@ -5300,6 +5532,10 @@ local messageQueueName = ARGV[2] local messageKeyValue = ARGV[3] local removeFromWorkerQueue = ARGV[4] local ckWildcardName = ARGV[5] +local keyPrefix = ARGV[6] +${QUEUE_GATES_LUA_HELPERS} + +local rawPayload = redis.call('GET', messageKey) local function decrFloored(key) if tonumber(redis.call('GET', key) or '0') > 0 then @@ -5357,6 +5593,7 @@ redis.call('SREM', envCurrentDequeuedKey, messageId) if removedFromDequeued == 1 then decrFloored(runningCounterKey) end +__gatesRelease(keyPrefix, rawPayload, messageId) -- Remove the message from the worker queue if removeFromWorkerQueue == '1' then @@ -5395,6 +5632,7 @@ local ckWildcardName = ARGV[5] local keyPrefix = ARGV[6] -- TTL (seconds) applied to counter lazy-init SETs local counterTtl = ARGV[7] +${QUEUE_GATES_LUA_HELPERS} local function decrFloored(key) if tonumber(redis.call('GET', key) or '0') > 0 then @@ -5420,6 +5658,7 @@ redis.call('SREM', envCurrentDequeuedKey, messageId) if removedFromDequeued == 1 then decrFloored(runningCounterKey) end +__gatesRelease(keyPrefix, messageData, messageId) -- Lazy-init lengthCounter if missing (e.g. expired via 24h TTL). nack re-queues a -- message, which means lengthCounter must be present before we INCR. Without this, @@ -5490,6 +5729,10 @@ local groupConcurrencyKey = KEYS[13] local messageId = ARGV[1] local messageQueueName = ARGV[2] local ckWildcardName = ARGV[3] +local keyPrefix = ARGV[4] +${QUEUE_GATES_LUA_HELPERS} + +local rawPayload = redis.call('GET', messageKey) local function decrFloored(key) if tonumber(redis.call('GET', key) or '0') > 0 then @@ -5546,26 +5789,31 @@ redis.call('SREM', envCurrentDequeuedKey, messageId) if removedFromDequeued == 1 then decrFloored(runningCounterKey) end +__gatesRelease(keyPrefix, rawPayload, messageId) `, }); this.redis.defineCommand("releaseConcurrency", { - numberOfKeys: 4, + numberOfKeys: 5, lua: ` -- Keys: local queueCurrentConcurrencyKey = KEYS[1] local envCurrentConcurrencyKey = KEYS[2] local queueCurrentDequeuedKey = KEYS[3] local envCurrentDequeuedKey = KEYS[4] +local messageKey = KEYS[5] -- Args: local messageId = ARGV[1] +local keyPrefix = ARGV[2] +${QUEUE_GATES_LUA_HELPERS} -- Update the concurrency keys redis.call('SREM', queueCurrentConcurrencyKey, messageId) redis.call('SREM', envCurrentConcurrencyKey, messageId) redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) +__gatesRelease(keyPrefix, redis.call('GET', messageKey), messageId) `, }); @@ -5574,7 +5822,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: 7, + numberOfKeys: 8, lua: ` -- Keys: local queueCurrentConcurrencyKey = KEYS[1] @@ -5584,12 +5832,14 @@ local envCurrentDequeuedKey = KEYS[4] local runningCounterKey = KEYS[5] local ckIndexKey = KEYS[6] local groupConcurrencyKey = KEYS[7] +local messageKey = KEYS[8] -- Args: local messageId = ARGV[1] local keyPrefix = ARGV[2] -- TTL (seconds) applied to counter lazy-init SETs local counterTtl = ARGV[3] +${QUEUE_GATES_LUA_HELPERS} -- Lazy-init runningCounter if missing (e.g. expired via 24h TTL). Runs BEFORE -- the SREM so the seed captures pre-release state; the subsequent DECR accounts @@ -5618,6 +5868,7 @@ if removedFromDequeued == 1 then redis.call('DECR', runningCounterKey) end end +__gatesRelease(keyPrefix, redis.call('GET', messageKey), messageId) `, }); @@ -5701,29 +5952,33 @@ return results }); this.redis.defineCommand("clearMessageFromConcurrencySets", { - numberOfKeys: 4, + numberOfKeys: 5, lua: ` -- Keys: local queueCurrentConcurrencyKey = KEYS[1] local envCurrentConcurrencyKey = KEYS[2] local queueCurrentDequeuedKey = KEYS[3] local envCurrentDequeuedKey = KEYS[4] +local messageKey = KEYS[5] -- Args: local messageId = ARGV[1] +local keyPrefix = ARGV[2] +${QUEUE_GATES_LUA_HELPERS} -- Update the concurrency keys redis.call('SREM', queueCurrentConcurrencyKey, messageId) redis.call('SREM', envCurrentConcurrencyKey, messageId) redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) +__gatesRelease(keyPrefix, redis.call('GET', messageKey), messageId) `, }); // Tracked variant of clearMessageFromConcurrencySets — see releaseConcurrencyTracked // for the contract. Only invoke for CK queues. this.redis.defineCommand("clearMessageFromConcurrencySetsTracked", { - numberOfKeys: 7, + numberOfKeys: 8, lua: ` -- Keys: local queueCurrentConcurrencyKey = KEYS[1] @@ -5733,12 +5988,14 @@ local envCurrentDequeuedKey = KEYS[4] local runningCounterKey = KEYS[5] local ckIndexKey = KEYS[6] local groupConcurrencyKey = KEYS[7] +local messageKey = KEYS[8] -- Args: local messageId = ARGV[1] local keyPrefix = ARGV[2] -- TTL (seconds) applied to counter lazy-init SETs local counterTtl = ARGV[3] +${QUEUE_GATES_LUA_HELPERS} -- Lazy-init runningCounter if missing — see releaseConcurrencyTracked for rationale. if redis.call('EXISTS', runningCounterKey) == 0 then @@ -5763,6 +6020,7 @@ if removedFromDequeued == 1 then redis.call('DECR', runningCounterKey) end end +__gatesRelease(keyPrefix, redis.call('GET', messageKey), messageId) `, }); } @@ -5802,6 +6060,8 @@ declare module "@internal/redis" { defaultEnvConcurrencyBurstFactor: string, currentTime: string, enableFastPath: string, + keyPrefix: string, + gatesEnabled: string, metricsEnabled: string, callback?: Callback<[number, number[] | null]> ): Result<[number, number[] | null], Context>; @@ -5833,6 +6093,8 @@ declare module "@internal/redis" { defaultEnvConcurrencyBurstFactor: string, currentTime: string, enableFastPath: string, + keyPrefix: string, + gatesEnabled: string, metricsEnabled: string, callback?: Callback<[number, number[] | null]> ): Result<[number, number[] | null], Context>; @@ -5870,6 +6132,7 @@ declare module "@internal/redis" { defaultEnvConcurrencyBurstFactor: string, keyPrefix: string, maxCount: string, + gatesEnabled: string, metricsEnabled: string, callback?: Callback<[string[] | null, number[] | null]> ): Result<[string[] | null, number[] | null], Context>; @@ -5903,6 +6166,7 @@ declare module "@internal/redis" { messageQueueName: string, messageKeyValue: string, removeFromWorkerQueue: string, + keyPrefix: string, callback?: Callback ): Result; @@ -5912,8 +6176,10 @@ declare module "@internal/redis" { envCurrentConcurrencyKey: string, queueCurrentDequeuedKey: string, envCurrentDequeuedKey: string, + messageKey: string, // args messageId: string, + keyPrefix: string, callback?: Callback ): Result; @@ -5932,6 +6198,7 @@ declare module "@internal/redis" { messageQueueName: string, messageData: string, messageScore: string, + keyPrefix: string, callback?: Callback ): Result; @@ -5949,6 +6216,7 @@ declare module "@internal/redis" { // args messageId: string, messageQueueName: string, + keyPrefix: string, callback?: Callback ): Result; @@ -5958,8 +6226,10 @@ declare module "@internal/redis" { envCurrentConcurrencyKey: string, queueCurrentDequeuedKey: string, envCurrentDequeuedKey: string, + messageKey: string, // args messageId: string, + keyPrefix: string, callback?: Callback ): Result; @@ -6162,6 +6432,7 @@ declare module "@internal/redis" { keyPrefix: string, counterTtl: string, totalConcurrencyEnabled: string, + gatesEnabled: string, metricsEnabled: string, callback?: Callback<[number, number[] | null]> ): Result<[number, number[] | null], Context>; @@ -6200,6 +6471,7 @@ declare module "@internal/redis" { keyPrefix: string, counterTtl: string, totalConcurrencyEnabled: string, + gatesEnabled: string, metricsEnabled: string, callback?: Callback<[number, number[] | null]> ): Result<[number, number[] | null], Context>; @@ -6225,6 +6497,7 @@ declare module "@internal/redis" { keyPrefix: string, maxCount: string, totalConcurrencyEnabled: string, + gatesEnabled: string, metricsEnabled: string, callback?: Callback<[string[] | null, number[] | null]> ): Result<[string[] | null, number[] | null], Context>; @@ -6255,6 +6528,7 @@ declare module "@internal/redis" { messageKeyValue: string, removeFromWorkerQueue: string, ckWildcardName: string, + keyPrefix: string, callback?: Callback ): Result; @@ -6298,6 +6572,7 @@ declare module "@internal/redis" { messageId: string, messageQueueName: string, ckWildcardName: string, + keyPrefix: string, callback?: Callback ): Result; @@ -6321,6 +6596,7 @@ declare module "@internal/redis" { runningCounterKey: string, ckIndexKey: string, groupConcurrencyKey: string, + messageKey: string, messageId: string, keyPrefix: string, counterTtl: string, @@ -6335,6 +6611,7 @@ declare module "@internal/redis" { runningCounterKey: string, ckIndexKey: string, groupConcurrencyKey: string, + messageKey: string, messageId: string, keyPrefix: string, counterTtl: string, diff --git a/internal-packages/run-engine/src/run-queue/tests/queueGates.test.ts b/internal-packages/run-engine/src/run-queue/tests/queueGates.test.ts new file mode 100644 index 00000000000..eb6dbda41fc --- /dev/null +++ b/internal-packages/run-engine/src/run-queue/tests/queueGates.test.ts @@ -0,0 +1,395 @@ +import { 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, gatesEnabled: boolean) { + return new RunQueue({ + ...testOptions, + gatesEnabled, + 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(); +} + +async function popWorkerQueue(queue: RunQueue, expected: string): Promise { + const next = await queue.dequeueMessageFromWorkerQueue("consumer-1", "main", { + blockingPop: false, + }); + return next?.messageId === expected; +} + +vi.setConfig({ testTimeout: 60_000 }); + +describe("RunQueue gates", () => { + redisTest("an unkeyed gate caps runs across its holders", async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "shared-gate", 1); + + const now = Date.now(); + for (const i of [0, 1]) { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: `r${i}`, + timestamp: now - 1000 + i, + gates: [{ queue: "shared-gate" }], + }), + workerQueue: "main", + }); + } + + const oneAdmitted = await waitFor( + async () => + (await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "shared-gate")) === 1 + ); + expect(oneAdmitted).toBe(true); + + /** The second run must stay queued while the gate is full. */ + await setTimeout(2000); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "shared-gate")).toBe(1); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + + expect(await popWorkerQueue(queue, "r0")).toBe(true); + await queue.acknowledgeMessage(authenticatedEnvDev.organization.id, "r0"); + + /** Acking r0 frees the gate slot and its home slot; r1 is admitted. */ + const r1Admitted = await waitFor(() => popWorkerQueue(queue, "r1")); + expect(r1Admitted).toBe(true); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "shared-gate")).toBe(1); + } finally { + await queue.quit(); + } + }); + + redisTest( + "a keyed gate caps a tenant across home concurrency keys", + async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "tenant", 1); + + const now = Date.now(); + for (const [i, ck] of ["ck-a", "ck-b"].entries()) { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: `r${i}`, + concurrencyKey: ck, + timestamp: now - 1000 + i, + gates: [{ queue: "tenant", concurrencyKey: "acme" }], + }), + workerQueue: "main", + }); + } + + const oneAdmitted = await waitFor( + async () => + (await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "tenant", "acme")) === 1 + ); + expect(oneAdmitted).toBe(true); + + await setTimeout(2000); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "tenant", "acme")).toBe( + 1 + ); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "tenant")).toBe(1); + + expect(await popWorkerQueue(queue, "r0")).toBe(true); + await queue.acknowledgeMessage(authenticatedEnvDev.organization.id, "r0"); + + const r1Admitted = await waitFor(() => popWorkerQueue(queue, "r1")); + expect(r1Admitted).toBe(true); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "tenant", "acme")).toBe( + 1 + ); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "tenant")).toBe(1); + } finally { + await queue.quit(); + } + } + ); + + redisTest("ignores gates and holds no gate slots when disabled", async ({ redisContainer }) => { + const queue = createQueue(redisContainer, false); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "shared-gate", 1); + + const now = Date.now(); + for (const i of [0, 1]) { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: `r${i}`, + timestamp: now - 1000 + i, + gates: [{ queue: "shared-gate" }], + }), + workerQueue: "main", + }); + } + + const bothAdmitted = await waitFor( + async () => (await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")) === 0 + ); + expect(bothAdmitted).toBe(true); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "shared-gate")).toBe(0); + } finally { + await queue.quit(); + } + }); + + redisTest("nacking releases the gate slot", async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "shared-gate", 1); + + const now = Date.now(); + for (const i of [0, 1]) { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: `r${i}`, + timestamp: now - 1000 + i, + gates: [{ queue: "shared-gate" }], + }), + workerQueue: "main", + }); + } + + const r0Admitted = await waitFor(async () => { + if ((await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "shared-gate")) !== 1) { + return false; + } + return popWorkerQueue(queue, "r0"); + }); + expect(r0Admitted).toBe(true); + + await queue.nackMessage({ + orgId: authenticatedEnvDev.organization.id, + messageId: "r0", + retryAt: Date.now() + 120_000, + }); + + /** r0's gate slot must be released so r1 (the only eligible run) is admitted. */ + const r1Admitted = await waitFor(() => popWorkerQueue(queue, "r1")); + expect(r1Admitted).toBe(true); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "shared-gate")).toBe(1); + } finally { + await queue.quit(); + } + }); + + redisTest("reconciles a leaked gate member instead of blocking", async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + const keys = testOptions.keys; + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "shared-gate", 1); + + /** A dead member with no message key occupies the gate. */ + await queue.redis.sadd( + keys.queueCurrentConcurrencyKey(authenticatedEnvDev, "shared-gate"), + "dead-0" + ); + + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: "r0", + timestamp: Date.now() - 1000, + gates: [{ queue: "shared-gate" }], + }), + workerQueue: "main", + }); + + const r0Admitted = await waitFor(() => popWorkerQueue(queue, "r0"), 30_000); + expect(r0Admitted).toBe(true); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "shared-gate")).toBe(1); + } finally { + await queue.quit(); + } + }); + + redisTest( + "a run already holding its own gate slot is never deadlocked by it", + async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + const keys = testOptions.keys; + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "shared-gate", 1); + + /** + * An unmirrored release (an older build's nack) leaves the run's own + * membership behind while the run goes back to waiting in its queue. The + * gate is "full" with the run itself; admission must still let it through. + */ + await queue.redis.sadd( + keys.queueCurrentConcurrencyKey(authenticatedEnvDev, "shared-gate"), + "r0" + ); + + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: "r0", + timestamp: Date.now() - 1000, + gates: [{ queue: "shared-gate" }], + }), + workerQueue: "main", + }); + + const r0Admitted = await waitFor(() => popWorkerQueue(queue, "r0"), 30_000); + expect(r0Admitted).toBe(true); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "shared-gate")).toBe(1); + } finally { + await queue.quit(); + } + } + ); + + redisTest( + "a gate without a key inherits the run's concurrency key", + async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "tenant", 1); + + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: "r0", + concurrencyKey: "acme", + timestamp: Date.now() - 1000, + gates: [{ queue: "tenant" }], + }), + workerQueue: "main", + }); + + const admitted = await waitFor( + async () => + (await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "tenant", "acme")) === 1 + ); + expect(admitted).toBe(true); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "tenant")).toBe(0); + } finally { + await queue.quit(); + } + } + ); + + redisTest("enqueue fast path respects a full gate", async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "shared-gate", 1); + + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: "r0", + timestamp: Date.now() - 1000, + gates: [{ queue: "shared-gate" }], + }), + workerQueue: "main", + enableFastPath: true, + skipDequeueProcessing: true, + }); + + /** The fast path admits synchronously and takes the gate slot. */ + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "shared-gate")).toBe(1); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(0); + + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: "r1", + timestamp: Date.now() - 999, + gates: [{ queue: "shared-gate" }], + }), + workerQueue: "main", + enableFastPath: true, + skipDequeueProcessing: true, + }); + + /** At the gate's limit the fast path must fall back to a normal enqueue. */ + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "shared-gate")).toBe(1); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + } finally { + await queue.quit(); + } + }); +}); 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 index cd6d299f151..8af76226042 100644 --- a/internal-packages/run-engine/src/run-queue/tests/totalConcurrency.test.ts +++ b/internal-packages/run-engine/src/run-queue/tests/totalConcurrency.test.ts @@ -365,6 +365,64 @@ describe("RunQueue total concurrency limit", () => { } ); + redisTest( + "reconciles a member whose run went back to waiting in a 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); + + /** + * A run on ANOTHER queue whose message exists and is queued there, leaked + * into this queue's group set by an unmirrored release. It can never be a + * dequeue candidate here, so only the reconcile can free the slot: queued + * and in-flight are mutually exclusive, so a queued member holds nothing. + */ + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: "x0", + queue: "other-queue", + concurrencyKey: "ck-x", + timestamp: Date.now() + 3_600_000, + }), + workerQueue: "main", + }); + await queue.redis.sadd( + keys.queueGroupConcurrencyKey(authenticatedEnvDev, "task/my-task"), + "x0" + ); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: "r0", + concurrencyKey: "ck-a", + timestamp: Date.now() - 1000, + }), + workerQueue: "main", + }); + + const r0Admitted = await waitFor(async () => { + await queue.redis.del( + `${keys.queueGroupConcurrencyKey(authenticatedEnvDev, "task/my-task")}: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(); + } + } + ); + redisTest( "reconciles a large leaked backlog across bounded passes", async ({ redisContainer }) => { diff --git a/internal-packages/run-engine/src/run-queue/types.ts b/internal-packages/run-engine/src/run-queue/types.ts index 10575d8b5d3..97fb99417fc 100644 --- a/internal-packages/run-engine/src/run-queue/types.ts +++ b/internal-packages/run-engine/src/run-queue/types.ts @@ -2,6 +2,19 @@ import { z } from "zod"; import { RuntimeEnvironmentType } from "@trigger.dev/database"; import type { MinimalAuthenticatedEnvironment } from "../shared/index.js"; +/** + * A gate is another declared queue this run must also hold a concurrency slot in + * while it executes. The run waits in its own queue; each gate contributes an + * extra admit condition (the gate queue's per-key limit, and its total limit when + * the entry is keyed) and an extra slot held until release. `queue` is the bare + * queue name; the org/project/env scope comes from the run's own payload. + */ +const QueueGate = z.object({ + queue: z.string().min(1).max(128), + concurrencyKey: z.string().min(1).max(128).optional(), +}); +type QueueGate = z.infer; + export const InputPayload = z.object({ runId: z.string(), /** Deprecated: not read on the V2 dequeue path; will stop being written in a follow-up. Optional to keep new readers compatible with old payloads that still include it, and vice versa. */ @@ -19,6 +32,8 @@ export const InputPayload = z.object({ attempt: z.number(), /** TTL expiration timestamp (unix ms). If set, run will be expired when this time is reached. */ ttlExpiresAt: z.number().optional(), + /** Additional queues this run must also hold a slot in while executing. At most two. */ + gates: QueueGate.array().max(2).optional(), }); export type InputPayload = z.infer;