KAFKA-19667: Close ShareConsumer in ShareConsumerPerformance after metrics displayed (#20467)

Ensure that metrics are retrieved and displayed (when requested) before
ShareConsumer.close() is called. This is important because metrics are
technically supposed to be removed on ShareConsumer.close(), which means
retrieving them after close() would yield an empty map.

Related to #20267.

Reviewers: Apoorv Mittal <[email protected]>

Author: AndrewJSchofield

tools/src/main/java/org/apache/kafka/tools/ShareConsumerPerformance.java

@@ -20,6 +20,7 @@
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
 import org.apache.kafka.clients.consumer.KafkaShareConsumer;
+import org.apache.kafka.clients.consumer.ShareConsumer;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.Metric;
 import org.apache.kafka.common.MetricName;
@@ -47,8 +48,8 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Function;
 
 import joptsimple.OptionException;
 import joptsimple.OptionSpec;
@@ -59,6 +60,10 @@ public class ShareConsumerPerformance {
     private static final Logger LOG = LoggerFactory.getLogger(ShareConsumerPerformance.class);
 
     public static void main(String[] args) {
+        run(args, KafkaShareConsumer::new);
+    }
+
+    static void run(String[] args, Function<Properties, ShareConsumer<byte[], byte[]>> shareConsumerCreator) {
         try {
             LOG.info("Starting share consumer/consumers...");
             ShareConsumerPerfOptions options = new ShareConsumerPerfOptions(args);
@@ -68,9 +73,9 @@ public static void main(String[] args) {
             if (!options.hideHeader())
                 printHeader();
 
-            List<KafkaShareConsumer<byte[], byte[]>> shareConsumers = new ArrayList<>();
+            List<ShareConsumer<byte[], byte[]>> shareConsumers = new ArrayList<>();
             for (int i = 0; i < options.threads(); i++) {
-                shareConsumers.add(new KafkaShareConsumer<>(options.props()));
+                shareConsumers.add(shareConsumerCreator.apply(options.props()));
             }
             long startMs = System.currentTimeMillis();
             consume(shareConsumers, options, totalMessagesRead, totalBytesRead, startMs);
@@ -83,7 +88,6 @@ public static void main(String[] args) {
             shareConsumers.forEach(shareConsumer -> {
                 @SuppressWarnings("UnusedLocalVariable")
                 Map<TopicIdPartition, Optional<KafkaException>> ignored = shareConsumer.commitSync();
-                shareConsumer.close(Duration.ofMillis(500));
             });
 
             // Print final stats for share group.
@@ -94,6 +98,7 @@ public static void main(String[] args) {
 
             shareConsumersMetrics.forEach(ToolsUtils::printMetrics);
 
+            shareConsumers.forEach(shareConsumer -> shareConsumer.close(Duration.ofMillis(500)));
         } catch (Throwable e) {
             System.err.println(e.getMessage());
             System.err.println(Utils.stackTrace(e));
@@ -106,11 +111,11 @@ protected static void printHeader() {
         System.out.printf("start.time, end.time, data.consumed.in.MB, MB.sec, nMsg.sec, data.consumed.in.nMsg%s%n", newFieldsInHeader);
     }
 
-    private static void consume(List<KafkaShareConsumer<byte[], byte[]>> shareConsumers,
+    private static void consume(List<ShareConsumer<byte[], byte[]>> shareConsumers,
                                 ShareConsumerPerfOptions options,
                                 AtomicLong totalMessagesRead,
                                 AtomicLong totalBytesRead,
-                                long startMs) throws ExecutionException, InterruptedException, TimeoutException {
+                                long startMs) throws ExecutionException, InterruptedException {
         long numMessages = options.numMessages();
         long recordFetchTimeoutMs = options.recordFetchTimeoutMs();
         shareConsumers.forEach(shareConsumer -> shareConsumer.subscribe(options.topic()));
@@ -180,7 +185,7 @@ private static void consume(List<KafkaShareConsumer<byte[], byte[]>> shareConsum
         totalBytesRead.set(bytesRead.get());
     }
 
-    private static void consumeMessagesForSingleShareConsumer(KafkaShareConsumer<byte[], byte[]> shareConsumer,
+    private static void consumeMessagesForSingleShareConsumer(ShareConsumer<byte[], byte[]> shareConsumer,
                                                               AtomicLong totalMessagesRead,
                                                               AtomicLong totalBytesRead,
                                                               ShareConsumerPerfOptions options,

tools/src/test/java/org/apache/kafka/tools/ShareConsumerPerformanceTest.java

@@ -17,7 +17,10 @@
 package org.apache.kafka.tools;
 
 import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.MockShareConsumer;
+import org.apache.kafka.clients.consumer.ShareConsumer;
 import org.apache.kafka.common.utils.Exit;
+import org.apache.kafka.common.utils.Utils;
 
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
@@ -30,6 +33,8 @@
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.text.SimpleDateFormat;
+import java.util.Properties;
+import java.util.function.Function;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -120,6 +125,21 @@ public void testDefaultClientId() throws IOException {
         assertEquals("perf-share-consumer-client", config.props().getProperty(ConsumerConfig.CLIENT_ID_CONFIG));
     }
 
+    @Test
+    public void testMetricsRetrievedBeforeConsumerClosed() {
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--messages", "0",
+            "--print-metrics"
+        };
+
+        Function<Properties, ShareConsumer<byte[], byte[]>> shareConsumerCreator = properties -> new MockShareConsumer<>();
+
+        String err = ToolsTestUtils.captureStandardErr(() -> ShareConsumerPerformance.run(args, shareConsumerCreator));
+        assertTrue(Utils.isBlank(err), "Should be no stderr message, but was \"" + err + "\"");
+    }
+
     private void testHeaderMatchContent(int expectedOutputLineCount, Runnable runnable) {
         String out = ToolsTestUtils.captureStandardOut(() -> {
             ShareConsumerPerformance.printHeader();