Improved get() benchmarks from @alucarded

Benchmarks are a subset of facebook#7597

We just want to benchmark existing JNI methods with some potential performance improvements.

java/jmh/src/main/java/org/rocksdb/jmh/GetBenchmarks.java

@@ -1,23 +1,24 @@
 /**
  * Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
- *  This source code is licensed under both the GPLv2 (found in the
- *  COPYING file in the root directory) and Apache 2.0 License
- *  (found in the LICENSE.Apache file in the root directory).
+ * This source code is licensed under both the GPLv2 (found in the
+ * COPYING file in the root directory) and Apache 2.0 License
+ * (found in the LICENSE.Apache file in the root directory).
  */
 package org.rocksdb.jmh;
 
-import org.openjdk.jmh.annotations.*;
-import org.rocksdb.*;
-import org.rocksdb.util.FileUtils;
+import static org.rocksdb.util.KVUtils.ba;
 
 import java.io.IOException;
+import java.nio.ByteBuffer;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicInteger;
-
-import static org.rocksdb.util.KVUtils.ba;
+import org.openjdk.jmh.annotations.*;
+import org.rocksdb.*;
+import org.rocksdb.util.FileUtils;
 
 @State(Scope.Benchmark)
 public class GetBenchmarks {
@@ -30,16 +31,25 @@ public class GetBenchmarks {
   })
   String columnFamilyTestType;
 
-  @Param("100000")
+  @Param({"1000", "100000"})
   int keyCount;
 
+  @Param({"12", "64", "128"}) int keySize;
+
+  @Param({"64", "1024", "65536"}) int valueSize;
+
   Path dbDir;
   DBOptions options;
+  ReadOptions readOptions;
   int cfs = 0;  // number of column families
   private AtomicInteger cfHandlesIdx;
   ColumnFamilyHandle[] cfHandles;
   RocksDB db;
   private final AtomicInteger keyIndex = new AtomicInteger();
+  private ByteBuffer keyBuf;
+  private ByteBuffer valueBuf;
+  private byte[] keyArr;
+  private byte[] valueArr;
 
   @Setup(Level.Trial)
   public void setup() throws IOException, RocksDBException {
@@ -50,6 +60,7 @@ public void setup() throws IOException, RocksDBException {
     options = new DBOptions()
         .setCreateIfMissing(true)
         .setCreateMissingColumnFamilies(true);
+    readOptions = new ReadOptions();
 
     final List<ColumnFamilyDescriptor> cfDescriptors = new ArrayList<>();
     cfDescriptors.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY));
@@ -74,16 +85,32 @@ public void setup() throws IOException, RocksDBException {
     cfHandles = cfHandlesList.toArray(new ColumnFamilyHandle[0]);
 
     // store initial data for retrieving via get
-    for (int i = 0; i < cfs; i++) {
+    keyArr = new byte[keySize];
+    valueArr = new byte[valueSize];
+    Arrays.fill(keyArr, (byte) 0x30);
+    Arrays.fill(valueArr, (byte) 0x30);
+    for (int i = 0; i <= cfs; i++) {
       for (int j = 0; j < keyCount; j++) {
-        db.put(cfHandles[i], ba("key" + j), ba("value" + j));
+        final byte[] keyPrefix = ba("key" + j);
+        final byte[] valuePrefix = ba("value" + j);
+        System.arraycopy(keyPrefix, 0, keyArr, 0, keyPrefix.length);
+        System.arraycopy(valuePrefix, 0, valueArr, 0, valuePrefix.length);
+        db.put(cfHandles[i], keyArr, valueArr);
       }
     }
 
-    try (final FlushOptions flushOptions = new FlushOptions()
-            .setWaitForFlush(true)) {
+    try (final FlushOptions flushOptions = new FlushOptions().setWaitForFlush(true)) {
       db.flush(flushOptions);
     }
+
+    keyBuf = ByteBuffer.allocateDirect(keySize);
+    valueBuf = ByteBuffer.allocateDirect(valueSize);
+    Arrays.fill(keyArr, (byte) 0x30);
+    Arrays.fill(valueArr, (byte) 0x30);
+    keyBuf.put(keyArr);
+    keyBuf.flip();
+    valueBuf.put(valueArr);
+    valueBuf.flip();
   }
 
   @TearDown(Level.Trial)
@@ -93,13 +120,14 @@ public void cleanup() throws IOException {
     }
     db.close();
     options.close();
+    readOptions.close();
     FileUtils.delete(dbDir);
   }
 
   private ColumnFamilyHandle getColumnFamily() {
     if (cfs == 0) {
       return cfHandles[0];
-    }  else if (cfs == 1) {
+    } else if (cfs == 1) {
       return cfHandles[1];
     } else {
       int idx = cfHandlesIdx.getAndIncrement();
@@ -131,9 +159,58 @@ private int next() {
     return idx;
   }
 
-  @Benchmark
-  public byte[] get() throws RocksDBException {
+  // String -> byte[]
+  private byte[] getKeyArr() {
+    final int MAX_LEN = 9; // key100000
+    final int keyIdx = next();
+    final byte[] keyPrefix = ba("key" + keyIdx);
+    System.arraycopy(keyPrefix, 0, keyArr, 0, keyPrefix.length);
+    Arrays.fill(keyArr, keyPrefix.length, MAX_LEN, (byte) 0x30);
+    return keyArr;
+  }
+
+  // String -> ByteBuffer
+  private ByteBuffer getKeyBuf() {
+    final int MAX_LEN = 9; // key100000
     final int keyIdx = next();
-    return db.get(getColumnFamily(), ba("key" + keyIdx));
+    final String keyStr = "key" + keyIdx;
+    for (int i = 0; i < keyStr.length(); ++i) {
+      keyBuf.put(i, (byte) keyStr.charAt(i));
+    }
+    for (int i = keyStr.length(); i < MAX_LEN; ++i) {
+      keyBuf.put(i, (byte) 0x30);
+    }
+    // Reset position for future reading
+    keyBuf.position(0);
+    return keyBuf;
+  }
+
+  private byte[] getValueArr() {
+    return valueArr;
+  }
+
+  private ByteBuffer getValueBuf() {
+    return valueBuf;
+  }
+
+  @Benchmark
+  public void get() throws RocksDBException {
+    db.get(getColumnFamily(), getKeyArr());
+  }
+
+  @Benchmark
+  public void preallocatedGet() throws RocksDBException {
+    db.get(getColumnFamily(), getKeyArr(), getValueArr());
+  }
+
+  @Benchmark
+  public void preallocatedByteBufferGet() throws RocksDBException {
+    int res = db.get(getColumnFamily(), readOptions, getKeyBuf(), getValueBuf());
+    // For testing correctness:
+    //    assert res > 0;
+    //    final byte[] ret = new byte[valueSize];
+    //    valueBuf.get(ret);
+    //    System.out.println(str(ret));
+    //    valueBuf.flip();
   }
-}
+}