Add autotune and #getBytesPerSecond() to RocksJava RateLimiter

Summary: Closes https://github.com/facebook/rocksdb/pull/3332

Differential Revision: D6667680

Pulled By: ajkr

fbshipit-source-id: b2bb6889257850a4eb6f6cbd7106f62df7b82730

3 files changed, 199 insertions, 23 deletions

diff --git a/java/src/main/java/org/rocksdb/RateLimiter.java b/java/src/main/java/org/rocksdb/RateLimiter.java
index fc2388777..4d6cb2129 100644
--- a/java/src/main/java/org/rocksdb/RateLimiter.java
+++ b/java/src/main/java/org/rocksdb/RateLimiter.java
@@ -12,8 +12,11 @@ package org.rocksdb;
  * @since 3.10.0
  */
 public class RateLimiter extends RocksObject {
-  private static final long DEFAULT_REFILL_PERIOD_MICROS = (100 * 1000);
-  private static final int DEFAULT_FAIRNESS = 10;
+  public static final long DEFAULT_REFILL_PERIOD_MICROS = 100 * 1000;
+  public static final int DEFAULT_FAIRNESS = 10;
+  public static final RateLimiterMode DEFAULT_MODE =
+      RateLimiterMode.WRITES_ONLY;
+  public static final boolean DEFAULT_AUTOTUNE = false;
 
   /**
    * RateLimiter constructor
@@ -21,24 +24,62 @@ public class RateLimiter extends RocksObject {
    * @param rateBytesPerSecond this is the only parameter you want to set
    *     most of the time. It controls the total write rate of compaction
    *     and flush in bytes per second. Currently, RocksDB does not enforce
-   *     rate limit for anything other than flush and compaction, e.g. write to WAL.
-   * @param refillPeriodMicros this controls how often tokens are refilled. For example,
+   *     rate limit for anything other than flush and compaction, e.g. write to
+   *     WAL.
+   */
+  public RateLimiter(final long rateBytesPerSecond) {
+    this(rateBytesPerSecond, DEFAULT_REFILL_PERIOD_MICROS, DEFAULT_FAIRNESS,
+        DEFAULT_MODE, DEFAULT_AUTOTUNE);
+  }
+
+  /**
+   * RateLimiter constructor
+   *
+   * @param rateBytesPerSecond this is the only parameter you want to set
+   *     most of the time. It controls the total write rate of compaction
+   *     and flush in bytes per second. Currently, RocksDB does not enforce
+   *     rate limit for anything other than flush and compaction, e.g. write to
+   *     WAL.
+   * @param refillPeriodMicros this controls how often tokens are refilled. For
+   *     example,
+   *     when rate_bytes_per_sec is set to 10MB/s and refill_period_us is set to
+   *     100ms, then 1MB is refilled every 100ms internally. Larger value can
+   *     lead to burstier writes while smaller value introduces more CPU
+   *     overhead. The default of 100,000ms should work for most cases.
+   */
+  public RateLimiter(final long rateBytesPerSecond,
+      final long refillPeriodMicros) {
+    this(rateBytesPerSecond, refillPeriodMicros, DEFAULT_FAIRNESS, DEFAULT_MODE,
+        DEFAULT_AUTOTUNE);
+  }
+
+  /**
+   * RateLimiter constructor
+   *
+   * @param rateBytesPerSecond this is the only parameter you want to set
+   *     most of the time. It controls the total write rate of compaction
+   *     and flush in bytes per second. Currently, RocksDB does not enforce
+   *     rate limit for anything other than flush and compaction, e.g. write to
+   *     WAL.
+   * @param refillPeriodMicros this controls how often tokens are refilled. For
+   *     example,
    *     when rate_bytes_per_sec is set to 10MB/s and refill_period_us is set to
-   *     100ms, then 1MB is refilled every 100ms internally. Larger value can lead to
-   *     burstier writes while smaller value introduces more CPU overhead.
-   *     The default should work for most cases.
+   *     100ms, then 1MB is refilled every 100ms internally. Larger value can
+   *     lead to burstier writes while smaller value introduces more CPU
+   *     overhead. The default of 100,000ms should work for most cases.
    * @param fairness RateLimiter accepts high-pri requests and low-pri requests.
-   *     A low-pri request is usually blocked in favor of hi-pri request. Currently,
-   *     RocksDB assigns low-pri to request from compaction and high-pri to request
-   *     from flush. Low-pri requests can get blocked if flush requests come in
-   *     continuously. This fairness parameter grants low-pri requests permission by
-   *     fairness chance even though high-pri requests exist to avoid starvation.
+   *     A low-pri request is usually blocked in favor of hi-pri request.
+   *     Currently, RocksDB assigns low-pri to request from compaction and
+   *     high-pri to request from flush. Low-pri requests can get blocked if
+   *     flush requests come in continuously. This fairness parameter grants
+   *     low-pri requests permission by fairness chance even though high-pri
+   *     requests exist to avoid starvation.
    *     You should be good by leaving it at default 10.
    */
   public RateLimiter(final long rateBytesPerSecond,
       final long refillPeriodMicros, final int fairness) {
-    super(newRateLimiterHandle(rateBytesPerSecond,
-        refillPeriodMicros, fairness));
+    this(rateBytesPerSecond, refillPeriodMicros, fairness, DEFAULT_MODE,
+        DEFAULT_AUTOTUNE);
   }
 
   /**
@@ -47,10 +88,65 @@ public class RateLimiter extends RocksObject {
    * @param rateBytesPerSecond this is the only parameter you want to set
    *     most of the time. It controls the total write rate of compaction
    *     and flush in bytes per second. Currently, RocksDB does not enforce
-   *     rate limit for anything other than flush and compaction, e.g. write to WAL.
+   *     rate limit for anything other than flush and compaction, e.g. write to
+   *     WAL.
+   * @param refillPeriodMicros this controls how often tokens are refilled. For
+   *     example,
+   *     when rate_bytes_per_sec is set to 10MB/s and refill_period_us is set to
+   *     100ms, then 1MB is refilled every 100ms internally. Larger value can
+   *     lead to burstier writes while smaller value introduces more CPU
+   *     overhead. The default of 100,000ms should work for most cases.
+   * @param fairness RateLimiter accepts high-pri requests and low-pri requests.
+   *     A low-pri request is usually blocked in favor of hi-pri request.
+   *     Currently, RocksDB assigns low-pri to request from compaction and
+   *     high-pri to request from flush. Low-pri requests can get blocked if
+   *     flush requests come in continuously. This fairness parameter grants
+   *     low-pri requests permission by fairness chance even though high-pri
+   *     requests exist to avoid starvation.
+   *     You should be good by leaving it at default 10.
+   * @param rateLimiterMode indicates which types of operations count against
+   *     the limit.
    */
-  public RateLimiter(final long rateBytesPerSecond) {
-    this(rateBytesPerSecond, DEFAULT_REFILL_PERIOD_MICROS, DEFAULT_FAIRNESS);
+  public RateLimiter(final long rateBytesPerSecond,
+      final long refillPeriodMicros, final int fairness,
+      final RateLimiterMode rateLimiterMode) {
+    this(rateBytesPerSecond, refillPeriodMicros, fairness, rateLimiterMode,
+        DEFAULT_AUTOTUNE);
+  }
+
+  /**
+   * RateLimiter constructor
+   *
+   * @param rateBytesPerSecond this is the only parameter you want to set
+   *     most of the time. It controls the total write rate of compaction
+   *     and flush in bytes per second. Currently, RocksDB does not enforce
+   *     rate limit for anything other than flush and compaction, e.g. write to
+   *     WAL.
+   * @param refillPeriodMicros this controls how often tokens are refilled. For
+   *     example,
+   *     when rate_bytes_per_sec is set to 10MB/s and refill_period_us is set to
+   *     100ms, then 1MB is refilled every 100ms internally. Larger value can
+   *     lead to burstier writes while smaller value introduces more CPU
+   *     overhead. The default of 100,000ms should work for most cases.
+   * @param fairness RateLimiter accepts high-pri requests and low-pri requests.
+   *     A low-pri request is usually blocked in favor of hi-pri request.
+   *     Currently, RocksDB assigns low-pri to request from compaction and
+   *     high-pri to request from flush. Low-pri requests can get blocked if
+   *     flush requests come in continuously. This fairness parameter grants
+   *     low-pri requests permission by fairness chance even though high-pri
+   *     requests exist to avoid starvation.
+   *     You should be good by leaving it at default 10.
+   * @param rateLimiterMode indicates which types of operations count against
+   *     the limit.
+   * @param autoTune Enables dynamic adjustment of rate limit within the range
+   *     {@code [rate_bytes_per_sec / 20, rate_bytes_per_sec]}, according to
+   *     the recent demand for background I/O.
+   */
+  public RateLimiter(final long rateBytesPerSecond,
+      final long refillPeriodMicros, final int fairness,
+      final RateLimiterMode rateLimiterMode, final boolean autoTune) {
+    super(newRateLimiterHandle(rateBytesPerSecond,
+        refillPeriodMicros, fairness, rateLimiterMode.getValue(), autoTune));
   }
 
   /**
@@ -65,6 +161,16 @@ public class RateLimiter extends RocksObject {
   }
 
   /**
+   * Returns the bytes per second.
+   *
+   * @return bytes per second.
+   */
+  public long getBytesPerSecond() {
+    assert(isOwningHandle());
+    return getBytesPerSecond(nativeHandle_);
+  }
+
+  /**
    * <p>Request for token to write bytes. If this request can not be satisfied,
    * the call is blocked. Caller is responsible to make sure
    * {@code bytes &lt; GetSingleBurstBytes()}.</p>
@@ -107,11 +213,13 @@ public class RateLimiter extends RocksObject {
   }
 
   private static native long newRateLimiterHandle(final long rateBytesPerSecond,
-      final long refillPeriodMicros, final int fairness);
+      final long refillPeriodMicros, final int fairness,
+      final byte rateLimiterMode, final boolean autoTune);
   @Override protected final native void disposeInternal(final long handle);
 
   private native void setBytesPerSecond(final long handle,
       final long bytesPerSecond);
+  private native long getBytesPerSecond(final long handle);
   private native void request(final long handle, final long bytes);
   private native long getSingleBurstBytes(final long handle);
   private native long getTotalBytesThrough(final long handle);
diff --git a/java/src/main/java/org/rocksdb/RateLimiterMode.java b/java/src/main/java/org/rocksdb/RateLimiterMode.java
new file mode 100644
index 000000000..4b029d816
--- /dev/null
+++ b/java/src/main/java/org/rocksdb/RateLimiterMode.java
@@ -0,0 +1,52 @@
+// 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).
+
+package org.rocksdb;
+
+/**
+ * Mode for {@link RateLimiter#RateLimiter(long, long, int, RateLimiterMode)}.
+ */
+public enum RateLimiterMode {
+  READS_ONLY((byte)0x0),
+  WRITES_ONLY((byte)0x1),
+  ALL_IO((byte)0x2);
+
+  private final byte value;
+
+  RateLimiterMode(final byte value) {
+    this.value = value;
+  }
+
+  /**
+   * <p>Returns the byte value of the enumerations value.</p>
+   *
+   * @return byte representation
+   */
+  public byte getValue() {
+    return value;
+  }
+
+  /**
+   * <p>Get the RateLimiterMode enumeration value by
+   * passing the byte identifier to this method.</p>
+   *
+   * @param byteIdentifier of RateLimiterMode.
+   *
+   * @return AccessHint instance.
+   *
+   * @throws IllegalArgumentException if the access hint for the byteIdentifier
+   *     cannot be found
+   */
+  public static RateLimiterMode getRateLimiterMode(final byte byteIdentifier) {
+    for (final RateLimiterMode rateLimiterMode : RateLimiterMode.values()) {
+      if (rateLimiterMode.getValue() == byteIdentifier) {
+        return rateLimiterMode;
+      }
+    }
+
+    throw new IllegalArgumentException(
+        "Illegal value provided for RateLimiterMode.");
+  }
+}
diff --git a/java/src/test/java/org/rocksdb/RateLimiterTest.java b/java/src/test/java/org/rocksdb/RateLimiterTest.java
index 27567e89d..c78f9876e 100644
--- a/java/src/test/java/org/rocksdb/RateLimiterTest.java
+++ b/java/src/test/java/org/rocksdb/RateLimiterTest.java
@@ -8,6 +8,7 @@ import org.junit.ClassRule;
 import org.junit.Test;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.rocksdb.RateLimiter.*;
 
 public class RateLimiterTest {
 
@@ -16,17 +17,21 @@ public class RateLimiterTest {
       new RocksMemoryResource();
 
   @Test
-  public void setBytesPerSecond() {
+  public void bytesPerSecond() {
     try(final RateLimiter rateLimiter =
-            new RateLimiter(1000, 100 * 1000, 1)) {
+            new RateLimiter(1000, DEFAULT_REFILL_PERIOD_MICROS,
+                DEFAULT_FAIRNESS, DEFAULT_MODE, DEFAULT_AUTOTUNE)) {
+      assertThat(rateLimiter.getBytesPerSecond()).isGreaterThan(0);
       rateLimiter.setBytesPerSecond(2000);
+      assertThat(rateLimiter.getBytesPerSecond()).isGreaterThan(0);
     }
   }
 
   @Test
   public void getSingleBurstBytes() {
     try(final RateLimiter rateLimiter =
-            new RateLimiter(1000, 100 * 1000, 1)) {
+            new RateLimiter(1000, DEFAULT_REFILL_PERIOD_MICROS,
+                DEFAULT_FAIRNESS, DEFAULT_MODE, DEFAULT_AUTOTUNE)) {
       assertThat(rateLimiter.getSingleBurstBytes()).isEqualTo(100);
     }
   }
@@ -34,7 +39,8 @@ public class RateLimiterTest {
   @Test
   public void getTotalBytesThrough() {
     try(final RateLimiter rateLimiter =
-            new RateLimiter(1000, 100 * 1000, 1)) {
+            new RateLimiter(1000, DEFAULT_REFILL_PERIOD_MICROS,
+                DEFAULT_FAIRNESS, DEFAULT_MODE, DEFAULT_AUTOTUNE)) {
       assertThat(rateLimiter.getTotalBytesThrough()).isEqualTo(0);
     }
   }
@@ -42,8 +48,18 @@ public class RateLimiterTest {
   @Test
   public void getTotalRequests() {
     try(final RateLimiter rateLimiter =
-            new RateLimiter(1000, 100 * 1000, 1)) {
+            new RateLimiter(1000, DEFAULT_REFILL_PERIOD_MICROS,
+                DEFAULT_FAIRNESS, DEFAULT_MODE, DEFAULT_AUTOTUNE)) {
       assertThat(rateLimiter.getTotalRequests()).isEqualTo(0);
     }
   }
+
+  @Test
+  public void autoTune() {
+    try(final RateLimiter rateLimiter =
+            new RateLimiter(1000, DEFAULT_REFILL_PERIOD_MICROS,
+                DEFAULT_FAIRNESS, DEFAULT_MODE, true)) {
+      assertThat(rateLimiter.getBytesPerSecond()).isGreaterThan(0);
+    }
+  }
 }