juplo.de Git - demos/kafka/chat/blob - src/main/java/de/juplo/kafka/chat/backend/persistence/kafka/ChatMessageChannel.java

   1 package de.juplo.kafka.chat.backend.persistence.kafka;
   2
   3 import de.juplo.kafka.chat.backend.domain.ChatRoom;
   4 import de.juplo.kafka.chat.backend.domain.Message;
   5 import de.juplo.kafka.chat.backend.persistence.KafkaLikeShardingStrategy;
   6 import lombok.Getter;
   7 import lombok.extern.slf4j.Slf4j;
   8 import org.apache.kafka.clients.consumer.Consumer;
   9 import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
  10 import org.apache.kafka.clients.consumer.ConsumerRecord;
  11 import org.apache.kafka.clients.consumer.ConsumerRecords;
  12 import org.apache.kafka.clients.producer.Producer;
  13 import org.apache.kafka.clients.producer.ProducerRecord;
  14 import org.apache.kafka.common.TopicPartition;
  15 import org.apache.kafka.common.errors.RecordDeserializationException;
  16 import org.apache.kafka.common.errors.WakeupException;
  17 import reactor.core.publisher.Flux;
  18 import reactor.core.publisher.Mono;
  19
  20 import java.time.*;
  21 import java.util.Collection;
  22 import java.util.List;
  23 import java.util.Map;
  24 import java.util.UUID;
  25 import java.util.stream.IntStream;
  26
  27
  28 @Slf4j
  29 public class ChatMessageChannel implements Runnable, ConsumerRebalanceListener
  30 {
  31   private final String topic;
  32   private final Consumer<String, MessageTo> consumer;
  33   private final Producer<String, MessageTo> producer;
  34   private final ZoneId zoneId;
  35   private final int numShards;
  36   private final boolean[] isShardOwned;
  37   private final long[] currentOffset;
  38   private final long[] nextOffset;
  39   private final Map<UUID, ChatRoom>[] chatrooms;
  40   private final KafkaLikeShardingStrategy shardingStrategy;
  41
  42   private boolean running;
  43   @Getter
  44   private volatile boolean loadInProgress;
  45
  46
  47   public ChatMessageChannel(
  48     String topic,
  49     Consumer<String, MessageTo> consumer,
  50     Producer<String, MessageTo> producer,
  51     ZoneId zoneId,
  52     int numShards)
  53   {
  54     log.debug(
  55         "Creating ChatMessageChannel for topic {} with {} partitions",
  56         topic,
  57         numShards);
  58     this.topic = topic;
  59     this.consumer = consumer;
  60     this.producer = producer;
  61     this.zoneId = zoneId;
  62     this.numShards = numShards;
  63     this.isShardOwned = new boolean[numShards];
  64     this.currentOffset = new long[numShards];
  65     this.nextOffset = new long[numShards];
  66     this.chatrooms = new Map[numShards];
  67     this.shardingStrategy = new KafkaLikeShardingStrategy(numShards);
  68   }
  69
  70
  71   Mono<Message> sendMessage(
  72       UUID chatRoomId,
  73       Message.MessageKey key,
  74       LocalDateTime timestamp,
  75       String text)
  76   {
  77     int shard = this.shardingStrategy.selectShard(chatRoomId);
  78     TopicPartition tp = new TopicPartition(topic, shard);
  79     ZonedDateTime zdt = ZonedDateTime.of(timestamp, zoneId);
  80     return Mono.create(sink ->
  81     {
  82       ProducerRecord<String, MessageTo> record =
  83           new ProducerRecord<>(
  84               tp.topic(),
  85               tp.partition(),
  86               zdt.toEpochSecond(),
  87               chatRoomId.toString(),
  88               MessageTo.of(key.getUsername(), key.getMessageId(), text));
  89
  90       producer.send(record, ((metadata, exception) ->
  91       {
  92         if (metadata != null)
  93         {
  94           // On successful send
  95           Message message = new Message(key, metadata.offset(), timestamp, text);
  96           log.info("Successfully send message {}", message);
  97           sink.success(message);
  98         }
  99         else
 100         {
 101           // On send-failure
 102           log.error(
 103               "Could not send message for chat-room={}, key={}, timestamp={}, text={}: {}",
 104               chatRoomId,
 105               key,
 106               timestamp,
 107               text,
 108               exception);
 109           sink.error(exception);
 110         }
 111       }));
 112     });
 113   }
 114
 115   @Override
 116   public void onPartitionsAssigned(Collection<TopicPartition> partitions)
 117   {
 118     log.info("Newly assigned partitions! Pausing normal operations...");
 119     loadInProgress = true;
 120
 121     consumer.endOffsets(partitions).forEach((topicPartition, currentOffset) ->
 122     {
 123       int partition = topicPartition.partition();
 124       isShardOwned[partition] =  true;
 125       this.currentOffset[partition] = currentOffset;
 126
 127       log.info(
 128           "Partition assigned: {} - loading messages: next={} -> current={}",
 129           partition,
 130           nextOffset[partition],
 131           currentOffset);
 132
 133       consumer.seek(topicPartition, nextOffset[partition]);
 134     });
 135
 136     consumer.resume(partitions);
 137   }
 138
 139   @Override
 140   public void onPartitionsRevoked(Collection<TopicPartition> partitions)
 141   {
 142     partitions.forEach(topicPartition ->
 143     {
 144       int partition = topicPartition.partition();
 145       isShardOwned[partition] = false;
 146       log.info("Partition revoked: {} - next={}", partition, nextOffset[partition]);
 147     });
 148   }
 149
 150   @Override
 151   public void onPartitionsLost(Collection<TopicPartition> partitions)
 152   {
 153     log.warn("Lost partitions: {}, partitions");
 154     // TODO: Muss auf den Verlust anders reagiert werden?
 155     onPartitionsRevoked(partitions);
 156   }
 157
 158   @Override
 159   public void run()
 160   {
 161     consumer.subscribe(List.of(topic));
 162
 163     running = true;
 164
 165     while (running)
 166     {
 167       try
 168       {
 169         ConsumerRecords<String, MessageTo> records = consumer.poll(Duration.ofMinutes(5));
 170         log.info("Fetched {} messages", records.count());
 171
 172         if (loadInProgress)
 173         {
 174           loadMessages(records);
 175
 176           if (isLoadingCompleted())
 177           {
 178             log.info("Loading of messages completed! Pausing all owned partitions...");
 179             pauseAllOwnedPartions();
 180             log.info("Resuming normal operations...");
 181             loadInProgress = false;
 182           }
 183         }
 184         else
 185         {
 186           if (!records.isEmpty())
 187           {
 188             throw new IllegalStateException("All owned partitions should be paused, when no load is in progress!");
 189           }
 190         }
 191       }
 192       catch (WakeupException e)
 193       {
 194         log.info("Received WakeupException, exiting!");
 195         running = false;
 196       }
 197     }
 198   }
 199
 200   void loadMessages(ConsumerRecords<String, MessageTo> records)
 201   {
 202     for (ConsumerRecord<String, MessageTo> record : records)
 203     {
 204       nextOffset[record.partition()] = record.offset() + 1;
 205       UUID chatRoomId = UUID.fromString(record.key());
 206       MessageTo messageTo = record.value();
 207
 208       Message.MessageKey key = Message.MessageKey.of(messageTo.getUser(), messageTo.getId());
 209
 210       Instant instant = Instant.ofEpochSecond(record.timestamp());
 211       LocalDateTime timestamp = LocalDateTime.ofInstant(instant, zoneId);
 212
 213       Message message = new Message(key, record.offset(), timestamp, messageTo.getText());
 214
 215       ChatRoom chatRoom = chatrooms[record.partition()].get(chatRoomId);
 216       if (chatRoom == null)
 217       {
 218         // Alles pausieren und erst von putChatRoom wieder resumen lassen!
 219       }
 220       KafkaChatRoomService kafkaChatRoomService =
 221           (KafkaChatRoomService) chatRoom.getChatRoomService();
 222
 223       kafkaChatRoomService.persistMessage(message);
 224     }
 225   }
 226
 227   boolean isLoadingCompleted()
 228   {
 229     return IntStream
 230         .range(0, numShards)
 231         .filter(shard -> isShardOwned[shard])
 232         .mapToObj(shard -> nextOffset[shard] >= currentOffset[shard])
 233         .collect(
 234             () -> Boolean.TRUE, // TODO: Boolean is immutable
 235             (acc, v) -> Boolean.valueOf(acc && v), // TODO: Boolean is immutable
 236             (a, b) -> Boolean.valueOf(a && b)); // TODO: Boolean is immutable
 237   }
 238
 239   void pauseAllOwnedPartions()
 240   {
 241     consumer.pause(IntStream
 242         .range(0, numShards)
 243         .filter(shard -> isShardOwned[shard])
 244         .mapToObj(shard -> new TopicPartition(topic, shard))
 245         .toList());
 246   }
 247
 248
 249   void putChatRoom(ChatRoom chatRoom)
 250   {
 251     Integer partition = chatRoom.getShard();
 252     UUID chatRoomId = chatRoom.getId();
 253     if (chatrooms[partition].containsKey(chatRoomId))
 254     {
 255       log.warn("Ignoring existing chat-room: " + chatRoom);
 256     }
 257     else
 258     {
 259       log.info(
 260           "Adding new chat-room to partition {}: {}",
 261           partition,
 262           chatRoom);
 263
 264       chatrooms[partition].put(chatRoomId, chatRoom);
 265     }
 266   }
 267
 268   Mono<ChatRoom> getChatRoom(int shard, UUID id)
 269   {
 270     return Mono.justOrEmpty(chatrooms[shard].get(id));
 271   }
 272
 273   Flux<ChatRoom> getChatRooms(int shard)
 274   {
 275     return Flux.fromStream(chatrooms[shard].values().stream());
 276   }
 277 }