dynamically configure auto_tuned mode of WriteAmpBasedRateLimiter (#217)

Add new APIs to `RateLimiter`, only implement it for `WriteAmpBasedRateLimiter`. Signed-off-by: tabokie <[email protected]>
tikv · May 11, 2022 · 4751586 · 4751586
1 parent f9aacb3
commit 4751586
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 5 deletions.
diff --git a/include/rocksdb/rate_limiter.h b/include/rocksdb/rate_limiter.h
@@ -51,6 +51,9 @@ class RateLimiter : public Customizable {
   // REQUIRED: bytes_per_second > 0
   virtual void SetBytesPerSecond(int64_t bytes_per_second) = 0;
 
+  // Dynamically change rate limiter's auto_tuned mode.
+  virtual void SetAutoTuned(bool /*auto_tuned*/) {}
+
   // Deprecated. New RateLimiter derived classes should override
   // Request(const int64_t, const Env::IOPriority, Statistics*) or
   // Request(const int64_t, const Env::IOPriority, Statistics*, OpType)
@@ -122,6 +125,8 @@ class RateLimiter : public Customizable {
 
   virtual int64_t GetBytesPerSecond() const = 0;
 
+  virtual bool GetAutoTuned() const { return false; }
+
   virtual bool IsRateLimited(OpType op_type) {
     if ((mode_ == RateLimiter::Mode::kWritesOnly &&
          op_type == RateLimiter::OpType::kRead) ||

diff --git a/utilities/rate_limiters/write_amp_based_rate_limiter.cc b/utilities/rate_limiters/write_amp_based_rate_limiter.cc
@@ -78,13 +78,33 @@ WriteAmpBasedRateLimiter::~WriteAmpBasedRateLimiter() {
 
 void WriteAmpBasedRateLimiter::SetBytesPerSecond(int64_t bytes_per_second) {
   assert(bytes_per_second > 0);
-  if (auto_tuned_) {
+  if (auto_tuned_.load(std::memory_order_acquire)) {
     max_bytes_per_sec_.store(bytes_per_second, std::memory_order_relaxed);
   } else {
     SetActualBytesPerSecond(bytes_per_second);
   }
 }
 
+void WriteAmpBasedRateLimiter::SetAutoTuned(bool auto_tuned) {
+  MutexLock g(&auto_tuned_mutex_);
+  if (auto_tuned_.load(std::memory_order_acquire) != auto_tuned) {
+    if (auto_tuned) {
+      max_bytes_per_sec_.store(rate_bytes_per_sec_, std::memory_order_relaxed);
+      refill_bytes_per_period_.store(
+          CalculateRefillBytesPerPeriod(rate_bytes_per_sec_),
+          std::memory_order_relaxed);
+    } else {
+      // must hold this lock to avoid tuner changing `rate_bytes_per_sec_`
+      MutexLock g2(&request_mutex_);
+      rate_bytes_per_sec_ = max_bytes_per_sec_.load(std::memory_order_relaxed);
+      refill_bytes_per_period_.store(
+          CalculateRefillBytesPerPeriod(rate_bytes_per_sec_),
+          std::memory_order_relaxed);
+    }
+    auto_tuned_.store(auto_tuned, std::memory_order_release);
+  }
+}
+
 void WriteAmpBasedRateLimiter::SetActualBytesPerSecond(
     int64_t bytes_per_second) {
   rate_bytes_per_sec_ = bytes_per_second;
@@ -100,7 +120,7 @@ void WriteAmpBasedRateLimiter::Request(int64_t bytes, const Env::IOPriority pri,
   TEST_SYNC_POINT("WriteAmpBasedRateLimiter::Request");
   TEST_SYNC_POINT_CALLBACK("WriteAmpBasedRateLimiter::Request:1",
                            &rate_bytes_per_sec_);
-  if (auto_tuned_ && pri == Env::IO_HIGH &&
+  if (auto_tuned_.load(std::memory_order_acquire) && pri == Env::IO_HIGH &&
       duration_highpri_bytes_through_ + duration_bytes_through_ + bytes <=
           max_bytes_per_sec_.load(std::memory_order_relaxed) *
               kSecondsPerTune) {
@@ -112,7 +132,7 @@ void WriteAmpBasedRateLimiter::Request(int64_t bytes, const Env::IOPriority pri,
   assert(bytes <= refill_bytes_per_period_.load(std::memory_order_relaxed));
   MutexLock g(&request_mutex_);
 
-  if (auto_tuned_) {
+  if (auto_tuned_.load(std::memory_order_acquire)) {
     std::chrono::microseconds now(NowMicrosMonotonic(env_));
     if (now - tuned_time_ >= std::chrono::microseconds(kMicrosPerTune)) {
       Tune();
@@ -352,7 +372,7 @@ Status WriteAmpBasedRateLimiter::Tune() {
 }
 
 void WriteAmpBasedRateLimiter::PaceUp() {
-  if (auto_tuned_) {
+  if (auto_tuned_.load(std::memory_order_acquire)) {
     should_pace_up_.store(true, std::memory_order_relaxed);
   }
 }

diff --git a/utilities/rate_limiters/write_amp_based_rate_limiter.h b/utilities/rate_limiters/write_amp_based_rate_limiter.h
@@ -34,6 +34,9 @@ class WriteAmpBasedRateLimiter : public RateLimiter {
   // When auto-tuned is on, this sets rate limit's upper bound instead.
   virtual void SetBytesPerSecond(int64_t bytes_per_second) override;
 
+  // Dynamically change rate limiter's auto_tuned mode.
+  virtual void SetAutoTuned(bool auto_tuned) override;
+
   // Request for token to write bytes. If this request can not be satisfied,
   // the call is blocked. Caller is responsible to make sure
   // bytes <= GetSingleBurstBytes()
@@ -68,6 +71,10 @@ class WriteAmpBasedRateLimiter : public RateLimiter {
     return rate_bytes_per_sec_;
   }
 
+  virtual bool GetAutoTuned() const override {
+    return auto_tuned_.load(std::memory_order_acquire);
+  }
+
   virtual void PaceUp() override;
 
  private:
@@ -108,7 +115,10 @@ class WriteAmpBasedRateLimiter : public RateLimiter {
   Req* leader_;
   std::deque<Req*> queue_[Env::IO_TOTAL];
 
-  bool auto_tuned_;
+  // only used to synchronize auto_tuned setters
+  port::Mutex auto_tuned_mutex_;
+
+  std::atomic<bool> auto_tuned_;
   std::atomic<int64_t> max_bytes_per_sec_;
   std::chrono::microseconds tuned_time_;
   int64_t duration_highpri_bytes_through_;