Fix alloc_bytes of ReadLimiter

dbms/src/Encryption/RateLimiter.cpp

@@ -23,6 +23,7 @@
 #include <boost/algorithm/string.hpp>
 #include <cassert>
 #include <fstream>
+#include <magic_enum.hpp>
 
 namespace CurrentMetrics
 {
@@ -107,6 +108,7 @@ WriteLimiter::WriteLimiter(Int64 rate_limit_per_sec_, LimiterType type_, UInt64
  , requests_to_wait{0}
  , type(type_)
  , alloc_bytes{0}
+ , log(Logger::get(std::string(magic_enum::enum_name(type))))
 {}
 
 WriteLimiter::~WriteLimiter()
@@ -131,7 +133,8 @@ void WriteLimiter::request(Int64 bytes)
  consumeBytes(bytes);
  return;
  }
-
+ Stopwatch sw_pending;
+ Int64 wait_times = 0;
  auto pending_request = pendingRequestMetrics(type);
 
  // request cannot be satisfied at this moment, enqueue
@@ -140,7 +143,7 @@ void WriteLimiter::request(Int64 bytes)
  while (!r.granted)
  {
  assert(!req_queue.empty());
-
+ wait_times++;
  bool timed_out = false;
  // if this request is in the front of req_queue,
  // then it is responsible to trigger the refill process.
@@ -191,6 +194,7 @@ void WriteLimiter::request(Int64 bytes)
  }
  }
  }
+ LOG_FMT_TRACE(log, "pending_us {} wait_times {} pending_count {} rate_limit_per_sec {}", sw_pending.elapsed() / 1000, wait_times, req_queue.size(), refill_balance_per_period * 1000 / refill_period_ms);
 }
 
 size_t WriteLimiter::setStop()
@@ -296,7 +300,7 @@ ReadLimiter::ReadLimiter(
  : WriteLimiter(rate_limit_per_sec_, type_, refill_period_ms_)
  , get_read_bytes(std::move(get_read_bytes_))
  , last_stat_bytes(get_read_bytes())
- , log(Logger::get("ReadLimiter"))
+ , last_refill_time(std::chrono::system_clock::now())
 {}
 
 Int64 ReadLimiter::getAvailableBalance()
@@ -317,17 +321,21 @@ Int64 ReadLimiter::getAvailableBalance()
  else
  {
  Int64 real_alloc_bytes = bytes - last_stat_bytes;
- metricAllocBytes(type, real_alloc_bytes);
+ // `alloc_bytes` is the number of byte that ReadLimiter has allocated.
+ if (available_balance > 0)
+ {
+ auto can_alloc_bytes = std::min(real_alloc_bytes, available_balance);
+ alloc_bytes += can_alloc_bytes;
+ metricAllocBytes(type, can_alloc_bytes);
+ }
  available_balance -= real_alloc_bytes;
- alloc_bytes += real_alloc_bytes;
  }
  last_stat_bytes = bytes;
  return available_balance;
 }
 
-void ReadLimiter::consumeBytes(Int64 bytes)
+void ReadLimiter::consumeBytes([[maybe_unused]] Int64 bytes)
 {
- metricRequestBytes(type, bytes);
  // Do nothing for read.
 }
 
@@ -338,10 +346,26 @@ bool ReadLimiter::canGrant([[maybe_unused]] Int64 bytes)
 
 void ReadLimiter::refillAndAlloc()
 {
- if (available_balance < refill_balance_per_period)
+ // `available_balance` of `ReadLimiter` may be overdrawn.
+ if (available_balance < 0)
+ {
+ // Limiter may not be called for a long time.
+ // During this time, limiter can be refilled at most `max_refill_times` times and covers some overdraft.
+ auto elapsed_duration = std::chrono::system_clock::now() - last_refill_time;
+ UInt64 elapsed_ms = std::chrono::duration_cast<std::chrono::milliseconds>(elapsed_duration).count();
+ // At least refill one time.
+ Int64 max_refill_times = std::max(elapsed_ms, refill_period_ms) / refill_period_ms;
+ Int64 max_refill_bytes = max_refill_times * refill_balance_per_period;
+ Int64 can_alloc_bytes = std::min(-available_balance, max_refill_bytes);
+ alloc_bytes += can_alloc_bytes;
+ metricAllocBytes(type, can_alloc_bytes);
+ available_balance = std::min(available_balance + max_refill_bytes, refill_balance_per_period);
+ }
+ else
  {
- available_balance += refill_balance_per_period;
+ available_balance = refill_balance_per_period;
  }
+ last_refill_time = std::chrono::system_clock::now();
 
  assert(!req_queue.empty());
  auto * head_req = req_queue.front();
@@ -710,26 +734,13 @@ IOLimitTuner::TuneResult IOLimitTuner::tune() const
  }
 
  auto [max_read_bytes_per_sec, max_write_bytes_per_sec, rw_tuned] = tuneReadWrite();
- LOG_FMT_INFO(
- log,
- "tuneReadWrite: max_read {} max_write {} rw_tuned {}",
- max_read_bytes_per_sec,
- max_write_bytes_per_sec,
- rw_tuned);
  auto [max_bg_read_bytes_per_sec, max_fg_read_bytes_per_sec, read_tuned] = tuneRead(max_read_bytes_per_sec);
- LOG_FMT_INFO(
- log,
- "tuneRead: bg_read {} fg_read {} read_tuned {}",
- max_bg_read_bytes_per_sec,
- max_fg_read_bytes_per_sec,
- read_tuned);
  auto [max_bg_write_bytes_per_sec, max_fg_write_bytes_per_sec, write_tuned] = tuneWrite(max_write_bytes_per_sec);
- LOG_FMT_INFO(
- log,
- "tuneWrite: bg_write {} fg_write {} write_tuned {}",
- max_bg_write_bytes_per_sec,
- max_fg_write_bytes_per_sec,
- write_tuned);
+ if (rw_tuned || read_tuned || write_tuned)
+ {
+ LOG_FMT_INFO(log, "tune_msg: bg_write {} => {} fg_write {} => {} bg_read {} => {} fg_read {} => {}", bg_write_stat != nullptr ? bg_write_stat->maxBytesPerSec() : 0, max_bg_write_bytes_per_sec, fg_write_stat != nullptr ? fg_write_stat->maxBytesPerSec() : 0, max_fg_write_bytes_per_sec, bg_read_stat != nullptr ? bg_read_stat->maxBytesPerSec() : 0, max_bg_read_bytes_per_sec, fg_read_stat != nullptr ? fg_read_stat->maxBytesPerSec() : 0, max_fg_read_bytes_per_sec);
+ }
+
  return {.max_bg_read_bytes_per_sec = max_bg_read_bytes_per_sec,
  .max_fg_read_bytes_per_sec = max_fg_read_bytes_per_sec,
  .read_tuned = read_tuned || rw_tuned,

dbms/src/Encryption/RateLimiter.h

@@ -93,7 +93,10 @@ class WriteLimiter
  void request(Int64 bytes);
 
  // just for test purpose
- inline UInt64 getTotalBytesThrough() const { return alloc_bytes; }
+ inline UInt64 getTotalBytesThrough() const
+ {
+ return available_balance < 0 ? alloc_bytes - available_balance : alloc_bytes;
+ }
 
  LimiterStat getStat();
 
@@ -153,6 +156,7 @@ class WriteLimiter
 
  Stopwatch stat_stop_watch;
  UInt64 alloc_bytes;
+ LoggerPtr log;
 };
 
 using WriteLimiterPtr = std::shared_ptr<WriteLimiter>;
@@ -194,7 +198,7 @@ class ReadLimiter : public WriteLimiter
 
  std::function<Int64()> get_read_bytes;
  Int64 last_stat_bytes;
- LoggerPtr log;
+ std::chrono::time_point<std::chrono::system_clock> last_refill_time;
 };
 
 using ReadLimiterPtr = std::shared_ptr<ReadLimiter>;

dbms/src/Encryption/tests/gtest_rate_limiter.cpp

@@ -22,8 +22,6 @@
 #include <random>
 #include <thread>
 
-#include "common/types.h"
-
 #ifdef __linux__
 #include <sys/syscall.h>
 #endif
@@ -333,23 +331,47 @@ TEST(ReadLimiterTest, LimiterStat)
  ASSERT_GT(stat.pct(), 100) << stat.toString();
  }
 
- static constexpr UInt64 alloc_bytes = 2047;
+ static constexpr UInt64 total_bytes = 2047;
  for (int i = 0; i < 11; i++)
  {
  request(read_limiter, 1 << i);
  }
-
  std::this_thread::sleep_for(100ms);
  ASSERT_EQ(read_limiter.getAvailableBalance(), -947);
 
  stat = read_limiter.getStat();
- ASSERT_EQ(stat.alloc_bytes, alloc_bytes);
- ASSERT_GE(stat.elapsed_ms, alloc_bytes / 100 + 1);
+ ASSERT_EQ(stat.alloc_bytes, total_bytes + read_limiter.getAvailableBalance());
+ ASSERT_GE(stat.elapsed_ms, stat.alloc_bytes / 100 + 1);
  ASSERT_EQ(stat.refill_period_ms, 100ul);
  ASSERT_EQ(stat.refill_bytes_per_period, 100);
  ASSERT_EQ(stat.maxBytesPerSec(), 1000);
- ASSERT_EQ(stat.avgBytesPerSec(), static_cast<Int64>(alloc_bytes * 1000 / stat.elapsed_ms)) << stat.toString();
- ASSERT_EQ(stat.pct(), static_cast<Int64>(alloc_bytes * 1000 / stat.elapsed_ms) * 100 / stat.maxBytesPerSec()) << stat.toString();
+ ASSERT_EQ(stat.avgBytesPerSec(), static_cast<Int64>(stat.alloc_bytes * 1000 / stat.elapsed_ms)) << stat.toString();
+ ASSERT_EQ(stat.pct(), static_cast<Int64>(stat.alloc_bytes * 1000 / stat.elapsed_ms) * 100 / stat.maxBytesPerSec()) << stat.toString();
+}
+
+TEST(ReadLimiterTest, ReadMany)
+{
+ Int64 real_read_bytes{0};
+ auto get_read_bytes = [&]() {
+ return real_read_bytes;
+ };
+ auto request = [&](ReadLimiter & limiter, Int64 bytes) {
+ limiter.request(bytes);
+ real_read_bytes += bytes;
+ };
+
+ constexpr Int64 bytes_per_sec = 1000;
+ constexpr UInt64 refill_period_ms = 100;
+ ReadLimiter read_limiter(get_read_bytes, bytes_per_sec, LimiterType::UNKNOW, refill_period_ms);
+ ASSERT_EQ(read_limiter.getAvailableBalance(), 100);
+ request(read_limiter, 1000);
+ ASSERT_EQ(read_limiter.getAvailableBalance(), -900);
+ ASSERT_EQ(read_limiter.alloc_bytes, 100);
+
+ std::this_thread::sleep_for(1200ms);
+ Stopwatch sw;
+ request(read_limiter, 100);
+ ASSERT_LE(sw.elapsedMilliseconds(), 1); // Not blocked.
 }
 
 #ifdef __linux__

dbms/src/TestUtils/MockReadLimiter.h

@@ -31,8 +31,8 @@ class MockReadLimiter final : public ReadLimiter
 protected:
  void consumeBytes(Int64 bytes) override
  {
- // Need soft limit here.
- WriteLimiter::consumeBytes(bytes); // NOLINT(bugprone-parent-virtual-call)
+ alloc_bytes += std::min(available_balance, bytes);
+ available_balance -= bytes;
  }
 };