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.*;
  12
  13
  14 @RequiredArgsConstructor
  15 @Slf4j
  16 public class ApplicationRebalanceListener implements PollIntervalAwareConsumerRebalanceListener
  17 {
  18   private final ApplicationRecordHandler recordHandler;
  19   private final AdderResults adderResults;
  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 final Set<Integer> partitions = new HashSet<>();
  28
  29   private Instant lastCommit = Instant.EPOCH;
  30   private boolean commitsEnabled = true;
  31
  32   @Override
  33   public void onPartitionsAssigned(Collection<TopicPartition> partitions)
  34   {
  35     partitions.forEach(tp ->
  36     {
  37       Integer partition = tp.partition();
  38       this.partitions.add(partition);
  39       StateDocument document =
  40           stateRepository
  41               .findById(Integer.toString(partition))
  42               .orElse(new StateDocument(partition));
  43       log.info("{} - adding partition: {}, offset={}", id, partition, document.offset);
  44       if (document.offset >= 0)
  45       {
  46         // Only seek, if a stored offset was found
  47         // Otherwise: Use initial offset, generated by Kafka
  48         consumer.seek(tp, document.offset);
  49       }
  50       recordHandler.addPartition(partition, document.state);
  51       adderResults.addPartition(partition, document.results);
  52     });
  53   }
  54
  55   @Override
  56   public void onPartitionsRevoked(Collection<TopicPartition> partitions)
  57   {
  58     partitions.forEach(tp ->
  59     {
  60       Integer partition = tp.partition();
  61       this.partitions.remove(partition);
  62       Long offset = consumer.position(tp);
  63       log.info(
  64           "{} - removing partition: {}, offset of next message {})",
  65           id,
  66           partition,
  67           offset);
  68       if (commitsEnabled)
  69       {
  70         Map<String, AdderResult> state = recordHandler.removePartition(partition);
  71         Map<String, List<AdderResult>> results = adderResults.removePartition(partition);
  72         stateRepository.save(new StateDocument(partition, state, results, offset));
  73       }
  74       else
  75       {
  76         log.info("Offset commits are disabled! Last commit: {}", lastCommit);
  77       }
  78     });
  79   }
  80
  81
  82   @Override
  83   public void beforeNextPoll()
  84   {
  85     if (!commitsEnabled)
  86     {
  87       log.info("Offset commits are disabled! Last commit: {}", lastCommit);
  88       return;
  89     }
  90
  91     if (lastCommit.plus(commitInterval).isBefore(clock.instant()))
  92     {
  93       log.debug("Storing data and offsets, last commit: {}", lastCommit);
  94       partitions.forEach(partition -> stateRepository.save(
  95           new StateDocument(
  96               partition,
  97               recordHandler.getState(partition).getState(),
  98               adderResults.getState(partition),
  99               consumer.position(new TopicPartition(topic, partition)))));
 100       lastCommit = clock.instant();
 101     }
 102   }
 103
 104   @Override
 105   public void enableCommits()
 106   {
 107     commitsEnabled = true;
 108   }
 109
 110   @Override
 111   public void disableCommits()
 112   {
 113     commitsEnabled = false;
 114   }
 115 }