juplo.de Git - demos/kafka/training/blob - src/main/java/de/juplo/kafka/EndlessConsumer.java

   1 package de.juplo.kafka;
   2
   3 import lombok.extern.slf4j.Slf4j;
   4 import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
   5 import org.apache.kafka.clients.consumer.ConsumerRecord;
   6 import org.apache.kafka.clients.consumer.ConsumerRecords;
   7 import org.apache.kafka.clients.consumer.KafkaConsumer;
   8 import org.apache.kafka.common.TopicPartition;
   9 import org.apache.kafka.common.errors.WakeupException;
  10 import org.apache.kafka.common.serialization.StringDeserializer;
  11
  12 import javax.annotation.PreDestroy;
  13 import java.time.Duration;
  14 import java.util.*;
  15 import java.util.concurrent.ExecutionException;
  16 import java.util.concurrent.ExecutorService;
  17 import java.util.concurrent.Future;
  18 import java.util.concurrent.atomic.AtomicBoolean;
  19
  20
  21 @Slf4j
  22 public class EndlessConsumer implements Runnable
  23 {
  24   private final ExecutorService executor;
  25   private final String bootstrapServer;
  26   private final String groupId;
  27   private final String id;
  28   private final String topic;
  29   private final String autoOffsetReset;
  30
  31   private AtomicBoolean running = new AtomicBoolean();
  32   private long consumed = 0;
  33   private KafkaConsumer<String, String> consumer = null;
  34   private Future<?> future = null;
  35
  36   private final Map<Integer, Map<String, Long>> seen = new HashMap<>();
  37   private final Map<Integer, Long> offsets = new HashMap<>();
  38
  39
  40   public EndlessConsumer(
  41       ExecutorService executor,
  42       String bootstrapServer,
  43       String groupId,
  44       String clientId,
  45       String topic,
  46       String autoOffsetReset)
  47   {
  48     this.executor = executor;
  49     this.bootstrapServer = bootstrapServer;
  50     this.groupId = groupId;
  51     this.id = clientId;
  52     this.topic = topic;
  53     this.autoOffsetReset = autoOffsetReset;
  54   }
  55
  56   @Override
  57   public void run()
  58   {
  59     try
  60     {
  61       Properties props = new Properties();
  62       props.put("bootstrap.servers", bootstrapServer);
  63       props.put("group.id", groupId);
  64       props.put("client.id", id);
  65       props.put("auto.offset.reset", autoOffsetReset);
  66       props.put("metadata.max.age.ms", "1000");
  67       props.put("key.deserializer", StringDeserializer.class.getName());
  68       props.put("value.deserializer", StringDeserializer.class.getName());
  69
  70       this.consumer = new KafkaConsumer<>(props);
  71
  72       log.info("{} - Subscribing to topic {}", id, topic);
  73       consumer.subscribe(Arrays.asList(topic), new ConsumerRebalanceListener()
  74       {
  75         @Override
  76         public void onPartitionsRevoked(Collection<TopicPartition> partitions)
  77         {
  78           partitions.forEach(tp ->
  79           {
  80             Integer partition = tp.partition();
  81             Long newOffset = consumer.position(tp);
  82             Long oldOffset = offsets.remove(partition);
  83             log.info(
  84                 "{} - removing partition: {}, consumed {} records (offset {} -> {})",
  85                 id,
  86                 partition,
  87                 newOffset - oldOffset,
  88                 oldOffset,
  89                 newOffset);
  90             Map<String, Long> removed = seen.remove(partition);
  91             for (String key : removed.keySet())
  92             {
  93               log.info(
  94                   "{} - Seen {} messages for partition={}|key={}",
  95                   id,
  96                   removed.get(key),
  97                   partition,
  98                   key);
  99             }
 100           });
 101         }
 102
 103         @Override
 104         public void onPartitionsAssigned(Collection<TopicPartition> partitions)
 105         {
 106           partitions.forEach(tp ->
 107           {
 108             Integer partition = tp.partition();
 109             Long offset = consumer.position(tp);
 110             log.info("{} - adding partition: {}, offset={}", id, partition, offset);
 111             offsets.put(partition, offset);
 112             seen.put(partition, new HashMap<>());
 113           });
 114         }
 115       });
 116
 117       while (true)
 118       {
 119         ConsumerRecords<String, String> records =
 120             consumer.poll(Duration.ofSeconds(1));
 121
 122         // Do something with the data...
 123         log.info("{} - Received {} messages", id, records.count());
 124         for (ConsumerRecord<String, String> record : records)
 125         {
 126           consumed++;
 127           log.info(
 128               "{} - {}: {}/{} - {}={}",
 129               id,
 130               record.offset(),
 131               record.topic(),
 132               record.partition(),
 133               record.key(),
 134               record.value()
 135           );
 136
 137           Integer partition = record.partition();
 138           String key = record.key() == null ? "NULL" : record.key();
 139           Map<String, Long> byKey = seen.get(partition);
 140
 141           if (!byKey.containsKey(key))
 142             byKey.put(key, 0l);
 143
 144           long seenByKey = byKey.get(key);
 145           seenByKey++;
 146           byKey.put(key, seenByKey);
 147         }
 148       }
 149     }
 150     catch(WakeupException e)
 151     {
 152       log.info("{} - RIIING!", id);
 153     }
 154     catch(Exception e)
 155     {
 156       log.error("{} - Unexpected error: {}", id, e.toString(), e);
 157       running.set(false); // Mark the instance as not running
 158     }
 159     finally
 160     {
 161       log.info("{} - Closing the KafkaConsumer", id);
 162       consumer.close();
 163       log.info("{} - Consumer-Thread exiting", id);
 164     }
 165   }
 166
 167   public Map<Integer, Map<String, Long>> getSeen()
 168   {
 169     return seen;
 170   }
 171
 172   public synchronized void start()
 173   {
 174     boolean stateChanged = running.compareAndSet(false, true);
 175     if (!stateChanged)
 176       throw new RuntimeException("Consumer instance " + id + " is already running!");
 177
 178     log.info("{} - Starting - consumed {} messages before", id, consumed);
 179     future = executor.submit(this);
 180   }
 181
 182   public Optional<Long> getOffset(Integer partition)
 183   {
 184     return Optional.ofNullable(offsets.get(partition));
 185   }
 186
 187   public synchronized void stop() throws ExecutionException, InterruptedException
 188   {
 189     boolean stateChanged = running.compareAndSet(true, false);
 190     if (!stateChanged)
 191       throw new RuntimeException("Consumer instance " + id + " is not running!");
 192
 193     log.info("{} - Stopping", id);
 194     consumer.wakeup();
 195     future.get();
 196     log.info("{} - Stopped - consumed {} messages so far", id, consumed);
 197   }
 198
 199   @PreDestroy
 200   public void destroy() throws ExecutionException, InterruptedException
 201   {
 202     log.info("{} - Destroy!", id);
 203     try
 204     {
 205       stop();
 206     }
 207     catch (IllegalStateException e)
 208     {
 209       log.info("{} - Was already stopped", id);
 210     }
 211     finally
 212     {
 213       log.info("{}: Consumed {} messages in total, exiting!", id, consumed);
 214     }
 215   }
 216 }