From 5d8efd0591ef3260189e777e13699b0412a94310 Mon Sep 17 00:00:00 2001
From: Antony Stubbs <antony.stubbs@gmail.com>
Date: Tue, 17 Oct 2017 16:32:06 +0100
Subject: [PATCH 1/2] MINOR: Adds an option to consume continuously

---
 .../kafka/tools/ConsumerPerformance.scala     | 44 +++++++++++++------
 1 file changed, 31 insertions(+), 13 deletions(-)

diff --git a/core/src/main/scala/kafka/tools/ConsumerPerformance.scala b/core/src/main/scala/kafka/tools/ConsumerPerformance.scala
index bdec41f6029aa..355426082cb02 100644
--- a/core/src/main/scala/kafka/tools/ConsumerPerformance.scala
+++ b/core/src/main/scala/kafka/tools/ConsumerPerformance.scala
@@ -65,7 +65,7 @@ object ConsumerPerformance {
       val consumer = new KafkaConsumer[Array[Byte], Array[Byte]](config.props)
       consumer.subscribe(Collections.singletonList(config.topic))
       startMs = System.currentTimeMillis
-      consume(consumer, List(config.topic), config.numMessages, 1000, config, totalMessagesRead, totalBytesRead, joinGroupTimeInMs, startMs)
+      consume(consumer, List(config.topic), config.numMessages, 1000, config, totalMessagesRead, totalBytesRead, joinGroupTimeInMs, startMs, config.runContinuously)
       endMs = System.currentTimeMillis
 
       if (config.printMetrics) {
@@ -141,7 +141,9 @@ object ConsumerPerformance {
               totalMessagesRead: AtomicLong,
               totalBytesRead: AtomicLong,
               joinTime: AtomicLong,
-              testStartTime: Long) {
+              testStartTime: Long,
+              readContinuously: Boolean
+             ) {
     var bytesRead = 0L
     var messagesRead = 0L
     var lastBytesRead = 0L
@@ -166,7 +168,23 @@ object ConsumerPerformance {
     var lastConsumedTime = System.currentTimeMillis
     var currentTimeMillis = lastConsumedTime
 
-    while (messagesRead < count && currentTimeMillis - lastConsumedTime <= timeout) {
+    def maybeReport = {
+      if (currentTimeMillis - lastReportTime >= config.reportingInterval) {
+        if (config.showDetailedStats)
+          printNewConsumerProgress(0, bytesRead, lastBytesRead, messagesRead, lastMessagesRead,
+            lastReportTime, currentTimeMillis, config.dateFormat, joinTimeMsInSingleRound)
+        joinTimeMsInSingleRound = 0L
+        lastReportTime = currentTimeMillis
+        lastMessagesRead = messagesRead
+        lastBytesRead = bytesRead
+      }
+    }
+
+    var moreToRead = true
+    var timeSinceLastConsume = 0L
+    var withinTimeout = true
+
+    while (moreToRead && withinTimeout || readContinuously) {
       val records = consumer.poll(100).asScala
       currentTimeMillis = System.currentTimeMillis
       if (records.nonEmpty)
@@ -178,16 +196,13 @@ object ConsumerPerformance {
         if (record.value != null)
           bytesRead += record.value.size
 
-        if (currentTimeMillis - lastReportTime >= config.reportingInterval) {
-          if (config.showDetailedStats)
-            printNewConsumerProgress(0, bytesRead, lastBytesRead, messagesRead, lastMessagesRead,
-              lastReportTime, currentTimeMillis, config.dateFormat, joinTimeMsInSingleRound)
-          joinTimeMsInSingleRound = 0L
-          lastReportTime = currentTimeMillis
-          lastMessagesRead = messagesRead
-          lastBytesRead = bytesRead
-        }
+        maybeReport
       }
+      moreToRead = messagesRead < count
+      timeSinceLastConsume = currentTimeMillis - lastConsumedTime
+      withinTimeout = timeSinceLastConsume <= timeout
+
+      maybeReport
     }
 
     totalMessagesRead.set(messagesRead)
@@ -304,6 +319,7 @@ object ConsumerPerformance {
     val printMetricsOpt = parser.accepts("print-metrics", "Print out the metrics. This only applies to new consumer.")
     val showDetailedStatsOpt = parser.accepts("show-detailed-stats", "If set, stats are reported for each reporting " +
       "interval as configured by reporting-interval")
+    val runContinuouslyOpt = parser.accepts("run-continuously", "Consume messages continuously, without end. Overrides message-count")
 
     val options = parser.parse(args: _*)
 
@@ -311,6 +327,7 @@ object ConsumerPerformance {
 
     val useOldConsumer = options.has(zkConnectOpt)
     val printMetrics = options.has(printMetricsOpt)
+    val runContinuously = options.has(runContinuouslyOpt)
 
     val props = if (options.has(consumerConfigOpt))
       Utils.loadProps(options.valueOf(consumerConfigOpt))
@@ -374,10 +391,11 @@ object ConsumerPerformance {
       var lastReportTime: Long = startMs
       var lastBytesRead = 0L
       var lastMessagesRead = 0L
+      var runContinuously = true
 
       try {
         val iter = stream.iterator
-        while (iter.hasNext && messagesRead < config.numMessages) {
+        while ((iter.hasNext && messagesRead < config.numMessages) || runContinuously) {
           val messageAndMetadata = iter.next()
           messagesRead += 1
           bytesRead += messageAndMetadata.message.length

From 08f8a616bb44b5439eaf75166c9bb221ec728591 Mon Sep 17 00:00:00 2001
From: Antony Stubbs <antony.stubbs@gmail.com>
Date: Tue, 17 Oct 2017 18:18:55 +0100
Subject: [PATCH 2/2] Explicitly set continuous to false to trick findbugs

---
 core/src/main/scala/kafka/tools/ConsumerPerformance.scala | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/core/src/main/scala/kafka/tools/ConsumerPerformance.scala b/core/src/main/scala/kafka/tools/ConsumerPerformance.scala
index 355426082cb02..5b42e26e37fbf 100644
--- a/core/src/main/scala/kafka/tools/ConsumerPerformance.scala
+++ b/core/src/main/scala/kafka/tools/ConsumerPerformance.scala
@@ -142,7 +142,7 @@ object ConsumerPerformance {
               totalBytesRead: AtomicLong,
               joinTime: AtomicLong,
               testStartTime: Long,
-              readContinuously: Boolean
+              readContinuously: Boolean = false
              ) {
     var bytesRead = 0L
     var messagesRead = 0L