Add Wrap for auto-tuned Rate Limiter (#257)

librocksdb_sys/crocksdb/c.cc

@@ -2780,6 +2780,31 @@ crocksdb_ratelimiter_t* crocksdb_ratelimiter_create(
   return rate_limiter;
 }
 
+crocksdb_ratelimiter_t* crocksdb_ratelimiter_create_with_auto_tuned(
+    int64_t rate_bytes_per_sec,
+    int64_t refill_period_us,
+    int32_t fairness,
+    crocksdb_ratelimiter_mode_t mode,
+    bool auto_tuned) {
+    crocksdb_ratelimiter_t* rate_limiter = new crocksdb_ratelimiter_t;
+    RateLimiter::Mode m = RateLimiter::Mode::kWritesOnly;
+    switch (mode) {
+      case kReadsOnly:
+        m = RateLimiter::Mode::kReadsOnly;
+        break;
+      case kWritesOnly:
+        m = RateLimiter::Mode::kWritesOnly;
+        break;
+      case kAllIo:
+        m = RateLimiter::Mode::kAllIo;
+        break;
+    }
+    rate_limiter->rep = NewGenericRateLimiter(rate_bytes_per_sec,
+                                              refill_period_us, fairness,
+                                              m, auto_tuned);
+    return rate_limiter;
+}
+
 void crocksdb_ratelimiter_destroy(crocksdb_ratelimiter_t *limiter) {
   if (limiter->rep) {
       delete limiter->rep;

librocksdb_sys/crocksdb/crocksdb/c.h

@@ -163,6 +163,12 @@ typedef enum crocksdb_table_property_t {
   kCompressionName = 17,
 } crocksdb_table_property_t;
 
+typedef enum crocksdb_ratelimiter_mode_t {
+  kReadsOnly = 1,
+  kWritesOnly = 2,
+  kAllIo = 3,
+} crocksdb_ratelimiter_mode_t;
+
 /* DB operations */
 
 extern C_ROCKSDB_LIBRARY_API crocksdb_t* crocksdb_open(
@@ -1099,6 +1105,9 @@ extern C_ROCKSDB_LIBRARY_API void crocksdb_options_set_delayed_write_rate(
 /* RateLimiter */
 extern C_ROCKSDB_LIBRARY_API crocksdb_ratelimiter_t* crocksdb_ratelimiter_create(
     int64_t rate_bytes_per_sec, int64_t refill_period_us, int32_t fairness);
+extern C_ROCKSDB_LIBRARY_API crocksdb_ratelimiter_t* crocksdb_ratelimiter_create_with_auto_tuned(
+    int64_t rate_bytes_per_sec, int64_t refill_period_us, int32_t fairness,
+    crocksdb_ratelimiter_mode_t mode, bool auto_tuned);
 extern C_ROCKSDB_LIBRARY_API void crocksdb_ratelimiter_destroy(crocksdb_ratelimiter_t*);
 extern C_ROCKSDB_LIBRARY_API void crocksdb_ratelimiter_set_bytes_per_second(
     crocksdb_ratelimiter_t *limiter, int64_t rate_bytes_per_sec);

librocksdb_sys/src/lib.rs

@@ -268,6 +268,14 @@ pub enum DBBottommostLevelCompaction {
     Force = 2,
 }
 
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+#[repr(C)]
+pub enum DBRateLimiterMode {
+    ReadOnly = 1,
+    WriteOnly = 2,
+    AllIo = 3,
+}
+
 pub fn error_message(ptr: *mut c_char) -> String {
     let c_str = unsafe { CStr::from_ptr(ptr) };
     let s = format!("{}", c_str.to_string_lossy());
@@ -548,6 +556,13 @@ extern "C" {
         refill_period_us: i64,
         fairness: i32,
     ) -> *mut DBRateLimiter;
+    pub fn crocksdb_ratelimiter_create_with_auto_tuned(
+        rate_bytes_per_sec: i64,
+        refill_period_us: i64,
+        fairness: i32,
+        mode: DBRateLimiterMode,
+        auto_tuned: bool,
+    ) -> *mut DBRateLimiter;
     pub fn crocksdb_ratelimiter_destroy(limiter: *mut DBRateLimiter);
     pub fn crocksdb_ratelimiter_set_bytes_per_second(
         limiter: *mut DBRateLimiter,

src/lib.rs

@@ -27,7 +27,8 @@ pub use event_listener::{
 pub use librocksdb_sys::{
     self as crocksdb_ffi, new_bloom_filter, CompactionPriority, CompactionReason,
     DBBottommostLevelCompaction, DBCompactionStyle, DBCompressionType, DBEntryType, DBInfoLogLevel,
-    DBRecoveryMode, DBStatisticsHistogramType, DBStatisticsTickerType, WriteStallCondition,
+    DBRateLimiterMode, DBRecoveryMode, DBStatisticsHistogramType, DBStatisticsTickerType,
+    WriteStallCondition,
 };
 pub use merge_operator::MergeOperands;
 pub use metadata::{ColumnFamilyMetaData, LevelMetaData, SstFileMetaData};

src/rocksdb_options.rs

@@ -18,9 +18,9 @@ use comparator::{self, compare_callback, ComparatorCallback};
 use crocksdb_ffi::{
     self, DBBlockBasedTableOptions, DBBottommostLevelCompaction, DBCompactOptions,
     DBCompactionOptions, DBCompressionType, DBFifoCompactionOptions, DBFlushOptions,
-    DBInfoLogLevel, DBInstance, DBRateLimiter, DBReadOptions, DBRecoveryMode, DBRestoreOptions,
-    DBSnapshot, DBStatisticsHistogramType, DBStatisticsTickerType, DBTitanDBOptions,
-    DBWriteOptions, Options,
+    DBInfoLogLevel, DBInstance, DBRateLimiter, DBRateLimiterMode, DBReadOptions, DBRecoveryMode,
+    DBRestoreOptions, DBSnapshot, DBStatisticsHistogramType, DBStatisticsTickerType,
+    DBTitanDBOptions, DBWriteOptions, Options,
 };
 use event_listener::{new_event_listener, EventListener};
 use libc::{self, c_double, c_int, c_uchar, c_void, size_t};
@@ -181,6 +181,25 @@ impl RateLimiter {
         RateLimiter { inner: limiter }
     }
 
+    pub fn new_with_auto_tuned(
+        rate_bytes_per_sec: i64,
+        refill_period_us: i64,
+        fairness: i32,
+        mode: DBRateLimiterMode,
+        auto_tuned: bool,
+    ) -> RateLimiter {
+        let limiter = unsafe {
+            crocksdb_ffi::crocksdb_ratelimiter_create_with_auto_tuned(
+                rate_bytes_per_sec,
+                refill_period_us,
+                fairness,
+                mode,
+                auto_tuned,
+            )
+        };
+        RateLimiter { inner: limiter }
+    }
+
     pub fn set_bytes_per_second(&self, bytes_per_sec: i64) {
         unsafe {
             crocksdb_ffi::crocksdb_ratelimiter_set_bytes_per_second(self.inner, bytes_per_sec);
@@ -899,6 +918,24 @@ impl DBOptions {
         }
     }
 
+    pub fn set_ratelimiter_with_auto_tuned(
+        &mut self,
+        rate_bytes_per_sec: i64,
+        mode: DBRateLimiterMode,
+        auto_tuned: bool,
+    ) {
+        let rate_limiter = RateLimiter::new_with_auto_tuned(
+            rate_bytes_per_sec,
+            DEFAULT_REFILL_PERIOD_US,
+            DEFAULT_FAIRNESS,
+            mode,
+            auto_tuned,
+        );
+        unsafe {
+            crocksdb_ffi::crocksdb_options_set_ratelimiter(self.inner, rate_limiter.inner);
+        }
+    }
+
     // Create a info log with `path` and save to options logger field directly.
     // TODO: export more logger options like level, roll size, time, etc...
     pub fn create_info_log(&self, path: &str) -> Result<(), String> {

tests/cases/test_rocksdb_options.rs

@@ -12,7 +12,7 @@
 // limitations under the License.
 
 use rocksdb::crocksdb_ffi::{
-    CompactionPriority, DBCompressionType, DBInfoLogLevel as InfoLogLevel,
+    CompactionPriority, DBCompressionType, DBInfoLogLevel as InfoLogLevel, DBRateLimiterMode,
     DBStatisticsHistogramType as HistogramType, DBStatisticsTickerType as TickerType,
 };
 use rocksdb::{
@@ -162,6 +162,16 @@ fn test_set_ratelimiter() {
     drop(db);
 }
 
+#[test]
+fn test_set_ratelimiter_with_auto_tuned() {
+    let path = TempDir::new("_rust_rocksdb_test_set_rate_limiter_with_auto_tuned").expect("");
+    let mut opts = DBOptions::new();
+    opts.create_if_missing(true);
+    opts.set_ratelimiter_with_auto_tuned(100 * 1024 * 1024, DBRateLimiterMode::AllIo, true);
+    let db = DB::open(opts, path.path().to_str().unwrap()).unwrap();
+    drop(db);
+}
+
 #[test]
 fn test_set_wal_opt() {
     let path = TempDir::new("_rust_rocksdb_test_set_wal_opt").expect("");