invoke consume in pollRecords

Author: Rancho-7

tools/src/main/java/org/apache/kafka/tools/automq/perf/ConsumerService.java

@@ -59,7 +59,6 @@
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
-import java.util.concurrent.atomic.LongAdder;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -105,17 +104,9 @@ public int createConsumers(List<Topic> topics, ConsumersConfig config) {
 
     public void start(ConsumerCallback callback, int pollRate) {
         BlockingBucket bucket = rateLimitBucket(pollRate);
-        LongAdder counter = new LongAdder();
-        ConsumerCallback callbackWithRateLimit = (tp, p, st) -> {
-            callback.messageReceived(tp, p, st);
-            counter.increment();
-            if ((counter.sum()) % 1000 == 0) {
-                bucket.consume(1);
-            }
-        };
         CompletableFuture.allOf(
             groups.stream()
-                .map(group -> group.start(callbackWithRateLimit))
+                .map(group -> group.start(callback, bucket))
                 .toArray(CompletableFuture[]::new)
         ).join();
     }
@@ -282,8 +273,8 @@ public Group(int index, int consumersPerGroup, List<Topic> topics, ConsumersConf
             }
         }
 
-        public CompletableFuture<Void> start(ConsumerCallback callback) {
-            consumers().forEach(consumer -> consumer.start(callback));
+        public CompletableFuture<Void> start(ConsumerCallback callback, BlockingBucket bucket) {
+            consumers().forEach(consumer -> consumer.start(callback, bucket));
 
             // wait for all consumers to join the group
             return CompletableFuture.allOf(consumers()
@@ -382,8 +373,8 @@ public Consumer(Properties properties, String topic) {
             consumer.subscribe(List.of(topic), subscribeListener());
         }
 
-        public void start(ConsumerCallback callback) {
-            this.task = this.executor.submit(() -> pollRecords(consumer, callback));
+        public void start(ConsumerCallback callback, BlockingBucket bucket) {
+            this.task = this.executor.submit(() -> pollRecords(consumer, callback, bucket));
         }
 
         public CompletableFuture<Void> started() {
@@ -413,7 +404,7 @@ public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
             };
         }
 
-        private void pollRecords(KafkaConsumer<String, byte[]> consumer, ConsumerCallback callback) {
+        private void pollRecords(KafkaConsumer<String, byte[]> consumer, ConsumerCallback callback, BlockingBucket bucket) {
             while (!closing) {
                 try {
                     while (paused) {
@@ -425,6 +416,7 @@ private void pollRecords(KafkaConsumer<String, byte[]> consumer, ConsumerCallbac
                         TopicPartition topicPartition = new TopicPartition(record.topic(), record.partition());
                         callback.messageReceived(topicPartition, record.value(), sendTimeNanos);
                     }
+                    bucket.consume(records.count());
                 } catch (InterruptException | InterruptedException e) {
                     // ignore, as we are closing
                 } catch (Exception e) {