Add assign/unassign within rebalance callbacks

Also remove onPartitionsAssigned and onPartitionsRevoked.

Author: milindl

examples/kafkajs/consumer.js

@@ -1,5 +1,5 @@
 // require('kafkajs') is replaced with require('@confluentinc/kafka-javascript').KafkaJS.
-const { Kafka } = require('@confluentinc/kafka-javascript').KafkaJS;
+const { Kafka, ErrorCodes } = require('@confluentinc/kafka-javascript').KafkaJS;
 
 async function consumerStart() {
   let consumer;
@@ -22,22 +22,17 @@ async function consumerStart() {
     kafkaJS: {
       groupId: 'test-group',
       autoCommit: false,
-      rebalanceListener: {
-        onPartitionsAssigned: async (assignment) => {
-          console.log(`Assigned partitions ${JSON.stringify(assignment)}`);
-        },
-        onPartitionsRevoked: async (assignment) => {
-          console.log(`Revoked partitions ${JSON.stringify(assignment)}`);
-          if (!stopped) {
-            await consumer.commitOffsets().catch((e) => {
-              console.error(`Failed to commit ${e}`);
-            })
-          }
-        }
-      },
     },
-
     /* Properties from librdkafka can also be used */
+    rebalance_cb: (err, assignment) => {
+      if (err.code === ErrorCodes.ERR__ASSIGN_PARTITIONS) {
+        console.log(`Assigned partitions ${JSON.stringify(assignment)}`);
+      } else if (err.code === ErrorCodes.ERR__REVOKE_PARTITIONS) {
+        console.log(`Revoked partitions ${JSON.stringify(assignment)}`);
+      } else {
+        console.error(`Rebalance error ${err}`);
+      }
+    },
     'auto.commit.interval.ms': 6000,
   });
 

lib/kafkajs/_common.js

@@ -52,7 +52,6 @@ const kafkaJSProperties = {
     'autoCommit',
     'autoCommitInterval',
     'autoCommitThreshold',
-    'rebalanceListener',
   ],
   admin: [],
 };

lib/kafkajs/_consumer.js

@@ -248,45 +248,67 @@ class Consumer {
     err = LibrdKafkaError.create(err);
     const userSpecifiedRebalanceCb = this.#userConfig['rebalance_cb'];
 
-    let call;
+    let assignmentFnCalled = false;
+    function assignmentFn(userAssignment) {
+      if (assignmentFnCalled)
+        return;
+      assignmentFnCalled = true;
+
+      if (this.#internalClient.rebalanceProtocol() === "EAGER") {
+        this.#internalClient.assign(userAssignment);
+        this.#partitionCount = userAssignment.length;
+      } else {
+        this.#internalClient.incrementalAssign(userAssignment);
+        this.#partitionCount += userAssignment.length;
+      }
+    }
 
-    /* Since we don't expose assign() or incremental_assign() methods, we allow the user
-     * to modify the assignment by returning it. If a truthy value is returned, we use that
-     * and do not apply any pending seeks to it either. */
+    function unassignmentFn(userAssignment) {
+      if (assignmentFnCalled)
+        return;
+
+      assignmentFnCalled = true;
+      if (this.#internalClient.rebalanceProtocol() === "EAGER") {
+        this.#internalClient.unassign();
+        this.#messageCache.removeTopicPartitions();
+        this.#partitionCount = 0;
+      } else {
+        this.#internalClient.incrementalUnassign(userAssignment);
+        this.#messageCache.removeTopicPartitions(userAssignment);
+        this.#partitionCount -= userAssignment.length;
+      }
+    }
+
+    let call = Promise.resolve();
+
+    /* We allow the user to modify the assignment by returning it. If a truthy
+     * value is returned, we use that and do not apply any pending seeks to it either.
+     * The user can alternatively use the assignmentFns argument.
+     * Precedence is given to the calling of functions within assignmentFns. */
     let assignmentModified = false;
     if (typeof userSpecifiedRebalanceCb === 'function') {
       call = new Promise((resolve, reject) => {
-        try {
-          const alternateAssignment = userSpecifiedRebalanceCb(err, assignment);
+        const assignmentFns = {
+          assign: assignmentFn.bind(this),
+          unassign: unassignmentFn.bind(this),
+        };
+
+        /* The user specified callback may be async, or sync. Wrapping it in a
+         * Promise.resolve ensures that we always get a promise back. */
+        return Promise.resolve(
+          userSpecifiedRebalanceCb(err, assignment, assignmentFns)
+        ).then(alternateAssignment => {
           if (alternateAssignment) {
             assignment = alternateAssignment;
             assignmentModified = true;
           }
           resolve();
-        } catch (e) {
-          reject(e);
-        }
+        }).catch(reject);
       });
-    } else {
-      switch (err.code) {
-        // TODO: is this the right way to handle this error?
-        // We might just be able to throw, because the error is something the user has caused.
-        case LibrdKafkaError.codes.ERR__ASSIGN_PARTITIONS:
-          call = (this.#userConfig.rebalanceListener.onPartitionsAssigned ?
-            this.#userConfig.rebalanceListener.onPartitionsAssigned(assignment) :
-            Promise.resolve()).catch(e => this.#logger.error(e));
-          break;
-        case LibrdKafkaError.codes.ERR__REVOKE_PARTITIONS:
-          call = (this.#userConfig.rebalanceListener.onPartitionsRevoked ?
-            this.#userConfig.rebalanceListener.onPartitionsRevoked(assignment) :
-            Promise.resolve()).catch(e => this.#logger.error(e));
-          break;
-        default:
-          call = Promise.reject(`Unexpected rebalanceListener error code ${err.code}`).catch((e) => {
-            this.#logger.error(e);
-          });
-          break;
-      }
+    } else if (err.code !== LibrdKafkaError.codes.ERR__ASSIGN_PARTITIONS && err.code !== LibrdKafkaError.codes.ERR__REVOKE_PARTITIONS) {
+        call = Promise.reject(`Unexpected rebalance_cb error code ${err.code}`).catch((e) => {
+          this.#logger.error(e);
+        });
     }
 
     call
@@ -311,16 +333,10 @@ class Consumer {
           if (err.code === LibrdKafkaError.codes.ERR__ASSIGN_PARTITIONS) {
 
             const checkPendingSeeks = this.#pendingSeeks.size !== 0;
-            if (checkPendingSeeks && !assignmentModified)
+            if (checkPendingSeeks && !assignmentModified && !assignmentFnCalled)
               assignment = this.#assignAsPerSeekedOffsets(assignment);
 
-            if (this.#internalClient.rebalanceProtocol() === "EAGER") {
-              this.#internalClient.assign(assignment);
-              this.#partitionCount = assignment.length;
-            } else {
-              this.#internalClient.incrementalAssign(assignment);
-              this.#partitionCount += assignment.length;
-            }
+            assignmentFn.call(this, assignment);
 
             if (checkPendingSeeks) {
               const offsetsToCommit = assignment
@@ -342,15 +358,7 @@ class Consumer {
             this.#messageCache.addTopicPartitions(assignment);
 
           } else {
-            if (this.#internalClient.rebalanceProtocol() === "EAGER") {
-              this.#internalClient.unassign();
-              this.#messageCache.removeTopicPartitions();
-              this.#partitionCount = 0;
-            } else {
-              this.#internalClient.incrementalUnassign(assignment);
-              this.#messageCache.removeTopicPartitions(assignment);
-              this.#partitionCount -= assignment.length;
-            }
+            unassignmentFn.call(this, assignment);
           }
         } catch (e) {
           // Ignore exceptions if we are not connected
@@ -522,16 +530,10 @@ class Consumer {
 
     /* Delete properties which are already processed, or cannot be passed to node-rdkafka */
     delete rdKafkaConfig.kafkaJS;
-    delete rdKafkaConfig.rebalanceListener;
 
     /* Certain properties that the user has set are overridden. We use trampolines to accommodate the user's callbacks.
      * TODO: add trampoline method for offset commit callback. */
     rdKafkaConfig['offset_commit_cb'] = true;
-
-    if (!Object.hasOwn(this.#userConfig, 'rebalanceListener')) {
-      /* We might want to do certain things to maintain internal state in rebalance listener, so we need to set it to an empty object. */
-      this.#userConfig.rebalanceListener = {};
-    }
     rdKafkaConfig['rebalance_cb'] = this.#rebalanceCallback.bind(this);
 
     /* Offset management is different from case to case.
@@ -1587,6 +1589,7 @@ class Consumer {
       if (!topic.partitions) {
         toppar.partitions = this.#getAllAssignedPartition(topic.topic);
       } else {
+        /* TODO: add a check here to make sure we own each partition */
         toppar.partitions = [...topic.partitions];
       }
 
@@ -1597,6 +1600,8 @@ class Consumer {
     if (flattenedToppars.length === 0) {
       return;
     }
+
+    /* TODO: error handling is lacking for pause, including partition level errors. */
     this.#internalClient.pause(flattenedToppars);
 
     /* Mark the messages in the cache as stale, runInternal* will deal with
@@ -1608,7 +1613,7 @@ class Consumer {
       .filter(key => this.#topicPartitionToBatchPayload.has(key))
       .forEach(key => this.#topicPartitionToBatchPayload.get(key)._stale = true);
 
-      flattenedToppars.map(JSON.stringify).forEach(topicPartition => this.#pausedPartitions.add(topicPartition));
+    flattenedToppars.map(JSON.stringify).forEach(topicPartition => this.#pausedPartitions.add(topicPartition));
 
     /* Note: we don't use flattenedToppars here because resume flattens them again. */
     return () => this.resume(toppars);

test/promisified/consumer/incrementalRebalance.spec.js

@@ -148,12 +148,11 @@ describe('Consumer > incremental rebalance', () => {
         let revokes = 0;
 
         consumer = createConsumer(consumerConfig, {
-            rebalanceListener: {
-                onPartitionsAssigned: async (assignment) => {
+            rebalance_cb: async (err, assignment) => {
+                if (err.code === ErrorCodes.ERR__ASSIGN_PARTITIONS) {
                     assigns++;
                     expect(assignment.length).toBe(2);
-                },
-                onPartitionsRevoked: async (assignment) => {
+                } else if (err.code === ErrorCodes.ERR__REVOKE_PARTITIONS) {
                     revokes++;
                     expect(assignment.length).toBe(2);
                 }

test/promisified/consumer/rebalanceCallback.spec.js

@@ -0,0 +1,143 @@
+jest.setTimeout(30000);
+
+const { waitFor,
+    secureRandom,
+    createTopic,
+    createConsumer,
+    createProducer,
+    sleep, } = require("../testhelpers");
+const { ErrorCodes } = require('../../../lib').KafkaJS;
+
+describe('Consumer', () => {
+    let consumer;
+    let groupId, topicName;
+    let consumerConfig;
+
+    beforeEach(async () => {
+        topicName = `test-topic-${secureRandom()}`;
+        groupId = `consumer-group-id-${secureRandom()}`;
+        consumerConfig = {
+            groupId,
+        };
+        consumer = null;
+        await createTopic({ topic: topicName, partitions: 3 });
+    });
+
+    afterEach(async () => {
+        consumer && (await consumer.disconnect());
+    });
+
+    it('calls rebalance callback', async () => {
+        let calls = 0;
+        consumer = createConsumer(consumerConfig, {
+            rebalance_cb: function () {
+                calls++;
+            }
+        });
+
+        await consumer.connect();
+        await consumer.subscribe({ topic: topicName });
+        consumer.run({ eachMessage: async () => { } });
+        await waitFor(() => consumer.assignment().length > 0, () => null, 1000);
+        expect(calls).toBe(1); /* assign */
+        await consumer.disconnect();
+        expect(calls).toBe(2); /* assign + unassign */
+        consumer = null;
+    });
+
+    it('allows modifying the assignment via returns', async () => {
+        consumer = createConsumer(consumerConfig, {
+            rebalance_cb: function (err, assignment) {
+                if (err.code === ErrorCodes.ERR__ASSIGN_PARTITIONS) {
+                    expect(assignment.length).toBe(3);
+                    return assignment.filter(a => a.partition !== 0);
+                }
+            }
+        });
+
+        await consumer.connect();
+        await consumer.subscribe({ topic: topicName });
+        consumer.run({ eachMessage: async () => { } });
+        await waitFor(() => consumer.assignment().length > 0, () => null, 1000);
+        expect(consumer.assignment().length).toBe(2);
+        expect(consumer.assignment()).toEqual(
+            expect.arrayContaining([
+                { topic: topicName, partition: 1 },
+                { topic: topicName, partition: 2 }]));
+    });
+
+    it('allows modifying the assigment via assignment functions', async () => {
+        let calls = 0;
+        consumer = createConsumer(consumerConfig, {
+            rebalance_cb: function (err, assignment, assignmentFns) {
+                calls++;
+                if (err.code === ErrorCodes.ERR__ASSIGN_PARTITIONS) {
+                    expect(assignment.length).toBe(3);
+                    assignmentFns.assign(assignment.filter(a => a.partition !== 0));
+                } else {
+                    assignmentFns.unassign(assignment);
+                }
+            }
+        });
+
+        await consumer.connect();
+        await consumer.subscribe({ topic: topicName });
+        consumer.run({ eachMessage: async () => { } });
+        await waitFor(() => consumer.assignment().length > 0, () => null, 1000);
+        expect(consumer.assignment().length).toBe(2);
+        expect(consumer.assignment()).toEqual(
+            expect.arrayContaining([
+                { topic: topicName, partition: 1 },
+                { topic: topicName, partition: 2 }]));
+        await consumer.disconnect();
+        expect(calls).toBe(2);
+        consumer = null;
+    });
+
+    it('pauses correctly from the rebalance callback after assign', async () => {
+        consumer = createConsumer(consumerConfig, {
+            rebalance_cb: function (err, assignment, assignmentFns) {
+                if (err.code === ErrorCodes.ERR__ASSIGN_PARTITIONS) {
+                    expect(assignment.length).toBe(3);
+
+                    /* Assign first so we can pause. */
+                    assignmentFns.assign(assignment);
+
+                    /* Convert the assignment into format suitable for pause argument. */
+                    const pausablePartitions = [{ topic: topicName, partitions: [0, 1, 2] }];
+                    consumer.pause(pausablePartitions);
+                } else {
+                    assignmentFns.unassign(assignment);
+                }
+            }
+        });
+
+        let messagesConsumed = [];
+        await consumer.connect();
+        await consumer.subscribe({ topic: topicName });
+        consumer.run({ eachMessage: async (e) => { messagesConsumed.push(e); } });
+        await waitFor(() => consumer.assignment().length > 0, () => null, 1000);
+
+        const producer = createProducer({});
+        await producer.connect();
+        const key1 = secureRandom();
+        for (const partition of [0, 1, 2]) {
+            const message = { key: `key-${key1}`, value: `value-${key1}`, partition };
+            await producer.send({
+                topic: topicName,
+                messages: [message],
+            });
+        }
+        await producer.disconnect();
+
+        expect(consumer.paused()).toEqual([{ topic: topicName, partitions: [0, 1, 2] }]);
+
+        /* Give it some extra time just in case - should be enough to get the messages if a partition isn't paused. */
+        await sleep(1000);
+        expect(messagesConsumed.length).toBe(0);
+
+        consumer.resume([ { topic: topicName } ]);
+        await waitFor(() => messagesConsumed.length === 3, () => null, 1000);
+        expect(messagesConsumed.length).toBe(3);
+    });
+});