From f62fbd2c8575f6f14dabe1567f99a62a8ebe07c0 Mon Sep 17 00:00:00 2001
From: sdong <siying.d@fb.com>
Date: Fri, 27 May 2016 16:14:24 -0700
Subject: [PATCH] Handle overflow case of rate limiter's paramters

Summary: When rate_bytes_per_sec * refill_period_us_ overflows, the actual limited rate is very low. Handle this case so the rate will be large.

Test Plan: Add a unit test for it.

Reviewers: IslamAbdelRahman, andrewkr

Reviewed By: andrewkr

Subscribers: yiwu, lightmark, leveldb, andrewkr, dhruba

Differential Revision: https://reviews.facebook.net/D58929
---
 port/port_posix.h         |  1 +
 port/win/port_win.h       |  2 ++
 util/rate_limiter.cc      | 13 +++++++++++++
 util/rate_limiter.h       |  5 +----
 util/rate_limiter_test.cc |  5 +++++
 5 files changed, 22 insertions(+), 4 deletions(-)

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<int32_t>::max();
 const uint64_t kMaxUint64 = std::numeric_limits<uint64_t>::max();
+const int64_t kMaxInt64 = std::numeric_limits<int64_t>::max();
 const size_t kMaxSizet = std::numeric_limits<size_t>::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<int>::max();
 const uint64_t kMaxUint64 = std::numeric_limits<uint64_t>::max();
+const int64_t kMaxInt64 = std::numeric_limits<int64_t>::max();
 
 const size_t kMaxSizet = std::numeric_limits<size_t>::max();
 
@@ -94,6 +95,7 @@ const size_t kMaxSizet = std::numeric_limits<size_t>::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<RateLimiter> limiter(new GenericRateLimiter(100, 100, 10));
 }