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.ConsumerRebalanceListener;
   6 import org.apache.kafka.common.TopicPartition;
   7 import org.springframework.kafka.listener.ConsumerAwareRebalanceListener;
   8
   9 import java.util.*;
  10
  11
  12 @RequiredArgsConstructor
  13 @Slf4j
  14 public class ApplicationRebalanceListener implements ConsumerAwareRebalanceListener
  15 {
  16   private final ApplicationRecordHandler recordHandler;
  17   private final AdderResults adderResults;
  18   private final StateRepository stateRepository;
  19   private final String id;
  20
  21   private final Set<Integer> partitions = new HashSet<>();
  22
  23   @Override
  24   public void onPartitionsAssigned(Collection<TopicPartition> partitions)
  25   {
  26     partitions.forEach(tp ->
  27     {
  28       Integer partition = tp.partition();
  29       log.info("{} - adding partition: {}", id, partition);
  30       this.partitions.add(partition);
  31       StateDocument document =
  32           stateRepository
  33               .findById(Integer.toString(partition))
  34               .orElse(new StateDocument(partition));
  35       recordHandler.addPartition(partition, document.state);
  36       for (String user : document.state.keySet())
  37       {
  38         log.info(
  39             "{} - Restored state for partition={}|user={}: {}",
  40             id,
  41             partition,
  42             user,
  43             document.state.get(user));
  44       }
  45       adderResults.addPartition(partition, document.results);
  46     });
  47   }
  48
  49   @Override
  50   public void onPartitionsRevoked(Collection<TopicPartition> partitions)
  51   {
  52     partitions.forEach(tp ->
  53     {
  54       Integer partition = tp.partition();
  55       log.info("{} - removing partition: {}", id, partition);
  56       this.partitions.remove(partition);
  57       Map<String, AdderResult> state = recordHandler.removePartition(partition);
  58       for (String user : state.keySet())
  59       {
  60         log.info(
  61             "{} - Saved state for partition={}|user={}: {}",
  62             id,
  63             partition,
  64             user,
  65             state.get(user));
  66       }
  67       Map<String, List<AdderResult>> results = adderResults.removePartition(partition);
  68       stateRepository.save(new StateDocument(partition, state, results));
  69     });
  70   }
  71 }