eachBatch should respect partition restriction despite concurrency (#224)

* Add failing test

* Do not remove partition unless no messages are fetched

* Update CHANGELOG.md

Author: milindl

CHANGELOG.md

@@ -6,6 +6,10 @@ v1.2.0 is a feature release. It is supported for all usage.
 
 1. Add support for an Admin API to fetch topic offsets by timestamp (#206).
 
+## Fixes
+
+1. Fixes an issue where the `eachBatch` callback was being called for the same partition concurrently (#224).
+
 
 # confluent-kafka-javascript v1.1.0
 

lib/kafkajs/_consumer_cache.js

@@ -241,11 +241,10 @@ class MessageCache {
 
         let nextN = ppc._nextN(size);
 
-        if (size === -1 || nextN.length < size) {
+        if (!nextN.length) {
             this.#removeEmptyPartition(ppc);
-        }
-        if (!nextN.length)
             return this.nextN(null, size);
+        }
 
         this.#size -= nextN.length;
         return [nextN, ppc];

test/promisified/consumer/consumeMessages.spec.js

@@ -812,4 +812,50 @@ describe.each(cases)('Consumer - partitionsConsumedConcurrently = %s -', (partit
         expect(lost).toEqual(1);
         expect(errors).toEqual(false);
     }, 60000);
+
+    it('consumes no more than one batch from the same partition at at time', async () => {
+        if (partitionsConsumedConcurrently === 1) {
+            return;
+        }
+
+        await producer.connect();
+
+        const messages = Array(200).fill().map(() => {
+            const value = secureRandom();
+            return { value: `value-${value}` };
+        });
+
+        await producer.send({
+            topic: topicName,
+            messages: messages,
+        });
+
+        await consumer.connect();
+        await consumer.subscribe({ topic: topicName });
+
+        const batchesHandled = new Map();
+        for (let i = 0; i < partitions; i++) {
+            batchesHandled.set(i, 0);
+        }
+
+        let messagesConsumed = 0;
+        let batchesCountExceeds1 = false;
+
+        consumer.run({
+            partitionsConsumedConcurrently,
+            eachBatch: async ({ batch }) => {
+                messagesConsumed += batch.messages.length;
+                batchesHandled.set(batch.partition, batchesHandled.get(batch.partition) + 1);
+                if (batchesHandled.get(batch.partition) > 1) {
+                    batchesCountExceeds1 = true;
+                } else {
+                    await sleep(100);
+                }
+                batchesHandled.set(batch.partition, batchesHandled.get(batch.partition) - 1);
+            }
+        });
+
+        await waitFor(() => (messagesConsumed === messages.length || batchesCountExceeds1), () => { }, 100);
+        expect(batchesCountExceeds1).toBe(false);
+    });
 });

test/promisified/unit/cache.spec.js

@@ -126,6 +126,27 @@ describe('MessageCache', () => {
         expect(receivedMessages.filter(msg => msg.partition === 1).length).toBe(30);
     });
 
+    it('caches messages and does not return messages from the same partition unless PPC is returned', () => {
+        // Just partition 0 has messages.
+        const msgs = messages.slice(0, 10).filter(msg => msg.partition === 0);
+        cache.addMessages(msgs);
+
+        let next = [null, null];
+        let nextPpc = [null, null];
+
+        // Fetch messages all the way to the end for partition 0 with size = 10 (can also use -1 instead of 10).
+        next[0] = cache.nextN(nextPpc[0], 10);
+        expect(next[0]).not.toBeNull();
+
+        // More messages get added in the meanwhile to the same partition.
+        cache.addMessages(msgs);
+
+        // This call should not return anything - since the partition is still
+        // owned by the first nextN caller, they have not returned it yet.
+        next[1] = cache.nextN(nextPpc[1], -1);
+        expect(next[1]).toBeNull();
+    });
+
     it('does not allow fetching messages more than available partitions at a time', () => {
         const msgs = messages.slice(0, 90);
         cache.addMessages(msgs);