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, AbstractTo> producer;
  29   private final Consumer<String, AbstractTo> 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, AbstractTo> producer,
  46     Consumer<String, AbstractTo> 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     IntStream
  64         .range(0, numShards)
  65         .forEach(shard -> this.chatrooms[shard] = new HashMap<>());
  66     this.shardingStrategy = new KafkaLikeShardingStrategy(numShards);
  67   }
  68
  69
  70   Mono<Message> sendMessage(
  71       UUID chatRoomId,
  72       Message.MessageKey key,
  73       LocalDateTime timestamp,
  74       String text)
  75   {
  76     int shard = this.shardingStrategy.selectShard(chatRoomId);
  77     TopicPartition tp = new TopicPartition(topic, shard);
  78     ZonedDateTime zdt = ZonedDateTime.of(timestamp, zoneId);
  79     return Mono.create(sink ->
  80     {
  81       ProducerRecord<String, ChatMessageTo> record =
  82           new ProducerRecord<>(
  83               tp.topic(),
  84               tp.partition(),
  85               zdt.toEpochSecond(),
  86               chatRoomId.toString(),
  87               ChatMessageTo.of(key.getUsername(), key.getMessageId(), text));
  88
  89       producer.send(record, ((metadata, exception) ->
  90       {
  91         if (metadata != null)
  92         {
  93           // On successful send
  94           Message message = new Message(key, metadata.offset(), timestamp, text);
  95           log.info("Successfully send message {}", message);
  96           sink.success(message);
  97         }
  98         else
  99         {
 100           // On send-failure
 101           log.error(
 102               "Could not send message for chat-room={}, key={}, timestamp={}, text={}: {}",
 103               chatRoomId,
 104               key,
 105               timestamp,
 106               text,
 107               exception);
 108           sink.error(exception);
 109         }
 110       }));
 111     });
 112   }
 113
 114   @Override
 115   public void onPartitionsAssigned(Collection<TopicPartition> partitions)
 116   {
 117     log.info("Newly assigned partitions! Pausing normal operations...");
 118     loadInProgress = true;
 119
 120     consumer.endOffsets(partitions).forEach((topicPartition, currentOffset) ->
 121     {
 122       int partition = topicPartition.partition();
 123       isShardOwned[partition] =  true;
 124       this.currentOffset[partition] = currentOffset;
 125
 126       log.info(
 127           "Partition assigned: {} - loading messages: next={} -> current={}",
 128           partition,
 129           nextOffset[partition],
 130           currentOffset);
 131
 132       consumer.seek(topicPartition, nextOffset[partition]);
 133     });
 134
 135     consumer.resume(partitions);
 136   }
 137
 138   @Override
 139   public void onPartitionsRevoked(Collection<TopicPartition> partitions)
 140   {
 141     partitions.forEach(topicPartition ->
 142     {
 143       int partition = topicPartition.partition();
 144       isShardOwned[partition] = false;
 145       log.info("Partition revoked: {} - next={}", partition, nextOffset[partition]);
 146     });
 147   }
 148
 149   @Override
 150   public void onPartitionsLost(Collection<TopicPartition> partitions)
 151   {
 152     log.warn("Lost partitions: {}, partitions");
 153     // TODO: Muss auf den Verlust anders reagiert werden?
 154     onPartitionsRevoked(partitions);
 155   }
 156
 157   @Override
 158   public void run()
 159   {
 160     consumer.subscribe(List.of(topic), this);
 161
 162     running = true;
 163
 164     while (running)
 165     {
 166       try
 167       {
 168         ConsumerRecords<String, AbstractTo> records = consumer.poll(Duration.ofMinutes(5));
 169         log.info("Fetched {} messages", records.count());
 170
 171         if (loadInProgress)
 172         {
 173           loadChatRoom(records);
 174
 175           if (isLoadingCompleted())
 176           {
 177             log.info("Loading of messages completed! Pausing all owned partitions...");
 178             pauseAllOwnedPartions();
 179             log.info("Resuming normal operations...");
 180             loadInProgress = false;
 181           }
 182         }
 183         else
 184         {
 185           if (!records.isEmpty())
 186           {
 187             throw new IllegalStateException("All owned partitions should be paused, when no load is in progress!");
 188           }
 189         }
 190       }
 191       catch (WakeupException e)
 192       {
 193         log.info("Received WakeupException, exiting!");
 194         running = false;
 195       }
 196     }
 197
 198     log.info("Exiting normally");
 199   }
 200
 201   void loadChatRoom(ConsumerRecords<String, AbstractTo> records)
 202   {
 203     for (ConsumerRecord<String, AbstractTo> record : records)
 204     {
 205       switch (record.value().getType())
 206       {
 207         case CREATE_CHATROOM_REQUEST:
 208           createChatRoom((CreateChatRoomRequestTo) record.value());
 209           break;
 210
 211         case MESSAGE_SENT:
 212           UUID chatRoomId = UUID.fromString(record.key());
 213           Instant instant = Instant.ofEpochSecond(record.timestamp());
 214           LocalDateTime timestamp = LocalDateTime.ofInstant(instant, zoneId);
 215           loadChatMessage(
 216               chatRoomId,
 217               timestamp,
 218               record.offset(),
 219               (ChatMessageTo) record.value(),
 220               record.partition());
 221           break;
 222       }
 223
 224       nextOffset[record.partition()] = record.offset() + 1;
 225     }
 226   }
 227
 228   void createChatRoom(
 229       CreateChatRoomRequestTo createChatRoomRequestTo,
 230       int partition)
 231   {
 232     chatrooms[partition].put
 233   }
 234
 235   void loadChatMessage(
 236       UUID chatRoomId,
 237       LocalDateTime timestamp,
 238       long offset,
 239       ChatMessageTo chatMessageTo,
 240       int partition)
 241   {
 242     Message.MessageKey key = Message.MessageKey.of(chatMessageTo.getUser(), chatMessageTo.getId());
 243     Message message = new Message(key, offset, timestamp, chatMessageTo.getText());
 244
 245     ChatRoom chatRoom = chatrooms[partition].get(chatRoomId);
 246     KafkaChatRoomService kafkaChatRoomService =
 247         (KafkaChatRoomService) chatRoom.getChatRoomService();
 248
 249     kafkaChatRoomService.persistMessage(message);
 250   }
 251
 252   boolean isLoadingCompleted()
 253   {
 254     return IntStream
 255         .range(0, numShards)
 256         .filter(shard -> isShardOwned[shard])
 257         .allMatch(shard -> nextOffset[shard] >= currentOffset[shard]);
 258   }
 259
 260   void pauseAllOwnedPartions()
 261   {
 262     consumer.pause(IntStream
 263         .range(0, numShards)
 264         .filter(shard -> isShardOwned[shard])
 265         .mapToObj(shard -> new TopicPartition(topic, shard))
 266         .toList());
 267   }
 268
 269
 270   void putChatRoom(ChatRoom chatRoom)
 271   {
 272     Integer partition = chatRoom.getShard();
 273     UUID chatRoomId = chatRoom.getId();
 274     if (chatrooms[partition].containsKey(chatRoomId))
 275     {
 276       log.warn("Ignoring existing chat-room: " + chatRoom);
 277     }
 278     else
 279     {
 280       log.info(
 281           "Adding new chat-room to partition {}: {}",
 282           partition,
 283           chatRoom);
 284
 285       chatrooms[partition].put(chatRoomId, chatRoom);
 286     }
 287   }
 288
 289   Mono<ChatRoom> getChatRoom(int shard, UUID id)
 290   {
 291     return Mono.justOrEmpty(chatrooms[shard].get(id));
 292   }
 293
 294   Flux<ChatRoom> getChatRooms()
 295   {
 296     return Flux.fromStream(IntStream
 297         .range(0, numShards)
 298         .filter(shard -> isShardOwned[shard])
 299         .mapToObj(shard -> Integer.valueOf(shard))
 300         .flatMap(shard -> chatrooms[shard].values().stream()));
 301   }
 302 }