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