Introduce allowStall option for write buffer manager constructor (#9076)

Summary:
https://github.com/facebook/rocksdb/pull/7898 enable write buffer manager to stall write when memory_usage exceeds buffer_size, this is really useful for container running case to limit the memory usage. However, this feature is not visiable for rocksJava yet.

This PR targets to introduce this feature for rocksJava.

Pull Request resolved: https://github.com/facebook/rocksdb/pull/9076

Reviewed By: akankshamahajan15

Differential Revision: D31931092

Pulled By: anand1976

fbshipit-source-id: 5531c16a87598663a02368c07b5e13a503164578

java/rocksjni/write_buffer_manager.cc

@@ -16,14 +16,15 @@
  * Signature: (JJ)J
  */
 jlong Java_org_rocksdb_WriteBufferManager_newWriteBufferManager(
-        JNIEnv* /*env*/, jclass /*jclazz*/, jlong jbuffer_size, jlong jcache_handle) {
+    JNIEnv* /*env*/, jclass /*jclazz*/, jlong jbuffer_size, jlong jcache_handle,
+    jboolean allow_stall) {
   auto* cache_ptr =
       reinterpret_cast<std::shared_ptr<ROCKSDB_NAMESPACE::Cache>*>(
           jcache_handle);
   auto* write_buffer_manager =
       new std::shared_ptr<ROCKSDB_NAMESPACE::WriteBufferManager>(
-          std::make_shared<ROCKSDB_NAMESPACE::WriteBufferManager>(jbuffer_size,
-                                                                  *cache_ptr));
+          std::make_shared<ROCKSDB_NAMESPACE::WriteBufferManager>(
+              jbuffer_size, *cache_ptr, allow_stall));
   return reinterpret_cast<jlong>(write_buffer_manager);
 }
 

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

@@ -22,12 +22,29 @@ public class WriteBufferManager extends RocksObject {
    *
    * @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
+   * @param allowStall if set true, it will enable stalling of writes when memory_usage() exceeds
+   *     buffer_size.
+   *        It will wait for flush to complete and memory usage to drop down.
    */
+  public WriteBufferManager(
+      final long bufferSizeBytes, final Cache cache, final boolean allowStall) {
+    super(newWriteBufferManager(bufferSizeBytes, cache.nativeHandle_, allowStall));
+    this.allowStall_ = allowStall;
+  }
+
   public WriteBufferManager(final long bufferSizeBytes, final Cache cache){
-    super(newWriteBufferManager(bufferSizeBytes, cache.nativeHandle_));
+    this(bufferSizeBytes, cache, false);
+  }
+
+  public boolean allowStall() {
+    return allowStall_;
   }
 
-  private native static long newWriteBufferManager(final long bufferSizeBytes, final long cacheHandle);
+  private native static long newWriteBufferManager(
+      final long bufferSizeBytes, final long cacheHandle, final boolean allowStall);
+
   @Override
   protected native void disposeInternal(final long handle);
+
+  private boolean allowStall_;
 }

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

@@ -686,6 +686,16 @@ public class OptionsTest {
     }
   }
 
+  @Test
+  public void setWriteBufferManagerWithAllowStall() throws RocksDBException {
+    try (final Options opt = new Options(); final Cache cache = new LRUCache(1 * 1024 * 1024);
+         final WriteBufferManager writeBufferManager = new WriteBufferManager(2000l, cache, true)) {
+      opt.setWriteBufferManager(writeBufferManager);
+      assertThat(opt.writeBufferManager()).isEqualTo(writeBufferManager);
+      assertThat(opt.writeBufferManager().allowStall()).isEqualTo(true);
+    }
+  }
+
   @Test
   public void accessHintOnCompactionStart() {
     try (final Options opt = new Options()) {