Skip to content

Commit 771dc1b

Browse files
committed
fix(run-engine): a gate's total limit counts keyed and keyless holders together
1 parent 3c0f968 commit 771dc1b

2 files changed

Lines changed: 66 additions & 17 deletions

File tree

‎internal-packages/run-engine/src/run-queue/index.ts‎

Lines changed: 15 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -67,8 +67,10 @@ const SemanticAttributes = {
6767
* __gatesHaveCapacity / __gatesAcquire run on admit paths behind the gatesEnabled
6868
* flag. __gatesRelease runs on every release path UNCONDITIONALLY and is payload-
6969
* driven (a cheap substring probe before decoding), so slots acquired while the flag
70-
* was on always drain, and gateless messages pay near zero. __gateReconcile is the
71-
* same bounded self-heal as the total-cap gate. Group and gate sets are strict
70+
* was on always drain, and gateless messages pay near zero. Every holder joins the
71+
* gate's groupConcurrency set — keyed or keyless — so a totalConcurrency limit caps
72+
* across everything, matching the SDK's documented "total" semantics. __gateReconcile
73+
* is the same bounded self-heal as the total-cap gate. Group and gate sets are strict
7274
* mirrors of the member's home-queue currentConcurrency set (admits populate both
7375
* in one script), so a member whose message key is gone, or who is absent from its
7476
* home currentConcurrency set, holds no legitimate slot and is pruned. The home-set
@@ -124,15 +126,13 @@ local function __gatesHaveCapacity(gatesKeyPrefix, msg, messageId, envLimit, msg
124126
__gateReconcile(variant .. ':currentConcurrency', msgKeyPrefix, gatesKeyPrefix)
125127
return false
126128
end
127-
if gateKey and gateKey ~= '' then
128-
local rawTotal = redis.call('GET', base .. ':totalConcurrency')
129-
if rawTotal then
130-
local totalLimit = math.min(tonumber(rawTotal), envLimit)
131-
local groupKey = base .. ':groupConcurrency'
132-
if tonumber(redis.call('SCARD', groupKey) or '0') >= totalLimit and redis.call('SISMEMBER', groupKey, messageId) == 0 then
133-
__gateReconcile(groupKey, msgKeyPrefix, gatesKeyPrefix)
134-
return false
135-
end
129+
local rawTotal = redis.call('GET', base .. ':totalConcurrency')
130+
if rawTotal then
131+
local totalLimit = math.min(tonumber(rawTotal), envLimit)
132+
local groupKey = base .. ':groupConcurrency'
133+
if tonumber(redis.call('SCARD', groupKey) or '0') >= totalLimit and redis.call('SISMEMBER', groupKey, messageId) == 0 then
134+
__gateReconcile(groupKey, msgKeyPrefix, gatesKeyPrefix)
135+
return false
136136
end
137137
end
138138
end
@@ -142,11 +142,9 @@ end
142142
local function __gatesAcquire(gatesKeyPrefix, msg, messageId)
143143
if not msg.gates then return end
144144
for _, gate in ipairs(msg.gates) do
145-
local base, variant, gateKey = __gateKeys(gatesKeyPrefix, msg, gate)
145+
local base, variant = __gateKeys(gatesKeyPrefix, msg, gate)
146146
redis.call('SADD', variant .. ':currentConcurrency', messageId)
147-
if gateKey and gateKey ~= '' then
148-
redis.call('SADD', base .. ':groupConcurrency', messageId)
149-
end
147+
redis.call('SADD', base .. ':groupConcurrency', messageId)
150148
end
151149
end
152150
@@ -156,9 +154,9 @@ local function __gatesRelease(gatesKeyPrefix, rawPayload, messageId)
156154
local ok, msg = pcall(cjson.decode, rawPayload)
157155
if not ok or type(msg) ~= 'table' or not msg.gates then return end
158156
for _, gate in ipairs(msg.gates) do
159-
local base, variant, gateKey = __gateKeys(gatesKeyPrefix, msg, gate)
157+
local base, variant = __gateKeys(gatesKeyPrefix, msg, gate)
160158
local removed = redis.call('SREM', variant .. ':currentConcurrency', messageId)
161-
if removed == 1 and gateKey and gateKey ~= '' then
159+
if removed == 1 then
162160
redis.call('SREM', base .. ':groupConcurrency', messageId)
163161
end
164162
end

‎internal-packages/run-engine/src/run-queue/tests/queueGates.test.ts‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -180,6 +180,57 @@ describe("RunQueue gates", () => {
180180
}
181181
);
182182

183+
redisTest(
184+
"a gate's total limit caps keyed and keyless holders together",
185+
async ({ redisContainer }) => {
186+
const queue = createQueue(redisContainer, true);
187+
try {
188+
await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5);
189+
await queue.updateQueueTotalConcurrencyLimits(authenticatedEnvDev, "db-limit", 1);
190+
191+
const now = Date.now();
192+
await queue.enqueueMessage({
193+
env: authenticatedEnvDev,
194+
message: makeMessage({
195+
runId: "r0",
196+
timestamp: now - 1000,
197+
gates: [{ queue: "db-limit" }],
198+
}),
199+
workerQueue: "main",
200+
});
201+
await queue.enqueueMessage({
202+
env: authenticatedEnvDev,
203+
message: makeMessage({
204+
runId: "r1",
205+
timestamp: now - 999,
206+
gates: [{ queue: "db-limit", concurrencyKey: "acme" }],
207+
}),
208+
workerQueue: "main",
209+
});
210+
211+
const oneAdmitted = await waitFor(
212+
async () => (await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "db-limit")) === 1
213+
);
214+
expect(oneAdmitted).toBe(true);
215+
216+
/** The keyed run must stay queued: the keyless holder occupies the total pool. */
217+
await setTimeout(2000);
218+
expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "db-limit")).toBe(1);
219+
expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1);
220+
221+
expect(await popWorkerQueue(queue, "r0")).toBe(true);
222+
await queue.acknowledgeMessage(authenticatedEnvDev.organization.id, "r0");
223+
224+
/** Acking the keyless holder frees the total pool; the keyed run is admitted. */
225+
const r1Admitted = await waitFor(() => popWorkerQueue(queue, "r1"));
226+
expect(r1Admitted).toBe(true);
227+
expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "db-limit")).toBe(1);
228+
} finally {
229+
await queue.quit();
230+
}
231+
}
232+
);
233+
183234
redisTest("ignores gates and holds no gate slots when disabled", async ({ redisContainer }) => {
184235
const queue = createQueue(redisContainer, false);
185236
try {

0 commit comments

Comments
 (0)