Fixed correct node selection for CLUSTER batch command

Author: pendula95

src/main/java/io/vertx/redis/client/impl/RedisClusterConnection.java

@@ -329,7 +329,8 @@ public Future<List<Response>> batch(List<Request> requests) {
       LOG.debug("Empty batch");
       promise.complete(Collections.emptyList());
     } else {
-      int currentSlot = -1;
+      int correctSlot = -1;
+      String currentEndpoint = null;
       boolean readOnly = false;
       boolean forceMasterEndpoint = false;
 
@@ -353,6 +354,7 @@ public Future<List<Response>> batch(List<Request> requests) {
         final List<byte[]> keys = req.keys();
         forceMasterEndpoint |= MASTER_ONLY_COMMANDS.contains(cmd);
         int slot;
+        String endpoint;
 
         // process slots, need to verify if we can run this batch
         switch (keys.size()) {
@@ -362,10 +364,14 @@ public Future<List<Response>> batch(List<Request> requests) {
           case 1:
             // command is single key, as long as we're on the same slot, it's OK
             slot = ZModem.generate(keys.get(0));
+            // as cluster server serves range of slots we need to compare to server range and not exact slot
+            //always take master to make sure we have same endpoint
+            endpoint = slots.endpointsForKey(slot)[0];
             // we are checking the first request key
-            if (currentSlot == -1) {
-              currentSlot = slot;
-            } else if (currentSlot != slot) {
+            if (currentEndpoint == null) {
+              currentEndpoint = endpoint;
+              correctSlot = slot;
+            } else if (!currentEndpoint.equals(endpoint)) {
               // in cluster mode we currently do not handle batching commands which keys are not on the same slot
               promise.fail(buildCrossslotFailureMsg(req));
               return promise.future();
@@ -375,9 +381,11 @@ public Future<List<Response>> batch(List<Request> requests) {
             // multiple keys on the command
             for (byte[] key : keys) {
               slot = ZModem.generate(key);
-              if (currentSlot == -1) {
-                currentSlot = slot;
-              } else if (currentSlot != slot) {
+              endpoint = slots.endpointsForKey(slot)[0];
+              if (currentEndpoint == null) {
+                correctSlot = slot;
+                currentEndpoint = endpoint;
+              } else if (!currentEndpoint.equals(endpoint)) {
                 // in cluster mode we currently do not handle batching commands which keys are not on the same slot
                 promise.fail(buildCrossslotFailureMsg(req));
                 return promise.future();
@@ -388,7 +396,8 @@ public Future<List<Response>> batch(List<Request> requests) {
       }
 
       // all keys are on the same slot!
-      batch(selectEndpoint(currentSlot, readOnly, forceMasterEndpoint), RETRIES, requests, promise);
+      //we just need to decide which endpoint to use based on additional options
+      batch(selectEndpoint(correctSlot, readOnly, forceMasterEndpoint), RETRIES, requests, promise);
     }
 
     return promise.future();

src/test/java/io/vertx/redis/client/test/RedisClusterTest.java

@@ -9,14 +9,15 @@
 import io.vertx.ext.unit.junit.RunTestOnContext;
 import io.vertx.ext.unit.junit.VertxUnitRunner;
 import io.vertx.redis.client.*;
-import java.util.stream.Collectors;
+import io.vertx.redis.client.impl.ZModem;
 import org.junit.*;
 import org.junit.runner.RunWith;
 import org.testcontainers.containers.FixedHostPortGenericContainer;
 import org.testcontainers.containers.GenericContainer;
 
 import java.util.*;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
 
 import static io.vertx.redis.client.Command.*;
 import static io.vertx.redis.client.Request.cmd;
@@ -1163,6 +1164,70 @@ public void setAndWaitEmptyBatch(TestContext should) {
     ));
   }
 
+  @Test(timeout = 30_000)
+  @SuppressWarnings("unchecked")
+  public void batchSameSlotGroupByMultipleSlotsCommands(TestContext should) {
+    final Async test = should.async();
+
+    Map<Integer, List<Request>> mapCommands = new HashMap<>();
+    for (int i = 0; i < 100000; i++) {
+      String key = "key-" + i;
+      String value = "value-" + i;
+      int endpoint;
+      int slot = ZModem.generate(key);
+      if (slot < 16384 / (options.getEndpoints().size() / 2)) {
+        endpoint = 0;
+      } else if (slot <= 2 * 16384 / (options.getEndpoints().size() / 2)) {
+        endpoint = 1;
+      } else {
+        endpoint = 2;
+      }
+      List<Request> commands = mapCommands.computeIfAbsent(endpoint, k -> new ArrayList<>());
+      commands.add(Request.cmd(Command.SET).arg(key).arg(value));
+    }
+
+    Redis.createClient(rule.vertx(), options)
+      .connect()
+      .onComplete(should.asyncAssertSuccess(cluster -> {
+          List<Future<List<Response>>> futures = new ArrayList<>();
+          for (Map.Entry<Integer, List<Request>> entry : mapCommands.entrySet()) {
+            futures.add(cluster.batch(entry.getValue()));
+          }
+          Future.all(futures).onComplete(should.asyncAssertSuccess(responses -> {
+            should.assertEquals(mapCommands.values().stream().map(List::size).reduce(0 , Integer::sum), responses.result().list().stream().map(item -> ((List<Request>) item).size()).reduce(0, Integer::sum));
+
+            test.countDown();
+          })).onFailure(should::fail);
+        }
+      ));
+  }
+
+  @Test(timeout = 30_000)
+  public void batchSameSlotsCommands(TestContext should) {
+    final Async test = should.async();
+
+    List<Request> rawCommands = new ArrayList<>();
+    for (int i = 0; i < 100000; i++) {
+      String key = "key-" + i;
+      String value = "value-" + i;
+      int endpoint;
+      int slot = ZModem.generate(key);
+      if (slot < 16384 / (options.getEndpoints().size() / 2)) {
+        rawCommands.add(Request.cmd(Command.SET).arg(key).arg(value));
+      }
+    }
+
+    Redis.createClient(rule.vertx(), options)
+      .connect()
+      .onComplete(should.asyncAssertSuccess(cluster -> {
+          cluster.batch(rawCommands).onComplete(should.asyncAssertSuccess(responses -> {
+            should.assertEquals(rawCommands.size(), responses.size());
+            test.countDown();
+          })).onFailure(should::fail);
+        }
+      ));
+  }
+
   @Test(timeout = 30_000)
   public void clusterInfoReturnsVerbatimString(TestContext should) {
     final Async test = should.async();