Skip to content
11 changes: 10 additions & 1 deletion core/src/main/scala/kafka/tools/ConsumerPerformance.scala
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ object ConsumerPerformance extends LazyLogging {
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, config.recordFetchTimeoutMs, config, totalMessagesRead, totalBytesRead, joinGroupTimeInMs, startMs)
endMs = System.currentTimeMillis

if (config.printMetrics) {
Expand Down Expand Up @@ -188,6 +188,9 @@ object ConsumerPerformance extends LazyLogging {
}
}

if (messagesRead < count)
println(s"WARNING: Exiting before consuming the expected number of messages: timeout ($timeout ms) exceeded. " +
"You can use the --timeout option to increase the timeout."))
totalMessagesRead.set(messagesRead)
totalBytesRead.set(bytesRead)
}
Expand Down Expand Up @@ -302,6 +305,11 @@ object ConsumerPerformance extends LazyLogging {
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 recordFetchTimeoutOpt = parser.accepts("timeout", "The maximum allowed time in milliseconds between returned records.")
.withOptionalArg()
.describedAs("milliseconds")
.ofType(classOf[Long])
.defaultsTo(10000)

val options = parser.parse(args: _*)

Expand Down Expand Up @@ -354,6 +362,7 @@ object ConsumerPerformance extends LazyLogging {
val showDetailedStats = options.has(showDetailedStatsOpt)
val dateFormat = new SimpleDateFormat(options.valueOf(dateFormatOpt))
val hideHeader = options.has(hideHeaderOpt)
val recordFetchTimeoutMs = options.valueOf(recordFetchTimeoutOpt).longValue()
}

class ConsumerPerfThread(threadId: Int,
Expand Down