Pause, resume and assignment

Author: emasab

examples/kafkajs/consumer.js

@@ -3,6 +3,8 @@ const { Kafka } = require('../..').KafkaJS
 
 async function consumerStart() {
     let consumer;
+    var stopped = false;
+
     const kafka = new Kafka({
         brokers: ['<fill>'],
         ssl: true,
@@ -12,24 +14,28 @@ async function consumerStart() {
             username: '<fill>',
             password: '<fill>',
         },
-        rebalanceListener: {
-          onPartitionsAssigned: async (assignment) => {
-            console.log(`Assigned partitions ${JSON.stringify(assignment)}`);
-          },
-          onPartitionsRevoked: async (assignment) => {
-            console.log(`Revoked partitions ${JSON.stringify(assignment)}`);
+    });
+
+    consumer = kafka.consumer({
+      groupId: 'test-group',
+      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}`);
             })
           }
-        },
-        rdKafka: {
-          'enable.auto.commit': false
         }
+      },
+      rdKafka: {
+        'enable.auto.commit': false
+      }
     });
 
-    consumer = kafka.consumer({ groupId: 'test-group' });
-
     await consumer.connect();
     console.log("Connected successfully");
 
@@ -39,6 +45,7 @@ async function consumerStart() {
       ]
     })
 
+    // Batch consumer, commit and seek example
     var batch = 0;
     consumer.run({
       eachMessage: async ({ topic, partition, message }) => {
@@ -62,12 +69,48 @@ async function consumerStart() {
       },
     });
 
+    // Pause/Resume example
+    const pauseResumeLoop = async () => {
+      let paused = false;
+      let ticks = 0;
+      while (!stopped) {
+        await new Promise((resolve) => setTimeout(resolve, 100));
+        if (stopped)
+          break;
+
+        ticks++;
+        if (ticks == 200) {
+          ticks = 0;
+          const assignment = consumer.assignment();
+          if (paused) {
+            console.log(`Resuming partitions ${JSON.stringify(assignment)}`)
+            consumer.resume(assignment);
+          } else {
+            console.log(`Pausing partitions ${JSON.stringify(assignment)}`);
+            consumer.pause(assignment);
+          }
+          paused = !paused;
+        }
+      }
+    }
+
+    if (consumer.assignment) {
+      // KafkaJS doesn't have assignment()
+      pauseResumeLoop()
+    }
+
+    // Disconnect example
     const disconnect = () => {
       process.off('SIGINT', disconnect);
       process.off('SIGTERM', disconnect);
-      consumer.disconnect().finally(() => {
-        console.log("Disconnected successfully");
-      });
+      stopped = true;
+      consumer.commitOffsets()
+      .finally(() =>
+        consumer.disconnect()
+      )
+      .finally(() =>
+        console.log("Disconnected successfully")
+      );
     }
     process.on('SIGINT', disconnect);
     process.on('SIGTERM', disconnect);

lib/kafkajs/_consumer.js

@@ -166,6 +166,23 @@ class Consumer {
     }
   }
 
+  #flattenTopicPartitions(topics) {
+    const ret = [];
+    for (let topic of topics) {
+      if (topic.partition != null)
+        ret.push({
+          topic: topic.topic,
+          partition: topic.partition
+        });
+      else {
+        for (let partition of topic.partitions) {
+          ret.push({topic: topic.topic, partition});
+        }
+      }
+    }
+    return ret;
+  }
+
   async connect() {
       if (this.#state !== ConsumerState.INIT) {
           return Promise.reject('Connect has already been called elsewhere.');
@@ -243,15 +260,21 @@ class Consumer {
   }
 
   pause(topics) {
-    this.#notImplemented();
+    topics = this.#flattenTopicPartitions(topics);
+    this.#internalClient.pause(topics);
   }
 
   paused() {
     this.#notImplemented();
   }
 
+  assignment() {
+    return this.#flattenTopicPartitions(this.#internalClient.assignments());
+  }
+
   resume(topics) {
-    this.#notImplemented();
+    topics = this.#flattenTopicPartitions(topics);
+    this.#internalClient.resume(topics);
   }
 
   on(eventName, listener) {

lib/kafkajs/_kafka.js

@@ -9,22 +9,27 @@ class Kafka {
         this.#commonClientConfig = config;
     }
 
-    producer(config) {
-        if (config === null || !(config instanceof Object)) {
-            config = {};
+    #mergeConfiguration(config) {
+      let baseConfig = Object.assign({}, this.#commonClientConfig);
+      config = Object.assign({}, config);
+
+      let rdKafka = baseConfig.rdKafka;
+      Object.assign(baseConfig, config);
+      if (rdKafka && config.rdKafka) {
+        baseConfig.rdKafka = {
+          ...rdkafka,
+          ...config.rdKafka
         }
-
-        config = Object.assign(config, this.#commonClientConfig);
-        return new Producer(config);
+      }
+      return baseConfig;
     }
 
-    consumer(config) {
-      if (config === null || !(config instanceof Object)) {
-        config = {};
+    producer(config) {
+        return new Producer(this.#mergeConfiguration(config));
     }
 
-      config = Object.assign(config, this.#commonClientConfig);
-      return new Consumer(config);
+    consumer(config) {
+      return new Consumer(this.#mergeConfiguration(config));
     }
 }
 

types/kafkajs.d.ts

@@ -417,6 +417,7 @@ export type Consumer = {
   describeGroup(): Promise<GroupDescription>
   pause(topics: Array<{ topic: string; partitions?: number[] }>): void
   paused(): TopicPartitions[]
+  assignment(): TopicPartitions[]
   resume(topics: Array<{ topic: string; partitions?: number[] }>): void
   on(
     eventName: ConsumerEvents['HEARTBEAT'],