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.locks.Condition;
  18 import java.util.concurrent.locks.Lock;
  19 import java.util.concurrent.locks.ReentrantLock;
  20
  21
  22 @Slf4j
  23 public class EndlessConsumer implements Runnable
  24 {
  25   private final ExecutorService executor;
  26   private final PartitionStatisticsRepository repository;
  27   private final String bootstrapServer;
  28   private final String groupId;
  29   private final String id;
  30   private final String topic;
  31   private final String autoOffsetReset;
  32
  33   private final Lock lock = new ReentrantLock();
  34   private final Condition condition = lock.newCondition();
  35   private boolean running = false;
  36   private Exception exception;
  37   private long consumed = 0;
  38   private KafkaConsumer<String, String> consumer = null;
  39
  40
  41   private final Map<Integer, Map<String, Integer>> seen = new HashMap<>();
  42
  43
  44   public EndlessConsumer(
  45       ExecutorService executor,
  46       PartitionStatisticsRepository repository,
  47       String bootstrapServer,
  48       String groupId,
  49       String clientId,
  50       String topic,
  51       String autoOffsetReset)
  52   {
  53     this.executor = executor;
  54     this.repository = repository;
  55     this.bootstrapServer = bootstrapServer;
  56     this.groupId = groupId;
  57     this.id = clientId;
  58     this.topic = topic;
  59     this.autoOffsetReset = autoOffsetReset;
  60   }
  61
  62   @Override
  63   public void run()
  64   {
  65     try
  66     {
  67       Properties props = new Properties();
  68       props.put("bootstrap.servers", bootstrapServer);
  69       props.put("group.id", groupId);
  70       props.put("client.id", id);
  71       props.put("enable.auto.commit", false);
  72       props.put("auto.offset.reset", autoOffsetReset);
  73       props.put("metadata.max.age.ms", "1000");
  74       props.put("key.deserializer", StringDeserializer.class.getName());
  75       props.put("value.deserializer", StringDeserializer.class.getName());
  76
  77       this.consumer = new KafkaConsumer<>(props);
  78
  79       log.info("{} - Subscribing to topic {}", id, topic);
  80       consumer.subscribe(Arrays.asList(topic), new ConsumerRebalanceListener()
  81       {
  82         @Override
  83         public void onPartitionsRevoked(Collection<TopicPartition> partitions)
  84         {
  85           partitions.forEach(tp ->
  86           {
  87             log.info("{} - removing partition: {}", id, tp);
  88             Map<String, Integer> removed = seen.remove(tp.partition());
  89             for (String key : removed.keySet())
  90             {
  91               log.info(
  92                   "{} - Seen {} messages for partition={}|key={}",
  93                   id,
  94                   removed.get(key),
  95                   tp.partition(),
  96                   key);
  97             }
  98             repository.save(new StatisticsDocument(tp.partition(), removed, consumer.position(tp)));
  99           });
 100         }
 101
 102         @Override
 103         public void onPartitionsAssigned(Collection<TopicPartition> partitions)
 104         {
 105           partitions.forEach(tp ->
 106           {
 107             log.info("{} - adding partition: {}", id, tp);
 108             StatisticsDocument document =
 109                 repository
 110                     .findById(Integer.toString(tp.partition()))
 111                     .orElse(new StatisticsDocument(tp.partition()));
 112             consumer.seek(tp, document.offset);
 113             seen.put(tp.partition(), document.statistics);
 114           });
 115         }
 116       });
 117
 118       while (true)
 119       {
 120         ConsumerRecords<String, String> records =
 121             consumer.poll(Duration.ofSeconds(1));
 122
 123         // Do something with the data...
 124         log.info("{} - Received {} messages", id, records.count());
 125         for (ConsumerRecord<String, String> record : records)
 126         {
 127           consumed++;
 128           log.info(
 129               "{} - {}: {}/{} - {}={}",
 130               id,
 131               record.offset(),
 132               record.topic(),
 133               record.partition(),
 134               record.key(),
 135               record.value()
 136           );
 137
 138           Integer partition = record.partition();
 139           String key = record.key() == null ? "NULL" : record.key();
 140           Map<String, Integer> byKey = seen.get(partition);
 141
 142           if (!byKey.containsKey(key))
 143             byKey.put(key, 0);
 144
 145           int seenByKey = byKey.get(key);
 146           seenByKey++;
 147           byKey.put(key, seenByKey);
 148         }
 149
 150         seen.forEach((partiton, statistics) -> repository.save(
 151             new StatisticsDocument(
 152                 partiton,
 153                 statistics,
 154                 consumer.position(new TopicPartition(topic, partiton)))));
 155       }
 156     }
 157     catch(WakeupException e)
 158     {
 159       log.info("{} - RIIING!", id);
 160       shutdown();
 161     }
 162     catch(Exception e)
 163     {
 164       log.error("{} - Unexpected error: {}", id, e.toString(), e);
 165       shutdown(e);
 166     }
 167     finally
 168     {
 169       log.info("{} - Closing the KafkaConsumer", id);
 170       consumer.close();
 171       log.info("{} - Consumer-Thread exiting", id);
 172     }
 173   }
 174
 175   private void shutdown()
 176   {
 177     shutdown(null);
 178   }
 179
 180   private void shutdown(Exception e)
 181   {
 182     lock.lock();
 183     try
 184     {
 185       running = false;
 186       exception = e;
 187       condition.signal();
 188     }
 189     finally
 190     {
 191       lock.unlock();
 192     }
 193   }
 194
 195   public Map<Integer, Map<String, Integer>> getSeen()
 196   {
 197     return seen;
 198   }
 199
 200   public void start()
 201   {
 202     lock.lock();
 203     try
 204     {
 205       if (running)
 206         throw new IllegalStateException("Consumer instance " + id + " is already running!");
 207
 208       log.info("{} - Starting - consumed {} messages before", id, consumed);
 209       running = true;
 210       exception = null;
 211       executor.submit(this);
 212     }
 213     finally
 214     {
 215       lock.unlock();
 216     }
 217   }
 218
 219   public synchronized void stop() throws ExecutionException, InterruptedException
 220   {
 221     lock.lock();
 222     try
 223     {
 224       if (!running)
 225         throw new IllegalStateException("Consumer instance " + id + " is not running!");
 226
 227       log.info("{} - Stopping", id);
 228       consumer.wakeup();
 229       condition.await();
 230       log.info("{} - Stopped - consumed {} messages so far", id, consumed);
 231     }
 232     finally
 233     {
 234       lock.unlock();
 235     }
 236   }
 237
 238   @PreDestroy
 239   public void destroy() throws ExecutionException, InterruptedException
 240   {
 241     log.info("{} - Destroy!", id);
 242     try
 243     {
 244       stop();
 245     }
 246     catch (IllegalStateException e)
 247     {
 248       log.info("{} - Was already stopped", id);
 249     }
 250     catch (Exception e)
 251     {
 252       log.error("{} - Unexpected exception while trying to stop the consumer", id, e);
 253     }
 254     finally
 255     {
 256       log.info("{}: Consumed {} messages in total, exiting!", id, consumed);
 257     }
 258   }
 259
 260   public boolean running()
 261   {
 262     lock.lock();
 263     try
 264     {
 265       return running;
 266     }
 267     finally
 268     {
 269       lock.unlock();
 270     }
 271   }
 272
 273   public Optional<Exception> exitStatus()
 274   {
 275     lock.lock();
 276     try
 277     {
 278       if (running)
 279         throw new IllegalStateException("No exit-status available: Consumer instance " + id + " is running!");
 280
 281       return Optional.ofNullable(exception);
 282     }
 283     finally
 284     {
 285       lock.unlock();
 286     }
 287   }
 288 }