WriteBufferManager JNI fixes (#4579)

Summary:
1. `WriteBufferManager` should have a reference alive in Java side through `Options`/`DBOptions` otherwise, if it's GC'ed at java side, native side can seg fault.
2. native method `setWriteBufferManager()` in `DBOptions.java` doesn't have it's jni method invocation in rocksdbjni which is added in this PR
3. `DBOptionsTest.java` is referencing object of `Options`. Instead it should be testing against `DBOptions`. Seems like a copy paste error.
4. Add a getter for WriteBufferManager.
Pull Request resolved: #4579

Differential Revision: D10561150

Pulled By: sagar0

fbshipit-source-id: 139a15c7f051a9f77b4200215b88267b48fbc487

java/rocksjni/options.cc

@@ -5532,6 +5532,20 @@ void Java_org_rocksdb_DBOptions_setDbWriteBufferSize(
   opt->db_write_buffer_size = static_cast<size_t>(jdb_write_buffer_size);
 }
 
+/*
+ * Class:     org_rocksdb_DBOptions
+ * Method:    setWriteBufferManager
+ * Signature: (JJ)V
+ */
+void Java_org_rocksdb_DBOptions_setWriteBufferManager(JNIEnv* /*env*/, jobject /*jobj*/,
+                                                      jlong jdb_options_handle,
+                                                      jlong jwrite_buffer_manager_handle) {
+  auto* write_buffer_manager =
+      reinterpret_cast<std::shared_ptr<rocksdb::WriteBufferManager> *>(jwrite_buffer_manager_handle);
+  reinterpret_cast<rocksdb::DBOptions*>(jdb_options_handle)->write_buffer_manager =
+      *write_buffer_manager;
+}
+
 /*
  * Class:     org_rocksdb_DBOptions
  * Method:    dbWriteBufferSize

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

@@ -46,6 +46,7 @@ public DBOptions(DBOptions other) {
     this.numShardBits_ = other.numShardBits_;
     this.rateLimiter_ = other.rateLimiter_;
     this.rowCache_ = other.rowCache_;
+    this.writeBufferManager_ = other.writeBufferManager_;
   }
 
   /**
@@ -671,10 +672,17 @@ public DBOptions setDbWriteBufferSize(final long dbWriteBufferSize) {
   public DBOptions setWriteBufferManager(final WriteBufferManager writeBufferManager) {
     assert(isOwningHandle());
     setWriteBufferManager(nativeHandle_, writeBufferManager.nativeHandle_);
+    this.writeBufferManager_ = writeBufferManager;
     return this;
   }
 
   @Override
+  public WriteBufferManager writeBufferManager() {
+    assert(isOwningHandle());
+    return this.writeBufferManager_;
+  }
+
+    @Override
   public long dbWriteBufferSize() {
     assert(isOwningHandle());
     return dbWriteBufferSize(nativeHandle_);
@@ -1167,4 +1175,5 @@ private native void setAvoidFlushDuringShutdown(final long handle,
   private int numShardBits_;
   private RateLimiter rateLimiter_;
   private Cache rowCache_;
+  private WriteBufferManager writeBufferManager_;
 }

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

@@ -1004,6 +1004,15 @@ public interface DBOptionsInterface<T extends DBOptionsInterface> {
    */
   T setWriteBufferManager(final WriteBufferManager writeBufferManager);
 
+  /**
+   * Reference to {@link WriteBufferManager} used by it. <br>
+   *
+   * Default: null (Disabled)
+   *
+   * @return a reference to WriteBufferManager
+   */
+  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

@@ -70,6 +70,7 @@ public Options(Options other) {
     this.compactionOptionsFIFO_ = other.compactionOptionsFIFO_;
     this.compressionOptions_ = other.compressionOptions_;
     this.rowCache_ = other.rowCache_;
+    this.writeBufferManager_ = other.writeBufferManager_;
   }
 
   @Override
@@ -727,10 +728,17 @@ public Options setDbWriteBufferSize(final long dbWriteBufferSize) {
   public Options setWriteBufferManager(final WriteBufferManager writeBufferManager) {
     assert(isOwningHandle());
     setWriteBufferManager(nativeHandle_, writeBufferManager.nativeHandle_);
+    this.writeBufferManager_ = writeBufferManager;
     return this;
   }
 
   @Override
+  public WriteBufferManager writeBufferManager() {
+    assert(isOwningHandle());
+    return this.writeBufferManager_;
+  }
+
+    @Override
   public long dbWriteBufferSize() {
     assert(isOwningHandle());
     return dbWriteBufferSize(nativeHandle_);
@@ -1918,4 +1926,5 @@ private native void setForceConsistencyChecks(final long handle,
   private CompactionOptionsFIFO compactionOptionsFIFO_;
   private CompressionOptions compressionOptions_;
   private Cache rowCache_;
+  private WriteBufferManager writeBufferManager_;
 }

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

@@ -426,19 +426,21 @@ public void dbWriteBufferSize() {
 
   @Test
   public void setWriteBufferManager() throws RocksDBException {
-    try (final Options opt = new Options();
+    try (final DBOptions opt = new DBOptions();
          final Cache cache = new LRUCache(1 * 1024 * 1024);
          final WriteBufferManager writeBufferManager = new WriteBufferManager(2000l, cache)) {
       opt.setWriteBufferManager(writeBufferManager);
+      assertThat(opt.writeBufferManager()).isEqualTo(writeBufferManager);
     }
   }
 
   @Test
   public void setWriteBufferManagerWithZeroBufferSize() throws RocksDBException {
-    try (final Options opt = new Options();
+    try (final DBOptions opt = new DBOptions();
          final Cache cache = new LRUCache(1 * 1024 * 1024);
          final WriteBufferManager writeBufferManager = new WriteBufferManager(0l, cache)) {
       opt.setWriteBufferManager(writeBufferManager);
+      assertThat(opt.writeBufferManager()).isEqualTo(writeBufferManager);
     }
   }
 

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

@@ -651,6 +651,7 @@ public void setWriteBufferManager() throws RocksDBException {
          final Cache cache = new LRUCache(1 * 1024 * 1024);
          final WriteBufferManager writeBufferManager = new WriteBufferManager(2000l, cache)) {
       opt.setWriteBufferManager(writeBufferManager);
+      assertThat(opt.writeBufferManager()).isEqualTo(writeBufferManager);
     }
   }
 
@@ -660,6 +661,7 @@ public void setWriteBufferManagerWithZeroBufferSize() throws RocksDBException {
          final Cache cache = new LRUCache(1 * 1024 * 1024);
          final WriteBufferManager writeBufferManager = new WriteBufferManager(0l, cache)) {
       opt.setWriteBufferManager(writeBufferManager);
+      assertThat(opt.writeBufferManager()).isEqualTo(writeBufferManager);
     }
   }