Plumb WriteBufferManager through JNI (#4492)

Summary:
Allow rocks java to explicitly create WriteBufferManager by plumbing it to the native code through JNI.
Pull Request resolved: https://github.com/facebook/rocksdb/pull/4492

Differential Revision: D10428506

Pulled By: sagar0

fbshipit-source-id: cd9dd8c2ef745a0303416b44e2080547bdcca1fd

java/CMakeLists.txt

@@ -58,6 +58,7 @@ set(JNI_NATIVE_SOURCES
         rocksjni/writebatchhandlerjnicallback.cc
         rocksjni/write_batch_test.cc
         rocksjni/write_batch_with_index.cc
+        rocksjni/write_buffer_manager.cc
 )
 
 set(NATIVE_JAVA_CLASSES
@@ -145,6 +146,7 @@ set(NATIVE_JAVA_CLASSES
         org.rocksdb.SnapshotTest
         org.rocksdb.WriteBatchTest
         org.rocksdb.WriteBatchTestInternalHelper
+        org.rocksdb.WriteBufferManager
 )
 
 include(FindJava)
@@ -283,6 +285,7 @@ add_jar(
   src/main/java/org/rocksdb/WriteBatch.java
   src/main/java/org/rocksdb/WriteBatchWithIndex.java
   src/main/java/org/rocksdb/WriteOptions.java
+  src/main/java/org/rocksdb/WriteBufferManager.java
   src/main/java/org/rocksdb/util/BytewiseComparator.java
   src/main/java/org/rocksdb/util/DirectBytewiseComparator.java
   src/main/java/org/rocksdb/util/Environment.java

java/Makefile

@@ -67,6 +67,7 @@ NATIVE_JAVA_CLASSES = org.rocksdb.AbstractCompactionFilter\
 	org.rocksdb.WriteBatch.Handler\
 	org.rocksdb.WriteOptions\
 	org.rocksdb.WriteBatchWithIndex\
+	org.rocksdb.WriteBufferManager\
 	org.rocksdb.WBWIRocksIterator
 
 NATIVE_JAVA_TEST_CLASSES = org.rocksdb.RocksDBExceptionTest\

java/rocksjni/options.cc

@@ -250,6 +250,20 @@ void Java_org_rocksdb_Options_setWriteBufferSize(JNIEnv* env, jobject /*jobj*/,
   }
 }
 
+/*
+ * Class:     org_rocksdb_Options
+ * Method:    setWriteBufferManager
+ * Signature: (JJ)V
+ */
+void Java_org_rocksdb_Options_setWriteBufferManager(JNIEnv* /*env*/, jobject /*jobj*/,
+                                                    jlong joptions_handle,
+                                                    jlong jwrite_buffer_manager_handle) {
+  auto* write_buffer_manager =
+          reinterpret_cast<std::shared_ptr<rocksdb::WriteBufferManager> *>(jwrite_buffer_manager_handle);
+  reinterpret_cast<rocksdb::Options*>(joptions_handle)->write_buffer_manager =
+          *write_buffer_manager;
+}
+
 /*
  * Class:     org_rocksdb_Options
  * Method:    writeBufferSize

java/rocksjni/write_buffer_manager.cc

@@ -0,0 +1,38 @@
+// 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).
+
+#include <jni.h>
+
+#include "include/org_rocksdb_WriteBufferManager.h"
+
+#include "rocksdb/cache.h"
+#include "rocksdb/write_buffer_manager.h"
+
+/*
+ * Class:     org_rocksdb_WriteBufferManager
+ * Method:    newWriteBufferManager
+ * Signature: (JJ)J
+ */
+jlong Java_org_rocksdb_WriteBufferManager_newWriteBufferManager(
+        JNIEnv* /*env*/, jclass /*jclazz*/, jlong jbuffer_size, jlong jcache_handle) {
+    auto* cache_ptr =
+            reinterpret_cast<std::shared_ptr<rocksdb::Cache> *>(jcache_handle);
+    auto* write_buffer_manager = new std::shared_ptr<rocksdb::WriteBufferManager>(
+            std::make_shared<rocksdb::WriteBufferManager>(jbuffer_size, *cache_ptr));
+    return reinterpret_cast<jlong>(write_buffer_manager);
+}
+
+/*
+ * Class:     org_rocksdb_WriteBufferManager
+ * Method:    disposeInternal
+ * Signature: (J)V
+ */
+void Java_org_rocksdb_WriteBufferManager_disposeInternal(
+        JNIEnv* /*env*/, jobject /*jobj*/, jlong jhandle) {
+    auto* write_buffer_manager =
+            reinterpret_cast<std::shared_ptr<rocksdb::WriteBufferManager> *>(jhandle);
+    assert(write_buffer_manager != nullptr);
+    delete write_buffer_manager;
+}

java/src/main/java/org/rocksdb/DBOptions.java

@@ -667,6 +667,13 @@ public class DBOptions
     return this;
   }
 
+  @Override
+  public DBOptions setWriteBufferManager(final WriteBufferManager writeBufferManager) {
+    assert(isOwningHandle());
+    setWriteBufferManager(nativeHandle_, writeBufferManager.nativeHandle_);
+    return this;
+  }
+
   @Override
   public long dbWriteBufferSize() {
     assert(isOwningHandle());
@@ -1087,6 +1094,8 @@ public class DBOptions
   private native boolean adviseRandomOnOpen(long handle);
   private native void setDbWriteBufferSize(final long handle,
       final long dbWriteBufferSize);
+  private native void setWriteBufferManager(final long dbOptionsHandle,
+      final long writeBufferManagerHandle);
   private native long dbWriteBufferSize(final long handle);
   private native void setAccessHintOnCompactionStart(final long handle,
       final byte accessHintOnCompactionStart);

java/src/main/java/org/rocksdb/DBOptionsInterface.java

@@ -991,6 +991,19 @@ public interface DBOptionsInterface<T extends DBOptionsInterface> {
    */
   T setDbWriteBufferSize(long dbWriteBufferSize);
 
+  /**
+   * Use passed {@link WriteBufferManager} to control memory usage across
+   * multiple column families and/or DB instances.
+   *
+   * Check <a href="https://github.com/facebook/rocksdb/wiki/Write-Buffer-Manager">
+   *     https://github.com/facebook/rocksdb/wiki/Write-Buffer-Manager</a>
+   * for more details on when to use it
+   *
+   * @param writeBufferManager The WriteBufferManager to use
+   * @return the reference of the current options.
+   */
+  T setWriteBufferManager(final WriteBufferManager writeBufferManager);
+
   /**
    * Amount of data to build up in memtables across all column
    * families before writing to disk.

java/src/main/java/org/rocksdb/Options.java

@@ -723,6 +723,13 @@ public class Options extends RocksObject
     return this;
   }
 
+  @Override
+  public Options setWriteBufferManager(final WriteBufferManager writeBufferManager) {
+    assert(isOwningHandle());
+    setWriteBufferManager(nativeHandle_, writeBufferManager.nativeHandle_);
+    return this;
+  }
+
   @Override
   public long dbWriteBufferSize() {
     assert(isOwningHandle());
@@ -1690,6 +1697,8 @@ public class Options extends RocksObject
   private native boolean adviseRandomOnOpen(long handle);
   private native void setDbWriteBufferSize(final long handle,
       final long dbWriteBufferSize);
+  private native void setWriteBufferManager(final long handle,
+      final long writeBufferManagerHandle);
   private native long dbWriteBufferSize(final long handle);
   private native void setAccessHintOnCompactionStart(final long handle,
       final byte accessHintOnCompactionStart);

java/src/main/java/org/rocksdb/WriteBufferManager.java

@@ -0,0 +1,30 @@
+package org.rocksdb;
+
+import org.rocksdb.Cache;
+
+/**
+ * Java wrapper over native write_buffer_manager class
+ */
+public class WriteBufferManager extends RocksObject {
+  static {
+    RocksDB.loadLibrary();
+  }
+
+  /**
+   * Construct a new instance of WriteBufferManager.
+   *
+   * Check <a href="https://github.com/facebook/rocksdb/wiki/Write-Buffer-Manager">
+   *     https://github.com/facebook/rocksdb/wiki/Write-Buffer-Manager</a>
+   * for more details on when to use it
+   *
+   * @param bufferSizeBytes buffer size(in bytes) to use for native write_buffer_manager
+   * @param cache cache whose memory should be bounded by this write buffer manager
+   */
+  public WriteBufferManager(final long bufferSizeBytes, final Cache cache){
+    super(newWriteBufferManager(bufferSizeBytes, cache.nativeHandle_));
+  }
+
+  private native static long newWriteBufferManager(final long bufferSizeBytes, final long cacheHandle);
+  @Override
+  protected native void disposeInternal(final long handle);
+}

java/src/test/java/org/rocksdb/DBOptionsTest.java

@@ -424,6 +424,24 @@ public class DBOptionsTest {
     }
   }
 
+  @Test
+  public void setWriteBufferManager() throws RocksDBException {
+    try (final Options opt = new Options();
+         final Cache cache = new LRUCache(1 * 1024 * 1024);
+         final WriteBufferManager writeBufferManager = new WriteBufferManager(2000l, cache)) {
+      opt.setWriteBufferManager(writeBufferManager);
+    }
+  }
+
+  @Test
+  public void setWriteBufferManagerWithZeroBufferSize() throws RocksDBException {
+    try (final Options opt = new Options();
+         final Cache cache = new LRUCache(1 * 1024 * 1024);
+         final WriteBufferManager writeBufferManager = new WriteBufferManager(0l, cache)) {
+      opt.setWriteBufferManager(writeBufferManager);
+    }
+  }
+
   @Test
   public void accessHintOnCompactionStart() {
     try(final DBOptions opt = new DBOptions()) {

java/src/test/java/org/rocksdb/OptionsTest.java

@@ -645,6 +645,24 @@ public class OptionsTest {
     }
   }
 
+  @Test
+  public void setWriteBufferManager() throws RocksDBException {
+    try (final Options opt = new Options();
+         final Cache cache = new LRUCache(1 * 1024 * 1024);
+         final WriteBufferManager writeBufferManager = new WriteBufferManager(2000l, cache)) {
+      opt.setWriteBufferManager(writeBufferManager);
+    }
+  }
+
+  @Test
+  public void setWriteBufferManagerWithZeroBufferSize() throws RocksDBException {
+    try (final Options opt = new Options();
+         final Cache cache = new LRUCache(1 * 1024 * 1024);
+         final WriteBufferManager writeBufferManager = new WriteBufferManager(0l, cache)) {
+      opt.setWriteBufferManager(writeBufferManager);
+    }
+  }
+
   @Test
   public void accessHintOnCompactionStart() {
     try (final Options opt = new Options()) {

src.mk

@@ -470,4 +470,5 @@ JNI_NATIVE_SOURCES =                                          \
   java/rocksjni/write_batch.cc                                \
   java/rocksjni/writebatchhandlerjnicallback.cc               \
   java/rocksjni/write_batch_test.cc                           \
-  java/rocksjni/write_batch_with_index.cc
+  java/rocksjni/write_batch_with_index.cc                     \
+  java/rocksjni/write_buffer_manager.cc