WIP
[demos/kafka/training] / src / main / java / de / juplo / kafka / SumRebalanceListener.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
13
14 @RequiredArgsConstructor
15 @Slf4j
16 public class SumRebalanceListener implements PollIntervalAwareConsumerRebalanceListener
17 {
18   private final SumRecordHandler handler;
19   private final PartitionStatisticsRepository repository;
20   private final String id;
21   private final String topic;
22   private final Clock clock;
23   private final Duration commitInterval;
24   private final Consumer<String, String> consumer;
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       Long offset = consumer.position(tp);
35       log.info("{} - adding partition: {}, offset={}", id, partition, offset);
36       StateDocument document =
37           repository
38               .findById(Integer.toString(partition))
39               .orElse(new StateDocument(partition));
40       if (document.offset >= 0)
41       {
42         // Only seek, if a stored offset was found
43         // Otherwise: Use initial offset, generated by Kafka
44         consumer.seek(tp, document.offset);
45       }
46       handler.addPartition(partition, document.state);
47     });
48   }
49
50   @Override
51   public void onPartitionsRevoked(Collection<TopicPartition> partitions)
52   {
53     partitions.forEach(tp ->
54     {
55       Integer partition = tp.partition();
56       Long newOffset = consumer.position(tp);
57       log.info(
58           "{} - removing partition: {}, offset of next message {})",
59           id,
60           partition,
61           newOffset);
62       repository.save(new StateDocument(partition, handler.removePartition(partition), newOffset));
63     });
64   }
65
66
67   @Override
68   public void beforeNextPoll()
69   {
70     if (lastCommit.plus(commitInterval).isBefore(clock.instant()))
71     {
72       log.debug("Storing data and offsets, last commit: {}", lastCommit);
73       handler.getState().forEach((partiton, sumBusinessLogic) -> repository.save(
74           new StateDocument(
75               partiton,
76               sumBusinessLogic.getState(),
77               consumer.position(new TopicPartition(topic, partiton)))));
78       lastCommit = clock.instant();
79     }
80   }
81 }