JNI get_helper code sharing / multiGet() use efficient batch C++ support [replaced by https://github.com/facebook/rocksdb/pull/12344]

include/rocksdb/utilities/transaction.h

@@ -417,6 +417,23 @@ class Transaction {
                               std::string* value, bool exclusive = true,
                               const bool do_validate = true) = 0;
 
+  // An overload of the above method that receives a PinnableSlice
+  // For backward compatibility a default implementation is provided
+  virtual Status GetForUpdate(const ReadOptions& options, const Slice& key,
+                              PinnableSlice* pinnable_val,
+                              bool exclusive = true,
+                              const bool do_validate = true) {
+    if (pinnable_val == nullptr) {
+      std::string* null_str = nullptr;
+      return GetForUpdate(options, key, null_str, exclusive, do_validate);
+    } else {
+      auto s = GetForUpdate(options, key, pinnable_val->GetSelf(), exclusive,
+                            do_validate);
+      pinnable_val->PinSelf();
+      return s;
+    }
+  }
+
   virtual std::vector<Status> MultiGetForUpdate(
       const ReadOptions& options,
       const std::vector<ColumnFamilyHandle*>& column_family,

java/CMakeLists.txt

@@ -46,6 +46,7 @@ set(JNI_NATIVE_SOURCES
         rocksjni/hyper_clock_cache.cc
         rocksjni/ingest_external_file_options.cc
         rocksjni/iterator.cc
+        rocksjni/jni_multiget_helpers.cc
         rocksjni/jnicallback.cc
         rocksjni/loggerjnicallback.cc
         rocksjni/lru_cache.cc

java/jmh/README.md

@@ -6,10 +6,10 @@ 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.
 
-For instance, this is how to install the OSX jar you just built for 6.26.0
+For instance, this is how to install the OSX jar you just built for 8.11.0
 
 ```bash
-$ mvn install:install-file -Dfile=./java/target/rocksdbjni-6.26.0-SNAPSHOT-osx.jar -DgroupId=org.rocksdb -DartifactId=rocksdbjni -Dversion=6.26.0-SNAPSHOT -Dpackaging=jar
+$ mvn install:install-file -Dfile=./java/target/rocksdbjni-8.11.0-SNAPSHOT-osx.jar -DgroupId=org.rocksdb -DartifactId=rocksdbjni -Dversion=8.11.0-SNAPSHOT -Dpackaging=jar
 ```
 
 ```bash

java/jmh/pom.xml

@@ -38,8 +38,8 @@
     </scm>
 
     <properties>
-        <project.build.source>1.7</project.build.source>
-        <project.build.target>1.7</project.build.target>
+        <project.build.source>17</project.build.source>
+        <project.build.target>17</project.build.target>
         <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
 
         <jmh.version>1.22</jmh.version>
@@ -50,7 +50,7 @@
         <dependency>
             <groupId>org.rocksdb</groupId>
             <artifactId>rocksdbjni</artifactId>
-            <version>7.9.0-SNAPSHOT</version>
+            <version>8.11.0</version>
         </dependency>
 
         <dependency>

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

@@ -55,6 +55,9 @@ public class MultiGetBenchmarks {
   RocksDB db;
   private final AtomicInteger keyIndex = new AtomicInteger();
 
+  private List<ColumnFamilyHandle> defaultCFHandles = new ArrayList<>();
+  private List<ColumnFamilyHandle> randomCFHandles = new ArrayList<>();
+
   @Setup(Level.Trial)
   public void setup() throws IOException, RocksDBException {
     RocksDB.loadLibrary();
@@ -88,13 +91,30 @@ public void setup() throws IOException, RocksDBException {
     cfHandles = cfHandlesList.toArray(new ColumnFamilyHandle[0]);
 
     // store initial data for retrieving via get
+    for (int j = 0; j < keyCount; j++) {
+      final byte[] paddedValue = Arrays.copyOf(ba("value" + j), valueSize);
+      db.put(ba("key" + j), paddedValue);
+    }
+
+    // store initial data for retrieving via get - column families
     for (int i = 0; i < cfs; i++) {
       for (int j = 0; j < keyCount; j++) {
         final byte[] paddedValue = Arrays.copyOf(ba("value" + j), valueSize);
         db.put(cfHandles[i], ba("key" + j), paddedValue);
       }
     }
 
+    // build a big list of default column families for efficient passing
+    final ColumnFamilyHandle defaultCFH = db.getDefaultColumnFamily();
+    for (int i = 0; i < keyCount; i++) {
+      defaultCFHandles.add(defaultCFH);
+    }
+
+    // list of random cfs
+    for (int i = 0; i < keyCount; i++) {
+      randomCFHandles.add(cfHandlesList.get((int) (Math.random() * cfs)));
+    }
+
     try (final FlushOptions flushOptions = new FlushOptions()
             .setWaitForFlush(true)) {
       db.flush(flushOptions);
@@ -163,15 +183,13 @@ private int next(final int inc, final int limit) {
 
   @Setup
   public void allocateSliceBuffers() {
-    keysBuffer = ByteBuffer.allocateDirect(keyCount * valueSize);
+    keysBuffer = ByteBuffer.allocateDirect(keyCount * keySize);
     valuesBuffer = ByteBuffer.allocateDirect(keyCount * valueSize);
     valueBuffersList = new ArrayList<>();
     keyBuffersList = new ArrayList<>();
     for (int i = 0; i < keyCount; i++) {
-      valueBuffersList.add(valuesBuffer.slice());
-      valuesBuffer.position(i * valueSize);
-      keyBuffersList.add(keysBuffer.slice());
-      keysBuffer.position(i * keySize);
+      valueBuffersList.add(valuesBuffer.slice(i * valueSize, valueSize));
+      keyBuffersList.add(keysBuffer.slice(i * keySize, keySize));
     }
   }
 
@@ -181,7 +199,7 @@ public void freeSliceBuffers() {
   }
 
   @Benchmark
-  public List<byte[]> multiGet10() throws RocksDBException {
+  public void multiGetList10() throws RocksDBException {
     final int fromKeyIdx = next(multiGetSize, keyCount);
     if (fromKeyIdx >= 0) {
       final List<byte[]> keys = keys(fromKeyIdx, fromKeyIdx + multiGetSize);
@@ -191,6 +209,53 @@ public List<byte[]> multiGet10() throws RocksDBException {
           throw new RuntimeException("Test valueSize assumption wrong");
       }
     }
+  }
+
+  @Benchmark
+  public void multiGetListExplicitCF20() throws RocksDBException {
+    final int fromKeyIdx = next(multiGetSize, keyCount);
+    if (fromKeyIdx >= 0) {
+      final List<byte[]> keys = keys(fromKeyIdx, fromKeyIdx + multiGetSize);
+      final List<ColumnFamilyHandle> columnFamilyHandles =
+          defaultCFHandles.subList(fromKeyIdx, fromKeyIdx + multiGetSize);
+      final List<byte[]> valueResults = db.multiGetAsList(columnFamilyHandles, keys);
+      for (final byte[] result : valueResults) {
+        if (result.length != valueSize)
+          throw new RuntimeException("Test valueSize assumption wrong");
+      }
+    }
+  }
+
+  @Benchmark
+  public void multiGetListRandomCF30() throws RocksDBException {
+    final int fromKeyIdx = next(multiGetSize, keyCount);
+    if (fromKeyIdx >= 0) {
+      final List<byte[]> keys = keys(fromKeyIdx, fromKeyIdx + multiGetSize);
+      final List<ColumnFamilyHandle> columnFamilyHandles =
+          randomCFHandles.subList(fromKeyIdx, fromKeyIdx + multiGetSize);
+      final List<byte[]> valueResults = db.multiGetAsList(columnFamilyHandles, keys);
+      for (final byte[] result : valueResults) {
+        if (result.length != valueSize)
+          throw new RuntimeException("Test valueSize assumption wrong");
+      }
+    }
+  }
+
+  @Benchmark
+  public List<byte[]> multiGetBB200() throws RocksDBException {
+    final int fromKeyIdx = next(multiGetSize, keyCount);
+    if (fromKeyIdx >= 0) {
+      final List<ByteBuffer> keys = keys(keyBuffersList, fromKeyIdx, fromKeyIdx + multiGetSize);
+      final List<ByteBuffer> values =
+          valueBuffersList.subList(fromKeyIdx, fromKeyIdx + multiGetSize);
+      final List<ByteBufferGetStatus> statusResults = db.multiGetByteBuffers(keys, values);
+      for (final ByteBufferGetStatus result : statusResults) {
+        if (result.status.getCode() != Status.Code.Ok)
+          throw new RuntimeException("Test status not OK: " + result.status);
+        if (result.value.limit() != valueSize)
+          throw new RuntimeException("Test valueSize assumption wrong");
+      }
+    }
     return new ArrayList<>();
   }