From b409091a36a51376081ea456e0351b08438d3079 Mon Sep 17 00:00:00 2001
From: jinhelin <linjinhe33@gmail.com>
Date: Mon, 10 Oct 2022 21:54:54 +0800
Subject: [PATCH] Fix I/O limiter auto-tuner.

---
 dbms/src/Encryption/RateLimiter.cpp           | 28 +++++++++++++++++
 dbms/src/Encryption/RateLimiter.h             |  5 +--
 .../Encryption/tests/gtest_rate_limiter.cpp   | 31 +++++++++++++++++++
 3 files changed, 62 insertions(+), 2 deletions(-)

diff --git a/dbms/src/Encryption/RateLimiter.cpp b/dbms/src/Encryption/RateLimiter.cpp
index afd24b72ad3..b8695dd631a 100644
--- a/dbms/src/Encryption/RateLimiter.cpp
+++ b/dbms/src/Encryption/RateLimiter.cpp
@@ -897,4 +897,32 @@ IOLimitTuner::Watermark IOLimitTuner::getWatermark(int pct) const
         return Watermark::Low;
     }
 }
+
+IOLimitTuner::Watermark IOLimitTuner::getWatermark(const LimiterStatUPtr & fg, const LimiterStatUPtr & bg, int pct) const
+{
+    // There is a cornar case:
+    // 1. Both `fg_mbps` and `bg_mbps` are less than `min_bytes_per_sec`.
+    // 2. `bg` runs out of the quota, but `fg` has not been used.
+    // 3. In this case, `pct` can be less than `medium_pct` and watermark is LOW.
+    // 4. Low watermark means it is unnecessary to increase the quota by descrease the quota of others.
+    // 5. However, because `fg_mbps` is less than `min_bytes_per_sec`, `bg_mbps` cannot be increased by decreasing `fg_mbps`.
+    if (fg != nullptr && bg != nullptr)
+    {
+        auto fg_wm = getWatermark(fg->pct());
+        auto bg_wm = getWatermark(bg->pct());
+        auto fg_mbps = fg->maxBytesPerSec();
+        auto bg_mbps = bg->maxBytesPerSec();
+        // `fg` needs more bandwidth, but `bg`'s bandwidth is small.
+        if (fg_wm > bg_wm && bg_mbps <= io_config.min_bytes_per_sec * 2)
+        {
+            return fg_wm;
+        }
+        // `bg_read` needs more bandwidth, but `fg_read`'s bandwidth is small.
+        else if (bg_wm > fg_wm && fg_mbps <= io_config.min_bytes_per_sec * 2)
+        {
+            return bg_wm;
+        }
+    }
+    return getWatermark(pct);
+}
 } // namespace DB
diff --git a/dbms/src/Encryption/RateLimiter.h b/dbms/src/Encryption/RateLimiter.h
index b0578433d58..25f82212602 100644
--- a/dbms/src/Encryption/RateLimiter.h
+++ b/dbms/src/Encryption/RateLimiter.h
@@ -412,13 +412,14 @@ class IOLimitTuner
     };
     Watermark writeWatermark() const
     {
-        return getWatermark(writePct());
+        return getWatermark(fg_write_stat, bg_write_stat, writePct());
     }
     Watermark readWatermark() const
     {
-        return getWatermark(readPct());
+        return getWatermark(fg_read_stat, bg_read_stat, readPct());
     }
     Watermark getWatermark(int pct) const;
+    Watermark getWatermark(const LimiterStatUPtr & fg, const LimiterStatUPtr & bg, int pct) const;
 
     // Returns <max_read_bytes_per_sec, max_write_bytes_per_sec, has_tuned>
     std::tuple<Int64, Int64, bool> tuneReadWrite() const;
diff --git a/dbms/src/Encryption/tests/gtest_rate_limiter.cpp b/dbms/src/Encryption/tests/gtest_rate_limiter.cpp
index 9913aab9dd7..613e552def5 100644
--- a/dbms/src/Encryption/tests/gtest_rate_limiter.cpp
+++ b/dbms/src/Encryption/tests/gtest_rate_limiter.cpp
@@ -731,5 +731,36 @@ TEST(IOLimitTunerTest, Tune)
     }
 }
 
+TEST(IOLimitTunerTest, Tune2)
+{
+    StorageIORateLimitConfig io_config;
+    io_config.max_bytes_per_sec = 2000;
+    io_config.min_bytes_per_sec = 10;
+
+    auto bg_write_stat = createLimiterStat(0, 1000, 1000, 990);
+    auto fg_write_stat = createLimiterStat(0, 1000, 1000, 990);
+    auto bg_read_stat = createLimiterStat(10, 1000, 1000, 10);
+    auto fg_read_stat = createLimiterStat(0, 1000, 1000, 10);
+
+    ASSERT_EQ(bg_write_stat->pct(), 0);
+    ASSERT_EQ(fg_write_stat->pct(), 0);
+    ASSERT_EQ(bg_read_stat->pct(), 100);
+    ASSERT_EQ(fg_read_stat->pct(), 0);
+    ASSERT_EQ(bg_read_stat->maxBytesPerSec(), 10);
+
+    IOLimitTuner tuner(
+        std::move(bg_write_stat),
+        std::move(fg_write_stat),
+        std::move(bg_read_stat),
+        std::move(fg_read_stat),
+        io_config);
+    ASSERT_EQ(tuner.readWatermark(), Emergency);
+
+    auto res = tuner.tune();
+    ASSERT_TRUE(res.write_tuned);
+    ASSERT_TRUE(res.read_tuned);
+    ASSERT_GT(res.max_bg_read_bytes_per_sec, 10);
+}
+
 } // namespace tests
 } // namespace DB