juplo.de Git - demos/kafka/training/blob - src/main/java/de/juplo/kafka/WordcountRebalanceListener.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.common.TopicPartition;
   7
   8 import java.time.Clock;
   9 import java.time.Duration;
  10 import java.time.Instant;
  11 import java.util.Collection;
  12 import java.util.Map;
  13
  14
  15 @RequiredArgsConstructor
  16 @Slf4j
  17 public class WordcountRebalanceListener implements PollIntervalAwareConsumerRebalanceListener
  18 {
  19   private final WordcountRecordHandler handler;
  20   private final PartitionStatisticsRepository repository;
  21   private final String id;
  22   private final String topic;
  23   private final Clock clock;
  24   private final Duration commitInterval;
  25   private final Consumer<String, String> consumer;
  26
  27   private Instant lastCommit = Instant.EPOCH;
  28
  29   @Override
  30   public void onPartitionsAssigned(Collection<TopicPartition> partitions)
  31   {
  32     partitions.forEach(tp ->
  33     {
  34       Integer partition = tp.partition();
  35       Long offset = consumer.position(tp);
  36       log.info("{} - adding partition: {}, offset={}", id, partition, offset);
  37       StatisticsDocument document =
  38           repository
  39               .findById(Integer.toString(partition))
  40               .orElse(new StatisticsDocument(partition));
  41       if (document.offset >= 0)
  42       {
  43         // Only seek, if a stored offset was found
  44         // Otherwise: Use initial offset, generated by Kafka
  45         consumer.seek(tp, document.offset);
  46       }
  47       handler.addPartition(partition, document.statistics);
  48     });
  49   }
  50
  51   @Override
  52   public void onPartitionsRevoked(Collection<TopicPartition> partitions)
  53   {
  54     partitions.forEach(tp ->
  55     {
  56       Integer partition = tp.partition();
  57       Long newOffset = consumer.position(tp);
  58       log.info(
  59           "{} - removing partition: {}, offset of next message {})",
  60           id,
  61           partition,
  62           newOffset);
  63       Map<String, Map<String, Long>> removed = handler.removePartition(partition);
  64       repository.save(new StatisticsDocument(partition, removed, consumer.position(tp)));
  65     });
  66   }
  67
  68
  69   @Override
  70   public void beforeNextPoll()
  71   {
  72     if (lastCommit.plus(commitInterval).isBefore(clock.instant()))
  73     {
  74       log.debug("Storing data and offsets, last commit: {}", lastCommit);
  75       handler.getSeen().forEach((partiton, statistics) -> repository.save(
  76           new StatisticsDocument(
  77               partiton,
  78               statistics,
  79               consumer.position(new TopicPartition(topic, partiton)))));
  80       lastCommit = clock.instant();
  81     }
  82   }
  83 }