Tests: Umbau für einen Commit im Fehlerfall und Anpassung des Tests
[demos/kafka/training] / src / main / java / de / juplo / kafka / EndlessConsumer.java
1 package de.juplo.kafka;
2
3 import lombok.RequiredArgsConstructor;
4 import lombok.extern.slf4j.Slf4j;
5 import org.apache.kafka.clients.consumer.*;
6 import org.apache.kafka.common.TopicPartition;
7 import org.apache.kafka.common.errors.RecordDeserializationException;
8 import org.apache.kafka.common.errors.WakeupException;
9
10 import javax.annotation.PreDestroy;
11 import java.time.Duration;
12 import java.util.*;
13 import java.util.concurrent.ExecutionException;
14 import java.util.concurrent.ExecutorService;
15 import java.util.concurrent.locks.Condition;
16 import java.util.concurrent.locks.Lock;
17 import java.util.concurrent.locks.ReentrantLock;
18
19
20 @Slf4j
21 @RequiredArgsConstructor
22 public class EndlessConsumer<K, V> implements Runnable
23 {
24   private final ExecutorService executor;
25   private final String id;
26   private final String topic;
27   private final Consumer<K, V> consumer;
28   private final java.util.function.Consumer<ConsumerRecord<K, V>> handler;
29
30   private final Lock lock = new ReentrantLock();
31   private final Condition condition = lock.newCondition();
32   private boolean running = false;
33   private Exception exception;
34   private long consumed = 0;
35
36   private final Map<Integer, Map<String, Long>> seen = new HashMap<>();
37   private final Map<Integer, Long> offsets = new HashMap<>();
38
39
40   @Override
41   public void run()
42   {
43     try
44     {
45       log.info("{} - Subscribing to topic {}", id, topic);
46       consumer.subscribe(Arrays.asList(topic), new ConsumerRebalanceListener()
47       {
48         @Override
49         public void onPartitionsRevoked(Collection<TopicPartition> partitions)
50         {
51           partitions.forEach(tp ->
52           {
53             Integer partition = tp.partition();
54             Long newOffset = consumer.position(tp);
55             Long oldOffset = offsets.remove(partition);
56             log.info(
57                 "{} - removing partition: {}, consumed {} records (offset {} -> {})",
58                 id,
59                 partition,
60                 newOffset - oldOffset,
61                 oldOffset,
62                 newOffset);
63             Map<String, Long> removed = seen.remove(partition);
64             for (String key : removed.keySet())
65             {
66               log.info(
67                   "{} - Seen {} messages for partition={}|key={}",
68                   id,
69                   removed.get(key),
70                   partition,
71                   key);
72             }
73           });
74         }
75
76         @Override
77         public void onPartitionsAssigned(Collection<TopicPartition> partitions)
78         {
79           partitions.forEach(tp ->
80           {
81             Integer partition = tp.partition();
82             Long offset = consumer.position(tp);
83             log.info("{} - adding partition: {}, offset={}", id, partition, offset);
84             offsets.put(partition, offset);
85             seen.put(partition, new HashMap<>());
86           });
87         }
88       });
89
90       while (true)
91       {
92         ConsumerRecords<K, V> records =
93             consumer.poll(Duration.ofSeconds(1));
94
95         // Do something with the data...
96         log.info("{} - Received {} messages", id, records.count());
97         for (ConsumerRecord<K, V> record : records)
98         {
99           log.info(
100               "{} - {}: {}/{} - {}={}",
101               id,
102               record.offset(),
103               record.topic(),
104               record.partition(),
105               record.key(),
106               record.value()
107           );
108
109           handler.accept(record);
110
111           consumed++;
112
113           Integer partition = record.partition();
114           String key = record.key() == null ? "NULL" : record.key().toString();
115           Map<String, Long> byKey = seen.get(partition);
116
117           if (!byKey.containsKey(key))
118             byKey.put(key, 0l);
119
120           long seenByKey = byKey.get(key);
121           seenByKey++;
122           byKey.put(key, seenByKey);
123         }
124       }
125     }
126     catch(WakeupException e)
127     {
128       log.info("{} - RIIING! Request to stop consumption - commiting current offsets!", id);
129       consumer.commitSync();
130       shutdown();
131     }
132     catch(RecordDeserializationException e)
133     {
134       TopicPartition tp = e.topicPartition();
135       long offset = e.offset();
136       log.error(
137           "{} - Could not deserialize  message on topic {} with offset={}: {}",
138           id,
139           tp,
140           offset,
141           e.getCause().toString());
142
143       consumer.commitSync();
144       shutdown(e);
145     }
146     catch(Exception e)
147     {
148       log.error("{} - Unexpected error: {}", id, e.toString(), e);
149       shutdown(e);
150     }
151     finally
152     {
153       log.info("{} - Consumer-Thread exiting", id);
154     }
155   }
156
157   private void shutdown()
158   {
159     shutdown(null);
160   }
161
162   private void shutdown(Exception e)
163   {
164     lock.lock();
165     try
166     {
167       try
168       {
169         log.info("{} - Unsubscribing from topic {}", id, topic);
170         consumer.unsubscribe();
171       }
172       catch (Exception ue)
173       {
174         log.error(
175             "{} - Error while unsubscribing from topic {}: {}",
176             id,
177             topic,
178             ue.toString());
179       }
180       finally
181       {
182         running = false;
183         exception = e;
184         condition.signal();
185       }
186     }
187     finally
188     {
189       lock.unlock();
190     }
191   }
192
193   public Map<Integer, Map<String, Long>> getSeen()
194   {
195     return seen;
196   }
197
198   public void start()
199   {
200     lock.lock();
201     try
202     {
203       if (running)
204         throw new IllegalStateException("Consumer instance " + id + " is already running!");
205
206       log.info("{} - Starting - consumed {} messages before", id, consumed);
207       running = true;
208       exception = null;
209       executor.submit(this);
210     }
211     finally
212     {
213       lock.unlock();
214     }
215   }
216
217   public synchronized void stop() throws ExecutionException, InterruptedException
218   {
219     lock.lock();
220     try
221     {
222       if (!running)
223         throw new IllegalStateException("Consumer instance " + id + " is not running!");
224
225       log.info("{} - Stopping", id);
226       consumer.wakeup();
227       condition.await();
228       log.info("{} - Stopped - consumed {} messages so far", id, consumed);
229     }
230     finally
231     {
232       lock.unlock();
233     }
234   }
235
236   @PreDestroy
237   public void destroy() throws ExecutionException, InterruptedException
238   {
239     log.info("{} - Destroy!", id);
240     try
241     {
242       stop();
243     }
244     catch (IllegalStateException e)
245     {
246       log.info("{} - Was already stopped", id);
247     }
248     catch (Exception e)
249     {
250       log.error("{} - Unexpected exception while trying to stop the consumer", id, e);
251     }
252     finally
253     {
254       log.info("{}: Consumed {} messages in total, exiting!", id, consumed);
255     }
256   }
257
258   public boolean running()
259   {
260     lock.lock();
261     try
262     {
263       return running;
264     }
265     finally
266     {
267       lock.unlock();
268     }
269   }
270
271   public Optional<Exception> exitStatus()
272   {
273     lock.lock();
274     try
275     {
276       if (running)
277         throw new IllegalStateException("No exit-status available: Consumer instance " + id + " is running!");
278
279       return Optional.ofNullable(exception);
280     }
281     finally
282     {
283       lock.unlock();
284     }
285   }
286 }