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

   1 package de.juplo.kafka;
   2
   3 import lombok.RequiredArgsConstructor;
   4 import lombok.extern.slf4j.Slf4j;
   5 import org.apache.kafka.clients.consumer.Consumer;
   6 import org.apache.kafka.clients.consumer.ConsumerRecord;
   7 import org.apache.kafka.common.TopicPartition;
   8
   9 import java.time.Clock;
  10 import java.time.Duration;
  11 import java.time.Instant;
  12 import java.util.HashMap;
  13 import java.util.Map;
  14 import java.util.regex.Pattern;
  15
  16
  17 @RequiredArgsConstructor
  18 @Slf4j
  19 public class WordcountRecordHandler implements RecordHandler<String, String>
  20 {
  21   final static Pattern PATTERN = Pattern.compile("\\W+");
  22
  23
  24   private final PartitionStatisticsRepository repository;
  25   private final String id;
  26   private final String topic;
  27   private final Clock clock;
  28   private final Duration commitInterval;
  29   private final Consumer<String, String> consumer;
  30
  31   private final Map<Integer, Map<String, Map<String, Long>>> seen = new HashMap<>();
  32
  33   private Instant lastCommit = Instant.EPOCH;
  34
  35
  36   @Override
  37   public void accept(ConsumerRecord<String, String> record)
  38   {
  39     Integer partition = record.partition();
  40     String user = record.key();
  41     Map<String, Map<String, Long>> users = seen.get(partition);
  42
  43     Map<String, Long> words = users.get(user);
  44     if (words == null)
  45     {
  46       words = new HashMap<>();
  47       users.put(user, words);
  48     }
  49
  50     for (String word : PATTERN.split(record.value()))
  51     {
  52       Long num = words.get(word);
  53       if (num == null)
  54       {
  55         num = 1l;
  56       }
  57       else
  58       {
  59         num++;
  60       }
  61       words.put(word, num);
  62     }
  63   }
  64
  65
  66   @Override
  67   public void beforeNextPoll()
  68   {
  69     if (lastCommit.plus(commitInterval).isBefore(clock.instant()))
  70     {
  71       log.debug("Storing data and offsets, last commit: {}", lastCommit);
  72       seen.forEach((partiton, statistics) -> repository.save(
  73           new StatisticsDocument(
  74               partiton,
  75               statistics,
  76               consumer.position(new TopicPartition(topic, partiton)))));
  77       lastCommit = clock.instant();
  78     }
  79   }
  80
  81   @Override
  82   public void onPartitionAssigned(TopicPartition tp)
  83   {
  84     Integer partition = tp.partition();
  85     Long offset = consumer.position(tp);
  86     log.info("{} - adding partition: {}, offset={}", id, partition, offset);
  87     StatisticsDocument document =
  88         repository
  89             .findById(Integer.toString(partition))
  90             .orElse(new StatisticsDocument(partition));
  91     if (document.offset >= 0)
  92     {
  93       // Only seek, if a stored offset was found
  94       // Otherwise: Use initial offset, generated by Kafka
  95       consumer.seek(tp, document.offset);
  96     }
  97     seen.put(partition, document.statistics);
  98   }
  99
 100   @Override
 101   public void onPartitionRevoked(TopicPartition tp)
 102   {
 103     Integer partition = tp.partition();
 104     Long newOffset = consumer.position(tp);
 105     log.info(
 106         "{} - removing partition: {}, offset of next message {})",
 107         id,
 108         partition,
 109         newOffset);
 110     Map<String, Map<String, Long>> removed = seen.remove(partition);
 111     repository.save(new StatisticsDocument(partition, removed, consumer.position(tp)));
 112   }
 113
 114
 115   public Map<Integer, Map<String, Map<String, Long>>> getSeen()
 116   {
 117     return seen;
 118   }
 119 }