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

java/rocksjni/portal.h

@@ -21,6 +21,7 @@
 
 #include "rocksdb/db.h"
 #include "rocksdb/filter_policy.h"
+#include "rocksdb/rate_limiter.h"
 #include "rocksdb/status.h"
 #include "rocksdb/utilities/backupable_db.h"
 #include "rocksdb/utilities/write_batch_with_index.h"
@@ -2934,6 +2935,45 @@ class StatsLevelJni {
   }
 };
 
+// The portal class for org.rocksdb.RateLimiterMode
+class RateLimiterModeJni {
+ public:
+  // Returns the equivalent org.rocksdb.RateLimiterMode for the provided
+  // C++ rocksdb::RateLimiter::Mode enum
+  static jbyte toJavaRateLimiterMode(
+      const rocksdb::RateLimiter::Mode& rate_limiter_mode) {
+    switch(rate_limiter_mode) {
+      case rocksdb::RateLimiter::Mode::kReadsOnly:
+        return 0x0;
+      case rocksdb::RateLimiter::Mode::kWritesOnly:
+        return 0x1;
+      case rocksdb::RateLimiter::Mode::kAllIo:
+        return 0x2;
+
+      default:
+        // undefined/default
+        return 0x1;
+    }
+  }
+
+  // Returns the equivalent C++ rocksdb::RateLimiter::Mode enum for the
+  // provided Java org.rocksdb.RateLimiterMode
+  static rocksdb::RateLimiter::Mode toCppRateLimiterMode(jbyte jrate_limiter_mode) {
+    switch(jrate_limiter_mode) {
+      case 0x0:
+        return rocksdb::RateLimiter::Mode::kReadsOnly;
+      case 0x1:
+        return rocksdb::RateLimiter::Mode::kWritesOnly;
+      case 0x2:
+        return rocksdb::RateLimiter::Mode::kAllIo;
+
+      default:
+        // undefined/default
+        return rocksdb::RateLimiter::Mode::kWritesOnly;
+    }
+  }
+};
+
 // various utility functions for working with RocksDB and JNI
 class JniUtil {
  public:

java/rocksjni/ratelimiterjni.cc

@@ -12,16 +12,21 @@
 /*
  * Class:     org_rocksdb_RateLimiter
  * Method:    newRateLimiterHandle
- * Signature: (JJI)J
+ * Signature: (JJIBZ)J
  */
 jlong Java_org_rocksdb_RateLimiter_newRateLimiterHandle(
     JNIEnv* env, jclass jclazz, jlong jrate_bytes_per_second,
-    jlong jrefill_period_micros, jint jfairness) {
+    jlong jrefill_period_micros, jint jfairness, jbyte jrate_limiter_mode,
+    jboolean jauto_tune) {
+  auto rate_limiter_mode = rocksdb::RateLimiterModeJni::toCppRateLimiterMode(
+      jrate_limiter_mode);
   auto * sptr_rate_limiter =
       new std::shared_ptr<rocksdb::RateLimiter>(rocksdb::NewGenericRateLimiter(
           static_cast<int64_t>(jrate_bytes_per_second),
           static_cast<int64_t>(jrefill_period_micros),
-          static_cast<int32_t>(jfairness)));
+          static_cast<int32_t>(jfairness),
+          rate_limiter_mode,
+          jauto_tune));
 
   return reinterpret_cast<jlong>(sptr_rate_limiter);
 }
@@ -50,6 +55,17 @@ void Java_org_rocksdb_RateLimiter_setBytesPerSecond(
       SetBytesPerSecond(jbytes_per_second);
 }
 
+/*
+ * Class:     org_rocksdb_RateLimiter
+ * Method:    getBytesPerSecond
+ * Signature: (J)J
+ */
+jlong Java_org_rocksdb_RateLimiter_getBytesPerSecond(
+    JNIEnv* env, jobject jobj, jlong handle) {
+  return reinterpret_cast<std::shared_ptr<rocksdb::RateLimiter> *>(handle)->get()->
+      GetBytesPerSecond();
+}
+
 /*
  * Class:     org_rocksdb_RateLimiter
  * Method:    request

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);
+  public static final long DEFAULT_REFILL_PERIOD_MICROS = 100 * 1000;
-  private static final int DEFAULT_FAIRNESS = 10;
+  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.
+   *     rate limit for anything other than flush and compaction, e.g. write to
-   * @param refillPeriodMicros this controls how often tokens are refilled. For example,
+   *     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
+   *     100ms, then 1MB is refilled every 100ms internally. Larger value can
-   *     burstier writes while smaller value introduces more CPU overhead.
+   *     lead to burstier writes while smaller value introduces more CPU
-   *     The default should work for most cases.
+   *     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,
+   *     A low-pri request is usually blocked in favor of hi-pri request.
-   *     RocksDB assigns low-pri to request from compaction and high-pri to request
+   *     Currently, RocksDB assigns low-pri to request from compaction and
-   *     from flush. Low-pri requests can get blocked if flush requests come in
+   *     high-pri to request from flush. Low-pri requests can get blocked if
-   *     continuously. This fairness parameter grants low-pri requests permission by
+   *     flush requests come in continuously. This fairness parameter grants
-   *     fairness chance even though high-pri requests exist to avoid starvation.
+   *     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,
+    this(rateBytesPerSecond, refillPeriodMicros, fairness, DEFAULT_MODE,
-        refillPeriodMicros, fairness));
+        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) {
+  public RateLimiter(final long rateBytesPerSecond,
-    this(rateBytesPerSecond, DEFAULT_REFILL_PERIOD_MICROS, DEFAULT_FAIRNESS);
+      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));
   }
 
   /**
@@ -64,6 +160,16 @@ public class RateLimiter extends RocksObject {
     setBytesPerSecond(nativeHandle_, bytesPerSecond);
   }
 
+  /**
+   * 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
@@ -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);

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.");
+  }
+}

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);
+    }
+  }
 }