ENHANCE: Change decode logic in BTreeGetBulk api.

naver · Sep 13, 2023 · edd4541 · edd4541
1 parent 5a3d4db
commit edd4541
Show file tree

Hide file tree

Showing 6 changed files with 67 additions and 10 deletions.
diff --git a/src/main/java/net/spy/memcached/ArcusClient.java b/src/main/java/net/spy/memcached/ArcusClient.java
@@ -4025,6 +4025,9 @@ private <T> CollectionGetBulkFuture<Map<String, BTreeGetResult<Long, T>>> btreeG
 
     for (BTreeGetBulk<T> getBulk : getBulkList) {
       Operation op = opFact.bopGetBulk(getBulk, new BTreeGetBulkOperation.Callback() {
+        private final Map<String, List<BTreeElement<Long, CachedData>>> cachedDataMap =
+                new HashMap<String, List<BTreeElement<Long, CachedData>>>();
+
         @Override
         public void receivedStatus(OperationStatus status) {
           // Nothing to do here because the user MUST search the result Map instance.
@@ -4047,9 +4050,23 @@ public void gotKey(String key, int elementCount, OperationStatus status) {
 
         @Override
         public void gotElement(String key, int flags, Object bkey, byte[] eflag, byte[] data) {
-          result.get(key).addElement(
-                  new BTreeElement<Long, T>((Long) bkey, eflag,
-                          tc.decode(new CachedData(flags, data, tc.getMaxSize()))));
+          List<BTreeElement<Long, CachedData>> elements = cachedDataMap.get(key);
+          if (elements == null) {
+            elements = new ArrayList<BTreeElement<Long, CachedData>>();
+            cachedDataMap.put(key, elements);
+          }
+          elements.add(new BTreeElement<Long, CachedData>((Long) bkey, eflag,
+                  new CachedData(flags, data, tc.getMaxSize())));
+        }
+
+        @Override
+        public void addResult() {
+          if (!cachedDataMap.isEmpty()) {
+            for (Entry<String, List<BTreeElement<Long, CachedData>>> entry : cachedDataMap.entrySet()) {
+              result.get(entry.getKey()).addElements(entry.getValue(), tc);
+            }
+            cachedDataMap.clear();
+          }
         }
       });
       ops.add(op);
@@ -4081,6 +4098,9 @@ public void gotElement(String key, int flags, Object bkey, byte[] eflag, byte[]
 
     for (BTreeGetBulk<T> getBulk : getBulkList) {
       Operation op = opFact.bopGetBulk(getBulk, new BTreeGetBulkOperation.Callback() {
+        private final Map<String, List<BTreeElement<ByteArrayBKey, CachedData>>> cachedDataMap =
+                new HashMap<String, List<BTreeElement<ByteArrayBKey, CachedData>>>();
+
         @Override
         public void receivedStatus(OperationStatus status) {
         }
@@ -4103,10 +4123,24 @@ public void gotKey(String key, int elementCount, OperationStatus status) {
 
         @Override
         public void gotElement(String key, int flags, Object bkey, byte[] eflag, byte[] data) {
-          result.get(key).addElement(
-                  new BTreeElement<ByteArrayBKey, T>(
-                          new ByteArrayBKey((byte[]) bkey),
-                          eflag, tc.decode(new CachedData(flags, data, tc.getMaxSize()))));
+          List<BTreeElement<ByteArrayBKey, CachedData>> elements = cachedDataMap.get(key);
+          if (elements == null) {
+            elements = new ArrayList<BTreeElement<ByteArrayBKey, CachedData>>();
+            cachedDataMap.put(key, elements);
+          }
+          elements.add(new BTreeElement<ByteArrayBKey, CachedData>(
+                  new ByteArrayBKey((byte[]) bkey), eflag,
+                  new CachedData(flags, data, tc.getMaxSize())));
+        }
+
+        @Override
+        public void addResult() {
+          if (!cachedDataMap.isEmpty()) {
+            for (Entry<String, List<BTreeElement<ByteArrayBKey, CachedData>>> entry : cachedDataMap.entrySet()) {
+              result.get(entry.getKey()).addElements(entry.getValue(), tc);
+            }
+            cachedDataMap.clear();
+          }
         }
       });
       ops.add(op);

diff --git a/src/main/java/net/spy/memcached/collection/BTreeGetResult.java b/src/main/java/net/spy/memcached/collection/BTreeGetResult.java
@@ -16,10 +16,13 @@
  */
 package net.spy.memcached.collection;
 
+import java.util.List;
 import java.util.Map;
 import java.util.SortedMap;
 
+import net.spy.memcached.CachedData;
 import net.spy.memcached.ops.CollectionOperationStatus;
+import net.spy.memcached.transcoders.Transcoder;
 
 public class BTreeGetResult<K, V> {
 
@@ -40,7 +43,13 @@ public CollectionOperationStatus getCollectionResponse() {
     return opStatus;
   }
 
-  public void addElement(BTreeElement<K, V> element) {
-    this.elements.put(element.getBkey(), element);
+  public void addElements(List<BTreeElement<K, CachedData>> cachedData, Transcoder<V> tc) {
+    if (elements != null && elements.isEmpty()) {
+      for (BTreeElement<K, CachedData> elem : cachedData) {
+        BTreeElement<K, V> decodedElem =
+                new BTreeElement<K, V>(elem.getBkey(), elem.getEflag(), tc.decode(elem.getValue()));
+        elements.put(decodedElem.getBkey(), decodedElem);
+      }
+    }
   }
 }
diff --git a/src/main/java/net/spy/memcached/internal/CollectionGetBulkFuture.java b/src/main/java/net/spy/memcached/internal/CollectionGetBulkFuture.java
@@ -26,6 +26,7 @@
 
 import net.spy.memcached.MemcachedConnection;
 import net.spy.memcached.OperationTimeoutException;
+import net.spy.memcached.ops.CollectionGetOpCallback;
 import net.spy.memcached.ops.CollectionOperationStatus;
 import net.spy.memcached.ops.Operation;
 import net.spy.memcached.ops.OperationState;
@@ -84,6 +85,10 @@ public T get(long duration, TimeUnit units)
         throw new ExecutionException(new RuntimeException(op.getCancelCause()));
       }
     }
+    for (Operation op : ops) {
+      CollectionGetOpCallback callback = (CollectionGetOpCallback) op.getCallback();
+      callback.addResult();
+    }
 
     return result;
   }

diff --git a/src/main/java/net/spy/memcached/ops/BTreeGetBulkOperation.java b/src/main/java/net/spy/memcached/ops/BTreeGetBulkOperation.java
@@ -21,7 +21,7 @@
 public interface BTreeGetBulkOperation extends KeyedOperation {
   BTreeGetBulk<?> getBulk();
 
-  interface Callback extends OperationCallback {
+  interface Callback extends CollectionGetOpCallback {
     void gotElement(String key, int flags, Object subkey, byte[] eflag, byte[] data);
 
     void gotKey(String key, int elementCount, OperationStatus status);

diff --git a/src/main/java/net/spy/memcached/ops/MultiBTreeGetBulkOperationCallback.java b/src/main/java/net/spy/memcached/ops/MultiBTreeGetBulkOperationCallback.java
@@ -33,5 +33,10 @@ public void gotElement(String key, int flags, Object subkey, byte[] eflag, byte[
   public void gotKey(String key, int elementCount, OperationStatus status) {
     ((BTreeGetBulkOperation.Callback) originalCallback).gotKey(key, elementCount, status);
   }
+
+  @Override
+  public void addResult() {
+    ((BTreeGetBulkOperation.Callback) originalCallback).addResult();
+  }
 }
 
diff --git a/src/test/manual/net/spy/memcached/MultibyteKeyTest.java b/src/test/manual/net/spy/memcached/MultibyteKeyTest.java
@@ -540,6 +540,10 @@ public void receivedStatus(OperationStatus status) {
             @Override
             public void complete() {
             }
+
+            @Override
+            public void addResult() {
+            }
           }).initialize();
     } catch (java.nio.BufferOverflowException e) {
       Assert.fail();