juplo.de Git - demos/kafka/training/blob - src/main/java/de/juplo/kafka/ApplicationRebalanceListener.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 ApplicationRebalanceListener implements PollIntervalAwareConsumerRebalanceListener
  18 {
  19   private final ApplicationRecordHandler recordHandler;
  20   private final StateRepository stateRepository;
  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   private boolean commitsEnabled = true;
  29
  30   @Override
  31   public void onPartitionsAssigned(Collection<TopicPartition> partitions)
  32   {
  33     partitions.forEach(tp ->
  34     {
  35       Integer partition = tp.partition();
  36       StateDocument document =
  37           stateRepository
  38               .findById(Integer.toString(partition))
  39               .orElse(new StateDocument(partition));
  40       log.info("{} - adding partition: {}, offset={}", id, partition, document.offset);
  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       recordHandler.addPartition(partition, document.state);
  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 offset = consumer.position(tp);
  58       log.info(
  59           "{} - removing partition: {}, offset of next message {})",
  60           id,
  61           partition,
  62           offset);
  63       if (commitsEnabled)
  64       {
  65         Map<String, Long> removed = recordHandler.removePartition(partition);
  66         stateRepository.save(new StateDocument(partition, removed, offset));
  67       }
  68       else
  69       {
  70         log.info("Offset commits are disabled! Last commit: {}", lastCommit);
  71       }
  72     });
  73   }
  74
  75
  76   @Override
  77   public void beforeNextPoll()
  78   {
  79     if (!commitsEnabled)
  80     {
  81       log.info("Offset commits are disabled! Last commit: {}", lastCommit);
  82       return;
  83     }
  84
  85     if (lastCommit.plus(commitInterval).isBefore(clock.instant()))
  86     {
  87       log.debug("Storing data and offsets, last commit: {}", lastCommit);
  88       recordHandler.getState().forEach((partiton, adder) -> stateRepository.save(
  89           new StateDocument(
  90               partiton,
  91               adder.getState(),
  92               consumer.position(new TopicPartition(topic, partiton)))));
  93       lastCommit = clock.instant();
  94     }
  95   }
  96
  97   @Override
  98   public void enableCommits()
  99   {
 100     commitsEnabled = true;
 101   }
 102
 103   @Override
 104   public void disableCommits()
 105   {
 106     commitsEnabled = false;
 107   }
 108 }