Add optimization where the size of messages in our subscription is less than the consumer cache size

Author: milindl

lib/kafka-consumer.js

@@ -20,6 +20,7 @@ var TopicPartition = require('./topic-partition');
 var shallowCopy = require('./util').shallowCopy;
 var DEFAULT_CONSUME_LOOP_TIMEOUT_DELAY = 500;
 var DEFAULT_CONSUME_TIME_OUT = 1000;
+const DEFAULT_IS_TIMEOUT_ONLY_FOR_FIRST_MESSAGE = false;
 util.inherits(KafkaConsumer, Client);
 
 /**
@@ -142,6 +143,7 @@ function KafkaConsumer(conf, topicConf) {
 
   this._consumeTimeout = DEFAULT_CONSUME_TIME_OUT;
   this._consumeLoopTimeoutDelay = DEFAULT_CONSUME_LOOP_TIMEOUT_DELAY;
+  this._consumeIsTimeoutOnlyForFirstMessage = DEFAULT_IS_TIMEOUT_ONLY_FOR_FIRST_MESSAGE;
 }
 
 /**
@@ -160,6 +162,20 @@ KafkaConsumer.prototype.setDefaultConsumeLoopTimeoutDelay = function(intervalMs)
   this._consumeLoopTimeoutDelay = intervalMs;
 };
 
+/**
+ * If true:
+ *  In consume(number, cb), we will wait for `timeoutMs` for the first message to be fetched.
+ *  Subsequent messages will not be waited for and will be fetched (upto `number`) if already ready.
+ *
+ * If false:
+ *  In consume(number, cb), we will wait for upto `timeoutMs` for each message to be fetched.
+ *
+ * @param {boolean} isTimeoutOnlyForFirstMessage
+ */
+KafkaConsumer.prototype.setDefaultIsTimeoutOnlyForFirstMessage = function(isTimeoutOnlyForFirstMessage) {
+  this._consumeIsTimeoutOnlyForFirstMessage = isTimeoutOnlyForFirstMessage;
+};
+
 /**
  * Get a stream representation of this KafkaConsumer
  *
@@ -512,7 +528,7 @@ KafkaConsumer.prototype._consumeLoop = function(timeoutMs, cb) {
 KafkaConsumer.prototype._consumeNum = function(timeoutMs, numMessages, cb) {
   var self = this;
 
-  this._client.consume(timeoutMs, numMessages, function(err, messages, eofEvents) {
+  this._client.consume(timeoutMs, numMessages, this._consumeIsTimeoutOnlyForFirstMessage, function(err, messages, eofEvents) {
     if (err) {
       err = LibrdKafkaError.create(err);
       if (cb) {

lib/kafkajs/_consumer.js

@@ -573,6 +573,9 @@ class Consumer {
     }
     this.#state = ConsumerState.CONNECTED;
 
+    /* Slight optimization for cases where the size of messages in our subscription is less than the cache size. */
+    this.#internalClient.setDefaultIsTimeoutOnlyForFirstMessage(false);
+
     // Resolve the promise.
     this.#connectPromiseFunc['resolve']();
   }

src/kafka-consumer.cc

@@ -1330,7 +1330,11 @@ NAN_METHOD(KafkaConsumer::NodeConsume) {
   }
 
   if (info[1]->IsNumber()) {
-    if (!info[2]->IsFunction()) {
+    if (!info[2]->IsBoolean()) {
+      return Nan::ThrowError("Need to specify a boolean");
+    }
+
+    if (!info[3]->IsFunction()) {
       return Nan::ThrowError("Need to specify a callback");
     }
 
@@ -1344,12 +1348,22 @@ NAN_METHOD(KafkaConsumer::NodeConsume) {
       numMessages = numMessagesMaybe.FromJust();
     }
 
+    v8::Local<v8::Boolean> isTimeoutOnlyForFirstMessageBoolean = info[2].As<v8::Boolean>();
+    Nan::Maybe<bool> isTimeoutOnlyForFirstMessageMaybe = Nan::To<bool>(isTimeoutOnlyForFirstMessageBoolean);  // NOLINT
+
+    bool isTimeoutOnlyForFirstMessage;
+    if (isTimeoutOnlyForFirstMessageMaybe.IsNothing()) {
+      return Nan::ThrowError("Parameter must be a boolean");
+    } else {
+      isTimeoutOnlyForFirstMessage = isTimeoutOnlyForFirstMessageMaybe.FromJust();
+    }
+
     KafkaConsumer* consumer = ObjectWrap::Unwrap<KafkaConsumer>(info.This());
 
-    v8::Local<v8::Function> cb = info[2].As<v8::Function>();
+    v8::Local<v8::Function> cb = info[3].As<v8::Function>();
     Nan::Callback *callback = new Nan::Callback(cb);
     Nan::AsyncQueueWorker(
-      new Workers::KafkaConsumerConsumeNum(callback, consumer, numMessages, timeout_ms));  // NOLINT
+      new Workers::KafkaConsumerConsumeNum(callback, consumer, numMessages, timeout_ms, isTimeoutOnlyForFirstMessage));  // NOLINT
 
   } else {
     if (!info[1]->IsFunction()) {

src/workers.cc

@@ -808,11 +808,13 @@ void KafkaConsumerConsumeLoop::HandleErrorCallback() {
 KafkaConsumerConsumeNum::KafkaConsumerConsumeNum(Nan::Callback *callback,
                                      KafkaConsumer* consumer,
                                      const uint32_t & num_messages,
-                                     const int & timeout_ms) :
+                                     const int & timeout_ms,
+                                     bool timeout_only_for_first_message) :
   ErrorAwareWorker(callback),
   m_consumer(consumer),
   m_num_messages(num_messages),
-  m_timeout_ms(timeout_ms) {}
+  m_timeout_ms(timeout_ms),
+  m_timeout_only_for_first_message(timeout_only_for_first_message) {}
 
 KafkaConsumerConsumeNum::~KafkaConsumerConsumeNum() {}
 
@@ -849,6 +851,14 @@ void KafkaConsumerConsumeNum::Execute() {
           break;
         case RdKafka::ERR_NO_ERROR:
           m_messages.push_back(b.data<RdKafka::Message*>());
+
+          // This allows getting ready messages, while not waiting for new ones.
+          // This is useful when we want to get the as many messages as possible
+          // within the timeout but not wait if we already have one or more messages.
+          if (m_timeout_only_for_first_message) {
+            timeout_ms = 1;
+          }
+
           break;
         default:
           // Set the error for any other errors and break

src/workers.h

@@ -441,7 +441,7 @@ class KafkaConsumerSeek : public ErrorAwareWorker {
 class KafkaConsumerConsumeNum : public ErrorAwareWorker {
  public:
   KafkaConsumerConsumeNum(Nan::Callback*, NodeKafka::KafkaConsumer*,
-    const uint32_t &, const int &);
+    const uint32_t &, const int &, bool);
   ~KafkaConsumerConsumeNum();
 
   void Execute();
@@ -451,6 +451,7 @@ class KafkaConsumerConsumeNum : public ErrorAwareWorker {
   NodeKafka::KafkaConsumer * m_consumer;
   const uint32_t m_num_messages;
   const int m_timeout_ms;
+  const bool m_timeout_only_for_first_message;
   std::vector<RdKafka::Message*> m_messages;
 };