From 53dcdb2ea0e3cfae4e3eea0e69616e2efe691465 Mon Sep 17 00:00:00 2001
From: brido4125 <hcs4125@gmail.com>
Date: Mon, 17 Jul 2023 18:15:55 +0900
Subject: [PATCH] INTERNAL: refactoring groupingKeys method logic in
 ArcusClient.

---
 .../java/net/spy/memcached/ArcusClient.java   | 58 +++++++------------
 1 file changed, 22 insertions(+), 36 deletions(-)

diff --git a/src/main/java/net/spy/memcached/ArcusClient.java b/src/main/java/net/spy/memcached/ArcusClient.java
index 4d51384a1..4126206ee 100644
--- a/src/main/java/net/spy/memcached/ArcusClient.java
+++ b/src/main/java/net/spy/memcached/ArcusClient.java
@@ -2146,48 +2146,34 @@ public SMGetFuture<List<SMGetElement<Object>>> asyncBopSortMergeGet(
    * @return list of grouped (memcached node + keys) in the group
    */
   private Collection<Entry<MemcachedNode, List<String>>> groupingKeys(List<String> keyList, int groupSize) {
-    Map<String, Integer> chunkCount = new HashMap<String, Integer>();
-    Map<String, Entry<MemcachedNode, List<String>>> result = new HashMap<String,
-        Entry<MemcachedNode, List<String>>>();
-
+    Map<Integer, Entry<MemcachedNode, List<String>>> result =
+            new HashMap<Integer, Entry<MemcachedNode, List<String>>>();
+    Map<MemcachedNode, List<String>> temp = new HashMap<MemcachedNode, List<String>>();
     MemcachedConnection conn = getMemcachedConnection();
+    Integer index = 0;
+
+    for (String key : keyList) {
+      MemcachedNode qa = conn.findNodeByKey(key);
 
-    for (String k : keyList) {
-      MemcachedNode qa = conn.findNodeByKey(k);
-      String node;
       if (qa == null) {
-        node = "fake_node";
-      } else {
-        node = qa.getSocketAddress().toString();
+        throw new IllegalStateException("Could not find node for key: " + key);
       }
-      int cc;
-      if (chunkCount.containsKey(node)) {
-        cc = chunkCount.get(node);
-      } else {
-        cc = 0;
-        chunkCount.put(node, 0);
+      if (!temp.containsKey(qa)) {
+        temp.put(qa, new ArrayList<String>());
       }
-
-      String resultKey = node + cc;
-
-      List<String> arrangedKeyList = null;
-
-      if (result.containsKey(resultKey)) {
-        if (result.get(resultKey).getValue().size() >= groupSize) {
-          arrangedKeyList = new ArrayList<String>();
-          cc++;
-          result.put(node + cc, new AbstractMap.SimpleEntry<MemcachedNode,
-              List<String>>(qa, arrangedKeyList));
-          chunkCount.put(node, cc);
-        } else {
-          arrangedKeyList = result.get(resultKey).getValue();
-        }
-      } else {
-        arrangedKeyList = new ArrayList<String>();
-        result.put(resultKey, new AbstractMap.SimpleEntry<MemcachedNode,
-            List<String>>(qa, arrangedKeyList));
+      if (temp.get(qa).size() >= groupSize) {
+        result.put(index, new AbstractMap.SimpleEntry<MemcachedNode, List<String>>(qa, temp.get(qa)));
+        temp.put(qa, new ArrayList<String>());
+        index++;
       }
-      arrangedKeyList.add(k);
+      temp.get(qa).add(key);
+    }
+    /*
+    * Add the new Entry instance which is not full(smaller than groupSize) to the result.
+    * */
+    for (Entry<MemcachedNode, List<String>> entry : temp.entrySet()) {
+      result.put(index, new AbstractMap.SimpleEntry<MemcachedNode, List<String>>(entry.getKey(), entry.getValue()));
+      index++;
     }
     return result.values();
   }