KAFKA-14752: Kafka examples improvements - consumer changes

examples/src/main/java/kafka/examples/Consumer.java

@@ -22,96 +22,110 @@
 import org.apache.kafka.clients.consumer.ConsumerRecords;
 import org.apache.kafka.clients.consumer.KafkaConsumer;
 import org.apache.kafka.common.TopicPartition;
-import org.apache.kafka.common.errors.WakeupException;
+import org.apache.kafka.common.serialization.IntegerDeserializer;
+import org.apache.kafka.common.serialization.StringDeserializer;
 
 import java.time.Duration;
 import java.util.Collection;
-import java.util.Collections;
 import java.util.Optional;
 import java.util.Properties;
+import java.util.UUID;
 import java.util.concurrent.CountDownLatch;
 
+import static java.util.Collections.singleton;
+
 /**
- * A simple consumer thread that demonstrate subscribe and poll use case. The thread subscribes to a topic,
- * then runs a loop to poll new messages, and print the message out. The thread closes until the target {@code
- * numMessageToConsume} is hit or catching an exception.
+ * A simple consumer thread that subscribes to a topic, fetches new records and prints them.
+ * The thread does not stop until all records are completed or an exception is raised.
  */
 public class Consumer extends Thread implements ConsumerRebalanceListener {
-    private final KafkaConsumer<Integer, String> consumer;
+    private final String bootstrapServers;
     private final String topic;
     private final String groupId;
-    private final int numMessageToConsume;
-    private int messageRemaining;
+    private final Optional<String> instanceId;
+    private final boolean readCommitted;
+    private final int numRecords;
     private final CountDownLatch latch;
+    private volatile boolean closed;
+    private int remainingRecords;
 
-    public Consumer(final String topic,
-                    final String groupId,
-                    final Optional<String> instanceId,
-                    final boolean readCommitted,
-                    final int numMessageToConsume,
-                    final CountDownLatch latch) {
-        super("KafkaConsumerExample");
-        this.groupId = groupId;
-        Properties props = new Properties();
-        props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, KafkaProperties.KAFKA_SERVER_URL + ":" + KafkaProperties.KAFKA_SERVER_PORT);
-        props.put(ConsumerConfig.GROUP_ID_CONFIG, groupId);
-        instanceId.ifPresent(id -> props.put(ConsumerConfig.GROUP_INSTANCE_ID_CONFIG, id));
-        props.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true");
-        props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.IntegerDeserializer");
-        props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.StringDeserializer");
-        if (readCommitted) {
-            props.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG, "read_committed");
-        }
-        props.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
-
-        consumer = new KafkaConsumer<>(props);
+    public Consumer(String threadName,
+                    String bootstrapServers,
+                    String topic,
+                    String groupId,
+                    Optional<String> instanceId,
+                    boolean readCommitted,
+                    int numRecords,
+                    CountDownLatch latch) {
+        super(threadName);
+        this.bootstrapServers = bootstrapServers;
         this.topic = topic;
-        this.numMessageToConsume = numMessageToConsume;
-        this.messageRemaining = numMessageToConsume;
+        this.groupId = groupId;
+        this.instanceId = instanceId;
+        this.readCommitted = readCommitted;
+        this.numRecords = numRecords;
+        this.remainingRecords = numRecords;
         this.latch = latch;
     }
 
-    KafkaConsumer<Integer, String> get() {
-        return consumer;
-    }
-
     @Override
     public void run() {
-        try {
-            System.out.println("Subscribe to:" + this.topic);
-            consumer.subscribe(Collections.singletonList(this.topic), this);
-            do {
-                doWork();
-            } while (messageRemaining > 0);
-            System.out.println(groupId + " finished reading " + numMessageToConsume + " messages");
-        } catch (WakeupException e) {
-            // swallow the wakeup
-        } catch (Exception e) {
-            System.out.println("Unexpected termination, exception thrown:" + e);
-        } finally {
-            shutdown();
+        // the consumer instance is NOT thread safe
+        try (KafkaConsumer<Integer, String> consumer = createKafkaConsumer()) {
+            consumer.subscribe(singleton(topic), this);
+            Utils.printOut("Subscribed to %s", topic);
+            while (!closed && remainingRecords > 0) {
+                try {
+                    // next poll must be called within session.timeout.ms to avoid rebalance
+                    ConsumerRecords<Integer, String> records = consumer.poll(Duration.ofSeconds(1));
+                    for (ConsumerRecord<Integer, String> record : records) {
+                        Utils.maybePrintRecord(numRecords, record);
+                    }
+                    remainingRecords -= records.count();
+                } catch (Throwable e) {
+                    // add your application retry strategy here
+                    Utils.printErr(e.getMessage());
+                    break;
+                }
+            }
+        } catch (Throwable e) {
+            Utils.printOut("Fatal error");
+            e.printStackTrace();
         }
+        Utils.printOut("Fetched %d records", numRecords - remainingRecords);
+        shutdown();
     }
-    public void doWork() {
-        ConsumerRecords<Integer, String> records = consumer.poll(Duration.ofSeconds(1));
-        for (ConsumerRecord<Integer, String> record : records) {
-            System.out.println(groupId + " received message : from partition " + record.partition() + ", (" + record.key() + ", " + record.value() + ") at offset " + record.offset());
+
+    public void shutdown() {
+        if (!closed) {
+            closed = true;
+            latch.countDown();
         }
-        messageRemaining -= records.count();
     }
 
-    public void shutdown() {
-        this.consumer.close();
-        latch.countDown();
+    public KafkaConsumer<Integer, String> createKafkaConsumer() {
+        Properties props = new Properties();
+        props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers);
+        props.put(ConsumerConfig.CLIENT_ID_CONFIG, "client-" + UUID.randomUUID());
+        props.put(ConsumerConfig.GROUP_ID_CONFIG, groupId);
+        instanceId.ifPresent(id -> props.put(ConsumerConfig.GROUP_INSTANCE_ID_CONFIG, id));
+        props.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, readCommitted ? "false" : "true");
+        props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, IntegerDeserializer.class);
+        props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
+        if (readCommitted) {
+            props.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG, "read_committed");
+        }
+        props.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
+        return new KafkaConsumer<>(props);
     }
 
     @Override
     public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
-        System.out.println("Revoking partitions:" + partitions);
+        Utils.printOut("Revoked partitions: %s", partitions);
     }
 
     @Override
     public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
-        System.out.println("Assigning partitions:" + partitions);
+        Utils.printOut("Assigned partitions: %s", partitions);
     }
 }

examples/src/main/java/kafka/examples/ExactlyOnceMessageProcessor.java

@@ -70,8 +70,10 @@ public ExactlyOnceMessageProcessor(final String inputTopic,
         // Consumer must be in read_committed mode, which means it won't be able to read uncommitted data.
         // Consumer could optionally configure groupInstanceId to avoid unnecessary rebalances.
         this.groupInstanceId = "Txn-consumer-" + instanceIdx;
-        consumer = new Consumer(inputTopic, "Eos-consumer",
-            Optional.of(groupInstanceId), READ_COMMITTED, -1, null).get();
+        boolean readCommitted = true;
+        consumer = new Consumer(
+            "processor-consumer", KafkaProperties.KAFKA_SERVER_URL + ":" + KafkaProperties.KAFKA_SERVER_PORT, inputTopic, "processor-group", Optional.of(groupInstanceId), readCommitted, -1, null)
+                .createKafkaConsumer();
         this.latch = latch;
     }
 

examples/src/main/java/kafka/examples/KafkaConsumerProducerDemo.java

@@ -29,7 +29,8 @@ public static void main(String[] args) throws InterruptedException {
         Producer producerThread = new Producer(KafkaProperties.TOPIC, isAsync, null, false, 10000, -1, latch);
         producerThread.start();
 
-        Consumer consumerThread = new Consumer(KafkaProperties.TOPIC, "DemoConsumer", Optional.empty(), false, 10000, latch);
+        Consumer consumerThread = new Consumer(
+            "consumer", KafkaProperties.KAFKA_SERVER_URL + ":" + KafkaProperties.KAFKA_SERVER_PORT, KafkaProperties.TOPIC, "DemoConsumer", Optional.empty(), false, 10000, latch);
         consumerThread.start();
 
         if (!latch.await(5, TimeUnit.MINUTES)) {

examples/src/main/java/kafka/examples/KafkaExactlyOnceDemo.java

@@ -114,7 +114,8 @@ public static void main(String[] args) throws InterruptedException, ExecutionExc
         CountDownLatch consumeLatch = new CountDownLatch(1);
 
         /* Stage 4: consume all processed messages to verify exactly once */
-        Consumer consumerThread = new Consumer(OUTPUT_TOPIC, "Verify-consumer", Optional.empty(), true, numRecords, consumeLatch);
+        Consumer consumerThread = new Consumer(
+            "consumer", "DemoConsumer", OUTPUT_TOPIC, "Verify-consumer", Optional.empty(), true, numRecords, consumeLatch);
         consumerThread.start();
 
         if (!consumeLatch.await(5, TimeUnit.MINUTES)) {

examples/src/main/java/kafka/examples/Utils.java

@@ -0,0 +1,106 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package kafka.examples;
+
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.AdminClientConfig;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.producer.RecordMetadata;
+import org.apache.kafka.common.errors.TopicExistsException;
+import org.apache.kafka.common.errors.UnknownTopicOrPartitionException;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.Properties;
+import java.util.UUID;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
+
+import static java.lang.String.format;
+
+public class Utils {
+    private Utils() {
+    }
+
+    public static void printHelp(String message, Object... args) {
+        System.out.println(format(message, args));
+    }
+
+    public static void printOut(String message, Object... args) {
+        System.out.printf("%s - %s%n", Thread.currentThread().getName(), format(message, args));
+    }
+
+    public static void printErr(String message, Object... args) {
+        System.err.printf("%s - %s%n", Thread.currentThread().getName(), format(message, args));
+    }
+
+    public static void maybePrintRecord(long numRecords, ConsumerRecord<Integer, String> record) {
+        maybePrintRecord(numRecords, record.key(), record.value(), record.topic(), record.partition(), record.offset());
+    }
+
+    public static void maybePrintRecord(long numRecords, int key, String value, RecordMetadata metadata) {
+        maybePrintRecord(numRecords, key, value, metadata.topic(), metadata.partition(), metadata.offset());
+    }
+
+    private static void maybePrintRecord(long numRecords, int key, String value, String topic, int partition, long offset) {
+        // we only print 10 records when there are 20 or more to send
+        if (key % Math.max(1, numRecords / 10) == 0) {
+            printOut("Sample: record(%d, %s), partition(%s-%d), offset(%d)", key, value, topic, partition, offset);
+        }
+    }
+
+    public static void recreateTopics(String bootstrapServers, int numPartitions, String... topicNames) {
+        Properties props = new Properties();
+        props.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers);
+        props.put(AdminClientConfig.CLIENT_ID_CONFIG, "client-" + UUID.randomUUID());
+        try (Admin admin = Admin.create(props)) {
+            // delete topics if present
+            try {
+                admin.deleteTopics(Arrays.asList(topicNames)).all().get();
+            } catch (ExecutionException e) {
+                if (!(e.getCause() instanceof UnknownTopicOrPartitionException)) {
+                    throw e;
+                }
+                printErr("Topics deletion error: %s", e.getCause());
+            }
+            printOut("Deleted topics: %s", Arrays.toString(topicNames));
+            // create topics in a retry loop
+            while (true) {
+                // use default RF to avoid NOT_ENOUGH_REPLICAS error with minISR>1
+                short replicationFactor = -1;
+                List<NewTopic> newTopics = Arrays.stream(topicNames)
+                    .map(name -> new NewTopic(name, numPartitions, replicationFactor))
+                    .collect(Collectors.toList());
+                try {
+                    admin.createTopics(newTopics).all().get();
+                    printOut("Created topics: %s", Arrays.toString(topicNames));
+                    break;
+                } catch (ExecutionException e) {
+                    if (!(e.getCause() instanceof TopicExistsException)) {
+                        throw e;
+                    }
+                    printOut("Waiting for topics metadata cleanup");
+                    TimeUnit.MILLISECONDS.sleep(1_000);
+                }
+            }
+        } catch (Throwable e) {
+            throw new RuntimeException("Topics creation error", e);
+        }
+    }
+}