support changing rate limiter dynamically (#510) (#516) (#515)

cherry-pick #510 and #516 

librocksdb_sys/crocksdb/c.cc

@@ -235,7 +235,7 @@ struct crocksdb_ingestexternalfileoptions_t  { IngestExternalFileOptions rep; };
 struct crocksdb_sstfilereader_t   { SstFileReader*    rep; };
 struct crocksdb_sstfilewriter_t   { SstFileWriter*    rep; };
 struct crocksdb_externalsstfileinfo_t   { ExternalSstFileInfo rep; };
-struct crocksdb_ratelimiter_t     { RateLimiter*      rep; };
+struct crocksdb_ratelimiter_t     { std::shared_ptr<RateLimiter> rep; };
 struct crocksdb_histogramdata_t   { HistogramData     rep; };
 struct crocksdb_pinnableslice_t   { PinnableSlice     rep; };
 struct crocksdb_flushjobinfo_t {
@@ -2939,12 +2939,24 @@ unsigned char crocksdb_options_statistics_get_histogram(
   return 0;
 }
 
-void crocksdb_options_set_ratelimiter(crocksdb_options_t *opt, crocksdb_ratelimiter_t *limiter) {
-  opt->rep.rate_limiter.reset(limiter->rep);
+void crocksdb_options_set_ratelimiter(crocksdb_options_t* opt,
+                                      crocksdb_ratelimiter_t* limiter) {
+  opt->rep.rate_limiter = limiter->rep;
   limiter->rep = nullptr;
 }
 
-void crocksdb_options_set_vector_memtable_factory(crocksdb_options_t* opt, uint64_t reserved_bytes) {
+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;
+    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));
 }
 
@@ -2977,8 +2989,8 @@ crocksdb_ratelimiter_t* crocksdb_ratelimiter_create(
     int64_t refill_period_us,
     int32_t fairness) {
   crocksdb_ratelimiter_t* rate_limiter = new crocksdb_ratelimiter_t;
-  rate_limiter->rep = NewGenericRateLimiter(rate_bytes_per_sec,
-                                            refill_period_us, fairness);
+  rate_limiter->rep = std::shared_ptr<RateLimiter>(
+      NewGenericRateLimiter(rate_bytes_per_sec, refill_period_us, fairness));
   return rate_limiter;
 }
 
@@ -3001,15 +3013,14 @@ crocksdb_ratelimiter_t* crocksdb_ratelimiter_create_with_auto_tuned(
         m = RateLimiter::Mode::kAllIo;
         break;
     }
-    rate_limiter->rep = NewGenericRateLimiter(rate_bytes_per_sec,
-                                              refill_period_us, fairness,
-                                              m, auto_tuned);
+    rate_limiter->rep = std::shared_ptr<RateLimiter>(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;
+    limiter->rep.reset();
   }
   delete limiter;
 }

librocksdb_sys/crocksdb/crocksdb/c.h

@@ -1168,6 +1168,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);
 

librocksdb_sys/src/lib.rs

@@ -733,6 +733,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(

scripts/format-diff.sh

+#!/bin/bash
+
+git diff `git merge-base master HEAD` ./librocksdb_sys/crocksdb | clang-format-diff -style=google -p1 -i

src/rocksdb_options.rs

@@ -1005,6 +1005,35 @@ 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());
+        }
+
+        let rate_limiter = RateLimiter { inner: limiter };
+
+        unsafe {
+            crocksdb_ffi::crocksdb_ratelimiter_set_bytes_per_second(
+                rate_limiter.inner,
+                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_limiter = RateLimiter { inner: limiter };
+        let rate =
+            unsafe { crocksdb_ffi::crocksdb_ratelimiter_get_bytes_per_second(rate_limiter.inner) };
+        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");