Add rate limiter

java/Makefile

@@ -1,4 +1,4 @@
-NATIVE_JAVA_CLASSES = org.rocksdb.RocksDB org.rocksdb.Options org.rocksdb.WriteBatch org.rocksdb.WriteBatchInternal org.rocksdb.WriteBatchTest org.rocksdb.WriteOptions org.rocksdb.BackupableDB org.rocksdb.BackupableDBOptions org.rocksdb.Statistics org.rocksdb.RocksIterator org.rocksdb.VectorMemTableConfig org.rocksdb.SkipListMemTableConfig org.rocksdb.HashLinkedListMemTableConfig org.rocksdb.HashSkipListMemTableConfig org.rocksdb.PlainTableConfig org.rocksdb.BlockBasedTableConfig org.rocksdb.ReadOptions org.rocksdb.Filter org.rocksdb.BloomFilter org.rocksdb.RestoreOptions org.rocksdb.RestoreBackupableDB org.rocksdb.RocksEnv
+NATIVE_JAVA_CLASSES = org.rocksdb.RocksDB org.rocksdb.Options org.rocksdb.WriteBatch org.rocksdb.WriteBatchInternal org.rocksdb.WriteBatchTest org.rocksdb.WriteOptions org.rocksdb.BackupableDB org.rocksdb.BackupableDBOptions org.rocksdb.Statistics org.rocksdb.RocksIterator org.rocksdb.VectorMemTableConfig org.rocksdb.SkipListMemTableConfig org.rocksdb.HashLinkedListMemTableConfig org.rocksdb.HashSkipListMemTableConfig org.rocksdb.PlainTableConfig org.rocksdb.BlockBasedTableConfig org.rocksdb.ReadOptions org.rocksdb.Filter org.rocksdb.BloomFilter org.rocksdb.RestoreOptions org.rocksdb.RestoreBackupableDB org.rocksdb.RocksEnv org.rocksdb.GenericRateLimiterConfig
 
 NATIVE_INCLUDE = ./include
 ROCKSDB_JAR = rocksdbjni.jar

java/RocksDBSample.java

@@ -76,6 +76,10 @@ public class RocksDBSample {
     options.setAllowMmapReads(true);
     assert(options.tableFactoryName().equals("PlainTable"));
     
+    options.setRateLimiterConfig(new GenericRateLimiterConfig(10000000,
+            10000, 10));
+    options.setRateLimiterConfig(new GenericRateLimiterConfig(10000000));
+
     BlockBasedTableConfig table_options = new BlockBasedTableConfig();
     table_options.setBlockCacheSize(64 * SizeUnit.KB)
                  .setFilterBitsPerKey(10)

java/org/rocksdb/GenericRateLimiterConfig.java

@@ -0,0 +1,36 @@
+// Copyright (c) 2014, Facebook, Inc.  All rights reserved.
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree. An additional grant
+// of patent rights can be found in the PATENTS file in the same directory.
+package org.rocksdb;
+
+/**
+ * Config for rate limiter, which is used to control write rate of flush and
+ * compaction.
+ */
+public class GenericRateLimiterConfig extends RateLimiterConfig {
+  private static final long DEFAULT_REFILL_PERIOD_MICROS = (100 * 1000);
+  private static final int DEFAULT_FAIRNESS = 10;
+    
+  public GenericRateLimiterConfig(long rateBytesPerSecond,
+      long refillPeriodMicros, int fairness) {
+    rateBytesPerSecond_ = rateBytesPerSecond;
+    refillPeriodMicros_ = refillPeriodMicros;
+    fairness_ = fairness;
+  }
+  
+  public GenericRateLimiterConfig(long rateBytesPerSecond) {
+    this(rateBytesPerSecond, DEFAULT_REFILL_PERIOD_MICROS, DEFAULT_FAIRNESS);
+  }
+  
+  @Override protected long newRateLimiterHandle() {
+    return newRateLimiterHandle(rateBytesPerSecond_, refillPeriodMicros_,
+        fairness_);
+  }
+    
+  private native long newRateLimiterHandle(long rateBytesPerSecond,
+      long refillPeriodMicros, int fairness);
+  private final long rateBytesPerSecond_;
+  private final long refillPeriodMicros_;
+  private final int fairness_;
+}

java/org/rocksdb/Options.java

@@ -1105,6 +1105,19 @@ public class Options extends RocksObject {
     return this;
   }
   
+  /**
+   * Use to control write rate of flush and compaction. Flush has higher
+   * priority than compaction. Rate limiting is disabled if nullptr.
+   * Default: nullptr
+   *
+   * @param config rate limiter config.
+   * @return the instance of the current Options.
+   */
+  public Options setRateLimiterConfig(RateLimiterConfig config) {
+    setRateLimiter(nativeHandle_, config.newRateLimiterHandle());
+    return this;
+  }
+
   /**
    * Returns the name of the current mem table representation.
    * Memtable format can be set using setTableFormatConfig.
@@ -2192,6 +2205,8 @@ public class Options extends RocksObject {
   private native long statisticsPtr(long optHandle);
 
   private native void setMemTableFactory(long handle, long factoryHandle);
+  private native void setRateLimiter(long handle,
+      long rateLimiterHandle);
   private native String memTableFactoryName(long handle);
 
   private native void setTableFactory(long handle, long factoryHandle);

java/org/rocksdb/RateLimiterConfig.java

@@ -0,0 +1,20 @@
+// Copyright (c) 2014, Facebook, Inc.  All rights reserved.
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree. An additional grant
+// of patent rights can be found in the PATENTS file in the same directory.
+package org.rocksdb;
+
+/**
+ * Config for rate limiter, which is used to control write rate of flush and
+ * compaction.
+ */
+public abstract class RateLimiterConfig {
+  /**
+   * This function should only be called by Options.setRateLimiter(),
+   * which will create a c++ shared-pointer to the c++ RateLimiter
+   * that is associated with the Java RateLimtierConifg.
+   *
+   * @see Options.setRateLimiter()
+   */
+  abstract protected long newRateLimiterHandle();
+}

java/rocksjni/options.cc

@@ -21,6 +21,7 @@
 #include "rocksdb/memtablerep.h"
 #include "rocksdb/table.h"
 #include "rocksdb/slice_transform.h"
+#include "rocksdb/rate_limiter.h"
 
 /*
  * Class:     org_rocksdb_Options
@@ -459,6 +460,17 @@ void Java_org_rocksdb_Options_setMemTableFactory(
       reinterpret_cast<rocksdb::MemTableRepFactory*>(jfactory_handle));
 }
 
+/*
+ * Class:     org_rocksdb_Options
+ * Method:    setRateLimiter
+ * Signature: (JJ)V
+ */
+void Java_org_rocksdb_Options_setRateLimiter(
+    JNIEnv* env, jobject jobj, jlong jhandle, jlong jrate_limiter_handle) {
+  reinterpret_cast<rocksdb::Options*>(jhandle)->rate_limiter.reset(
+      reinterpret_cast<rocksdb::RateLimiter*>(jrate_limiter_handle));
+}
+
 /*
  * Class:     org_rocksdb_Options
  * Method:    tableCacheNumshardbits

java/rocksjni/ratelimiterjni.cc

@@ -0,0 +1,24 @@
+// Copyright (c) 2014, Facebook, Inc.  All rights reserved.
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree. An additional grant
+// of patent rights can be found in the PATENTS file in the same directory.
+//
+// This file implements the "bridge" between Java and C++ for RateLimiter.
+
+#include "rocksjni/portal.h"
+#include "include/org_rocksdb_GenericRateLimiterConfig.h"
+#include "rocksdb/rate_limiter.h"
+
+/*
+ * Class:     org_rocksdb_GenericRateLimiterConfig
+ * Method:    newRateLimiterHandle
+ * Signature: (JJI)J
+ */
+jlong Java_org_rocksdb_GenericRateLimiterConfig_newRateLimiterHandle(
+    JNIEnv* env, jobject jobj, jlong jrate_bytes_per_second,
+    jlong jrefill_period_micros, jint jfairness) {
+  return reinterpret_cast<jlong>(rocksdb::NewGenericRateLimiter(
+      rocksdb::jlong_to_size_t(jrate_bytes_per_second),
+      rocksdb::jlong_to_size_t(jrefill_period_micros),
+      static_cast<int32_t>(jfairness)));
+}