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.common.TopicPartition;
   6
   7 import java.time.Clock;
   8 import java.time.Duration;
   9 import java.time.Instant;
  10 import java.util.*;
  11
  12
  13 @RequiredArgsConstructor
  14 @Slf4j
  15 public class ApplicationRebalanceListener implements PollIntervalAwareConsumerRebalanceListener
  16 {
  17   private final ApplicationRecordHandler recordHandler;
  18   private final AdderResults adderResults;
  19   private final StateRepository stateRepository;
  20   private final String id;
  21   private final Clock clock;
  22   private final Duration commitInterval;
  23
  24   private final Set<Integer> partitions = new HashSet<>();
  25
  26   private Instant lastCommit = Instant.EPOCH;
  27
  28   @Override
  29   public void onPartitionsAssigned(Collection<TopicPartition> partitions)
  30   {
  31     partitions.forEach(tp ->
  32     {
  33       Integer partition = tp.partition();
  34       log.info("{} - adding partition: {}", id, partition);
  35       this.partitions.add(partition);
  36       StateDocument document =
  37           stateRepository
  38               .findById(Integer.toString(partition))
  39               .orElse(new StateDocument(partition));
  40       recordHandler.addPartition(partition, document.state);
  41       for (String user : document.state.keySet())
  42       {
  43         log.info(
  44             "{} - Restored state for partition={}|user={}: {}",
  45             id,
  46             partition,
  47             user,
  48             document.state.get(user));
  49       }
  50       adderResults.addPartition(partition, document.results);
  51     });
  52   }
  53
  54   @Override
  55   public void onPartitionsRevoked(Collection<TopicPartition> partitions)
  56   {
  57     partitions.forEach(tp ->
  58     {
  59       Integer partition = tp.partition();
  60       log.info("{} - removing partition: {}", id, partition);
  61       this.partitions.remove(partition);
  62       Map<String, AdderResult> state = recordHandler.removePartition(partition);
  63       for (String user : state.keySet())
  64       {
  65         log.info(
  66             "{} - Saved state for partition={}|user={}: {}",
  67             id,
  68             partition,
  69             user,
  70             state.get(user));
  71       }
  72       Map<String, List<AdderResult>> results = adderResults.removePartition(partition);
  73       stateRepository.save(new StateDocument(partition, state, results));
  74     });
  75   }
  76
  77
  78   @Override
  79   public void beforeNextPoll()
  80   {
  81     if (lastCommit.plus(commitInterval).isBefore(clock.instant()))
  82     {
  83       log.debug("Storing data, last commit: {}", lastCommit);
  84       partitions.forEach(partition -> stateRepository.save(
  85           new StateDocument(
  86               partition,
  87               recordHandler.getState(partition).getState(),
  88               adderResults.getState(partition))));
  89       lastCommit = clock.instant();
  90     }
  91   }
  92 }