support changing rate limiter dynamically (#510)

* support change rate limiter dynamically
Signed-off-by: Ryan Leung <rleungx@gmail.com>

librocksdb_sys/crocksdb/c.cc

@@ -2907,6 +2907,16 @@ void crocksdb_options_set_ratelimiter(crocksdb_options_t* opt,
   limiter->rep = nullptr;
 }
 
+crocksdb_ratelimiter_t* crocksdb_options_get_ratelimiter(
+    crocksdb_options_t* opt) {
+  if (opt->rep.rate_limiter != nullptr) {
+    crocksdb_ratelimiter_t* limiter = new crocksdb_ratelimiter_t;
+    limiter->rep = opt->rep.rate_limiter.get();
+    return limiter;
+  }
+  return nullptr;
+}
+
 void crocksdb_options_set_vector_memtable_factory(crocksdb_options_t* opt,
                                                   uint64_t reserved_bytes) {
   opt->rep.memtable_factory.reset(new VectorRepFactory(reserved_bytes));

librocksdb_sys/crocksdb/crocksdb/c.h

@@ -1188,6 +1188,8 @@ extern C_ROCKSDB_LIBRARY_API void crocksdb_options_set_fifo_compaction_options(
     crocksdb_options_t* opt, crocksdb_fifo_compaction_options_t* fifo);
 extern C_ROCKSDB_LIBRARY_API void crocksdb_options_set_ratelimiter(
     crocksdb_options_t* opt, crocksdb_ratelimiter_t* limiter);
+extern C_ROCKSDB_LIBRARY_API crocksdb_ratelimiter_t*
+crocksdb_options_get_ratelimiter(crocksdb_options_t* opt);
 extern C_ROCKSDB_LIBRARY_API void crocksdb_options_set_vector_memtable_factory(
     crocksdb_options_t* opt, uint64_t reserved_bytes);
 extern C_ROCKSDB_LIBRARY_API void crocksdb_options_set_atomic_flush(

librocksdb_sys/src/lib.rs

@@ -734,6 +734,7 @@ extern "C" {
     pub fn crocksdb_options_set_delayed_write_rate(options: *mut Options, rate: u64);
     pub fn crocksdb_options_set_force_consistency_checks(options: *mut Options, v: bool);
     pub fn crocksdb_options_set_ratelimiter(options: *mut Options, limiter: *mut DBRateLimiter);
+    pub fn crocksdb_options_get_ratelimiter(options: *mut Options) -> *mut DBRateLimiter;
     pub fn crocksdb_options_set_info_log(options: *mut Options, logger: *mut DBLogger);
     pub fn crocksdb_options_get_block_cache_usage(options: *const Options) -> usize;
     pub fn crocksdb_options_set_block_cache_capacity(

src/rocksdb_options.rs

@@ -1006,6 +1006,28 @@ impl DBOptions {
         }
     }
 
+    pub fn set_rate_bytes_per_sec(&mut self, rate_bytes_per_sec: i64) -> Result<(), String> {
+        let limiter = unsafe { crocksdb_ffi::crocksdb_options_get_ratelimiter(self.inner) };
+        if limiter.is_null() {
+            return Err("Failed to get rate limiter".to_owned());
+        }
+
+        unsafe {
+            crocksdb_ffi::crocksdb_ratelimiter_set_bytes_per_second(limiter, rate_bytes_per_sec);
+        }
+        Ok(())
+    }
+
+    pub fn get_rate_bytes_per_sec(&self) -> Option<i64> {
+        let limiter = unsafe { crocksdb_ffi::crocksdb_options_get_ratelimiter(self.inner) };
+        if limiter.is_null() {
+            return None
+        }
+
+        let rate = unsafe { crocksdb_ffi::crocksdb_ratelimiter_get_bytes_per_second(limiter) };
+        Some(rate)
+    }
+
     // 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

@@ -176,6 +176,20 @@ fn test_set_ratelimiter_with_auto_tuned() {
     drop(db);
 }
 
+#[test]
+fn test_set_ratelimiter_bytes_per_second() {
+    let path = tempdir_with_prefix("_rust_rocksdb_test_set_rate_limiter_bytes_per_second");
+    let mut opts = DBOptions::new();
+    opts.create_if_missing(true);
+    // compaction and flush rate limited below 100MB/sec
+    opts.set_ratelimiter(100 * 1024 * 1024);
+    let db = DB::open(opts, path.path().to_str().unwrap()).unwrap();
+    let mut opts = db.get_db_options();
+    assert!(opts.set_rate_bytes_per_sec(200 * 1024 * 1024).is_ok(), true);
+    assert_eq!(opts.get_rate_bytes_per_sec().unwrap(), 200 * 1024 * 1024);
+    drop(db);
+}
+
 #[test]
 fn test_set_wal_opt() {
     let path = tempdir_with_prefix("_rust_rocksdb_test_set_wal_opt");