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