diff --git a/.changeset/fair-queue-concurrency-slot-leak.md b/.changeset/fair-queue-concurrency-slot-leak.md new file mode 100644 index 00000000000..aa325118421 --- /dev/null +++ b/.changeset/fair-queue-concurrency-slot-leak.md @@ -0,0 +1,7 @@ +--- +"@trigger.dev/redis-worker": patch +--- + +Fair queue consumers no longer leak the concurrency slots that gate a tenant's throughput. Slots were held by messages that had already finished, were never reclaimed, and once enough of them accumulated every queue belonging to that tenant stopped being served. Slots are now freed on the paths that previously skipped them, freed before the record needed to recover them is discarded, and released before a reclaimed message goes back on the queue. A failed release is now surfaced instead of being silently treated as success. + +Concurrency groups keyed on queue metadata rather than the tenant can still resolve to the wrong group when a consumer completes a message it did not enqueue, so this does not yet cover that case. diff --git a/packages/redis-worker/src/fair-queue/concurrency.ts b/packages/redis-worker/src/fair-queue/concurrency.ts index 641899ccc8d..a06592b644a 100644 --- a/packages/redis-worker/src/fair-queue/concurrency.ts +++ b/packages/redis-worker/src/fair-queue/concurrency.ts @@ -103,7 +103,37 @@ export class ConcurrencyManager { pipeline.srem(key, messageId); } - await pipeline.exec(); + this.#assertPipelineSucceeded(await pipeline.exec(), 1); + } + + /** + * Throw if any command in a released pipeline failed. ioredis resolves `exec()` even when + * individual commands error, so an unchecked pipeline reports success while leaving the + * slot held, which strands it permanently once the caller drops the in-flight record. + */ + #assertPipelineSucceeded( + results: Array<[Error | null, unknown]> | null, + messageCount: number + ): void { + if (results === null) { + throw new Error( + `Concurrency release pipeline for ${messageCount} message(s) was discarded without executing` + ); + } + + const errors = results + .map(([error]) => error) + .filter((error): error is Error => Boolean(error)); + + if (errors.length > 0) { + throw new Error( + `Failed to release ${errors.length} of ${ + results?.length ?? 0 + } concurrency slot commands across ${messageCount} message(s): ${errors + .map((error) => error.message) + .join("; ")}` + ); + } } /** @@ -127,7 +157,7 @@ export class ConcurrencyManager { } } - await pipeline.exec(); + this.#assertPipelineSucceeded(await pipeline.exec(), messages.length); } /** diff --git a/packages/redis-worker/src/fair-queue/index.ts b/packages/redis-worker/src/fair-queue/index.ts index e24de876d85..749259479f6 100644 --- a/packages/redis-worker/src/fair-queue/index.ts +++ b/packages/redis-worker/src/fair-queue/index.ts @@ -25,6 +25,7 @@ import type { FairScheduler, QueueCooloffState, QueueDescriptor, + ReclaimedMessageInfo, SchedulerContext, StoredMessage, TenantQueues, @@ -1245,22 +1246,20 @@ export class FairQueue { } } - const descriptor: QueueDescriptor = storedMessage - ? (this.queueDescriptorCache.get(queueId) ?? { - id: queueId, - tenantId: storedMessage.tenantId, - metadata: storedMessage.metadata ?? {}, - }) - : { id: queueId, tenantId: this.keys.extractTenantId(queueId), metadata: {} }; - - // Complete in visibility manager - await this.visibilityManager.complete(messageId, queueId); + const descriptor: QueueDescriptor = this.queueDescriptorCache.get(queueId) ?? { + id: queueId, + tenantId: storedMessage?.tenantId ?? this.keys.extractTenantId(queueId), + metadata: storedMessage?.metadata ?? {}, + }; // Release concurrency - if (this.concurrencyManager && storedMessage) { + if (this.concurrencyManager) { await this.concurrencyManager.release(descriptor, messageId); } + // Complete in visibility manager + await this.visibilityManager.complete(messageId, queueId); + // Update both old and new indexes, clean up caches if queue is empty const { queueEmpty } = await this.#updateAllIndexesAfterDequeue(queueId, descriptor.tenantId); if (queueEmpty) { @@ -1300,13 +1299,16 @@ export class FairQueue { } } - const descriptor: QueueDescriptor = storedMessage - ? (this.queueDescriptorCache.get(queueId) ?? { - id: queueId, - tenantId: storedMessage.tenantId, - metadata: storedMessage.metadata ?? {}, - }) - : { id: queueId, tenantId: this.keys.extractTenantId(queueId), metadata: {} }; + const descriptor: QueueDescriptor = this.queueDescriptorCache.get(queueId) ?? { + id: queueId, + tenantId: storedMessage?.tenantId ?? this.keys.extractTenantId(queueId), + metadata: storedMessage?.metadata ?? {}, + }; + + // Release concurrency + if (this.concurrencyManager) { + await this.concurrencyManager.release(descriptor, messageId); + } // Release back to queue (visibility manager updates dispatch indexes atomically) // Dispatch shard is tenant-based, not queue-based @@ -1324,17 +1326,32 @@ export class FairQueue { Date.now() // Put at back of queue ); - // Release concurrency - if (this.concurrencyManager && storedMessage) { - await this.concurrencyManager.release(descriptor, messageId); - } - this.logger.debug("Message released", { messageId, queueId, }); } + /** + * Release a concurrency slot for a message we can no longer describe, deriving the + * group from the queue id alone. Used on paths that bail out before the stored + * message is available, where the slot would otherwise be held with nothing left + * to reclaim it. + */ + async #releaseOrphanedConcurrency(messageId: string, queueId: string): Promise { + if (!this.concurrencyManager) { + return; + } + + const descriptor: QueueDescriptor = this.queueDescriptorCache.get(queueId) ?? { + id: queueId, + tenantId: this.keys.extractTenantId(queueId), + metadata: {}, + }; + + await this.concurrencyManager.release(descriptor, messageId); + } + /** * Mark a message as failed. This will trigger retry logic if configured, * or move the message to the dead letter queue. @@ -1353,6 +1370,7 @@ export class FairQueue { const dataJson = await this.redis.hget(inflightDataKey, messageId); if (!dataJson) { this.logger.error("Cannot fail message: not found in in-flight data", { messageId, queueId }); + await this.#releaseOrphanedConcurrency(messageId, queueId); return; } @@ -1364,6 +1382,7 @@ export class FairQueue { messageId, queueId, }); + await this.#releaseOrphanedConcurrency(messageId, queueId); return; } @@ -1411,6 +1430,11 @@ export class FairQueue { attempt: storedMessage.attempt + 1, }; + // Release concurrency + if (this.concurrencyManager) { + await this.concurrencyManager.release(descriptor, storedMessage.id); + } + // Release with delay, passing the updated message data so the Lua script // atomically writes the incremented attempt count when re-queuing. const tenantQueueIndexKey = this.keys.tenantQueueIndexKey(descriptor.tenantId); @@ -1427,11 +1451,6 @@ export class FairQueue { JSON.stringify(updatedMessage) ); - // Release concurrency - if (this.concurrencyManager) { - await this.concurrencyManager.release(descriptor, storedMessage.id); - } - this.telemetry.recordRetry(); this.logger.debug("Message scheduled for retry", { @@ -1445,13 +1464,13 @@ export class FairQueue { } } - // Move to DLQ - await this.#moveToDeadLetterQueue(storedMessage, error?.message); - // Release concurrency if (this.concurrencyManager) { await this.concurrencyManager.release(descriptor, storedMessage.id); } + + // Move to DLQ + await this.#moveToDeadLetterQueue(storedMessage, error?.message); } async #moveToDeadLetterQueue( @@ -1525,51 +1544,81 @@ export class FairQueue { } } - async #reclaimTimedOutMessages(): Promise { - let totalReclaimed = 0; + /** + * Free every timed-out message's concurrency slot in one pipeline, before any of them + * is put back on the queue. Throwing here aborts the requeue for the whole batch, which + * leaves the messages in-flight for the next reclaim tick. That is the safe direction: + * requeuing a message whose slot is still held is what strands the slot permanently. + */ + async #releaseReclaimedConcurrency(messages: ReclaimedMessageInfo[]): Promise { + if (!this.concurrencyManager || messages.length === 0) { + return []; + } - for (let shardId = 0; shardId < this.shardCount; shardId++) { - const reclaimedMessages = await this.visibilityManager.reclaimTimedOut(shardId, (queueId) => { - const tenantId = this.keys.extractTenantId(queueId); - const dispatchShardId = this.tenantDispatch.getShardForTenant(tenantId); - return { - queueKey: this.keys.queueKey(queueId), - queueItemsKey: this.keys.queueItemsKey(queueId), - tenantQueueIndexKey: this.keys.tenantQueueIndexKey(tenantId), - dispatchKey: this.keys.dispatchKey(dispatchShardId), - tenantId, - }; + const descriptorFor = (message: ReclaimedMessageInfo) => ({ + id: message.queueId, + tenantId: message.tenantId, + metadata: message.metadata ?? {}, + }); + + try { + await this.concurrencyManager.releaseBatch( + messages.map((message) => ({ queue: descriptorFor(message), messageId: message.messageId })) + ); + return []; + } catch (error) { + this.logger.error("Batch concurrency release failed, retrying message by message", { + count: messages.length, + error: error instanceof Error ? error.message : String(error), }); + } - if (reclaimedMessages.length > 0) { - // Release concurrency for all reclaimed messages in a single batch - // This is critical: when a message times out, its concurrency slot must be freed - // so the message can be processed again when it's re-claimed from the queue - if (this.concurrencyManager) { - try { - await this.concurrencyManager.releaseBatch( - reclaimedMessages.map((msg) => ({ - queue: { - id: msg.queueId, - tenantId: msg.tenantId, - metadata: msg.metadata ?? {}, - }, - messageId: msg.messageId, - })) - ); - } catch (error) { - this.logger.error("Failed to release concurrency for reclaimed messages", { - count: reclaimedMessages.length, - error: error instanceof Error ? error.message : String(error), - }); - } - } + const failed: string[] = []; - // Dispatch indexes are updated atomically by the releaseMessage Lua script - // inside reclaimTimedOut, so no separate index update needed here. + for (const message of messages) { + try { + await this.concurrencyManager.release(descriptorFor(message), message.messageId); + } catch (error) { + failed.push(message.messageId); + this.logger.error("Failed to release concurrency for reclaimed message", { + messageId: message.messageId, + queueId: message.queueId, + error: error instanceof Error ? error.message : String(error), + }); } + } + + return failed; + } - totalReclaimed += reclaimedMessages.length; + async #reclaimTimedOutMessages(): Promise { + let totalReclaimed = 0; + + for (let shardId = 0; shardId < this.shardCount; shardId++) { + try { + const reclaimedMessages = await this.visibilityManager.reclaimTimedOut( + shardId, + (queueId) => { + const tenantId = this.keys.extractTenantId(queueId); + const dispatchShardId = this.tenantDispatch.getShardForTenant(tenantId); + return { + queueKey: this.keys.queueKey(queueId), + queueItemsKey: this.keys.queueItemsKey(queueId), + tenantQueueIndexKey: this.keys.tenantQueueIndexKey(tenantId), + dispatchKey: this.keys.dispatchKey(dispatchShardId), + tenantId, + }; + }, + this.#releaseReclaimedConcurrency.bind(this) + ); + + totalReclaimed += reclaimedMessages.length; + } catch (error) { + this.logger.error("Failed to reclaim shard, leaving messages in-flight for the next tick", { + shardId, + error: error instanceof Error ? error.message : String(error), + }); + } } if (totalReclaimed > 0) { diff --git a/packages/redis-worker/src/fair-queue/tests/fairQueue.test.ts b/packages/redis-worker/src/fair-queue/tests/fairQueue.test.ts index 5ae1f390f9f..5fbf2299fab 100644 --- a/packages/redis-worker/src/fair-queue/tests/fairQueue.test.ts +++ b/packages/redis-worker/src/fair-queue/tests/fairQueue.test.ts @@ -10,7 +10,7 @@ import { WorkerQueueManager, } from "../index.js"; import type { FairQueueKeyProducer, FairQueueOptions } from "../types.js"; -import type { RedisOptions } from "@internal/redis"; +import { createRedisClient, type RedisOptions } from "@internal/redis"; // Define a common payload schema for tests const TestPayloadSchema = z.object({ value: z.string() }); @@ -1370,4 +1370,390 @@ describe("FairQueue", () => { } ); }); + + describe("concurrency slot release", () => { + redisTest( + "should release the concurrency slot when the in-flight record is already gone", + { timeout: 15000 }, + async ({ redisOptions }) => { + const processed: string[] = []; + keys = new DefaultFairQueueKeyProducer({ prefix: "test" }); + + const scheduler = new DRRScheduler({ + redis: redisOptions, + keys, + quantum: 10, + maxDeficit: 100, + }); + + const queue = new TestFairQueueHelper(redisOptions, keys, { + scheduler, + payloadSchema: TestPayloadSchema, + shardCount: 1, + consumerCount: 1, + consumerIntervalMs: 20, + visibilityTimeoutMs: 60000, + concurrencyGroups: [ + { + name: "tenant", + extractGroupId: (q) => q.tenantId, + getLimit: async () => 1, + defaultLimit: 1, + }, + ], + startConsumers: false, + }); + + const redis = createRedisClient(redisOptions); + + try { + queue.onMessage(async (ctx) => { + if (ctx.message.payload.value === "msg-0") { + await redis.hdel(keys.inflightDataKey(0), ctx.message.id); + } + processed.push(ctx.message.payload.value); + await ctx.complete(); + }); + + for (let i = 0; i < 2; i++) { + await queue.enqueue({ + queueId: "tenant:t1:queue:q1", + tenantId: "t1", + payload: { value: `msg-${i}` }, + }); + } + + queue.start(); + + await vi.waitFor( + () => { + expect(processed).toHaveLength(2); + }, + { timeout: 10000 } + ); + + const held = await redis.scard(keys.concurrencyKey("tenant", "t1")); + expect(held).toBe(0); + } finally { + await redis.quit(); + await queue.close(); + } + } + ); + + redisTest( + "should release metadata-derived groups from the cached descriptor when the in-flight record is gone", + { timeout: 15000 }, + async ({ redisOptions }) => { + const processed: string[] = []; + keys = new DefaultFairQueueKeyProducer({ prefix: "test" }); + + const scheduler = new DRRScheduler({ + redis: redisOptions, + keys, + quantum: 10, + maxDeficit: 100, + }); + + const queue = new TestFairQueueHelper(redisOptions, keys, { + scheduler, + payloadSchema: TestPayloadSchema, + shardCount: 1, + consumerCount: 1, + consumerIntervalMs: 20, + visibilityTimeoutMs: 60000, + concurrencyGroups: [ + { + name: "tenant", + extractGroupId: (q) => q.tenantId, + getLimit: async () => 5, + defaultLimit: 5, + }, + { + name: "organization", + extractGroupId: (q) => (q.metadata.orgId as string) ?? "default", + getLimit: async () => 1, + defaultLimit: 1, + }, + ], + startConsumers: false, + }); + + const redis = createRedisClient(redisOptions); + + try { + queue.onMessage(async (ctx) => { + if (ctx.message.payload.value === "msg-0") { + await redis.hdel(keys.inflightDataKey(0), ctx.message.id); + } + processed.push(ctx.message.payload.value); + await ctx.complete(); + }); + + for (let i = 0; i < 2; i++) { + await queue.enqueue({ + queueId: "tenant:t1:queue:q1", + tenantId: "t1", + metadata: { orgId: "org-1" }, + payload: { value: `msg-${i}` }, + }); + } + + queue.start(); + + await vi.waitFor( + () => { + expect(processed).toHaveLength(2); + }, + { timeout: 10000 } + ); + + expect(await redis.scard(keys.concurrencyKey("organization", "org-1"))).toBe(0); + expect(await redis.scard(keys.concurrencyKey("organization", "default"))).toBe(0); + } finally { + await redis.quit(); + await queue.close(); + } + } + ); + + redisTest( + "should not requeue a reclaimed message when freeing its slot fails", + { timeout: 20000 }, + async ({ redisOptions }) => { + keys = new DefaultFairQueueKeyProducer({ prefix: "test" }); + + const scheduler = new DRRScheduler({ + redis: redisOptions, + keys, + quantum: 10, + maxDeficit: 100, + }); + + const queue = new TestFairQueueHelper(redisOptions, keys, { + scheduler, + payloadSchema: TestPayloadSchema, + shardCount: 1, + consumerCount: 1, + consumerIntervalMs: 20, + visibilityTimeoutMs: 500, + reclaimIntervalMs: 100, + concurrencyGroups: [ + { + name: "tenant", + extractGroupId: (q) => q.tenantId, + getLimit: async () => 5, + defaultLimit: 5, + }, + ], + startConsumers: false, + }); + + const redis = createRedisClient(redisOptions); + const queueId = "tenant:t1:queue:release-fails"; + const concurrencyKey = keys.concurrencyKey("tenant", "t1"); + + try { + queue.onMessage(async () => { + await new Promise((resolve) => setTimeout(resolve, 15000)); + }); + + await queue.enqueue({ + queueId, + tenantId: "t1", + payload: { value: "msg-0" }, + }); + + queue.start(); + + await vi.waitFor( + async () => { + expect(await redis.zcard(keys.inflightKey(0))).toBe(1); + }, + { timeout: 5000 } + ); + + await redis.del(concurrencyKey); + await redis.set(concurrencyKey, "not-a-set"); + + await new Promise((resolve) => setTimeout(resolve, 2000)); + + expect(await redis.zcard(keys.inflightKey(0))).toBe(1); + expect(await redis.zcard(keys.queueKey(queueId))).toBe(0); + + const [stuckMember] = await redis.zrange(keys.inflightKey(0), 0, 0); + const stuckDeadline = Number(await redis.zscore(keys.inflightKey(0), stuckMember!)); + + await redis.del(concurrencyKey); + + await vi.waitFor( + async () => { + const score = await redis.zscore(keys.inflightKey(0), stuckMember!); + expect(score === null || Number(score) > stuckDeadline).toBe(true); + }, + { timeout: 8000 } + ); + } finally { + await redis.del(concurrencyKey); + await redis.quit(); + await queue.close(); + } + } + ); + + redisTest( + "should release the concurrency slot even when the reclaim requeue fails", + { timeout: 15000 }, + async ({ redisOptions }) => { + keys = new DefaultFairQueueKeyProducer({ prefix: "test" }); + + const scheduler = new DRRScheduler({ + redis: redisOptions, + keys, + quantum: 10, + maxDeficit: 100, + }); + + const queue = new TestFairQueueHelper(redisOptions, keys, { + scheduler, + payloadSchema: TestPayloadSchema, + shardCount: 1, + consumerCount: 1, + consumerIntervalMs: 20, + visibilityTimeoutMs: 30000, + reclaimIntervalMs: 100, + concurrencyGroups: [ + { + name: "tenant", + extractGroupId: (q) => q.tenantId, + getLimit: async () => 1, + defaultLimit: 1, + }, + ], + startConsumers: false, + }); + + const redis = createRedisClient(redisOptions); + const queueId = "tenant:t1:queue:reclaim-fail"; + const inflightKey = keys.inflightKey(0); + + let unblockHandler: (() => void) | undefined; + const handlerBlocked = new Promise((resolve) => { + unblockHandler = resolve; + }); + + try { + queue.onMessage(async () => { + await handlerBlocked; + }); + + await queue.enqueue({ + queueId, + tenantId: "t1", + payload: { value: "msg-0" }, + }); + + queue.start(); + + await vi.waitFor( + async () => { + expect(await redis.scard(keys.concurrencyKey("tenant", "t1"))).toBe(1); + expect(await redis.zcard(inflightKey)).toBe(1); + }, + { timeout: 5000 } + ); + + await redis.del(keys.queueKey(queueId)); + await redis.set(keys.queueKey(queueId), "not-a-zset"); + + const [member] = await redis.zrange(inflightKey, 0, 0); + expect(member).toBeDefined(); + await redis.zadd(inflightKey, Date.now() - 60000, member!); + + await vi.waitFor( + async () => { + expect(await redis.scard(keys.concurrencyKey("tenant", "t1"))).toBe(0); + }, + { timeout: 8000 } + ); + + expect(await redis.zcard(inflightKey)).toBe(1); + } finally { + unblockHandler?.(); + await redis.del(keys.queueKey(queueId)); + await redis.quit(); + await queue.close(); + } + } + ); + + redisTest( + "should release the concurrency slot when failMessage cannot read the message", + { timeout: 15000 }, + async ({ redisOptions }) => { + const started: string[] = []; + keys = new DefaultFairQueueKeyProducer({ prefix: "test" }); + + const scheduler = new DRRScheduler({ + redis: redisOptions, + keys, + quantum: 10, + maxDeficit: 100, + }); + + const queue = new TestFairQueueHelper(redisOptions, keys, { + scheduler, + payloadSchema: TestPayloadSchema, + shardCount: 1, + consumerCount: 1, + consumerIntervalMs: 20, + visibilityTimeoutMs: 60000, + concurrencyGroups: [ + { + name: "tenant", + extractGroupId: (q) => q.tenantId, + getLimit: async () => 5, + defaultLimit: 5, + }, + ], + startConsumers: false, + }); + + const redis = createRedisClient(redisOptions); + + try { + queue.onMessage(async (ctx) => { + started.push(ctx.message.payload.value); + await redis.hdel(keys.inflightDataKey(0), ctx.message.id); + await ctx.fail(new Error("boom")); + }); + + await queue.enqueue({ + queueId: "tenant:t1:queue:q1", + tenantId: "t1", + payload: { value: "msg-0" }, + }); + + queue.start(); + + await vi.waitFor( + () => { + expect(started).toHaveLength(1); + }, + { timeout: 10000 } + ); + + await vi.waitFor( + async () => { + expect(await redis.scard(keys.concurrencyKey("tenant", "t1"))).toBe(0); + }, + { timeout: 5000 } + ); + } finally { + await redis.quit(); + await queue.close(); + } + } + ); + }); }); diff --git a/packages/redis-worker/src/fair-queue/tests/visibility.test.ts b/packages/redis-worker/src/fair-queue/tests/visibility.test.ts index e20b6671075..83d7e7f3e88 100644 --- a/packages/redis-worker/src/fair-queue/tests/visibility.test.ts +++ b/packages/redis-worker/src/fair-queue/tests/visibility.test.ts @@ -912,5 +912,70 @@ describe("VisibilityManager", () => { await redis.quit(); } ); + + redisTest( + "should drop a dangling in-flight entry whose payload is gone", + { timeout: 10000 }, + async ({ redisOptions }) => { + keys = new DefaultFairQueueKeyProducer({ prefix: "test" }); + + const manager = new VisibilityManager({ + redis: redisOptions, + keys, + shardCount: 1, + defaultTimeoutMs: 100, + }); + + const redis = createRedisClient(redisOptions); + const queueId = "tenant:t1:queue:dangling"; + const queueKey = keys.queueKey(queueId); + const queueItemsKey = keys.queueItemsKey(queueId); + const dispatchKey = keys.dispatchKey(0); + const inflightKey = keys.inflightKey(0); + const inflightDataKey = keys.inflightDataKey(0); + + try { + const messageId = "dangling-msg"; + const storedMessage = { + id: messageId, + queueId, + tenantId: "t1", + payload: { id: 1, value: "test" }, + timestamp: Date.now() - 1000, + attempt: 1, + }; + + await redis.zadd(queueKey, storedMessage.timestamp, messageId); + await redis.hset(queueItemsKey, messageId, JSON.stringify(storedMessage)); + + const claimResult = await manager.claim( + queueId, + queueKey, + queueItemsKey, + "consumer-1", + 100 + ); + expect(claimResult.claimed).toBe(true); + + await redis.hdel(inflightDataKey, messageId); + expect(await redis.zcard(inflightKey)).toBe(1); + + await new Promise((resolve) => setTimeout(resolve, 150)); + + await manager.reclaimTimedOut(0, (qId) => ({ + queueKey: keys.queueKey(qId), + queueItemsKey: keys.queueItemsKey(qId), + tenantQueueIndexKey: keys.tenantQueueIndexKey(keys.extractTenantId(qId)), + dispatchKey, + tenantId: keys.extractTenantId(qId), + })); + + expect(await redis.zcard(inflightKey)).toBe(0); + } finally { + await manager.close(); + await redis.quit(); + } + } + ); }); }); diff --git a/packages/redis-worker/src/fair-queue/visibility.ts b/packages/redis-worker/src/fair-queue/visibility.ts index 8ecc3d9927a..65d633a5fc1 100644 --- a/packages/redis-worker/src/fair-queue/visibility.ts +++ b/packages/redis-worker/src/fair-queue/visibility.ts @@ -398,7 +398,8 @@ export class VisibilityManager { tenantQueueIndexKey: string; dispatchKey: string; tenantId: string; - } + }, + onBeforeRequeue?: (messages: ReclaimedMessageInfo[]) => Promise ): Promise { const inflightKey = this.keys.inflightKey(shardId); const inflightDataKey = this.keys.inflightDataKey(shardId); @@ -415,30 +416,77 @@ export class VisibilityManager { 100 // Process in batches ); - const reclaimedMessages: ReclaimedMessageInfo[] = []; + const candidates: Array<{ + member: string; + deadlineScore: string; + info: ReclaimedMessageInfo; + storedMessage: StoredMessage | null; + }> = []; for (let i = 0; i < timedOut.length; i += 2) { const member = timedOut[i]; - const _deadlineScore = timedOut[i + 1]; // This is the visibility deadline, not the original timestamp - if (!member || !_deadlineScore) { + const deadlineScore = timedOut[i + 1]; + if (!member || !deadlineScore) { continue; } const { messageId, queueId } = this.#parseMember(member); + + const dataJson = await this.redis.hget(inflightDataKey, messageId); + let storedMessage: StoredMessage | null = null; + if (dataJson) { + try { + storedMessage = JSON.parse(dataJson); + } catch { + // Ignore parse error, proceed with reclaim + } + } + + if (!storedMessage) { + this.logger.error("Missing or corrupted message data during reclaim, using fallback", { + messageId, + queueId, + }); + } + + candidates.push({ + member, + deadlineScore, + storedMessage, + info: { + messageId, + queueId, + tenantId: storedMessage?.tenantId ?? this.keys.extractTenantId(queueId), + metadata: storedMessage?.metadata ?? {}, + }, + }); + } + + if (candidates.length === 0) { + return []; + } + + const notReleased = new Set( + (onBeforeRequeue + ? await onBeforeRequeue(candidates.map((candidate) => candidate.info)) + : undefined) ?? [] + ); + + const reclaimedMessages: ReclaimedMessageInfo[] = []; + + for (const { member, deadlineScore, storedMessage, info } of candidates) { + const { messageId, queueId } = info; + + if (notReleased.has(messageId)) { + this.logger.error("Skipping requeue, concurrency slot was not released", { + messageId, + queueId, + }); + continue; + } const { queueKey, queueItemsKey, tenantQueueIndexKey, dispatchKey, tenantId } = getQueueKeys(queueId); try { - // Get message data BEFORE releasing so we can extract tenantId for concurrency release - const dataJson = await this.redis.hget(inflightDataKey, messageId); - let storedMessage: StoredMessage | null = null; - if (dataJson) { - try { - storedMessage = JSON.parse(dataJson); - } catch { - // Ignore parse error, proceed with reclaim - } - } - // Re-add to queue with original timestamp to preserve priority // Fall back to now if we can't get the original timestamp const score = storedMessage?.timestamp ?? now; @@ -457,34 +505,12 @@ export class VisibilityManager { tenantId ); - // Track reclaimed message for concurrency release - // Always add to reclaimedMessages to avoid concurrency leaks - if (storedMessage) { - reclaimedMessages.push({ - messageId, - queueId, - tenantId: storedMessage.tenantId, - metadata: storedMessage.metadata, - }); - } else { - // Fallback: extract tenantId from queueId when message data is missing or corrupted - // This ensures concurrency is released even if we can't get the full metadata - this.logger.error("Missing or corrupted message data during reclaim, using fallback", { - messageId, - queueId, - }); - reclaimedMessages.push({ - messageId, - queueId, - tenantId: this.keys.extractTenantId(queueId), - metadata: {}, - }); - } + reclaimedMessages.push(info); this.logger.debug("Reclaimed timed-out message", { messageId, queueId, - deadline: _deadlineScore, + deadline: deadlineScore, }); } catch (error) { this.logger.error("Failed to reclaim message", { @@ -707,7 +733,9 @@ local tenantId = ARGV[6] -- Get message data from in-flight local payload = redis.call('HGET', inflightDataKey, messageId) if not payload then - -- Message not in in-flight or already released + -- Message not in in-flight or already released. Drop the dangling in-flight + -- entry so it stops consuming a slot in every future reclaim scan. + redis.call('ZREM', inflightKey, member) return 0 end @@ -717,14 +745,15 @@ if updatedData and updatedData ~= "" then payload = updatedData end +-- Add back to queue before dropping the in-flight record. Lua does not roll back on +-- error, so removing first would lose the message outright if the queue write failed. +redis.call('ZADD', queueKey, score, messageId) +redis.call('HSET', queueItemsKey, messageId, payload) + -- Remove from in-flight redis.call('ZREM', inflightKey, member) redis.call('HDEL', inflightDataKey, messageId) --- Add back to queue -redis.call('ZADD', queueKey, score, messageId) -redis.call('HSET', queueItemsKey, messageId, payload) - -- Update tenant queue index (Level 2) with queue's oldest message local oldest = redis.call('ZRANGE', queueKey, 0, 0, 'WITHSCORES') if #oldest >= 2 then @@ -771,15 +800,20 @@ for i = 0, numMessages - 1 do -- Get message data from in-flight local payload = redis.call('HGET', inflightDataKey, messageId) if payload then + -- Add back to queue before dropping the in-flight record, so a failed queue write + -- cannot lose the message: Lua does not roll back already-applied commands. + redis.call('ZADD', queueKey, score, messageId) + redis.call('HSET', queueItemsKey, messageId, payload) + -- Remove from in-flight redis.call('ZREM', inflightKey, member) redis.call('HDEL', inflightDataKey, messageId) - -- Add back to queue - redis.call('ZADD', queueKey, score, messageId) - redis.call('HSET', queueItemsKey, messageId, payload) - releasedCount = releasedCount + 1 + else + -- Dangling in-flight entry with no payload: drop it so it stops consuming + -- a slot in every future reclaim scan. + redis.call('ZREM', inflightKey, member) end end