Merge pull request #448 from vert-x3/447-cluster-batch-command-same-slot-calculation-is-wrong

Fixed correct node selection for CLUSTER batch command

Author: Ladicek

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

@@ -4,19 +4,21 @@
 import io.vertx.core.Context;
 import io.vertx.core.Future;
 import io.vertx.core.buffer.Buffer;
+import io.vertx.core.json.JsonObject;
 import io.vertx.ext.unit.Async;
 import io.vertx.ext.unit.TestContext;
 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 +1165,98 @@ public void setAndWaitEmptyBatch(TestContext should) {
     ));
   }
 
+  @Test(timeout = 30_000)
+  @SuppressWarnings("unchecked")
+  public void batchSameSlotGroupByMultipleSlotsCommands(TestContext should) {
+    final Async test = should.async();
+
+    Redis.createClient(rule.vertx(), options)
+      .connect()
+      .onComplete(should.asyncAssertSuccess(cluster -> {
+        cluster
+          .send(cmd(CLUSTER).arg("SLOTS"))
+          .compose(reply -> {
+            if (reply == null || reply.size() == 0) {
+              // no slots available we can't really proceed
+              return Future.failedFuture("CLUSTER SLOTS No slots available in the cluster.");
+            }
+
+            Map<Integer, JsonObject> slotRangeMap = new HashMap<>();
+            for (int i = 0; i < reply.size(); i++) {
+              Response s = reply.get(i);
+              JsonObject slotRange = new JsonObject()
+                .put("start", s.get(0).toInteger())
+                .put("end", s.get(1).toInteger());
+              slotRangeMap.put(i, slotRange);
+            }
+
+            Map<Integer, List<Request>> mapCommands = new HashMap<>();
+            for (int i = 0; i < 100000; i++) {
+              String key = "key-" + i;
+              String value = "value-" + i;
+              int group;
+              int slot = ZModem.generate(key);
+              for (Map.Entry<Integer, JsonObject> entry : slotRangeMap.entrySet()) {
+                if (slot >= entry.getValue().getInteger("start") && slot <= entry.getValue().getInteger("end")) {
+                  mapCommands.computeIfAbsent(entry.getKey(), k -> new ArrayList<>()).add(cmd(SET).arg(key).arg(value));
+                  break;
+                }
+              }
+            }
+
+            List<Future<List<Response>>> futures = new ArrayList<>();
+            for (Map.Entry<Integer, List<Request>> entry : mapCommands.entrySet()) {
+              futures.add(cluster.batch(entry.getValue()));
+            }
+
+            return 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.complete();
+              }));
+          }).onFailure(should::fail);
+      })).onFailure(should::fail);
+  }
+
+  @Test(timeout = 30_000)
+  public void batchSameSlotsCommands(TestContext should) {
+    final Async test = should.async();
+
+    Redis.createClient(rule.vertx(), options)
+      .connect()
+      .onComplete(should.asyncAssertSuccess(cluster -> {
+        cluster
+          .send(cmd(CLUSTER).arg("SLOTS"))
+          .compose(reply -> {
+            if (reply == null || reply.size() == 0) {
+              // no slots available we can't really proceed
+              return Future.failedFuture("CLUSTER SLOTS No slots available in the cluster.");
+            }
+
+            //take random slot
+            Response s = reply.get(0);
+            JsonObject slotRange = new JsonObject()
+              .put("start", s.get(0).toInteger())
+              .put("end", s.get(1).toInteger());
+
+            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 >= slotRange.getInteger("start") && slot <= slotRange.getInteger("end")) {
+                rawCommands.add(Request.cmd(Command.SET).arg(key).arg(value));
+              }
+            }
+            return cluster.batch(rawCommands).onComplete(should.asyncAssertSuccess(responses -> {
+              should.assertEquals(rawCommands.size(), responses.size());
+              test.complete();
+            }));
+          }).onFailure(should::fail);
+      })).onFailure(should::fail);
+  }
+
   @Test(timeout = 30_000)
   public void clusterInfoReturnsVerbatimString(TestContext should) {
     final Async test = should.async();