From 5a0cff2c9eeb2dc8272a0d78fc57ef60fbb31527 Mon Sep 17 00:00:00 2001
From: Ted Mostly <wanghenshui@qq.com>
Date: Mon, 8 May 2023 19:51:09 +0800
Subject: [PATCH] feat: ratelimiter support auto tune (#1374)

* feat: ratelimiter support auto tune

* add conf

* add commit

---------

Co-authored-by: Xin.Zh <dragoncharlie@foxmail.com>
---
 conf/pika.conf      |  6 ++++++
 include/pika_conf.h | 17 +++++++++++++++++
 src/pika_admin.cc   | 18 ++++++++++++++++++
 src/pika_conf.cc    | 16 ++++++++++++++++
 src/pika_server.cc  |  9 ++++++++-
 5 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/conf/pika.conf b/conf/pika.conf
index 1d36d1fab2..d97a4c809e 100644
--- a/conf/pika.conf
+++ b/conf/pika.conf
@@ -326,6 +326,12 @@ max-bytes-for-level-multiplier : 10
 # rate limiter bandwidth, default 200MB
 #rate-limiter-bandwidth : 209715200
 
+# rate limiter auto tune https://rocksdb.org/blog/2017/12/18/17-auto-tuned-rate-limiter.html
+# default: false
+#rate-limiter-refill-period-us : 100000
+#rate-limiter-fairness: 10
+#rate-limiter-auto-tuned : true
+
 ###########################
 # rocksdb blob configure
 # wiki https://github.com/facebook/rocksdb/wiki/BlobDB
diff --git a/include/pika_conf.h b/include/pika_conf.h
index 01d09062f0..64494664fc 100644
--- a/include/pika_conf.h
+++ b/include/pika_conf.h
@@ -266,6 +266,18 @@ class PikaConf : public pstd::BaseConf {
     RWLock l(&rwlock_, false);
     return rate_limiter_bandwidth_;
   }
+  int64_t rate_limiter_refill_period_us() {
+    RWLock l(&rwlock_, false);
+    return rate_limiter_refill_period_us_;
+  }
+  int64_t rate_limiter_fairness() {
+    RWLock l(&rwlock_, false);
+    return rate_limiter_fairness_;
+  }
+  bool rate_limiter_auto_tuned() {
+    RWLock l(&rwlock_, false);
+    return rate_limiter_auto_tuned_;
+  }
 
   bool enable_blob_files() { return enable_blob_files_; }
   int64_t min_blob_size() { return min_blob_size_; }
@@ -517,6 +529,7 @@ class PikaConf : public pstd::BaseConf {
   int expire_logs_nums_ = 0;
   bool slave_read_only_ = false;
   std::string conf_path_;
+
   int max_cache_statistic_keys_ = 0;
   int small_compaction_threshold_ = 0;
   int max_background_flushes_ = 0;
@@ -532,6 +545,10 @@ class PikaConf : public pstd::BaseConf {
   bool optimize_filters_for_hits_ = false;
   bool level_compaction_dynamic_level_bytes_ = false;
   int64_t rate_limiter_bandwidth_ = 200 * 1024 * 1024;  // 200M
+  int64_t rate_limiter_refill_period_us_ = 100 * 1000;
+  int64_t rate_limiter_fairness_ = 10;
+  bool rate_limiter_auto_tuned_ = true;
+
   std::atomic<int> sync_window_size_;
   std::atomic<int> max_conn_rbuf_size_;
   std::atomic<int> consensus_level_;
diff --git a/src/pika_admin.cc b/src/pika_admin.cc
index 5c486ad6d0..b24c4a2787 100644
--- a/src/pika_admin.cc
+++ b/src/pika_admin.cc
@@ -1641,6 +1641,24 @@ void ConfigCmd::ConfigGet(std::string& ret) {
     EncodeInt64(&config_body, g_pika_conf->rate_limiter_bandwidth());
   }
 
+  if (pstd::stringmatch(pattern.data(), "rate-limiter-refill-period-us", 1)) {
+    elements += 2;
+    EncodeString(&config_body, "rate-limiter-refill-period-us");
+    EncodeInt64(&config_body, g_pika_conf->rate_limiter_refill_period_us());
+  }
+
+  if (pstd::stringmatch(pattern.data(), "rate-limiter-fairness", 1)) {
+    elements += 2;
+    EncodeString(&config_body, "rate-limiter-fairness");
+    EncodeInt64(&config_body, g_pika_conf->rate_limiter_fairness());
+  }
+
+  if (pstd::stringmatch(pattern.data(), "rate-limiter-auto-tuned", 1)) {
+    elements += 2;
+    EncodeString(&config_body, "rate-limiter-auto-tuned");
+    EncodeString(&config_body, g_pika_conf->rate_limiter_auto_tuned() ? "yes" : "no");
+  }
+
   std::stringstream resp;
   resp << "*" << std::to_string(elements) << "\r\n" << config_body;
   ret = resp.str();
diff --git a/src/pika_conf.cc b/src/pika_conf.cc
index 81f4c6ebd8..408d5cfa9d 100644
--- a/src/pika_conf.cc
+++ b/src/pika_conf.cc
@@ -369,6 +369,22 @@ int PikaConf::Load() {
     rate_limiter_bandwidth_ = 200 * 1024 * 1024;  // 200MB
   }
 
+  // rate-limiter-refill-period-us
+  GetConfInt64("rate-limiter-refill-period-us", &rate_limiter_refill_period_us_);
+  if (rate_limiter_refill_period_us_ <= 0 ) {
+    rate_limiter_refill_period_us_ = 100 * 1000;
+  }
+
+  // rate-limiter-fairness
+  GetConfInt64("rate-limiter-fairness", &rate_limiter_fairness_);
+  if (rate_limiter_fairness_ <= 0 ) {
+    rate_limiter_fairness_ = 10;
+  }
+
+  std::string at;
+  GetConfStr("rate-limiter-auto-tuned", &at);
+  rate_limiter_auto_tuned_ = (at == "yes" || at.empty()) ? true : false;
+
   // max_write_buffer_num
   max_write_buffer_num_ = 2;
   GetConfInt("max-write-buffer-num", &max_write_buffer_num_);
diff --git a/src/pika_server.cc b/src/pika_server.cc
index 6c7c4a5073..a898ab5af2 100644
--- a/src/pika_server.cc
+++ b/src/pika_server.cc
@@ -1527,7 +1527,14 @@ void PikaServer::InitStorageOptions() {
   }
 
   storage_options_.options.rate_limiter =
-      std::shared_ptr<rocksdb::RateLimiter>(rocksdb::NewGenericRateLimiter(g_pika_conf->rate_limiter_bandwidth()));
+    std::shared_ptr<rocksdb::RateLimiter>(
+      rocksdb::NewGenericRateLimiter(
+        g_pika_conf->rate_limiter_bandwidth(),
+        g_pika_conf->rate_limiter_refill_period_us(),
+        g_pika_conf->rate_limiter_fairness(),
+        rocksdb::RateLimiter::Mode::kWritesOnly,
+        g_pika_conf->rate_limiter_auto_tuned()
+      ));
 
   // For Storage small compaction
   storage_options_.statistics_max_size = g_pika_conf->max_cache_statistic_keys();