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.Clock;
  12 import java.time.Duration;
  13 import java.time.Instant;
  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 import java.util.regex.Pattern;
  21
  22
  23 @Slf4j
  24 @RequiredArgsConstructor
  25 public class EndlessConsumer implements ConsumerRebalanceListener, Runnable
  26 {
  27   final static Pattern PATTERN = Pattern.compile("\\W+");
  28
  29
  30   private final ExecutorService executor;
  31   private final PartitionStatisticsRepository repository;
  32   private final String id;
  33   private final String topic;
  34   private final Clock clock;
  35   private final Duration commitInterval;
  36   private final Consumer<String, String> consumer;
  37
  38   private final Lock lock = new ReentrantLock();
  39   private final Condition condition = lock.newCondition();
  40   private boolean running = false;
  41   private Exception exception;
  42   private long consumed = 0;
  43
  44   private final Map<Integer, Map<String, Map<String, Long>>> seen = new HashMap<>();
  45
  46
  47   @Override
  48   public void onPartitionsRevoked(Collection<TopicPartition> partitions)
  49   {
  50     partitions.forEach(tp ->
  51     {
  52       Integer partition = tp.partition();
  53       Long newOffset = consumer.position(tp);
  54       log.info(
  55           "{} - removing partition: {}, offset of next message {})",
  56           id,
  57           partition,
  58           newOffset);
  59       Map<String, Map<String, Long>> removed = seen.remove(partition);
  60       repository.save(new StatisticsDocument(partition, removed, consumer.position(tp)));
  61     });
  62   }
  63
  64   @Override
  65   public void onPartitionsAssigned(Collection<TopicPartition> partitions)
  66   {
  67     partitions.forEach(tp ->
  68     {
  69       Integer partition = tp.partition();
  70       Long offset = consumer.position(tp);
  71       log.info("{} - adding partition: {}, offset={}", id, partition, offset);
  72       StatisticsDocument document =
  73           repository
  74               .findById(Integer.toString(partition))
  75               .orElse(new StatisticsDocument(partition));
  76       if (document.offset >= 0)
  77       {
  78         // Only seek, if a stored offset was found
  79         // Otherwise: Use initial offset, generated by Kafka
  80         consumer.seek(tp, document.offset);
  81       }
  82       seen.put(partition, document.statistics);
  83     });
  84   }
  85
  86
  87   @Override
  88   public void run()
  89   {
  90     try
  91     {
  92       log.info("{} - Subscribing to topic {}", id, topic);
  93       consumer.subscribe(Arrays.asList(topic), this);
  94
  95       Instant lastCommit = clock.instant();
  96
  97       while (true)
  98       {
  99         ConsumerRecords<String, String> 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<String, String> 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           consumed++;
 117
 118           Integer partition = record.partition();
 119           String user = record.key();
 120           Map<String, Map<String, Long>> users = seen.get(partition);
 121
 122           Map<String, Long> words = users.get(user);
 123           if (words == null)
 124           {
 125             words = new HashMap<>();
 126             users.put(user, words);
 127           }
 128
 129           for (String word : PATTERN.split(record.value()))
 130           {
 131             Long num = words.get(word);
 132             if (num == null)
 133             {
 134               num = 1l;
 135             }
 136             else
 137             {
 138               num++;
 139             }
 140             words.put(word, num);
 141           }
 142         }
 143
 144         if (lastCommit.plus(commitInterval).isBefore(clock.instant()))
 145         {
 146           log.debug("Storing data and offsets, last commit: {}", lastCommit);
 147           seen.forEach((partiton, statistics) -> repository.save(
 148               new StatisticsDocument(
 149                   partiton,
 150                   statistics,
 151                   consumer.position(new TopicPartition(topic, partiton)))));
 152           lastCommit = clock.instant();
 153         }
 154       }
 155     }
 156     catch(WakeupException e)
 157     {
 158       log.info("{} - RIIING! Request to stop consumption - commiting current offsets!", id);
 159       shutdown();
 160     }
 161     catch(RecordDeserializationException e)
 162     {
 163       TopicPartition tp = e.topicPartition();
 164       long offset = e.offset();
 165       log.error(
 166           "{} - Could not deserialize  message on topic {} with offset={}: {}",
 167           id,
 168           tp,
 169           offset,
 170           e.getCause().toString());
 171
 172       shutdown(e);
 173     }
 174     catch(Exception e)
 175     {
 176       log.error("{} - Unexpected error: {}", id, e.toString(), e);
 177       shutdown(e);
 178     }
 179     finally
 180     {
 181       log.info("{} - Consumer-Thread exiting", id);
 182     }
 183   }
 184
 185   private void shutdown()
 186   {
 187     shutdown(null);
 188   }
 189
 190   private void shutdown(Exception e)
 191   {
 192     lock.lock();
 193     try
 194     {
 195       try
 196       {
 197         log.info("{} - Unsubscribing from topic {}", id, topic);
 198         consumer.unsubscribe();
 199       }
 200       catch (Exception ue)
 201       {
 202         log.error(
 203             "{} - Error while unsubscribing from topic {}: {}",
 204             id,
 205             topic,
 206             ue.toString());
 207       }
 208       finally
 209       {
 210         running = false;
 211         exception = e;
 212         condition.signal();
 213       }
 214     }
 215     finally
 216     {
 217       lock.unlock();
 218     }
 219   }
 220
 221   public Map<Integer, Map<String, Map<String, Long>>> getSeen()
 222   {
 223     return seen;
 224   }
 225
 226   public void start()
 227   {
 228     lock.lock();
 229     try
 230     {
 231       if (running)
 232         throw new IllegalStateException("Consumer instance " + id + " is already running!");
 233
 234       log.info("{} - Starting - consumed {} messages before", id, consumed);
 235       running = true;
 236       exception = null;
 237       executor.submit(this);
 238     }
 239     finally
 240     {
 241       lock.unlock();
 242     }
 243   }
 244
 245   public synchronized void stop() throws InterruptedException
 246   {
 247     lock.lock();
 248     try
 249     {
 250       if (!running)
 251         throw new IllegalStateException("Consumer instance " + id + " is not running!");
 252
 253       log.info("{} - Stopping", id);
 254       consumer.wakeup();
 255       condition.await();
 256       log.info("{} - Stopped - consumed {} messages so far", id, consumed);
 257     }
 258     finally
 259     {
 260       lock.unlock();
 261     }
 262   }
 263
 264   @PreDestroy
 265   public void destroy() throws ExecutionException, InterruptedException
 266   {
 267     log.info("{} - Destroy!", id);
 268     log.info("{}: Consumed {} messages in total, exiting!", id, consumed);
 269   }
 270
 271   public boolean running()
 272   {
 273     lock.lock();
 274     try
 275     {
 276       return running;
 277     }
 278     finally
 279     {
 280       lock.unlock();
 281     }
 282   }
 283
 284   public Optional<Exception> exitStatus()
 285   {
 286     lock.lock();
 287     try
 288     {
 289       if (running)
 290         throw new IllegalStateException("No exit-status available: Consumer instance " + id + " is running!");
 291
 292       return Optional.ofNullable(exception);
 293     }
 294     finally
 295     {
 296       lock.unlock();
 297     }
 298   }
 299 }