From bc68bd5a13091ee03fcdb1896cbfda157746bccf Mon Sep 17 00:00:00 2001
From: sdong <siying.d@fb.com>
Date: Tue, 12 May 2015 17:40:27 -0700
Subject: [PATCH] db_bench to support rate limiter

Summary: Add --rate_limiter_bytes_per_sec to db_bench to allow rater limit to disk

Test Plan:
Run
./db_bench --benchmarks=fillseq --num=30000000 --rate_limiter_bytes_per_sec=3000000 --num_multi_db=8 -disable_wal
And see io_stats to have the rate limited.

Reviewers: yhchiang, rven, anthony, kradhakrishnan, igor

Reviewed By: igor

Subscribers: leveldb, dhruba

Differential Revision: https://reviews.facebook.net/D38385
---
 db/db_bench.cc | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/db/db_bench.cc b/db/db_bench.cc
index dee16878e..dd694d3ec 100644
--- a/db/db_bench.cc
+++ b/db/db_bench.cc
@@ -42,6 +42,7 @@ int main() {
 #include "rocksdb/write_batch.h"
 #include "rocksdb/slice.h"
 #include "rocksdb/filter_policy.h"
+#include "rocksdb/rate_limiter.h"
 #include "rocksdb/slice_transform.h"
 #include "rocksdb/perf_context.h"
 #include "rocksdb/utilities/flashcache.h"
@@ -509,6 +510,8 @@ DEFINE_int32(rate_limit_delay_max_milliseconds, 1000,
              "When hard_rate_limit is set then this is the max time a put will"
              " be stalled.");
 
+DEFINE_uint64(rate_limiter_bytes_per_sec, 0, "Set options.rate_limiter value.");
+
 DEFINE_int32(max_grandparent_overlap_factor, 10, "Control maximum bytes of "
              "overlaps in grandparent (i.e., level+2) before we stop building a"
              " single file in a level->level+1 compaction.");
@@ -2256,6 +2259,10 @@ class Benchmark {
     if (FLAGS_thread_status_per_interval > 0) {
       options.enable_thread_tracking = true;
     }
+    if (FLAGS_rate_limiter_bytes_per_sec > 0) {
+      options.rate_limiter.reset(
+          NewGenericRateLimiter(FLAGS_rate_limiter_bytes_per_sec));
+    }
 
     if (FLAGS_num_multi_db <= 1) {
       OpenDb(options, FLAGS_db, &db_);