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