fix two-stage processing with concurrency limit releasing in batches

Author: ericallam

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

@@ -977,14 +977,17 @@ export class FairQueue<TPayloadSchema extends z.ZodTypeAny = z.ZodUnknown> {
     let processedCount = 0;
 
     // Reserve concurrency and push each message to worker queue
-    for (const message of claimedMessages) {
+    for (let i = 0; i < claimedMessages.length; i++) {
+      const message = claimedMessages[i]!;
+
       // Reserve concurrency slot
       if (this.concurrencyManager) {
         const reserved = await this.concurrencyManager.reserve(descriptor, message.messageId);
         if (!reserved) {
-          // Release message back to queue (and ensure it's in master queue)
-          await this.visibilityManager.release(
-            message.messageId,
+          // Release ALL remaining messages (from index i onward) back to queue
+          // This prevents messages from being stranded in the in-flight set
+          await this.visibilityManager.releaseBatch(
+            claimedMessages.slice(i),
             queueId,
             queueKey,
             queueItemsKey,

packages/redis-worker/src/fair-queue/tests/fairQueue.test.ts

@@ -876,4 +876,106 @@ describe("FairQueue", () => {
       await queue.close();
     });
   });
+
+  describe("two-stage processing with concurrency limits", () => {
+    redisTest(
+      "should release remaining claimed messages when concurrency reservation fails",
+      { timeout: 30000 },
+      async ({ redisOptions }) => {
+        const processed: string[] = [];
+        const processingMessages = new Set<string>();
+        keys = new DefaultFairQueueKeyProducer({ prefix: "test" });
+
+        const scheduler = new DRRScheduler({
+          redis: redisOptions,
+          keys,
+          quantum: 10,
+          maxDeficit: 100,
+        });
+
+        // Create queue with:
+        // - Worker queue enabled (two-stage processing)
+        // - Concurrency limit of 2 per tenant
+        const queue = new FairQueue({
+          redis: redisOptions,
+          keys,
+          scheduler,
+          payloadSchema: TestPayloadSchema,
+          shardCount: 1,
+          consumerCount: 1,
+          consumerIntervalMs: 50,
+          visibilityTimeoutMs: 10000,
+          workerQueue: {
+            enabled: true,
+            blockingTimeoutSeconds: 1,
+          },
+          concurrency: {
+            groups: [
+              {
+                name: "tenant",
+                extractGroupId: (q) => q.tenantId,
+                getLimit: async () => 2, // Limit to 2 concurrent per tenant
+                defaultLimit: 2,
+              },
+            ],
+          },
+          startConsumers: false,
+        });
+
+        // Message handler that tracks what's being processed
+        queue.onMessage(async (ctx) => {
+          const value = ctx.message.payload.value;
+          processingMessages.add(value);
+
+          // Simulate some work
+          await new Promise((resolve) => setTimeout(resolve, 100));
+
+          processed.push(value);
+          processingMessages.delete(value);
+          await ctx.complete();
+        });
+
+        // Enqueue 5 messages to the same tenant queue
+        await queue.enqueueBatch({
+          queueId: "tenant:t1:queue:q1",
+          tenantId: "t1",
+          messages: [
+            { payload: { value: "msg-1" } },
+            { payload: { value: "msg-2" } },
+            { payload: { value: "msg-3" } },
+            { payload: { value: "msg-4" } },
+            { payload: { value: "msg-5" } },
+          ],
+        });
+
+        // Start processing
+        queue.start();
+
+        // Wait for all messages to be processed
+        // With concurrency limit of 2, it should process in batches
+        await vi.waitFor(
+          () => {
+            expect(processed.length).toBe(5);
+          },
+          { timeout: 20000 }
+        );
+
+        // All messages should have been processed
+        expect(processed).toContain("msg-1");
+        expect(processed).toContain("msg-2");
+        expect(processed).toContain("msg-3");
+        expect(processed).toContain("msg-4");
+        expect(processed).toContain("msg-5");
+
+        // Wait a bit for any cleanup to complete
+        await new Promise((resolve) => setTimeout(resolve, 100));
+
+        // No messages should be stuck in-flight
+        const inflightCount = await queue.getTotalInflightCount();
+        expect(inflightCount).toBe(0);
+
+        await queue.close();
+      }
+    );
+  });
 });

packages/redis-worker/src/fair-queue/tests/visibility.test.ts

@@ -244,5 +244,236 @@ describe("VisibilityManager", () => {
       }
     );
   });
+
+  describe("claimBatch", () => {
+    redisTest(
+      "should claim multiple messages atomically",
+      { timeout: 10000 },
+      async ({ redisOptions }) => {
+        keys = new DefaultFairQueueKeyProducer({ prefix: "test" });
+
+        const manager = new VisibilityManager({
+          redis: redisOptions,
+          keys,
+          shardCount: 1,
+          defaultTimeoutMs: 5000,
+        });
+
+        const redis = createRedisClient(redisOptions);
+        const queueId = "tenant:t1:queue:claim-batch";
+        const queueKey = keys.queueKey(queueId);
+        const queueItemsKey = keys.queueItemsKey(queueId);
+
+        // Add multiple messages to the queue
+        for (let i = 1; i <= 5; i++) {
+          const messageId = `msg-${i}`;
+          const storedMessage = {
+            id: messageId,
+            queueId,
+            tenantId: "t1",
+            payload: { value: `test-${i}` },
+            timestamp: Date.now() - (6 - i) * 1000,
+            attempt: 1,
+          };
+          await redis.zadd(queueKey, storedMessage.timestamp, messageId);
+          await redis.hset(queueItemsKey, messageId, JSON.stringify(storedMessage));
+        }
+
+        // Claim batch of 3 messages
+        const claimed = await manager.claimBatch(queueId, queueKey, queueItemsKey, "consumer-1", 3);
+
+        expect(claimed).toHaveLength(3);
+        expect(claimed[0]!.messageId).toBe("msg-1");
+        expect(claimed[1]!.messageId).toBe("msg-2");
+        expect(claimed[2]!.messageId).toBe("msg-3");
+
+        // Verify messages are in in-flight set
+        const inflightCount = await manager.getTotalInflightCount();
+        expect(inflightCount).toBe(3);
+
+        // Verify messages are removed from queue
+        const remainingCount = await redis.zcard(queueKey);
+        expect(remainingCount).toBe(2);
+
+        await manager.close();
+        await redis.quit();
+      }
+    );
+
+    redisTest(
+      "should return empty array when queue is empty",
+      { timeout: 10000 },
+      async ({ redisOptions }) => {
+        keys = new DefaultFairQueueKeyProducer({ prefix: "test" });
+
+        const manager = new VisibilityManager({
+          redis: redisOptions,
+          keys,
+          shardCount: 1,
+          defaultTimeoutMs: 5000,
+        });
+
+        const queueId = "tenant:t1:queue:empty";
+        const queueKey = keys.queueKey(queueId);
+        const queueItemsKey = keys.queueItemsKey(queueId);
+
+        const claimed = await manager.claimBatch(queueId, queueKey, queueItemsKey, "consumer-1", 5);
+        expect(claimed).toHaveLength(0);
+
+        await manager.close();
+      }
+    );
+  });
+
+  describe("releaseBatch", () => {
+    redisTest(
+      "should release multiple messages back to queue atomically",
+      { timeout: 10000 },
+      async ({ redisOptions }) => {
+        keys = new DefaultFairQueueKeyProducer({ prefix: "test" });
+
+        const manager = new VisibilityManager({
+          redis: redisOptions,
+          keys,
+          shardCount: 1,
+          defaultTimeoutMs: 5000,
+        });
+
+        const redis = createRedisClient(redisOptions);
+        const queueId = "tenant:t1:queue:release-batch";
+        const queueKey = keys.queueKey(queueId);
+        const queueItemsKey = keys.queueItemsKey(queueId);
+        const masterQueueKey = keys.masterQueueKey(0);
+
+        // Add messages to queue and claim them
+        for (let i = 1; i <= 5; i++) {
+          const messageId = `msg-${i}`;
+          const storedMessage = {
+            id: messageId,
+            queueId,
+            tenantId: "t1",
+            payload: { value: `test-${i}` },
+            timestamp: Date.now() - (6 - i) * 1000,
+            attempt: 1,
+          };
+          await redis.zadd(queueKey, storedMessage.timestamp, messageId);
+          await redis.hset(queueItemsKey, messageId, JSON.stringify(storedMessage));
+        }
+
+        // Claim all 5 messages
+        const claimed = await manager.claimBatch(queueId, queueKey, queueItemsKey, "consumer-1", 5);
+        expect(claimed).toHaveLength(5);
+
+        // Verify all messages are in-flight
+        let inflightCount = await manager.getTotalInflightCount();
+        expect(inflightCount).toBe(5);
+
+        // Queue should be empty
+        let queueCount = await redis.zcard(queueKey);
+        expect(queueCount).toBe(0);
+
+        // Release messages 3, 4, 5 back to queue (batch release)
+        const messagesToRelease = claimed.slice(2);
+        await manager.releaseBatch(
+          messagesToRelease,
+          queueId,
+          queueKey,
+          queueItemsKey,
+          masterQueueKey
+        );
+
+        // Verify 2 messages still in-flight
+        inflightCount = await manager.getTotalInflightCount();
+        expect(inflightCount).toBe(2);
+
+        // Verify 3 messages back in queue
+        queueCount = await redis.zcard(queueKey);
+        expect(queueCount).toBe(3);
+
+        // Verify the correct messages are back in queue
+        const queueMembers = await redis.zrange(queueKey, 0, -1);
+        expect(queueMembers).toContain("msg-3");
+        expect(queueMembers).toContain("msg-4");
+        expect(queueMembers).toContain("msg-5");
+
+        await manager.close();
+        await redis.quit();
+      }
+    );
+
+    redisTest(
+      "should handle empty messages array",
+      { timeout: 10000 },
+      async ({ redisOptions }) => {
+        keys = new DefaultFairQueueKeyProducer({ prefix: "test" });
+
+        const manager = new VisibilityManager({
+          redis: redisOptions,
+          keys,
+          shardCount: 1,
+          defaultTimeoutMs: 5000,
+        });
+
+        const queueId = "tenant:t1:queue:empty-release";
+        const queueKey = keys.queueKey(queueId);
+        const queueItemsKey = keys.queueItemsKey(queueId);
+        const masterQueueKey = keys.masterQueueKey(0);
+
+        // Should not throw when releasing empty array
+        await manager.releaseBatch([], queueId, queueKey, queueItemsKey, masterQueueKey);
+
+        await manager.close();
+      }
+    );
+
+    redisTest(
+      "should update master queue with oldest message timestamp",
+      { timeout: 10000 },
+      async ({ redisOptions }) => {
+        keys = new DefaultFairQueueKeyProducer({ prefix: "test" });
+
+        const manager = new VisibilityManager({
+          redis: redisOptions,
+          keys,
+          shardCount: 1,
+          defaultTimeoutMs: 5000,
+        });
+
+        const redis = createRedisClient(redisOptions);
+        const queueId = "tenant:t1:queue:master-update";
+        const queueKey = keys.queueKey(queueId);
+        const queueItemsKey = keys.queueItemsKey(queueId);
+        const masterQueueKey = keys.masterQueueKey(0);
+
+        // Add and claim messages
+        const baseTime = Date.now();
+        for (let i = 1; i <= 3; i++) {
+          const messageId = `msg-${i}`;
+          const storedMessage = {
+            id: messageId,
+            queueId,
+            tenantId: "t1",
+            payload: { value: `test-${i}` },
+            timestamp: baseTime + i * 1000, // Different timestamps
+            attempt: 1,
+          };
+          await redis.zadd(queueKey, storedMessage.timestamp, messageId);
+          await redis.hset(queueItemsKey, messageId, JSON.stringify(storedMessage));
+        }
+
+        const claimed = await manager.claimBatch(queueId, queueKey, queueItemsKey, "consumer-1", 3);
+
+        // Release all messages back
+        await manager.releaseBatch(claimed, queueId, queueKey, queueItemsKey, masterQueueKey);
+
+        // Master queue should have been updated
+        const masterScore = await redis.zscore(masterQueueKey, queueId);
+        expect(masterScore).not.toBeNull();
+
+        await manager.close();
+        await redis.quit();
+      }
+    );
+  });
 });