Provide a rebalance callback for cooperative rebalance strategy

Author: thynson

lib/kafka-consumer.js

@@ -21,6 +21,49 @@ var DEFAULT_CONSUME_LOOP_TIMEOUT_DELAY = 500;
 var DEFAULT_CONSUME_TIME_OUT = 1000;
 util.inherits(KafkaConsumer, Client);
 
+var eagerRebalanceCallback = function(err, assignment) {
+  // Create the librdkafka error
+  err = LibrdKafkaError.create(err);
+  // Emit the event
+  this.emit('rebalance', err, assignment);
+
+  // That's it
+  try {
+    if (err.code === -175 /*ERR__ASSIGN_PARTITIONS*/) {
+      this.assign(assignment);
+    } else if (err.code === -174 /*ERR__REVOKE_PARTITIONS*/) {
+      this.unassign();
+    }
+  } catch (e) {
+    // Ignore exceptions if we are not connected
+    if (this.isConnected()) {
+      this.emit('rebalance.error', e);
+    }
+  }
+}
+
+var cooperativeRebalanceCallback = function(err, assignment) {
+  // Create the librdkafka error
+  err = LibrdKafkaError.create(err);
+  // Emit the event
+  this.emit('rebalance', err, assignment);
+
+  // That's it
+  try {
+    if (err.code === -175 /*ERR__ASSIGN_PARTITIONS*/) {
+      this.incrementalAssign(assignment);
+    } else if (err.code === -174 /*ERR__REVOKE_PARTITIONS*/) {
+      this.incrementalUnassign(assignment);
+    }
+  } catch (e) {
+    // Ignore exceptions if we are not connected
+    if (this.isConnected()) {
+      this.emit('rebalance.error', e);
+    }
+  }
+}
+
+
 /**
  * KafkaConsumer class for reading messages from Kafka
  *
@@ -52,26 +95,9 @@ function KafkaConsumer(conf, topicConf) {
 
   // If rebalance is undefined we don't want any part of this
   if (onRebalance && typeof onRebalance === 'boolean') {
-    conf.rebalance_cb = function(err, assignment) {
-      // Create the librdkafka error
-      err = LibrdKafkaError.create(err);
-      // Emit the event
-      self.emit('rebalance', err, assignment);
-
-      // That's it
-      try {
-        if (err.code === -175 /*ERR__ASSIGN_PARTITIONS*/) {
-          self.assign(assignment);
-        } else if (err.code === -174 /*ERR__REVOKE_PARTITIONS*/) {
-          self.unassign();
-        }
-      } catch (e) {
-        // Ignore exceptions if we are not connected
-        if (self.isConnected()) {
-          self.emit('rebalance.error', e);
-        }
-      }
-    };
+      conf.rebalance_cb = conf['partition.assignment.strategy'] === 'cooperative-sticky'
+          ? cooperativeRebalanceCallback.bind(this)
+          : eagerRebalanceCallback.bind(this);
   } else if (onRebalance && typeof onRebalance === 'function') {
     /*
      * Once this is opted in to, that's it. It's going to manually rebalance