consumer: fix consumeNum to respect the consume timeout

Steps to reproduce:
* configure consumer with setDefaultConsumeTimeout(1000)
* produce 1 message every 500ms
* call consumer.consume(128, cb)

Actual outcome:
* consume returns 128 messages after 64 seconds

Expected outcome:
* consumer returns ~2 messages after 1 second

KafkaConsumerConsumeNum call underlaying c++ m_consumer->Consume
in cycle until
* either the accumulated batch is full
* or the call to c++ m_consumer->Consume times out on the total timeout

KafkaConsumerConsumeNum must enforce its timeout
over all m_consumer->Consume invocations altogether.

Author: katreniak

e2e/both.spec.js

@@ -272,6 +272,49 @@ describe('Consumer/Producer', function() {
     });
   });
 
+  it('should stop consuming batch after consume timeout', function(done) {
+    crypto.randomBytes(4096, function(ex, buffer) {
+      producer.setPollInterval(10);
+
+      producer.once('delivery-report', function(err, report) {
+        t.ifError(err);
+      });
+
+      consumer.subscribe([topic]);
+
+      var events = [];
+
+      consumer.once('data', function(msg) {
+        events.push("data");
+      });
+
+      consumer.on('partition.eof', function(eof) {
+        events.push("partition.eof");
+      });
+
+      let timeoutId;
+      let toProduce = 10;
+      produceLoop = () => {
+        producer.produce(topic, null, buffer, null);
+        if (--toProduce > 0) {
+          timeoutId = setTimeout(produceLoop, 500);
+        }
+      };
+      produceLoop();
+
+      consumer.setDefaultConsumeTimeout(2000);
+      const startedAt = Date.now();
+      consumer.consume(100, function(err, messages) {
+        t.ifError(err);
+        t(Date.now() - startedAt < 3000, 'Consume took longer than consume timeout');
+        t(messages.length > 2, 'Too few messages consumed within batch');
+        t(messages.length < 8, 'Too many messages consumed within batch');
+        clearTimeout(timeoutId);
+        done();
+      });
+    });
+  });
+
   it('should be able to produce and consume messages: consumeLoop', function(done) {
     var key = 'key';
 

src/workers.cc

@@ -7,6 +7,8 @@
  * of the MIT license.  See the LICENSE.txt file for details.
  */
 
+#include <algorithm>
+#include <chrono>
 #include <string>
 #include <vector>
 
@@ -810,6 +812,7 @@ void KafkaConsumerConsumeNum::Execute() {
   bool looping = true;
   int timeout_ms = m_timeout_ms;
   std::size_t eof_event_count = 0;
+  const auto end = std::chrono::steady_clock::now() + std::chrono::milliseconds(timeout_ms);
 
   while (m_messages.size() - eof_event_count < max && looping) {
     // Get a message
@@ -838,6 +841,8 @@ void KafkaConsumerConsumeNum::Execute() {
           break;
         case RdKafka::ERR_NO_ERROR:
           m_messages.push_back(b.data<RdKafka::Message*>());
+          timeout_ms = std::max(0, static_cast<int>(std::chrono::duration_cast<std::chrono::milliseconds>(
+            end - std::chrono::steady_clock::now()).count()));
           break;
         default:
           // Set the error for any other errors and break