Make route return multiple streams

Author: acogoluegnes

src/main/java/com/rabbitmq/stream/impl/Client.java

@@ -1240,7 +1240,7 @@ int getPort() {
     return port;
   }
 
-  public String route(String routingKey, String superStream) {
+  public List<String> route(String routingKey, String superStream) {
     if (routingKey == null || superStream == null) {
       throw new IllegalArgumentException("routing key and stream must not be null");
     }
@@ -1263,7 +1263,7 @@ public String route(String routingKey, String superStream) {
       bb.writeBytes(routingKey.getBytes(StandardCharsets.UTF_8));
       bb.writeShort(superStream.length());
       bb.writeBytes(superStream.getBytes(StandardCharsets.UTF_8));
-      OutstandingRequest<String> request = new OutstandingRequest<>(this.rpcTimeout);
+      OutstandingRequest<List<String>> request = new OutstandingRequest<>(this.rpcTimeout);
       outstandingRequests.put(correlationId, request);
       channel.writeAndFlush(bb);
       request.block();

src/main/java/com/rabbitmq/stream/impl/HashRoutingStrategy.java

@@ -14,20 +14,20 @@
 package com.rabbitmq.stream.impl;
 
 import com.rabbitmq.stream.Message;
+import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.function.Function;
 import java.util.function.ToIntFunction;
+import java.util.stream.Collectors;
 
 class HashRoutingStrategy implements RoutingStrategy {
 
   private final Function<Message, String> routingKeyExtractor;
 
   private final StreamEnvironment env;
 
-  private final String superStream;
-
-  private final List<String> partitions;
+  private final List<List<String>> partitions;
 
   private final ToIntFunction<String> hash;
 
@@ -38,15 +38,16 @@ class HashRoutingStrategy implements RoutingStrategy {
       ToIntFunction<String> hash) {
     this.routingKeyExtractor = routingKeyExtractor;
     this.env = env;
-    this.superStream = superStream;
     // TODO use async retry to get locator
     List<String> ps = this.env.locator().partitions(superStream);
-    this.partitions = new CopyOnWriteArrayList<>(ps);
+    this.partitions =
+        new CopyOnWriteArrayList<>(
+            ps.stream().map(Collections::singletonList).collect(Collectors.toList()));
     this.hash = hash;
   }
 
   @Override
-  public String route(Message message) {
+  public List<String> route(Message message) {
     String routingKey = routingKeyExtractor.apply(message);
     int hashValue = hash.applyAsInt(routingKey);
     return this.partitions.get((hashValue & 0x7FFFFFFF) % this.partitions.size());

src/main/java/com/rabbitmq/stream/impl/RoutingKeyRoutingStrategy.java

@@ -14,6 +14,7 @@
 package com.rabbitmq.stream.impl;
 
 import com.rabbitmq.stream.Message;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Function;
@@ -22,7 +23,7 @@ class RoutingKeyRoutingStrategy implements RoutingStrategy {
 
   private final Function<Message, String> routingKeyExtractor;
 
-  private final Map<String, String> routingKeysToStreams = new ConcurrentHashMap<>();
+  private final Map<String, List<String>> routingKeysToStreams = new ConcurrentHashMap<>();
 
   private final StreamEnvironment env;
 
@@ -36,15 +37,15 @@ class RoutingKeyRoutingStrategy implements RoutingStrategy {
   }
 
   @Override
-  public String route(Message message) {
+  public List<String> route(Message message) {
     String routingKey = this.routingKeyExtractor.apply(message);
-    String stream =
+    List<String> streams =
         routingKeysToStreams.computeIfAbsent(
             routingKey,
             routingKey1 -> {
               // TODO retry on locator lookup
               return env.locator().route(routingKey1, superStream);
             });
-    return stream;
+    return streams;
   }
 }

src/main/java/com/rabbitmq/stream/impl/RoutingStrategy.java

@@ -14,8 +14,9 @@
 package com.rabbitmq.stream.impl;
 
 import com.rabbitmq.stream.Message;
+import java.util.List;
 
-public interface RoutingStrategy {
+interface RoutingStrategy {
 
-  String route(Message message);
+  List<String> route(Message message);
 }

src/main/java/com/rabbitmq/stream/impl/ServerFrameHandler.java

@@ -879,28 +879,33 @@ int doHandle(Client client, ChannelHandlerContext ctx, ByteBuf message) {
       int read = 4;
       short responseCode = message.readShort();
       read += 2;
-      short streamSize = message.readShort();
-      read += 2;
-      String stream;
-      if (streamSize == -1) {
-        stream = null;
+      int streamCount = message.readInt();
+      read += 4;
+      List<String> streams;
+      if (streamCount == 0) {
+        streams = Collections.emptyList();
       } else {
-        byte[] bytes = new byte[streamSize];
-        message.readBytes(bytes);
-        stream = new String(bytes, StandardCharsets.UTF_8);
-        read += stream.length();
+        streams = new ArrayList<>(streamCount);
+        for (int i = 0; i < streamCount; i++) {
+          String stream = readString(message);
+          read += (2 + stream.length());
+          streams.add(stream);
+        }
       }
 
       if (responseCode != RESPONSE_CODE_OK) {
         LOGGER.info("Route returned error: {}", Utils.formatConstant(responseCode));
       }
 
-      OutstandingRequest<String> outstandingRequest =
-          remove(client.outstandingRequests, correlationId, String.class);
+      OutstandingRequest<List<String>> outstandingRequest =
+          remove(
+              client.outstandingRequests,
+              correlationId,
+              new ParameterizedTypeReference<List<String>>() {});
       if (outstandingRequest == null) {
         LOGGER.warn("Could not find outstanding request with correlation ID {}", correlationId);
       } else {
-        outstandingRequest.response().set(stream);
+        outstandingRequest.response().set(streams);
         outstandingRequest.countDown();
       }
       return read;

src/main/java/com/rabbitmq/stream/impl/SuperStreamProducer.java

@@ -85,11 +85,13 @@ public long getLastPublishingId() {
   public void send(Message message, ConfirmationHandler confirmationHandler) {
     // TODO handle when the stream is not found (no partition found for the message)
     // and call the confirmation handler with a failure
-    String stream = this.routingStrategy.route(message);
-    Producer producer =
-        producers.computeIfAbsent(
-            stream, stream1 -> producerBuilder.duplicate().stream(stream1).build());
-    producer.send(message, confirmationHandler);
+    List<String> streams = this.routingStrategy.route(message);
+    for (String stream : streams) {
+      Producer producer =
+          producers.computeIfAbsent(
+              stream, stream1 -> producerBuilder.duplicate().stream(stream1).build());
+      producer.send(message, confirmationHandler);
+    }
   }
 
   @Override

src/test/java/com/rabbitmq/stream/impl/RoutePartitionsTest.java

@@ -52,8 +52,8 @@ void tearDown() throws Exception {
   }
 
   @Test
-  void routeShouldReturnNullWhenExchangeDoesNotExist() {
-    assertThat(cf.get().route("", UUID.randomUUID().toString())).isNull();
+  void routeShouldReturnEmptyListWhenExchangeDoesNotExist() {
+    assertThat(cf.get().route("", UUID.randomUUID().toString())).isEmpty();
   }
 
   @Test
@@ -66,8 +66,8 @@ void routeShouldReturnNullWhenNoStreamForRoutingKey() throws Exception {
     declareSuperStreamTopology(connection, superStream, partitions);
 
     Client client = cf.get();
-    assertThat(client.route("0", superStream)).isEqualTo(superStream + "-0");
-    assertThat(client.route("42", superStream)).isNull();
+    assertThat(client.route("0", superStream)).hasSize(1).contains(superStream + "-0");
+    assertThat(client.route("42", superStream)).isEmpty();
   }
 
   @Test
@@ -88,8 +88,25 @@ void routeTopologyWithPartitionCount() throws Exception {
         .hasSize(partitions)
         .containsExactlyInAnyOrderElementsOf(
             IntStream.range(0, partitions).mapToObj(i -> superStream + "-" + i).collect(toList()));
-    assertThat(client.route("0", superStream)).isEqualTo(superStream + "-0");
-    assertThat(client.route("1", superStream)).isEqualTo(superStream + "-1");
-    assertThat(client.route("2", superStream)).isEqualTo(superStream + "-2");
+    assertThat(client.route("0", superStream)).hasSize(1).contains(superStream + "-0");
+    assertThat(client.route("1", superStream)).hasSize(1).contains(superStream + "-1");
+    assertThat(client.route("2", superStream)).hasSize(1).contains(superStream + "-2");
+  }
+
+  @Test
+  void routeReturnsMultipleStreamsIfMultipleBindingsForSameKey() throws Exception {
+    declareSuperStreamTopology(connection, superStream, 3);
+    connection.createChannel().queueBind(superStream + "-1", superStream, "0");
+    Client client = cf.get();
+    List<String> streams = client.partitions(superStream);
+    assertThat(streams)
+        .hasSize(partitions + 1)
+        .contains(
+            IntStream.range(0, partitions)
+                .mapToObj(i -> superStream + "-" + i)
+                .toArray(String[]::new));
+    assertThat(client.route("0", superStream))
+        .hasSize(2)
+        .contains(superStream + "-0", superStream + "-1");
   }
 }