NEU
[demos/kafka/chat] / src / main / java / de / juplo / kafka / chat / backend / persistence / kafka / ChatMessageChannel.java
index 7620461..69947a9 100644 (file)
@@ -12,7 +12,6 @@ import org.apache.kafka.clients.consumer.ConsumerRecords;
 import org.apache.kafka.clients.producer.Producer;
 import org.apache.kafka.clients.producer.ProducerRecord;
 import org.apache.kafka.common.TopicPartition;
-import org.apache.kafka.common.errors.RecordDeserializationException;
 import org.apache.kafka.common.errors.WakeupException;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
@@ -29,8 +28,8 @@ import java.util.stream.IntStream;
 public class ChatMessageChannel implements Runnable, ConsumerRebalanceListener
 {
   private final String topic;
-  private final Consumer<String, MessageTo> consumer;
   private final Producer<String, MessageTo> producer;
+  private final Consumer<String, MessageTo> consumer;
   private final ZoneId zoneId;
   private final int numShards;
   private final boolean[] isShardOwned;
@@ -46,8 +45,8 @@ public class ChatMessageChannel implements Runnable, ConsumerRebalanceListener
 
   public ChatMessageChannel(
     String topic,
-    Consumer<String, MessageTo> consumer,
     Producer<String, MessageTo> producer,
+    Consumer<String, MessageTo> consumer,
     ZoneId zoneId,
     int numShards)
   {
@@ -68,6 +67,50 @@ public class ChatMessageChannel implements Runnable, ConsumerRebalanceListener
   }
 
 
+  Mono<Message> sendMessage(
+      UUID chatRoomId,
+      Message.MessageKey key,
+      LocalDateTime timestamp,
+      String text)
+  {
+    int shard = this.shardingStrategy.selectShard(chatRoomId);
+    TopicPartition tp = new TopicPartition(topic, shard);
+    ZonedDateTime zdt = ZonedDateTime.of(timestamp, zoneId);
+    return Mono.create(sink ->
+    {
+      ProducerRecord<String, MessageTo> record =
+          new ProducerRecord<>(
+              tp.topic(),
+              tp.partition(),
+              zdt.toEpochSecond(),
+              chatRoomId.toString(),
+              MessageTo.of(key.getUsername(), key.getMessageId(), text));
+
+      producer.send(record, ((metadata, exception) ->
+      {
+        if (metadata != null)
+        {
+          // On successful send
+          Message message = new Message(key, metadata.offset(), timestamp, text);
+          log.info("Successfully send message {}", message);
+          sink.success(message);
+        }
+        else
+        {
+          // On send-failure
+          log.error(
+              "Could not send message for chat-room={}, key={}, timestamp={}, text={}: {}",
+              chatRoomId,
+              key,
+              timestamp,
+              text,
+              exception);
+          sink.error(exception);
+        }
+      }));
+    });
+  }
+
   @Override
   public void onPartitionsAssigned(Collection<TopicPartition> partitions)
   {
@@ -147,9 +190,8 @@ public class ChatMessageChannel implements Runnable, ConsumerRebalanceListener
       }
       catch (WakeupException e)
       {
-      }
-      catch (RecordDeserializationException e)
-      {
+        log.info("Received WakeupException, exiting!");
+        running = false;
       }
     }
   }
@@ -202,54 +244,24 @@ public class ChatMessageChannel implements Runnable, ConsumerRebalanceListener
         .toList());
   }
 
-  Mono<Message> sendMessage(
-      UUID chatRoomId,
-      Message.MessageKey key,
-      LocalDateTime timestamp,
-      String text)
-  {
-    int shard = this.shardingStrategy.selectShard(chatRoomId);
-    TopicPartition tp = new TopicPartition(topic, shard);
-    ZonedDateTime zdt = ZonedDateTime.of(timestamp, zoneId);
-    return Mono.create(sink ->
-    {
-      ProducerRecord<String, MessageTo> record =
-          new ProducerRecord<>(
-              tp.topic(),
-              tp.partition(),
-              zdt.toEpochSecond(),
-              chatRoomId.toString(),
-              MessageTo.of(key.getUsername(), key.getMessageId(), text));
-
-      producer.send(record, ((metadata, exception) ->
-      {
-        if (metadata != null)
-        {
-          // On successful send
-          Message message = new Message(key, metadata.offset(), timestamp, text);
-          log.info("Successfully send message {}", message);
-          sink.success(message);
-        }
-        else
-        {
-          // On send-failure
-          log.error(
-              "Could not send message for chat-room={}, key={}, timestamp={}, text={}: {}",
-              chatRoomId,
-              key,
-              timestamp,
-              text,
-              exception);
-          sink.error(exception);
-        }
-      }));
-    });
-  }
-
 
   void putChatRoom(ChatRoom chatRoom)
   {
-    chatrooms[chatRoom.getShard()].put(chatRoom.getId(), chatRoom);
+    Integer partition = chatRoom.getShard();
+    UUID chatRoomId = chatRoom.getId();
+    if (chatrooms[partition].containsKey(chatRoomId))
+    {
+      log.warn("Ignoring existing chat-room: " + chatRoom);
+    }
+    else
+    {
+      log.info(
+          "Adding new chat-room to partition {}: {}",
+          partition,
+          chatRoom);
+
+      chatrooms[partition].put(chatRoomId, chatRoom);
+    }
   }
 
   Mono<ChatRoom> getChatRoom(int shard, UUID id)
@@ -257,8 +269,12 @@ public class ChatMessageChannel implements Runnable, ConsumerRebalanceListener
     return Mono.justOrEmpty(chatrooms[shard].get(id));
   }
 
-  Flux<ChatRoom> getChatRooms(int shard)
+  Flux<ChatRoom> getChatRooms()
   {
-    return Flux.fromStream(chatrooms[shard].values().stream());
+    return Flux.fromStream(IntStream
+        .range(0, numShards)
+        .filter(shard -> isShardOwned[shard])
+        .mapToObj(shard -> Integer.valueOf(shard))
+        .flatMap(shard -> chatrooms[shard].values().stream()));
   }
 }