From dc00e4b1206cba9995fe042ec976df54d0f6e76f Mon Sep 17 00:00:00 2001
From: myasuka <myasuka@live.com>
Date: Tue, 26 Oct 2021 12:08:54 -0700
Subject: [PATCH] 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         |  7 ++++---
 .../java/org/rocksdb/WriteBufferManager.java  | 21 +++++++++++++++++--
 .../test/java/org/rocksdb/OptionsTest.java    | 10 +++++++++
 3 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/java/rocksjni/write_buffer_manager.cc b/java/rocksjni/write_buffer_manager.cc
index a52daed7d..e5c1e23f3 100644
--- a/java/rocksjni/write_buffer_manager.cc
+++ b/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);
 }
 
diff --git a/java/src/main/java/org/rocksdb/WriteBufferManager.java b/java/src/main/java/org/rocksdb/WriteBufferManager.java
index b244aa952..8ec963958 100644
--- a/java/src/main/java/org/rocksdb/WriteBufferManager.java
+++ b/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_;
 }
diff --git a/java/src/test/java/org/rocksdb/OptionsTest.java b/java/src/test/java/org/rocksdb/OptionsTest.java
index 5d732adc8..1bb055042 100644
--- a/java/src/test/java/org/rocksdb/OptionsTest.java
+++ b/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()) {