Add support for WriteAmpBasedRateLimiter (#547)

Add initializer for new WriteAmpBasedRateLimiter, also update rocksdb to include this plugin brought by tikv/rocksdb#193.
Signed-off-by: tabokie <xy.tao@outlook.com>

librocksdb_sys/bindings/aarch64-unknown-linux-gnu-bindings.rs

@@ -2458,6 +2458,15 @@ extern "C" {
         auto_tuned: libc::c_uchar,
     ) -> *mut crocksdb_ratelimiter_t;
 }
+extern "C" {
+    pub fn crocksdb_writeampbasedratelimiter_create_with_auto_tuned(
+        rate_bytes_per_sec: i64,
+        refill_period_us: i64,
+        fairness: i32,
+        mode: crocksdb_ratelimiter_mode_t,
+        auto_tuned: libc::c_uchar,
+    ) -> *mut crocksdb_ratelimiter_t;
+}
 extern "C" {
     pub fn crocksdb_ratelimiter_destroy(arg1: *mut crocksdb_ratelimiter_t);
 }

librocksdb_sys/bindings/x86_64-unknown-linux-gnu-bindings.rs

@@ -2561,6 +2561,15 @@ extern "C" {
         auto_tuned: libc::c_uchar,
     ) -> *mut crocksdb_ratelimiter_t;
 }
+extern "C" {
+    pub fn crocksdb_writeampbasedratelimiter_create_with_auto_tuned(
+        rate_bytes_per_sec: i64,
+        refill_period_us: i64,
+        fairness: i32,
+        mode: crocksdb_ratelimiter_mode_t,
+        auto_tuned: libc::c_uchar,
+    ) -> *mut crocksdb_ratelimiter_t;
+}
 extern "C" {
     pub fn crocksdb_ratelimiter_destroy(arg1: *mut crocksdb_ratelimiter_t);
 }

librocksdb_sys/crocksdb/c.cc

@@ -3151,6 +3151,28 @@ crocksdb_ratelimiter_t* crocksdb_ratelimiter_create_with_auto_tuned(
   return rate_limiter;
 }
 
+crocksdb_ratelimiter_t*
+crocksdb_writeampbasedratelimiter_create_with_auto_tuned(
+    int64_t rate_bytes_per_sec, int64_t refill_period_us, int32_t fairness,
+    crocksdb_ratelimiter_mode_t mode, unsigned char 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 = std::shared_ptr<RateLimiter>(NewWriteAmpBasedRateLimiter(
+      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) {
     limiter->rep.reset();

librocksdb_sys/crocksdb/crocksdb/c.h

@@ -1274,6 +1274,10 @@ crocksdb_ratelimiter_create_with_auto_tuned(int64_t rate_bytes_per_sec,
                                             int32_t fairness,
                                             crocksdb_ratelimiter_mode_t mode,
                                             unsigned char auto_tuned);
+extern C_ROCKSDB_LIBRARY_API crocksdb_ratelimiter_t*
+crocksdb_writeampbasedratelimiter_create_with_auto_tuned(
+    int64_t rate_bytes_per_sec, int64_t refill_period_us, int32_t fairness,
+    crocksdb_ratelimiter_mode_t mode, unsigned char 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(

rocksdb @ cf6a77d9

-Subproject commit 953f8b2f23eabdc0df9adb4fc50b566ae0c5a5e9
+Subproject commit cf6a77d9ba13de8df8f374d87c9fbfc7003c4f52

librocksdb_sys/src/lib.rs

@@ -807,6 +807,13 @@ extern "C" {
         mode: DBRateLimiterMode,
         auto_tuned: bool,
     ) -> *mut DBRateLimiter;
+    pub fn crocksdb_writeampbasedratelimiter_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/rocksdb_options.rs

@@ -228,6 +228,25 @@ impl RateLimiter {
         RateLimiter { inner: limiter }
     }
 
+    pub fn new_writeampbased_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_writeampbasedratelimiter_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);
@@ -1008,6 +1027,25 @@ impl DBOptions {
         }
     }
 
+    pub fn set_writeampbasedratelimiter_with_auto_tuned(
+        &mut self,
+        rate_bytes_per_sec: i64,
+        refill_period_us: i64,
+        mode: DBRateLimiterMode,
+        auto_tuned: bool,
+    ) {
+        let rate_limiter = RateLimiter::new_writeampbased_with_auto_tuned(
+            rate_bytes_per_sec,
+            refill_period_us,
+            DEFAULT_FAIRNESS,
+            mode,
+            auto_tuned,
+        );
+        unsafe {
+            crocksdb_ffi::crocksdb_options_set_ratelimiter(self.inner, rate_limiter.inner);
+        }
+    }
+
     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() {

tests/cases/test_rocksdb_options.rs

@@ -181,6 +181,22 @@ fn test_set_ratelimiter_with_auto_tuned() {
     drop(db);
 }
 
+#[test]
+fn test_set_writeampbasedratelimiter_with_auto_tuned() {
+    let path =
+        tempdir_with_prefix("_rust_rocksdb_test_set_write_amp_based_rate_limiter_with_auto_tuned");
+    let mut opts = DBOptions::new();
+    opts.create_if_missing(true);
+    opts.set_writeampbasedratelimiter_with_auto_tuned(
+        100 * 1024 * 1024,
+        10 * 100000,
+        DBRateLimiterMode::AllIo,
+        true,
+    );
+    let db = DB::open(opts, path.path().to_str().unwrap()).unwrap();
+    drop(db);
+}
+
 #[test]
 fn test_set_ratelimiter_bytes_per_second() {
     let path = tempdir_with_prefix("_rust_rocksdb_test_set_rate_limiter_bytes_per_second");