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.util.*;
   9
  10
  11 @RequiredArgsConstructor
  12 @Slf4j
  13 public class ApplicationRebalanceListener implements RebalanceListener
  14 {
  15   private final ApplicationRecordHandler recordHandler;
  16   private final AdderResults adderResults;
  17   private final StateRepository stateRepository;
  18   private final String id;
  19   private final String topic;
  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       if (document.offset >= 0)
  37       {
  38         // Only seek, if a stored offset was found
  39         // Otherwise: Use initial offset, generated by Kafka
  40         consumer.seek(tp, document.offset);
  41         log.info(
  42             "{} - Seeking to offset {} for partition {}",
  43             id,
  44             document.offset,
  45             partition);
  46       }
  47       recordHandler.addPartition(partition, document.state);
  48       for (String user : document.state.keySet())
  49       {
  50         log.info(
  51             "{} - Restored state for partition={}|user={}: {}",
  52             id,
  53             partition,
  54             user,
  55             document.state.get(user));
  56       }
  57       adderResults.addPartition(partition, document.results);
  58     });
  59   }
  60
  61   @Override
  62   public void onPartitionsRevoked(Collection<TopicPartition> partitions)
  63   {
  64     partitions.forEach(tp ->
  65     {
  66       Integer partition = tp.partition();
  67       log.info("{} - removing partition: {}", id, partition);
  68       this.partitions.remove(partition);
  69       Map<String, AdderResult> state = recordHandler.removePartition(partition);
  70       Long offset = consumer.position(tp);
  71       log.info(
  72           "{} - offset of next unseen message for partition {} is {}",
  73           id,
  74           partition,
  75           offset);
  76       for (String user : state.keySet())
  77       {
  78         log.info(
  79             "{} - Saved state for partition={}|user={}: {}",
  80             id,
  81             partition,
  82             user,
  83             state.get(user));
  84       }
  85       Map<String, List<AdderResult>> results = adderResults.removePartition(partition);
  86       stateRepository.save(new StateDocument(partition, state, results, offset));
  87     });
  88   }
  89
  90
  91   @Override
  92   public void beforeNextPoll()
  93   {
  94     partitions
  95       .stream()
  96       .forEach(partition ->
  97       {
  98         log.info("{} - persisting state & offset for partition: {}", id, partition);
  99         Map<String, AdderResult> state = recordHandler.getState(partition).getState();
 100         Long offset = consumer.position(new TopicPartition(topic, partition));
 101         log.info(
 102           "{} - offset of next unseen message for partition {} is {}",
 103           id,
 104           partition,
 105           offset);
 106         for (String user : state.keySet())
 107         {
 108           log.info(
 109             "{} - Saved state for partition={}|user={}: {}",
 110             id,
 111             partition,
 112             user,
 113             state.get(user));
 114         }
 115         Map<String, List<AdderResult>> results = adderResults.getState(partition);
 116         stateRepository.save(new StateDocument(partition, state, results, offset));
 117       });
 118   }
 119 }