INTERNAL: Limit bulk get keys size

naver · Jan 10, 2025 · 964b2c9 · 964b2c9
1 parent 5c13575
commit 964b2c9
Show file tree

Hide file tree

Showing 3 changed files with 89 additions and 38 deletions.
diff --git a/src/main/java/net/spy/memcached/protocol/TCPMemcachedNodeImpl.java b/src/main/java/net/spy/memcached/protocol/TCPMemcachedNodeImpl.java
@@ -40,6 +40,7 @@
 import net.spy.memcached.MemcachedNode;
 import net.spy.memcached.MemcachedReplicaGroup;
 import net.spy.memcached.compat.SpyObject;
+import net.spy.memcached.ops.GetOperation;
 import net.spy.memcached.ops.Operation;
 import net.spy.memcached.ops.OperationState;
 
@@ -265,7 +266,7 @@ public final void fillWriteBuffer(boolean shouldOptimize) {
           transitionWriteItem();
 
           preparePending();
-          if (shouldOptimize) {
+          if (shouldOptimize && o instanceof GetOperation) {
             optimize();
           }
 

diff --git a/src/main/java/net/spy/memcached/protocol/ascii/AsciiMemcachedNodeImpl.java b/src/main/java/net/spy/memcached/protocol/ascii/AsciiMemcachedNodeImpl.java
@@ -30,6 +30,9 @@
  * Memcached node for the ASCII protocol.
  */
 public final class AsciiMemcachedNodeImpl extends TCPMemcachedNodeImpl {
+
+  private static final int GET_BULK_CHUNK_SIZE = 200;
+
   public AsciiMemcachedNodeImpl(String name,
                                 SocketAddress sa,
                                 int bufSize, BlockingQueue<Operation> rq,
@@ -44,30 +47,39 @@ protected void optimize() {
     // make sure there are at least two get operations in a row before
     // attempting to optimize them.
     Operation nxtOp = writeQ.peek();
-    if (nxtOp instanceof GetOperation) {
-      optimizedOp = writeQ.remove();
-      nxtOp = writeQ.peek();
-      if (nxtOp instanceof GetOperation) {
-        OptimizedGetImpl og = new OptimizedGetImpl(
-                (GetOperation) optimizedOp, enabledMGetOp());
-        optimizedOp = og;
+    if (!(nxtOp instanceof GetOperation)
+            || ((GetOperation) nxtOp).getKeys().size() >= GET_BULK_CHUNK_SIZE) {
+      return;
+    }
 
-        do {
-          GetOperationImpl o = (GetOperationImpl) writeQ.remove();
-          if (!o.isCancelled()) {
-            og.addOperation(o);
-          }
-          nxtOp = writeQ.peek();
-        } while (nxtOp instanceof GetOperation);
+    int cnt = ((GetOperation) nxtOp).getKeys().size();
+    optimizedOp = writeQ.remove();
+    nxtOp = writeQ.peek();
+    OptimizedGetImpl og = null;
 
-        // Initialize the new mega get
-        optimizedOp.initialize();
-        assert optimizedOp.getState() == OperationState.WRITE_QUEUED;
-        ProxyCallback pcb = (ProxyCallback) og.getCallback();
-        getLogger().debug("Set up %s with %s keys and %s callbacks",
-                this, pcb.numKeys(), pcb.numCallbacks());
+    while (nxtOp instanceof GetOperation) {
+      cnt += ((GetOperation) nxtOp).getKeys().size();
+      if (cnt > GET_BULK_CHUNK_SIZE) {
+        break;
+      }
+      GetOperationImpl currentOp = (GetOperationImpl) writeQ.remove();
+      if (!currentOp.isCancelled()) {
+        if (og == null) {
+          og = new OptimizedGetImpl((GetOperation) optimizedOp, enabledMGetOp());
+          optimizedOp = og;
+        }
+        og.addOperation(currentOp);
       }
+      nxtOp = writeQ.peek();
     }
-  }
 
+    // Initialize the new mega get
+    if (og != null) {
+      optimizedOp.initialize();
+      assert optimizedOp.getState() == OperationState.WRITE_QUEUED;
+      ProxyCallback pcb = (ProxyCallback) optimizedOp.getCallback();
+      getLogger().debug("Set up %s with %s keys and %s callbacks",
+              this, pcb.numKeys(), pcb.numCallbacks());
+    }
+  }
 }
diff --git a/src/test/java/net/spy/memcached/protocol/ascii/OptimizeOperationTest.java b/src/test/java/net/spy/memcached/protocol/ascii/OptimizeOperationTest.java
@@ -1,10 +1,16 @@
 package net.spy.memcached.protocol.ascii;
 
+import java.net.InetSocketAddress;
 import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.List;
 
+import net.spy.memcached.ConnectionFactory;
+import net.spy.memcached.ConnectionFactoryBuilder;
 import net.spy.memcached.ops.GetOperation;
 import net.spy.memcached.ops.OperationStatus;
+import net.spy.memcached.protocol.TCPMemcachedNodeImpl;
 
 import org.junit.jupiter.api.Test;
 
@@ -13,24 +19,25 @@
 
 class OptimizeOperationTest {
 
+  private final GetOperation.Callback cb = new GetOperation.Callback() {
+    @Override
+    public void gotData(String key, int flags, byte[] data) {
+      // do nothing
+    }
+
+    @Override
+    public void receivedStatus(OperationStatus status) {
+      // do nothing
+    }
+
+    @Override
+    public void complete() {
+      // do nothing
+    }
+  };
+
   @Test
   void chooseGetOrMGet() {
-    GetOperation.Callback cb = new GetOperation.Callback() {
-      @Override
-      public void gotData(String key, int flags, byte[] data) {
-        // do nothing
-      }
-
-      @Override
-      public void receivedStatus(OperationStatus status) {
-        // do nothing
-      }
-
-      @Override
-      public void complete() {
-        // do nothing
-      }
-    };
     GetOperationImpl op1 = new GetOperationImpl("key", cb);
     GetOperationImpl op2 = new GetOperationImpl("key2", cb);
     OptimizedGetImpl optimizedOpWithMGet = new OptimizedGetImpl(op1, true);
@@ -53,4 +60,35 @@ public void complete() {
     String commandWithGet = new String(bytesWithGet, StandardCharsets.UTF_8);
     assertFalse(commandWithGet.contains("mget"));
   }
+
+
+  @Test
+  void doNotMergeTwoOperations() {
+    ConnectionFactoryBuilder builder = new ConnectionFactoryBuilder();
+    builder.setShouldOptimize(true);
+    ConnectionFactory cf = builder.build();
+    TCPMemcachedNodeImpl node = (TCPMemcachedNodeImpl) cf.createMemcachedNode("node1",
+            new InetSocketAddress("localhost", 11211), 4096);
+    node.setVersion("1.11.0");
+
+    List<String> keyList = new ArrayList<>();
+    for (int i = 0; i < 250; i++) {
+      keyList.add("k" + i);
+    }
+
+    GetOperationImpl op = new GetOperationImpl(keyList.get(0), cb);
+    GetOperationImpl op1 = new GetOperationImpl(keyList.subList(0, 190), cb, true);
+    GetOperationImpl op2 = new GetOperationImpl(keyList.subList(191, 205), cb, true);
+    node.addOpToWriteQ(op);
+    node.addOpToWriteQ(op1);
+    node.addOpToWriteQ(op2);
+    node.fillWriteBuffer(true);
+
+    ByteBuffer buffer = node.getWbuf();
+    byte[] bytesWithMGet = new byte[buffer.remaining()];
+    buffer.get(bytesWithMGet);
+    String commandWithMGet = new String(bytesWithMGet, StandardCharsets.UTF_8);
+    assertTrue(commandWithMGet.contains("mget 839 190"));
+    assertTrue(commandWithMGet.contains("mget 69 14"));
+  }
 }