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

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