Skip to content

Commit 3f8ed8b

Browse files
matt-aitkenmatt-aitken
committed
fix(redis-worker,run-engine): close the remaining ways a batch item is stranded or run twice
Releasing a slot could silently do nothing. The release pipeline never inspected its per-command errors, and ioredis resolves a pipeline whose commands failed, so a failed SREM reported success and the caller went on to destroy the in-flight record. Release now throws on a failed command. Reclaim then acted on that false success: it released per message, swallowed any error, and requeued regardless, which is exactly how a slot ends up held with the message gone. It now frees the whole timed-out batch in one pipeline before any of them is requeued, and a failure aborts the requeue so the messages stay in-flight for the next tick. That also restores the single round trip per shard. Requeuing removed the message from in-flight before writing it back to the queue. Redis does not roll back a script that fails partway, so a failed queue write lost the message outright. The order is now write-then-remove. Batch items were never heartbeated, so any item slower than the visibility timeout was reclaimed and handed to a second consumer while the first was still running it, executing the item twice. Items are now heartbeated for as long as their callback runs, and the visibility timeout is configurable so this is testable.
1 parent c6524ca commit 3f8ed8b

7 files changed

Lines changed: 259 additions & 139 deletions

File tree

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

Lines changed: 39 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,9 @@ const ENV_CONCURRENCY_KEY_PREFIX = "batch:env_concurrency";
5959
// then all messages are routed to this queue for BatchQueue's own consumer loop.
6060
const BATCH_WORKER_QUEUE_ID = "batch-worker-queue";
6161

62+
/** How long a claimed batch item stays invisible before the reclaim loop takes it back. */
63+
const BATCH_ITEM_VISIBILITY_TIMEOUT_MS = 60_000;
64+
6265
export class BatchQueue {
6366
private fairQueue: FairQueue<typeof BatchItemPayloadSchema>;
6467
private workerQueueManager: WorkerQueueManager;
@@ -67,6 +70,8 @@ export class BatchQueue {
6770
private tracer?: Tracer;
6871
private concurrencyRedis: Redis;
6972
private defaultConcurrency: number;
73+
private heartbeatIntervalMs: number;
74+
private visibilityTimeoutMs: number;
7075
private maxAttempts: number;
7176

7277
private processItemCallback?: ProcessBatchItemCallback;
@@ -95,6 +100,8 @@ export class BatchQueue {
95100
this.logger = options.logger ?? new Logger("BatchQueue", options.logLevel ?? "info");
96101
this.tracer = options.tracer;
97102
this.defaultConcurrency = options.defaultConcurrency ?? 10;
103+
this.visibilityTimeoutMs = options.visibilityTimeoutMs ?? BATCH_ITEM_VISIBILITY_TIMEOUT_MS;
104+
this.heartbeatIntervalMs = Math.max(50, Math.floor(this.visibilityTimeoutMs / 3));
98105
this.maxAttempts = options.retry?.maxAttempts ?? 1;
99106
this.abortController = new AbortController();
100107
this.workerQueueBlockingTimeoutSeconds = options.workerQueueBlockingTimeoutSeconds ?? 10;
@@ -154,7 +161,7 @@ export class BatchQueue {
154161
shardCount: options.shardCount ?? 1,
155162
consumerCount: options.consumerCount,
156163
consumerIntervalMs: options.consumerIntervalMs,
157-
visibilityTimeoutMs: 60_000, // 1 minute for batch item processing
164+
visibilityTimeoutMs: this.visibilityTimeoutMs,
158165
startConsumers: false, // We control when to start
159166
cooloff: {
160167
enabled: false,
@@ -752,6 +759,29 @@ export class BatchQueue {
752759
// Private - Message Handling
753760
// ============================================================================
754761

762+
/**
763+
* Keep extending a message's visibility deadline while its callback runs, and return a
764+
* function that stops doing so. Without this an item slower than the visibility timeout
765+
* is redelivered while still being processed, and the original consumer's completion
766+
* then destroys the redelivery's in-flight record, silently dropping the item and
767+
* leaving the batch short of its expected count forever.
768+
*/
769+
#startHeartbeat(messageId: string, queueId: string): () => void {
770+
const interval = setInterval(() => {
771+
this.fairQueue.heartbeatMessage(messageId, queueId).catch((error) => {
772+
this.logger.debug("Batch item heartbeat failed", {
773+
messageId,
774+
queueId,
775+
error: error instanceof Error ? error.message : String(error),
776+
});
777+
});
778+
}, this.heartbeatIntervalMs);
779+
780+
interval.unref?.();
781+
782+
return () => clearInterval(interval);
783+
}
784+
755785
async #handleMessage(consumerId: string, messageId: string, queueId: string): Promise<void> {
756786
// Get message data from FairQueue's in-flight storage
757787
const storedMessage = await this.fairQueue.getMessageData(messageId, queueId);
@@ -820,9 +850,10 @@ export class BatchQueue {
820850
let processedCount: number;
821851

822852
try {
823-
const result = await this.#startSpan(
824-
"BatchQueue.processItemCallback",
825-
async (innerSpan) => {
853+
const stopHeartbeat = this.#startHeartbeat(messageId, queueId);
854+
let result: Awaited<ReturnType<ProcessBatchItemCallback>>;
855+
try {
856+
result = await this.#startSpan("BatchQueue.processItemCallback", async (innerSpan) => {
826857
innerSpan?.setAttributes({
827858
"batch.id": batchId,
828859
"batch.itemIndex": itemIndex,
@@ -837,8 +868,10 @@ export class BatchQueue {
837868
attempt,
838869
isFinalAttempt,
839870
});
840-
}
841-
);
871+
});
872+
} finally {
873+
stopHeartbeat();
874+
}
842875

843876
if (result.success) {
844877
span?.setAttribute("batch.result", "success");

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

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -953,4 +953,56 @@ describe("BatchQueue", () => {
953953
}
954954
);
955955
});
956+
957+
describe("visibility heartbeat", () => {
958+
redisTest(
959+
"should not redeliver an item that takes longer than the visibility timeout",
960+
{ timeout: 60_000 },
961+
async ({ redisContainer }) => {
962+
const queue = new BatchQueue({
963+
redis: {
964+
host: redisContainer.getHost(),
965+
port: redisContainer.getPort(),
966+
keyPrefix: "test:",
967+
},
968+
drr: { quantum: 5, maxDeficit: 50 },
969+
consumerCount: 2,
970+
consumerIntervalMs: 50,
971+
visibilityTimeoutMs: 1_000,
972+
startConsumers: false,
973+
});
974+
975+
const invocations: number[] = [];
976+
977+
try {
978+
queue.onProcessItem(async ({ itemIndex }) => {
979+
const isFirst = invocations.length === 0;
980+
invocations.push(itemIndex);
981+
if (isFirst) {
982+
await new Promise((resolve) => setTimeout(resolve, 9_000));
983+
}
984+
return { success: true, runId: `run-${itemIndex}` };
985+
});
986+
987+
await queue.initializeBatch(createInitOptions("batch-hb", "env-hb", 1));
988+
await enqueueItems(queue, "batch-hb", "env-hb", createBatchItems(1));
989+
990+
queue.start();
991+
992+
await vi.waitFor(
993+
() => {
994+
expect(invocations.length).toBeGreaterThanOrEqual(1);
995+
},
996+
{ timeout: 10_000 }
997+
);
998+
999+
await new Promise((resolve) => setTimeout(resolve, 14_000));
1000+
1001+
expect(invocations).toEqual([0]);
1002+
} finally {
1003+
await queue.close();
1004+
}
1005+
}
1006+
);
1007+
});
9561008
});

‎internal-packages/run-engine/src/batch-queue/types.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -214,6 +214,12 @@ export type BatchQueueOptions = {
214214
* Items wait in queue until capacity frees up.
215215
*/
216216
defaultConcurrency?: number;
217+
/**
218+
* How long a claimed item stays invisible before the reclaim loop takes it back.
219+
* The item is heartbeated for as long as its callback runs, so this only bites when
220+
* a consumer stops making progress. Defaults to 60s.
221+
*/
222+
visibilityTimeoutMs?: number;
217223
/**
218224
* Optional global rate limiter to limit processing across all consumers.
219225
* When configured, limits the max items/second processed globally.

‎packages/redis-worker/src/fair-queue/concurrency.ts‎

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,31 @@ export class ConcurrencyManager {
103103
pipeline.srem(key, messageId);
104104
}
105105

106-
await pipeline.exec();
106+
this.#assertPipelineSucceeded(await pipeline.exec(), 1);
107+
}
108+
109+
/**
110+
* Throw if any command in a released pipeline failed. ioredis resolves `exec()` even when
111+
* individual commands error, so an unchecked pipeline reports success while leaving the
112+
* slot held, which strands it permanently once the caller drops the in-flight record.
113+
*/
114+
#assertPipelineSucceeded(
115+
results: Array<[Error | null, unknown]> | null,
116+
messageCount: number
117+
): void {
118+
const errors = (results ?? [])
119+
.map(([error]) => error)
120+
.filter((error): error is Error => Boolean(error));
121+
122+
if (errors.length > 0) {
123+
throw new Error(
124+
`Failed to release ${errors.length} of ${
125+
results?.length ?? 0
126+
} concurrency slot commands across ${messageCount} message(s): ${errors
127+
.map((error) => error.message)
128+
.join("; ")}`
129+
);
130+
}
107131
}
108132

109133
/**
@@ -127,7 +151,7 @@ export class ConcurrencyManager {
127151
}
128152
}
129153

130-
await pipeline.exec();
154+
this.#assertPipelineSucceeded(await pipeline.exec(), messages.length);
131155
}
132156

133157
/**

‎packages/redis-worker/src/fair-queue/index.ts‎

Lines changed: 38 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -1498,28 +1498,14 @@ export class FairQueue<TPayloadSchema extends z.ZodTypeAny = z.ZodUnknown> {
14981498
originalTimestamp: storedMessage.timestamp,
14991499
};
15001500

1501+
// Complete in visibility manager
1502+
await this.visibilityManager.complete(storedMessage.id, storedMessage.queueId);
1503+
15011504
// Add to DLQ
15021505
const pipeline = this.redis.pipeline();
15031506
pipeline.zadd(dlqKey, dlqMessage.deadLetteredAt, storedMessage.id);
15041507
pipeline.hset(dlqDataKey, storedMessage.id, JSON.stringify(dlqMessage));
1505-
const dlqResults = await pipeline.exec();
1506-
1507-
const dlqErrors = (dlqResults ?? [])
1508-
.map(([error]) => error)
1509-
.filter((error): error is Error => Boolean(error));
1510-
1511-
if (dlqErrors.length > 0) {
1512-
this.logger.error("Failed to write message to DLQ, leaving it in-flight to be reclaimed", {
1513-
messageId: storedMessage.id,
1514-
queueId: storedMessage.queueId,
1515-
tenantId: storedMessage.tenantId,
1516-
errors: dlqErrors.map((error) => error.message),
1517-
});
1518-
return;
1519-
}
1520-
1521-
// Complete in visibility manager
1522-
await this.visibilityManager.complete(storedMessage.id, storedMessage.queueId);
1508+
await pipeline.exec();
15231509

15241510
this.telemetry.recordDLQ();
15251511

@@ -1559,55 +1545,56 @@ export class FairQueue<TPayloadSchema extends z.ZodTypeAny = z.ZodUnknown> {
15591545
}
15601546

15611547
/**
1562-
* Free a reclaimed message's concurrency slot. Runs before the message is put back
1563-
* on the queue, so another consumer cannot reserve the same messageId and have its
1564-
* fresh reservation deleted by this release. Failures are logged rather than thrown
1565-
* so one bad message cannot stop the rest of the shard being reclaimed.
1548+
* Free every timed-out message's concurrency slot in one pipeline, before any of them
1549+
* is put back on the queue. Throwing here aborts the requeue for the whole batch, which
1550+
* leaves the messages in-flight for the next reclaim tick. That is the safe direction:
1551+
* requeuing a message whose slot is still held is what strands the slot permanently.
15661552
*/
1567-
async #releaseReclaimedConcurrency(message: ReclaimedMessageInfo): Promise<void> {
1568-
if (!this.concurrencyManager) {
1553+
async #releaseReclaimedConcurrency(messages: ReclaimedMessageInfo[]): Promise<void> {
1554+
if (!this.concurrencyManager || messages.length === 0) {
15691555
return;
15701556
}
15711557

1572-
try {
1573-
await this.concurrencyManager.release(
1574-
{
1558+
await this.concurrencyManager.releaseBatch(
1559+
messages.map((message) => ({
1560+
queue: {
15751561
id: message.queueId,
15761562
tenantId: message.tenantId,
15771563
metadata: message.metadata ?? {},
15781564
},
1579-
message.messageId
1580-
);
1581-
} catch (error) {
1582-
this.logger.error("Failed to release concurrency for reclaimed message", {
15831565
messageId: message.messageId,
1584-
queueId: message.queueId,
1585-
error: error instanceof Error ? error.message : String(error),
1586-
});
1587-
}
1566+
}))
1567+
);
15881568
}
15891569

15901570
async #reclaimTimedOutMessages(): Promise<void> {
15911571
let totalReclaimed = 0;
15921572

15931573
for (let shardId = 0; shardId < this.shardCount; shardId++) {
1594-
const reclaimedMessages = await this.visibilityManager.reclaimTimedOut(
1595-
shardId,
1596-
(queueId) => {
1597-
const tenantId = this.keys.extractTenantId(queueId);
1598-
const dispatchShardId = this.tenantDispatch.getShardForTenant(tenantId);
1599-
return {
1600-
queueKey: this.keys.queueKey(queueId),
1601-
queueItemsKey: this.keys.queueItemsKey(queueId),
1602-
tenantQueueIndexKey: this.keys.tenantQueueIndexKey(tenantId),
1603-
dispatchKey: this.keys.dispatchKey(dispatchShardId),
1604-
tenantId,
1605-
};
1606-
},
1607-
this.#releaseReclaimedConcurrency.bind(this)
1608-
);
1574+
try {
1575+
const reclaimedMessages = await this.visibilityManager.reclaimTimedOut(
1576+
shardId,
1577+
(queueId) => {
1578+
const tenantId = this.keys.extractTenantId(queueId);
1579+
const dispatchShardId = this.tenantDispatch.getShardForTenant(tenantId);
1580+
return {
1581+
queueKey: this.keys.queueKey(queueId),
1582+
queueItemsKey: this.keys.queueItemsKey(queueId),
1583+
tenantQueueIndexKey: this.keys.tenantQueueIndexKey(tenantId),
1584+
dispatchKey: this.keys.dispatchKey(dispatchShardId),
1585+
tenantId,
1586+
};
1587+
},
1588+
this.#releaseReclaimedConcurrency.bind(this)
1589+
);
16091590

1610-
totalReclaimed += reclaimedMessages.length;
1591+
totalReclaimed += reclaimedMessages.length;
1592+
} catch (error) {
1593+
this.logger.error("Failed to reclaim shard, leaving messages in-flight for the next tick", {
1594+
shardId,
1595+
error: error instanceof Error ? error.message : String(error),
1596+
});
1597+
}
16111598
}
16121599

16131600
if (totalReclaimed > 0) {

0 commit comments

Comments
 (0)