feat: Optimize MGET to GET for single-key partitions #2432

src/main/java/io/lettuce/core/cluster/RedisAdvancedClusterAsyncCommandsImpl.java

@@ -73,6 +73,7 @@
  * @author Mark Paluch
  * @author Jon Chambers
  * @author Tihomir Mateev
+ * @author devbini
  * @since 3.3
  */
 @SuppressWarnings("unchecked")
@@ -363,8 +364,22 @@ public RedisFuture<List<KeyValue<K, V>>> mget(Iterable<K> keys) {
         Map<Integer, RedisFuture<List<KeyValue<K, V>>>> executions = new HashMap<>(partitioned.size());
 
         for (Map.Entry<Integer, List<K>> entry : partitioned.entrySet()) {
-            RedisFuture<List<KeyValue<K, V>>> mget = super.mget(entry.getValue());
-            executions.put(entry.getKey(), mget);
+
+            List<K> keysInPartition = entry.getValue();
+
+            if (keysInPartition.size() == 1) {
+                K singleKey = keysInPartition.get(0);
+                RedisFuture<V> get = super.get(singleKey);
+
+                RedisFuture<List<KeyValue<K, V>>> convert = new PipelinedRedisFuture<>(
+                        get.thenApply(value -> java.util.Collections.singletonList(KeyValue.fromNullable(singleKey, value)))
+                                .toCompletableFuture());
+
+                executions.put(entry.getKey(), convert);
+            } else {
+                RedisFuture<List<KeyValue<K, V>>> mget = super.mget(keysInPartition);
+                executions.put(entry.getKey(), mget);
+            }
         }
 
         // restore order of key

src/test/java/io/lettuce/core/cluster/commands/StringClusterCommandIntegrationTests.java

@@ -4,13 +4,15 @@
 import static org.assertj.core.api.Assertions.*;
 
 import java.util.LinkedHashMap;
+import java.util.List;
 import java.util.Map;
 
 import javax.inject.Inject;
 
 import org.junit.jupiter.api.Tag;
 import org.junit.jupiter.api.Test;
 
+import io.lettuce.core.KeyValue;
 import io.lettuce.core.cluster.ClusterTestUtil;
 import io.lettuce.core.cluster.api.StatefulRedisClusterConnection;
 import io.lettuce.core.cluster.api.sync.RedisClusterCommands;
@@ -21,6 +23,7 @@
  * Integration tests for {@link io.lettuce.core.api.sync.RedisStringCommands} using Redis Cluster.
  *
  * @author Mark Paluch
+ * @author devbini
  */
 @Tag(INTEGRATION_TEST)
 class StringClusterCommandIntegrationTests extends StringCommandIntegrationTests {
@@ -58,4 +61,18 @@ void mgetStreaming() {
         assertThat(count.intValue()).isEqualTo(2);
     }
 
+    @Test
+    void mgetOptimization() {
+        redis.set("{slot-a}:test", "value-a");
+        redis.set("{slot-b}:test", "value-b");
+
+        List<KeyValue<String, String>> result = redis.mget("{slot-a}:test", "{slot-b}:test");
+
+        redis.del("{slot-a}:test", "{slot-b}:test");
+
+        assertThat(result).hasSize(2);
+        assertThat(result).containsExactly(KeyValue.just("{slot-a}:test", "value-a"),
+                KeyValue.just("{slot-b}:test", "value-b"));
+    }
+
 }