dynamically modify auto-tuned mode of rate limiter (#581)

Signed-off-by: tabokie <xy.tao@outlook.com>

librocksdb_sys/crocksdb/c.cc

@@ -3225,6 +3225,11 @@ void crocksdb_ratelimiter_set_bytes_per_second(crocksdb_ratelimiter_t* limiter,
   limiter->rep->SetBytesPerSecond(rate_bytes_per_sec);
 }
 
+void crocksdb_ratelimiter_set_auto_tuned(crocksdb_ratelimiter_t* limiter,
+                                         unsigned char auto_tuned) {
+  limiter->rep->SetAutoTuned(auto_tuned);
+}
+
 int64_t crocksdb_ratelimiter_get_singleburst_bytes(
     crocksdb_ratelimiter_t* limiter) {
   return limiter->rep->GetSingleBurstBytes();
@@ -3245,6 +3250,11 @@ int64_t crocksdb_ratelimiter_get_bytes_per_second(
   return limiter->rep->GetBytesPerSecond();
 }
 
+unsigned char crocksdb_ratelimiter_get_auto_tuned(
+    crocksdb_ratelimiter_t* limiter) {
+  return limiter->rep->GetAutoTuned();
+}
+
 int64_t crocksdb_ratelimiter_get_total_requests(crocksdb_ratelimiter_t* limiter,
                                                 unsigned char pri) {
   return limiter->rep->GetTotalRequests(static_cast<Env::IOPriority>(pri));

librocksdb_sys/crocksdb/crocksdb/c.h

@@ -1293,6 +1293,8 @@ 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);
+extern C_ROCKSDB_LIBRARY_API void crocksdb_ratelimiter_set_auto_tuned(
+    crocksdb_ratelimiter_t* limiter, unsigned char auto_tuned);
 extern C_ROCKSDB_LIBRARY_API int64_t
 crocksdb_ratelimiter_get_singleburst_bytes(crocksdb_ratelimiter_t* limiter);
 enum {
@@ -1307,6 +1309,8 @@ crocksdb_ratelimiter_get_total_bytes_through(crocksdb_ratelimiter_t* limiter,
                                              unsigned char pri);
 extern C_ROCKSDB_LIBRARY_API int64_t
 crocksdb_ratelimiter_get_bytes_per_second(crocksdb_ratelimiter_t* limiter);
+extern C_ROCKSDB_LIBRARY_API unsigned char crocksdb_ratelimiter_get_auto_tuned(
+    crocksdb_ratelimiter_t* limiter);
 extern C_ROCKSDB_LIBRARY_API int64_t crocksdb_ratelimiter_get_total_requests(
     crocksdb_ratelimiter_t* limiter, unsigned char pri);
 

rocksdb @ 4b0908b9

-Subproject commit fbbe2c873da2319c8986cb315f7ea8e75553a79c
+Subproject commit 4b0908b9d9620acbef4ed2003ac5868332c90897

librocksdb_sys/src/lib.rs

@@ -855,6 +855,7 @@ extern "C" {
         limiter: *mut DBRateLimiter,
         bytes_per_sec: i64,
     );
+    pub fn crocksdb_ratelimiter_set_auto_tuned(limiter: *mut DBRateLimiter, auto_tuned: bool);
     pub fn crocksdb_ratelimiter_get_singleburst_bytes(limiter: *mut DBRateLimiter) -> i64;
     pub fn crocksdb_ratelimiter_request(limiter: *mut DBRateLimiter, bytes: i64, pri: c_uchar);
     pub fn crocksdb_ratelimiter_get_total_bytes_through(
@@ -862,6 +863,7 @@ extern "C" {
         pri: c_uchar,
     ) -> i64;
     pub fn crocksdb_ratelimiter_get_bytes_per_second(limiter: *mut DBRateLimiter) -> i64;
+    pub fn crocksdb_ratelimiter_get_auto_tuned(limiter: *mut DBRateLimiter) -> bool;
     pub fn crocksdb_ratelimiter_get_total_requests(
         limiter: *mut DBRateLimiter,
         pri: c_uchar,

src/rocksdb_options.rs

@@ -1095,6 +1095,31 @@ impl DBOptions {
         Some(rate)
     }
 
+    pub fn set_auto_tuned(&mut self, auto_tuned: bool) -> 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_auto_tuned(rate_limiter.inner, auto_tuned);
+        }
+        Ok(())
+    }
+
+    pub fn get_auto_tuned(&self) -> Option<bool> {
+        let limiter = unsafe { crocksdb_ffi::crocksdb_options_get_ratelimiter(self.inner) };
+        if limiter.is_null() {
+            return None;
+        }
+
+        let rate_limiter = RateLimiter { inner: limiter };
+        let mode = unsafe { crocksdb_ffi::crocksdb_ratelimiter_get_auto_tuned(rate_limiter.inner) };
+        Some(mode)
+    }
+
     // 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

@@ -194,6 +194,9 @@ fn test_set_writeampbasedratelimiter_with_auto_tuned() {
         true,
     );
     let db = DB::open(opts, path.path().to_str().unwrap()).unwrap();
+    let mut opts = db.get_db_options();
+    assert!(opts.set_auto_tuned(false).is_ok(), true);
+    assert_eq!(opts.get_auto_tuned().unwrap(), false);
     drop(db);
 }