From 6251283442ff3a69ed7f9c4ab3ef8fc62197d1c0 Mon Sep 17 00:00:00 2001
From: Xinye Tao <xy.tao@outlook.com>
Date: Sat, 24 Oct 2020 11:42:10 +0800
Subject: [PATCH] Make auto-tuned rate limiter upper bound settable (#202)

Old implementation of auto-tuned rate limiter effectively ignore the
`SetBytesPerSecond` call from user because of its adaptive nature. Change this
semantics in `WriteAmpBasedRateLimiter` by setting `max-bytes-per-sec` in
auto-tuned mode.

Signed-off-by: tabokie <xy.tao@outlook.com>
---
 .../write_amp_based_rate_limiter.cc           | 26 +++++++++++++------
 .../write_amp_based_rate_limiter.h            |  4 ++-
 2 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/utilities/rate_limiters/write_amp_based_rate_limiter.cc b/utilities/rate_limiters/write_amp_based_rate_limiter.cc
index ee0dd47949c..18deb46b1a1 100644
--- a/utilities/rate_limiters/write_amp_based_rate_limiter.cc
+++ b/utilities/rate_limiters/write_amp_based_rate_limiter.cc
@@ -74,9 +74,17 @@ WriteAmpBasedRateLimiter::~WriteAmpBasedRateLimiter() {
   }
 }
 
-// This API allows user to dynamically change rate limiter's bytes per second.
 void WriteAmpBasedRateLimiter::SetBytesPerSecond(int64_t bytes_per_second) {
   assert(bytes_per_second > 0);
+  if (auto_tuned_) {
+    max_bytes_per_sec_.store(bytes_per_second, std::memory_order_relaxed);
+  } else {
+    SetActualBytesPerSecond(bytes_per_second);
+  }
+}
+
+void WriteAmpBasedRateLimiter::SetActualBytesPerSecond(
+    int64_t bytes_per_second) {
   rate_bytes_per_sec_ = bytes_per_second;
   refill_bytes_per_period_.store(
       CalculateRefillBytesPerPeriod(bytes_per_second),
@@ -91,8 +99,9 @@ void WriteAmpBasedRateLimiter::Request(int64_t bytes, const Env::IOPriority pri,
   TEST_SYNC_POINT_CALLBACK("WriteAmpBasedRateLimiter::Request:1",
                            &rate_bytes_per_sec_);
   if (auto_tuned_ && pri == Env::IO_HIGH &&
-      duration_highpri_bytes_through_ + bytes <=
-          max_bytes_per_sec_ * kSecondsPerTune) {
+      duration_highpri_bytes_through_ + duration_bytes_through_ + bytes <=
+          max_bytes_per_sec_.load(std::memory_order_relaxed) *
+              kSecondsPerTune) {
     total_bytes_through_[Env::IO_HIGH] += bytes;
     ++total_requests_[Env::IO_HIGH];
     duration_highpri_bytes_through_ += bytes;
@@ -321,12 +330,13 @@ Status WriteAmpBasedRateLimiter::Tune() {
   int64_t new_bytes_per_sec =
       (ratio + ratio_padding + ratio_delta_) *
       std::max(highpri_bytes_sampler_.GetRecentValue(), kHighBytesLower) / 10;
-  new_bytes_per_sec = std::max(
-      kMinBytesPerSec,
-      std::min(new_bytes_per_sec,
-               max_bytes_per_sec_ - highpri_bytes_sampler_.GetRecentValue()));
+  new_bytes_per_sec =
+      std::max(kMinBytesPerSec,
+               std::min(new_bytes_per_sec,
+                        max_bytes_per_sec_.load(std::memory_order_relaxed) -
+                            highpri_bytes_sampler_.GetRecentValue()));
   if (new_bytes_per_sec != prev_bytes_per_sec) {
-    SetBytesPerSecond(new_bytes_per_sec);
+    SetActualBytesPerSecond(new_bytes_per_sec);
   }
 
   duration_bytes_through_ = 0;
diff --git a/utilities/rate_limiters/write_amp_based_rate_limiter.h b/utilities/rate_limiters/write_amp_based_rate_limiter.h
index 02969b65db6..029925b7779 100644
--- a/utilities/rate_limiters/write_amp_based_rate_limiter.h
+++ b/utilities/rate_limiters/write_amp_based_rate_limiter.h
@@ -30,6 +30,7 @@ class WriteAmpBasedRateLimiter : public RateLimiter {
   virtual ~WriteAmpBasedRateLimiter();
 
   // This API allows user to dynamically change rate limiter's bytes per second.
+  // When auto-tuned is on, this sets rate limit's upper bound instead.
   virtual void SetBytesPerSecond(int64_t bytes_per_second) override;
 
   // Request for token to write bytes. If this request can not be satisfied,
@@ -69,6 +70,7 @@ class WriteAmpBasedRateLimiter : public RateLimiter {
  private:
   void Refill();
   int64_t CalculateRefillBytesPerPeriod(int64_t rate_bytes_per_sec);
+  void SetActualBytesPerSecond(int64_t bytes_per_second);
   Status Tune();
 
   uint64_t NowMicrosMonotonic(Env* env) {
@@ -104,7 +106,7 @@ class WriteAmpBasedRateLimiter : public RateLimiter {
   std::deque<Req*> queue_[Env::IO_TOTAL];
 
   bool auto_tuned_;
-  const int64_t max_bytes_per_sec_;
+  std::atomic<int64_t> max_bytes_per_sec_;
   std::chrono::microseconds tuned_time_;
   int64_t duration_highpri_bytes_through_;
   int64_t duration_bytes_through_;