diff --git a/port/port_posix.h b/port/port_posix.h index 454d6c1c3..15c4d0c0a 100644 --- a/port/port_posix.h +++ b/port/port_posix.h @@ -80,6 +80,7 @@ namespace port { // For use at db/file_indexer.h kLevelMaxIndex const int kMaxInt32 = std::numeric_limits::max(); const uint64_t kMaxUint64 = std::numeric_limits::max(); +const int64_t kMaxInt64 = std::numeric_limits::max(); const size_t kMaxSizet = std::numeric_limits::max(); static const bool kLittleEndian = PLATFORM_IS_LITTLE_ENDIAN; diff --git a/port/win/port_win.h b/port/win/port_win.h index a677bd662..54f10a24c 100644 --- a/port/win/port_win.h +++ b/port/win/port_win.h @@ -80,6 +80,7 @@ namespace port { // For use at db/file_indexer.h kLevelMaxIndex const int kMaxInt32 = std::numeric_limits::max(); const uint64_t kMaxUint64 = std::numeric_limits::max(); +const int64_t kMaxInt64 = std::numeric_limits::max(); const size_t kMaxSizet = std::numeric_limits::max(); @@ -94,6 +95,7 @@ const size_t kMaxSizet = std::numeric_limits::max(); // For use at db/file_indexer.h kLevelMaxIndex const int kMaxInt32 = INT32_MAX; +const int64_t kMaxInt64 = INT64_MAX; const uint64_t kMaxUint64 = UINT64_MAX; #ifdef _WIN64 diff --git a/util/rate_limiter.cc b/util/rate_limiter.cc index 352925bec..4e836d030 100644 --- a/util/rate_limiter.cc +++ b/util/rate_limiter.cc @@ -8,6 +8,7 @@ // found in the LICENSE file. See the AUTHORS file for names of contributors. #include "util/rate_limiter.h" +#include "port/port.h" #include "rocksdb/env.h" namespace rocksdb { @@ -204,6 +205,18 @@ void GenericRateLimiter::Refill() { } } +int64_t GenericRateLimiter::CalculateRefillBytesPerPeriod( + int64_t rate_bytes_per_sec) { + if (port::kMaxInt64 / rate_bytes_per_sec < refill_period_us_) { + // Avoid unexpected result in the overflow case. The result now is still + // inaccurate but is a number that is large enough. + return port::kMaxInt64 / 1000000; + } else { + return std::max(kMinRefillBytesPerPeriod, + rate_bytes_per_sec * refill_period_us_ / 1000000); + } +} + RateLimiter* NewGenericRateLimiter( int64_t rate_bytes_per_sec, int64_t refill_period_us, int32_t fairness) { assert(rate_bytes_per_sec > 0); diff --git a/util/rate_limiter.h b/util/rate_limiter.h index ea2975d64..ddeaeba10 100644 --- a/util/rate_limiter.h +++ b/util/rate_limiter.h @@ -60,10 +60,7 @@ class GenericRateLimiter : public RateLimiter { private: void Refill(); - int64_t CalculateRefillBytesPerPeriod(int64_t rate_bytes_per_sec) { - return std::max(kMinRefillBytesPerPeriod, - rate_bytes_per_sec * refill_period_us_ / 1000000); - } + int64_t CalculateRefillBytesPerPeriod(int64_t rate_bytes_per_sec); // This mutex guard all internal states mutable port::Mutex request_mutex_; diff --git a/util/rate_limiter_test.cc b/util/rate_limiter_test.cc index 9085835de..d1152ed56 100644 --- a/util/rate_limiter_test.cc +++ b/util/rate_limiter_test.cc @@ -22,6 +22,11 @@ namespace rocksdb { class RateLimiterTest : public testing::Test {}; +TEST_F(RateLimiterTest, OverflowRate) { + GenericRateLimiter limiter(port::kMaxInt64, 1000, 10); + ASSERT_GT(limiter.GetSingleBurstBytes(), 1000000000ll); +} + TEST_F(RateLimiterTest, StartStop) { std::unique_ptr limiter(new GenericRateLimiter(100, 100, 10)); }