juplo.de Git - demos/kafka/chat/blob - src/main/java/de/juplo/kafka/chat/backend/implementation/kafka/DataChannel.java

   1 package de.juplo.kafka.chat.backend.implementation.kafka;
   2
   3 import de.juplo.kafka.chat.backend.domain.*;
   4 import de.juplo.kafka.chat.backend.domain.exceptions.LoadInProgressException;
   5 import de.juplo.kafka.chat.backend.domain.exceptions.ShardNotOwnedException;
   6 import de.juplo.kafka.chat.backend.implementation.kafka.messages.AbstractMessageTo;
   7 import de.juplo.kafka.chat.backend.implementation.kafka.messages.data.EventChatMessageReceivedTo;
   8 import lombok.Getter;
   9 import lombok.extern.slf4j.Slf4j;
  10 import org.apache.kafka.clients.consumer.Consumer;
  11 import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
  12 import org.apache.kafka.clients.consumer.ConsumerRecord;
  13 import org.apache.kafka.clients.consumer.ConsumerRecords;
  14 import org.apache.kafka.clients.producer.Producer;
  15 import org.apache.kafka.clients.producer.ProducerRecord;
  16 import org.apache.kafka.common.TopicPartition;
  17 import org.apache.kafka.common.errors.WakeupException;
  18 import reactor.core.publisher.Mono;
  19
  20 import java.time.*;
  21 import java.util.*;
  22 import java.util.function.Function;
  23 import java.util.stream.IntStream;
  24
  25
  26 @Slf4j
  27 public class DataChannel implements Runnable, ConsumerRebalanceListener
  28 {
  29   private final String topic;
  30   private final Producer<String, AbstractMessageTo> producer;
  31   private final Consumer<String, AbstractMessageTo> consumer;
  32   private final ZoneId zoneId;
  33   private final int numShards;
  34   private final int bufferSize;
  35   private final Clock clock;
  36   private final boolean[] isShardOwned;
  37   private final long[] currentOffset;
  38   private final long[] nextOffset;
  39   private final Map<UUID, ChatRoomData>[] chatRoomData;
  40   private final InfoChannel infoChannel;
  41
  42   private boolean running;
  43   @Getter
  44   private volatile boolean loadInProgress;
  45
  46
  47   public DataChannel(
  48     String topic,
  49     Producer<String, AbstractMessageTo> producer,
  50     Consumer<String, AbstractMessageTo> dataChannelConsumer,
  51     ZoneId zoneId,
  52     int numShards,
  53     int bufferSize,
  54     Clock clock,
  55     InfoChannel infoChannel)
  56   {
  57     log.debug(
  58         "Creating DataChannel for topic {} with {} partitions",
  59         topic,
  60         numShards);
  61     this.topic = topic;
  62     this.consumer = dataChannelConsumer;
  63     this.producer = producer;
  64     this.zoneId = zoneId;
  65     this.numShards = numShards;
  66     this.bufferSize = bufferSize;
  67     this.clock = clock;
  68     this.isShardOwned = new boolean[numShards];
  69     this.currentOffset = new long[numShards];
  70     this.nextOffset = new long[numShards];
  71     this.chatRoomData = new Map[numShards];
  72     IntStream
  73         .range(0, numShards)
  74         .forEach(shard -> this.chatRoomData[shard] = new HashMap<>());
  75     this.infoChannel = infoChannel;
  76   }
  77
  78
  79
  80   Mono<Message> sendChatMessage(
  81       UUID chatRoomId,
  82       Message.MessageKey key,
  83       LocalDateTime timestamp,
  84       String text)
  85   {
  86     ZonedDateTime zdt = ZonedDateTime.of(timestamp, zoneId);
  87     return Mono.create(sink ->
  88     {
  89       ProducerRecord<String, AbstractMessageTo> record =
  90           new ProducerRecord<>(
  91               topic,
  92               null,
  93               zdt.toEpochSecond(),
  94               chatRoomId.toString(),
  95               EventChatMessageReceivedTo.of(key.getUsername(), key.getMessageId(), text));
  96
  97       producer.send(record, ((metadata, exception) ->
  98       {
  99         if (metadata != null)
 100         {
 101           // On successful send
 102           Message message = new Message(key, metadata.offset(), timestamp, text);
 103           log.info("Successfully send message {}", message);
 104           sink.success(message);
 105         }
 106         else
 107         {
 108           // On send-failure
 109           log.error(
 110               "Could not send message for chat-room={}, key={}, timestamp={}, text={}: {}",
 111               chatRoomId,
 112               key,
 113               timestamp,
 114               text,
 115               exception);
 116           sink.error(exception);
 117         }
 118       }));
 119     });
 120   }
 121
 122   @Override
 123   public void onPartitionsAssigned(Collection<TopicPartition> partitions)
 124   {
 125     log.info("Newly assigned partitions! Pausing normal operations...");
 126     loadInProgress = true;
 127
 128     consumer.endOffsets(partitions).forEach((topicPartition, currentOffset) ->
 129     {
 130       int partition = topicPartition.partition();
 131       isShardOwned[partition] =  true;
 132       this.currentOffset[partition] = currentOffset;
 133
 134       log.info(
 135           "Partition assigned: {} - loading messages: next={} -> current={}",
 136           partition,
 137           nextOffset[partition],
 138           currentOffset);
 139
 140       consumer.seek(topicPartition, nextOffset[partition]);
 141       infoChannel.sendShardAssignedEvent(partition);
 142     });
 143
 144     consumer.resume(partitions);
 145   }
 146
 147   @Override
 148   public void onPartitionsRevoked(Collection<TopicPartition> partitions)
 149   {
 150     partitions.forEach(topicPartition ->
 151     {
 152       int partition = topicPartition.partition();
 153       isShardOwned[partition] = false;
 154       log.info("Partition revoked: {} - next={}", partition, nextOffset[partition]);
 155       infoChannel.sendShardRevokedEvent(partition);
 156     });
 157   }
 158
 159   @Override
 160   public void onPartitionsLost(Collection<TopicPartition> partitions)
 161   {
 162     log.warn("Lost partitions: {}, partitions");
 163     // TODO: Muss auf den Verlust anders reagiert werden?
 164     onPartitionsRevoked(partitions);
 165   }
 166
 167   @Override
 168   public void run()
 169   {
 170     running = true;
 171
 172     while (running)
 173     {
 174       try
 175       {
 176         ConsumerRecords<String, AbstractMessageTo> records = consumer.poll(Duration.ofMinutes(1));
 177         log.info("Fetched {} messages", records.count());
 178
 179         if (loadInProgress)
 180         {
 181           loadChatRoomData(records);
 182
 183           if (isLoadingCompleted())
 184           {
 185             log.info("Loading of messages completed! Pausing all owned partitions...");
 186             pauseAllOwnedPartions();
 187             log.info("Resuming normal operations...");
 188             loadInProgress = false;
 189           }
 190         }
 191         else
 192         {
 193           if (!records.isEmpty())
 194           {
 195             throw new IllegalStateException("All owned partitions should be paused, when no load is in progress!");
 196           }
 197         }
 198       }
 199       catch (WakeupException e)
 200       {
 201         log.info("Received WakeupException, exiting!");
 202         running = false;
 203       }
 204     }
 205
 206     log.info("Exiting normally");
 207   }
 208
 209   private void loadChatRoomData(ConsumerRecords<String, AbstractMessageTo> records)
 210   {
 211     for (ConsumerRecord<String, AbstractMessageTo> record : records)
 212     {
 213       UUID chatRoomId = UUID.fromString(record.key());
 214
 215       switch (record.value().getType())
 216       {
 217         case EVENT_CHATMESSAGE_RECEIVED:
 218           Instant instant = Instant.ofEpochSecond(record.timestamp());
 219           LocalDateTime timestamp = LocalDateTime.ofInstant(instant, zoneId);
 220           loadChatMessage(
 221               chatRoomId,
 222               timestamp,
 223               record.offset(),
 224               (EventChatMessageReceivedTo) record.value(),
 225               record.partition());
 226           break;
 227
 228         default:
 229           log.debug(
 230               "Ignoring message for chat-room {} with offset {}: {}",
 231               chatRoomId,
 232               record.offset(),
 233               record.value());
 234       }
 235
 236       nextOffset[record.partition()] = record.offset() + 1;
 237     }
 238   }
 239
 240   private void loadChatMessage(
 241       UUID chatRoomId,
 242       LocalDateTime timestamp,
 243       long offset,
 244       EventChatMessageReceivedTo chatMessageTo,
 245       int partition)
 246   {
 247     Message.MessageKey key = Message.MessageKey.of(chatMessageTo.getUser(), chatMessageTo.getId());
 248     Message message = new Message(key, offset, timestamp, chatMessageTo.getText());
 249
 250     ChatRoomData chatRoomData = this
 251         .chatRoomData[partition]
 252         .computeIfAbsent(chatRoomId, this::computeChatRoomData);
 253     KafkaChatMessageService kafkaChatRoomService =
 254         (KafkaChatMessageService) chatRoomData.getChatRoomService();
 255
 256     kafkaChatRoomService.persistMessage(message);
 257   }
 258
 259   private boolean isLoadingCompleted()
 260   {
 261     return IntStream
 262         .range(0, numShards)
 263         .filter(shard -> isShardOwned[shard])
 264         .allMatch(shard -> nextOffset[shard] >= currentOffset[shard]);
 265   }
 266
 267   private void pauseAllOwnedPartions()
 268   {
 269     consumer.pause(IntStream
 270         .range(0, numShards)
 271         .filter(shard -> isShardOwned[shard])
 272         .mapToObj(shard -> new TopicPartition(topic, shard))
 273         .toList());
 274   }
 275
 276
 277   int[] getOwnedShards()
 278   {
 279     return IntStream
 280         .range(0, numShards)
 281         .filter(shard -> isShardOwned[shard])
 282         .toArray();
 283   }
 284
 285   Mono<ChatRoomData> getChatRoomData(int shard, UUID id)
 286   {
 287     if (loadInProgress)
 288     {
 289       return Mono.error(new LoadInProgressException());
 290     }
 291
 292     if (!isShardOwned[shard])
 293     {
 294       return Mono.error(new ShardNotOwnedException(shard));
 295     }
 296
 297     return infoChannel
 298         .getChatRoomInfo(id)
 299         .map(chatRoomInfo ->
 300             chatRoomData[shard].computeIfAbsent(id, this::computeChatRoomData));
 301   }
 302
 303   private ChatRoomData computeChatRoomData(UUID chatRoomId)
 304   {
 305     log.info("Creating ChatRoom {} with buffer-size {}", chatRoomId, bufferSize);
 306     KafkaChatMessageService service = new KafkaChatMessageService(this, chatRoomId);
 307     return new ChatRoomData(clock, service, bufferSize);
 308   }
 309 }