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("auto.offset.reset", autoOffsetReset);
  72       props.put("metadata.max.age.ms", "1000");
  73       props.put("key.deserializer", StringDeserializer.class.getName());
  74       props.put("value.deserializer", StringDeserializer.class.getName());
  75
  76       this.consumer = new KafkaConsumer<>(props);
  77
  78       log.info("{} - Subscribing to topic {}", id, topic);
  79       consumer.subscribe(Arrays.asList(topic), new ConsumerRebalanceListener()
  80       {
  81         @Override
  82         public void onPartitionsRevoked(Collection<TopicPartition> partitions)
  83         {
  84           partitions.forEach(tp ->
  85           {
  86             log.info("{} - removing partition: {}", id, tp);
  87             Map<String, Integer> removed = seen.remove(tp.partition());
  88             for (String key : removed.keySet())
  89             {
  90               log.info(
  91                   "{} - Seen {} messages for partition={}|key={}",
  92                   id,
  93                   removed.get(key),
  94                   tp.partition(),
  95                   key);
  96             }
  97             repository.save(new StatisticsDocument(tp.partition(), removed));
  98           });
  99         }
 100
 101         @Override
 102         public void onPartitionsAssigned(Collection<TopicPartition> partitions)
 103         {
 104           partitions.forEach(tp ->
 105           {
 106             log.info("{} - adding partition: {}", id, tp);
 107             seen.put(
 108                 tp.partition(),
 109                 repository
 110                     .findById(Integer.toString(tp.partition()))
 111                     .map(document -> document.statistics)
 112                     .orElse(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, Integer> byKey = seen.get(partition);
 140
 141           if (!byKey.containsKey(key))
 142             byKey.put(key, 0);
 143
 144           int 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       shutdown();
 154     }
 155     catch(Exception e)
 156     {
 157       log.error("{} - Unexpected error: {}", id, e.toString(), e);
 158       shutdown(e);
 159     }
 160     finally
 161     {
 162       log.info("{} - Closing the KafkaConsumer", id);
 163       consumer.close();
 164       log.info("{} - Consumer-Thread exiting", id);
 165     }
 166   }
 167
 168   private void shutdown()
 169   {
 170     shutdown(null);
 171   }
 172
 173   private void shutdown(Exception e)
 174   {
 175     lock.lock();
 176     try
 177     {
 178       running = false;
 179       exception = e;
 180       condition.signal();
 181     }
 182     finally
 183     {
 184       lock.unlock();
 185     }
 186   }
 187
 188   public Map<Integer, Map<String, Integer>> getSeen()
 189   {
 190     return seen;
 191   }
 192
 193   public void start()
 194   {
 195     lock.lock();
 196     try
 197     {
 198       if (running)
 199         throw new IllegalStateException("Consumer instance " + id + " is already running!");
 200
 201       log.info("{} - Starting - consumed {} messages before", id, consumed);
 202       running = true;
 203       exception = null;
 204       executor.submit(this);
 205     }
 206     finally
 207     {
 208       lock.unlock();
 209     }
 210   }
 211
 212   public synchronized void stop() throws ExecutionException, InterruptedException
 213   {
 214     lock.lock();
 215     try
 216     {
 217       if (!running)
 218         throw new IllegalStateException("Consumer instance " + id + " is not running!");
 219
 220       log.info("{} - Stopping", id);
 221       consumer.wakeup();
 222       condition.await();
 223       log.info("{} - Stopped - consumed {} messages so far", id, consumed);
 224     }
 225     finally
 226     {
 227       lock.unlock();
 228     }
 229   }
 230
 231   @PreDestroy
 232   public void destroy() throws ExecutionException, InterruptedException
 233   {
 234     log.info("{} - Destroy!", id);
 235     try
 236     {
 237       stop();
 238     }
 239     catch (IllegalStateException e)
 240     {
 241       log.info("{} - Was already stopped", id);
 242     }
 243     catch (Exception e)
 244     {
 245       log.error("{} - Unexpected exception while trying to stop the consumer", id, e);
 246     }
 247     finally
 248     {
 249       log.info("{}: Consumed {} messages in total, exiting!", id, consumed);
 250     }
 251   }
 252
 253   public boolean running()
 254   {
 255     lock.lock();
 256     try
 257     {
 258       return running;
 259     }
 260     finally
 261     {
 262       lock.unlock();
 263     }
 264   }
 265
 266   public Optional<Exception> exitStatus()
 267   {
 268     lock.lock();
 269     try
 270     {
 271       if (running)
 272         throw new IllegalStateException("No exit-status available: Consumer instance " + id + " is running!");
 273
 274       return Optional.ofNullable(exception);
 275     }
 276     finally
 277     {
 278       lock.unlock();
 279     }
 280   }
 281 }