RocksJava: Improve 'get' performance

java/CMakeLists.txt

@@ -47,6 +47,7 @@ set(JNI_NATIVE_SOURCES
         rocksjni/options.cc
         rocksjni/options_util.cc
         rocksjni/persistent_cache.cc
+        rocksjni/pinnable_slice.cc
         rocksjni/ratelimiterjni.cc
         rocksjni/remove_emptyvalue_compactionfilterjni.cc
         rocksjni/restorejni.cc
@@ -147,6 +148,7 @@ set(JAVA_MAIN_CLASSES
   src/main/java/org/rocksdb/EventListener.java
   src/main/java/org/rocksdb/Experimental.java
   src/main/java/org/rocksdb/ExternalFileIngestionInfo.java
+  src/main/java/org/rocksdb/FastBuffer.java
   src/main/java/org/rocksdb/Filter.java
   src/main/java/org/rocksdb/FileOperationInfo.java
   src/main/java/org/rocksdb/FlushJobInfo.java
@@ -189,6 +191,7 @@ set(JAVA_MAIN_CLASSES
   src/main/java/org/rocksdb/Options.java
   src/main/java/org/rocksdb/OptionsUtil.java
   src/main/java/org/rocksdb/PersistentCache.java
+  src/main/java/org/rocksdb/PinnableSlice.java
   src/main/java/org/rocksdb/PlainTableConfig.java
   src/main/java/org/rocksdb/Priority.java
   src/main/java/org/rocksdb/Range.java

java/jmh/README.md

@@ -6,6 +6,17 @@ These are micro-benchmarks for RocksJava functionality, using [JMH (Java Microbe
 
 **Note**: This uses a specific build of RocksDB that is set in the `<version>` element of the `dependencies` section of the `pom.xml` file. If you are testing local changes you should build and install a SNAPSHOT version of rocksdbjni, and update the `pom.xml` of rocksdbjni-jmh file to test with this.
 
+RocksJava jar can be installed to local Maven repository with command similar to
+```bash
+$ mvn install:install-file \
+   -Dfile=/path/to/target/rocksdbjni-[version]-[platform].jar \
+   -DgroupId=org.rocksdb \
+   -DartifactId=rocksdbjni \
+   -Dversion=[version]-SNAPSHOT \
+   -Dpackaging=jar \
+   -DgeneratePom=true
+```
+
 ```bash
 $ mvn package
 ```

java/jmh/pom.xml

@@ -50,7 +50,7 @@
         <dependency>
             <groupId>org.rocksdb</groupId>
             <artifactId>rocksdbjni</artifactId>
-            <version>6.6.0-SNAPSHOT</version>
+            <version>6.14.0-SNAPSHOT</version>
         </dependency>
 
         <dependency>

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 {
@@ -33,13 +34,22 @@ public class GetBenchmarks {
   @Param("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,70 @@ 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();
-    return db.get(getColumnFamily(), ba("key" + keyIdx));
+    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();
+    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 preallocatedCriticalGet() throws RocksDBException {
+    db.getCritical(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();
+  }
+
+  @Benchmark
+  public void unsafeGet() throws RocksDBException {
+    PinnableSlice pinnableSlice = db.getUnsafe(getColumnFamily(), readOptions, getKeyArr());
+    // assert pinnableSlice.capacity() > 0;
+    pinnableSlice.close();
   }
 }

java/rocksjni/pinnable_slice.cc

@@ -0,0 +1,16 @@
+#include "include/org_rocksdb_PinnableSlice.h"
+#include "rocksdb/slice.h"
+
+/*
+ * Class:     org_rocksdb_PinnableSlice
+ * Method:    deletePinnableSlice
+ * Signature: (J)V
+ */
+void Java_org_rocksdb_PinnableSlice_deletePinnableSlice(JNIEnv*, jclass,
+                                                        jlong handle) {
+  auto* pinnable_slice =
+      reinterpret_cast<ROCKSDB_NAMESPACE::PinnableSlice*>(handle);
+  assert(pinnable_slice != nullptr);
+  pinnable_slice->Reset();
+  delete pinnable_slice;
+}