From c7bea516cb0b4e2c35614995a6e032439df7b552 Mon Sep 17 00:00:00 2001 From: Matt Aitken Date: Sat, 29 Aug 2026 10:00:03 +0100 Subject: [PATCH 1/3] feat(run-engine): queue gates enforced from the message payload A gate is another declared queue a run must also hold a concurrency slot in while it executes, carried in the message payload as gates = [{queue, concurrencyKey?}] (at most two). Admit paths (both dequeue scripts and every enqueue fast path) check each gate's per-key limit and, for keyed entries, its total limit, then acquire a slot in each gate. Every release path removes the run from its gate sets, driven by the payload itself with a cheap substring probe before decoding, so gateless messages pay near zero and releases stay correct across flag toggles. Enforcement sits behind RUN_ENGINE_QUEUE_GATES_ENABLED (default off). Saturated gates self-heal with the same bounded message-key reconciliation as the total cap. The SDK and trigger surface that populates gates follows on top of this. --- apps/webapp/app/env.server.ts | 1 + apps/webapp/app/v3/runEngine.server.ts | 1 + .../run-engine/src/engine/index.ts | 1 + .../run-engine/src/engine/types.ts | 2 + .../run-engine/src/run-queue/index.ts | 375 +++++++++++++++--- .../src/run-queue/tests/queueGates.test.ts | 326 +++++++++++++++ .../run-engine/src/run-queue/types.ts | 15 + 7 files changed, 667 insertions(+), 54 deletions(-) create mode 100644 internal-packages/run-engine/src/run-queue/tests/queueGates.test.ts 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..2676ff602df 100644 --- a/internal-packages/run-engine/src/run-queue/index.ts +++ b/internal-packages/run-engine/src/run-queue/index.ts @@ -58,6 +58,93 @@ 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 variant = base + if gate.concurrencyKey and gate.concurrencyKey ~= '' then + variant = base .. ':ck:' .. gate.concurrencyKey + end + return base, variant +end + +local function __gateReconcile(setKey, msgKeyPrefix) + 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', '500') + redis.call('SET', cursorKey, scanResult[1], 'EX', '3600') + for _, memberId in ipairs(scanResult[2]) do + if redis.call('EXISTS', msgKeyPrefix .. memberId) == 0 then + redis.call('SREM', setKey, memberId) + end + end + end +end + +local function __gatesHaveCapacity(gatesKeyPrefix, msg, envLimit, msgKeyPrefix) + if not msg.gates then return true end + for _, gate in ipairs(msg.gates) do + local base, variant = __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 then + __gateReconcile(variant .. ':currentConcurrency', msgKeyPrefix) + return false + end + if gate.concurrencyKey and gate.concurrencyKey ~= '' then + local rawTotal = redis.call('GET', base .. ':totalConcurrency') + if rawTotal then + local totalLimit = math.min(tonumber(rawTotal), envLimit) + if tonumber(redis.call('SCARD', base .. ':groupConcurrency') or '0') >= totalLimit then + __gateReconcile(base .. ':groupConcurrency', msgKeyPrefix) + 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 = __gateKeys(gatesKeyPrefix, msg, gate) + redis.call('SADD', variant .. ':currentConcurrency', messageId) + if gate.concurrencyKey and gate.concurrencyKey ~= '' 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 = __gateKeys(gatesKeyPrefix, msg, gate) + local removed = redis.call('SREM', variant .. ':currentConcurrency', messageId) + if removed == 1 and gate.concurrencyKey and gate.concurrencyKey ~= '' 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 +286,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 +1389,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 +1401,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 +2384,7 @@ export class RunQueue { ckKeyPrefix, String(this.counterTtlSeconds), totalConcurrencyEnabledArg, + this.options.gatesEnabled ? "1" : "0", metricsGaugeArg ); } else { @@ -2321,6 +2421,7 @@ export class RunQueue { ckKeyPrefix, String(this.counterTtlSeconds), totalConcurrencyEnabledArg, + this.options.gatesEnabled ? "1" : "0", metricsGaugeArg ); } @@ -2353,6 +2454,8 @@ export class RunQueue { defaultEnvConcurrencyBurstFactor, currentTime, enableFastPathArg, + this.options.redis.keyPrefix ?? "", + this.options.gatesEnabled ? "1" : "0", metricsGaugeArg ); } else { @@ -2380,6 +2483,8 @@ export class RunQueue { defaultEnvConcurrencyBurstFactor, currentTime, enableFastPathArg, + this.options.redis.keyPrefix ?? "", + this.options.gatesEnabled ? "1" : "0", metricsGaugeArg ); } @@ -2460,6 +2565,7 @@ export class RunQueue { String(this.options.defaultEnvConcurrencyBurstFactor ?? 1), this.options.redis.keyPrefix ?? "", String(maxCount), + this.options.gatesEnabled ? "1" : "0", metricsGaugeArg ); @@ -2596,6 +2702,7 @@ export class RunQueue { this.options.redis.keyPrefix ?? "", String(maxCount), this.options.totalConcurrencyEnabled ? "1" : "0", + this.options.gatesEnabled ? "1" : "0", metricsGaugeArg ); @@ -2850,7 +2957,8 @@ export class RunQueue { messageQueue, messageKeyValue, removeFromWorkerQueue ? "1" : "0", - ckWildcardName + ckWildcardName, + this.options.redis.keyPrefix ?? "" ); } @@ -2867,7 +2975,8 @@ export class RunQueue { messageId, messageQueue, messageKeyValue, - removeFromWorkerQueue ? "1" : "0" + removeFromWorkerQueue ? "1" : "0", + this.options.redis.keyPrefix ?? "" ); } @@ -2910,6 +3019,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 +3031,9 @@ export class RunQueue { envCurrentConcurrencyKey, queueCurrentDequeuedKey, envCurrentDequeuedKey, - messageId + messageKey, + messageId, + this.options.redis.keyPrefix ?? "" ); } @@ -3001,7 +3113,8 @@ export class RunQueue { messageId, messageQueue, JSON.stringify(message), - String(messageScore) + String(messageScore), + this.options.redis.keyPrefix ?? "" ); } } @@ -3043,7 +3156,8 @@ export class RunQueue { this.keys.queueGroupConcurrencyKeyFromQueue(messageQueue), messageId, messageQueue, - ckWildcardName + ckWildcardName, + this.options.redis.keyPrefix ?? "" ); } else { await this.redis.moveToDeadLetterQueue( @@ -3057,7 +3171,8 @@ export class RunQueue { envQueueKey, deadLetterQueueKey, messageId, - messageQueue + messageQueue, + this.options.redis.keyPrefix ?? "" ); } } @@ -3445,8 +3560,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 +3583,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, 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 +3633,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 +3674,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 +3697,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, 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 +3751,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 +4021,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 +4057,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, 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 +4144,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 +4194,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 +4228,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, 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 +4308,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 +4444,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 +4476,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 +4492,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 +4563,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 +4634,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, 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 +4861,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() @@ -4780,25 +4979,35 @@ 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, envConcurrencyLimit, messageKeyPrefix) end - if ttlQueueKey and ttlQueueKey ~= '' and ttlExpiresAt then - local ttlMember = ckQueueName .. '|' .. messageId .. '|' .. (messageData.orgId or '') - redis.call('ZREM', ttlQueueKey, ttlMember) - end + if gatesAllow 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) + end + if gatesEnabled then + __gatesAcquire(keyPrefix, messageData, messageId) + end - table.insert(results, messageId) - table.insert(results, messageScore) - table.insert(results, messagePayload) + if ttlQueueKey and ttlQueueKey ~= '' and ttlExpiresAt then + local ttlMember = ckQueueName .. '|' .. messageId .. '|' .. (messageData.orgId or '') + redis.call('ZREM', ttlQueueKey, ttlMember) + end - dequeuedCount = dequeuedCount + 1 + 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 +5183,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 +5208,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 +5235,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 +5246,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 +5279,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 +5304,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 +5522,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 +5583,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 +5622,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 +5648,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 +5719,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 +5779,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 +5812,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 +5822,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 +5858,7 @@ if removedFromDequeued == 1 then redis.call('DECR', runningCounterKey) end end +__gatesRelease(keyPrefix, redis.call('GET', messageKey), messageId) `, }); @@ -5701,29 +5942,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 +5978,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 +6010,7 @@ if removedFromDequeued == 1 then redis.call('DECR', runningCounterKey) end end +__gatesRelease(keyPrefix, redis.call('GET', messageKey), messageId) `, }); } @@ -5802,6 +6050,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 +6083,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 +6122,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 +6156,7 @@ declare module "@internal/redis" { messageQueueName: string, messageKeyValue: string, removeFromWorkerQueue: string, + keyPrefix: string, callback?: Callback ): Result; @@ -5912,8 +6166,10 @@ declare module "@internal/redis" { envCurrentConcurrencyKey: string, queueCurrentDequeuedKey: string, envCurrentDequeuedKey: string, + messageKey: string, // args messageId: string, + keyPrefix: string, callback?: Callback ): Result; @@ -5932,6 +6188,7 @@ declare module "@internal/redis" { messageQueueName: string, messageData: string, messageScore: string, + keyPrefix: string, callback?: Callback ): Result; @@ -5949,6 +6206,7 @@ declare module "@internal/redis" { // args messageId: string, messageQueueName: string, + keyPrefix: string, callback?: Callback ): Result; @@ -5958,8 +6216,10 @@ declare module "@internal/redis" { envCurrentConcurrencyKey: string, queueCurrentDequeuedKey: string, envCurrentDequeuedKey: string, + messageKey: string, // args messageId: string, + keyPrefix: string, callback?: Callback ): Result; @@ -6162,6 +6422,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 +6461,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 +6487,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 +6518,7 @@ declare module "@internal/redis" { messageKeyValue: string, removeFromWorkerQueue: string, ckWildcardName: string, + keyPrefix: string, callback?: Callback ): Result; @@ -6298,6 +6562,7 @@ declare module "@internal/redis" { messageId: string, messageQueueName: string, ckWildcardName: string, + keyPrefix: string, callback?: Callback ): Result; @@ -6321,6 +6586,7 @@ declare module "@internal/redis" { runningCounterKey: string, ckIndexKey: string, groupConcurrencyKey: string, + messageKey: string, messageId: string, keyPrefix: string, counterTtl: string, @@ -6335,6 +6601,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..53efd420612 --- /dev/null +++ b/internal-packages/run-engine/src/run-queue/tests/queueGates.test.ts @@ -0,0 +1,326 @@ +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("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/types.ts b/internal-packages/run-engine/src/run-queue/types.ts index 10575d8b5d3..73feecb75c2 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(), + concurrencyKey: z.string().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; From f93520088df0b6b4829dcd597c8eec141d4c695d Mon Sep 17 00:00:00 2001 From: Matt Aitken Date: Sat, 29 Aug 2026 10:42:19 +0100 Subject: [PATCH 2/3] fix(run-engine): gates inherit the run key and never deadlock on their own slot A gate without a concurrencyKey now resolves to the run's own key inside the scripts, matching the documented semantics for every caller and for payloads written before a fix at the contract layer. Admission also passes a gate or total cap whose set already contains the message: an unmirrored release from an older build leaves the run's own membership behind, and blocking on it would deadlock the run against itself; since acquiring is an idempotent set add, an existing member simply keeps the slot it holds. Gate names and keys in the payload schema are bounded to 128 characters. --- .../run-engine/src/run-queue/index.ts | 69 ++++++++++++------- .../src/run-queue/tests/queueGates.test.ts | 69 +++++++++++++++++++ .../run-engine/src/run-queue/types.ts | 4 +- 3 files changed, 114 insertions(+), 28 deletions(-) diff --git a/internal-packages/run-engine/src/run-queue/index.ts b/internal-packages/run-engine/src/run-queue/index.ts index 2676ff602df..0f8a8fad9c5 100644 --- a/internal-packages/run-engine/src/run-queue/index.ts +++ b/internal-packages/run-engine/src/run-queue/index.ts @@ -74,11 +74,15 @@ const SemanticAttributes = { 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 gate.concurrencyKey and gate.concurrencyKey ~= '' then - variant = base .. ':ck:' .. gate.concurrencyKey + if gateKey and gateKey ~= '' then + variant = base .. ':ck:' .. gateKey end - return base, variant + return base, variant, gateKey end local function __gateReconcile(setKey, msgKeyPrefix) @@ -96,22 +100,23 @@ local function __gateReconcile(setKey, msgKeyPrefix) end end -local function __gatesHaveCapacity(gatesKeyPrefix, msg, envLimit, msgKeyPrefix) +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 = __gateKeys(gatesKeyPrefix, msg, gate) + 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 then + if occupancy >= perKeyLimit and redis.call('SISMEMBER', variant .. ':currentConcurrency', messageId) == 0 then __gateReconcile(variant .. ':currentConcurrency', msgKeyPrefix) return false end - if gate.concurrencyKey and gate.concurrencyKey ~= '' then + if gateKey and gateKey ~= '' then local rawTotal = redis.call('GET', base .. ':totalConcurrency') if rawTotal then local totalLimit = math.min(tonumber(rawTotal), envLimit) - if tonumber(redis.call('SCARD', base .. ':groupConcurrency') or '0') >= totalLimit then - __gateReconcile(base .. ':groupConcurrency', msgKeyPrefix) + local groupKey = base .. ':groupConcurrency' + if tonumber(redis.call('SCARD', groupKey) or '0') >= totalLimit and redis.call('SISMEMBER', groupKey, messageId) == 0 then + __gateReconcile(groupKey, msgKeyPrefix) return false end end @@ -123,9 +128,9 @@ end local function __gatesAcquire(gatesKeyPrefix, msg, messageId) if not msg.gates then return end for _, gate in ipairs(msg.gates) do - local base, variant = __gateKeys(gatesKeyPrefix, msg, gate) + local base, variant, gateKey = __gateKeys(gatesKeyPrefix, msg, gate) redis.call('SADD', variant .. ':currentConcurrency', messageId) - if gate.concurrencyKey and gate.concurrencyKey ~= '' then + if gateKey and gateKey ~= '' then redis.call('SADD', base .. ':groupConcurrency', messageId) end end @@ -137,9 +142,9 @@ local function __gatesRelease(gatesKeyPrefix, rawPayload, messageId) 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 = __gateKeys(gatesKeyPrefix, msg, gate) + local base, variant, gateKey = __gateKeys(gatesKeyPrefix, msg, gate) local removed = redis.call('SREM', variant .. ':currentConcurrency', messageId) - if removed == 1 and gate.concurrencyKey and gate.concurrencyKey ~= '' then + if removed == 1 and gateKey and gateKey ~= '' then redis.call('SREM', base .. ':groupConcurrency', messageId) end end @@ -3589,7 +3594,7 @@ if enableFastPath == '1' then local okDecode, decoded = pcall(cjson.decode, messageData) if okDecode and type(decoded) == 'table' and decoded.gates then gateMsg = decoded - gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, envLimit, nil) + gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, messageId, envLimit, nil) end end @@ -3703,7 +3708,7 @@ if enableFastPath == '1' then local okDecode, decoded = pcall(cjson.decode, messageData) if okDecode and type(decoded) == 'table' and decoded.gates then gateMsg = decoded - gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, envLimit, nil) + gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, messageId, envLimit, nil) end end @@ -4063,7 +4068,7 @@ if enableFastPath == '1' then local okDecode, decoded = pcall(cjson.decode, messageData) if okDecode and type(decoded) == 'table' and decoded.gates then gateMsg = decoded - gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, envLimit, nil) + gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, messageId, envLimit, nil) end end @@ -4234,7 +4239,7 @@ if enableFastPath == '1' then local okDecode, decoded = pcall(cjson.decode, messageData) if okDecode and type(decoded) == 'table' and decoded.gates then gateMsg = decoded - gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, envLimit, nil) + gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, messageId, envLimit, nil) end end @@ -4636,7 +4641,7 @@ for i = 1, #messages, 2 do else local gatesAllow = true if gatesEnabled then - gatesAllow = __gatesHaveCapacity(keyPrefix, messageData, envConcurrencyLimit, messageKeyPrefix) + gatesAllow = __gatesHaveCapacity(keyPrefix, messageData, messageId, envConcurrencyLimit, messageKeyPrefix) end if gatesAllow then @@ -4887,11 +4892,13 @@ 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 @@ -4923,7 +4930,7 @@ if totalConcurrencyEnabled then end end - actualMaxCount = math.min(actualMaxCount, totalConcurrencyLimit - groupCurrentConcurrency) + totalHeadroom = totalConcurrencyLimit - groupCurrentConcurrency end end @@ -4981,10 +4988,17 @@ for _, ckQueueName in ipairs(ckQueues) do else local gatesAllow = true if gatesEnabled then - gatesAllow = __gatesHaveCapacity(keyPrefix, messageData, envConcurrencyLimit, messageKeyPrefix) + gatesAllow = __gatesHaveCapacity(keyPrefix, messageData, messageId, envConcurrencyLimit, messageKeyPrefix) + end + + local alreadyInGroup = false + local totalAllows = true + if totalHeadroom ~= nil then + alreadyInGroup = redis.call('SISMEMBER', groupConcurrencyKey, messageId) == 1 + totalAllows = alreadyInGroup or totalHeadroom > 0 end - if gatesAllow then + if gatesAllow and totalAllows then redis.call('ZREM', fullQueueKey, messageId) redis.call('ZREM', envQueueKey, messageId) decrLengthCounter() @@ -4992,6 +5006,9 @@ for _, ckQueueName in ipairs(ckQueues) do 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) 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 index 53efd420612..eb6dbda41fc 100644 --- a/internal-packages/run-engine/src/run-queue/tests/queueGates.test.ts +++ b/internal-packages/run-engine/src/run-queue/tests/queueGates.test.ts @@ -282,6 +282,75 @@ describe("RunQueue gates", () => { } }); + 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 { diff --git a/internal-packages/run-engine/src/run-queue/types.ts b/internal-packages/run-engine/src/run-queue/types.ts index 73feecb75c2..97fb99417fc 100644 --- a/internal-packages/run-engine/src/run-queue/types.ts +++ b/internal-packages/run-engine/src/run-queue/types.ts @@ -10,8 +10,8 @@ import type { MinimalAuthenticatedEnvironment } from "../shared/index.js"; * queue name; the org/project/env scope comes from the run's own payload. */ const QueueGate = z.object({ - queue: z.string(), - concurrencyKey: z.string().optional(), + queue: z.string().min(1).max(128), + concurrencyKey: z.string().min(1).max(128).optional(), }); type QueueGate = z.infer; From 209641d1f7304db6ae00d76d911bc34839a2976b Mon Sep 17 00:00:00 2001 From: Matt Aitken Date: Sat, 29 Aug 2026 10:54:52 +0100 Subject: [PATCH 3/3] fix(run-engine): reconcile prunes members whose runs are queued A member whose message still exists but is sitting in a queue zset holds no legitimate slot: queued and in-flight are mutually exclusive on every path. The bounded reconcile now prunes those members too (the stored payload's queue field is the full variant key, so one ZSCORE answers it), which frees a saturated cap even when the self-holding run's variant lies beyond the dequeue candidate window and can never exempt itself. Covered by a test where the leaked member belongs to a different queue entirely. --- .../run-engine/src/run-queue/index.ts | 47 +++++++-------- .../run-queue/tests/totalConcurrency.test.ts | 58 +++++++++++++++++++ 2 files changed, 78 insertions(+), 27 deletions(-) diff --git a/internal-packages/run-engine/src/run-queue/index.ts b/internal-packages/run-engine/src/run-queue/index.ts index 0f8a8fad9c5..b6ef93a9ea3 100644 --- a/internal-packages/run-engine/src/run-queue/index.ts +++ b/internal-packages/run-engine/src/run-queue/index.ts @@ -85,16 +85,24 @@ local function __gateKeys(gatesKeyPrefix, msg, gate) return base, variant, gateKey end -local function __gateReconcile(setKey, msgKeyPrefix) +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', '500') + local scanResult = redis.call('SSCAN', setKey, cursor, 'COUNT', '100') redis.call('SET', cursorKey, scanResult[1], 'EX', '3600') for _, memberId in ipairs(scanResult[2]) do - if redis.call('EXISTS', msgKeyPrefix .. memberId) == 0 then + 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 @@ -107,7 +115,7 @@ local function __gatesHaveCapacity(gatesKeyPrefix, msg, messageId, envLimit, msg 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) + __gateReconcile(variant .. ':currentConcurrency', msgKeyPrefix, gatesKeyPrefix) return false end if gateKey and gateKey ~= '' then @@ -116,7 +124,7 @@ local function __gatesHaveCapacity(gatesKeyPrefix, msg, messageId, envLimit, msg 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) + __gateReconcile(groupKey, msgKeyPrefix, gatesKeyPrefix) return false end end @@ -4905,29 +4913,14 @@ if totalConcurrencyEnabled 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 totalHeadroom = totalConcurrencyLimit - groupCurrentConcurrency 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 }) => {