2fccb4f28209ff6ef4c268bcb7c9eb42cd444f08
[demos/kafka/training] / 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.Collection;
12 import java.util.Map;
13
14
15 @RequiredArgsConstructor
16 @Slf4j
17 public class ApplicationRebalanceListener implements PollIntervalAwareConsumerRebalanceListener
18 {
19   private final ApplicationRecordHandler recordHandler;
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, Long> consumer;
26
27   private Instant lastCommit = Instant.EPOCH;
28
29   @Override
30   public void onPartitionsAssigned(Collection<TopicPartition> partitions)
31   {
32     partitions.forEach(tp ->
33     {
34       Integer partition = tp.partition();
35       Long offset = consumer.position(tp);
36       log.info("{} - adding partition: {}, offset={}", id, partition, offset);
37       StateDocument document =
38           stateRepository
39               .findById(Integer.toString(partition))
40               .orElse(new StateDocument(partition));
41       if (document.offset >= 0)
42       {
43         // Only seek, if a stored offset was found
44         // Otherwise: Use initial offset, generated by Kafka
45         consumer.seek(tp, document.offset);
46       }
47       recordHandler.addPartition(partition, document.state);
48     });
49   }
50
51   @Override
52   public void onPartitionsRevoked(Collection<TopicPartition> partitions)
53   {
54     partitions.forEach(tp ->
55     {
56       Integer partition = tp.partition();
57       Long newOffset = consumer.position(tp);
58       log.info(
59           "{} - removing partition: {}, offset of next message {})",
60           id,
61           partition,
62           newOffset);
63       Map<String, Long> removed = recordHandler.removePartition(partition);
64       stateRepository.save(new StateDocument(partition, removed, consumer.position(tp)));
65     });
66   }
67
68
69   @Override
70   public void beforeNextPoll()
71   {
72     if (lastCommit.plus(commitInterval).isBefore(clock.instant()))
73     {
74       log.debug("Storing data and offsets, last commit: {}", lastCommit);
75       recordHandler.getState().forEach((partiton, state) -> stateRepository.save(
76           new StateDocument(
77               partiton,
78               state,
79               consumer.position(new TopicPartition(topic, partiton)))));
80       lastCommit = clock.instant();
81     }
82   }
83 }