1 //  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
2 //  This source code is licensed under both the GPLv2 (found in the
3 //  COPYING file in the root directory) and Apache 2.0 License
4 //  (found in the LICENSE.Apache file in the root directory).
5 //
6 // Copyright (c) 2011 The LevelDB Authors. All rights reserved.
7 // Use of this source code is governed by a BSD-style license that can be
8 // found in the LICENSE file. See the AUTHORS file for names of contributors.
9 
10 #pragma once
11 
12 #include <algorithm>
13 #include <atomic>
14 #include <chrono>
15 #include <deque>
16 #include "port/port.h"
17 #include "rocksdb/env.h"
18 #include "rocksdb/rate_limiter.h"
19 #include "util/mutexlock.h"
20 #include "util/random.h"
21 
22 namespace ROCKSDB_NAMESPACE {
23 
24 class GenericRateLimiter : public RateLimiter {
25  public:
26   GenericRateLimiter(int64_t refill_bytes, int64_t refill_period_us,
27                      int32_t fairness, RateLimiter::Mode mode, Env* env,
28                      bool auto_tuned);
29 
30   virtual ~GenericRateLimiter();
31 
32   // This API allows user to dynamically change rate limiter's bytes per second.
33   virtual void SetBytesPerSecond(int64_t bytes_per_second) override;
34 
35   // Request for token to write bytes. If this request can not be satisfied,
36   // the call is blocked. Caller is responsible to make sure
37   // bytes <= GetSingleBurstBytes()
38   using RateLimiter::Request;
39   virtual void Request(const int64_t bytes, const Env::IOPriority pri,
40                        Statistics* stats) override;
41 
GetSingleBurstBytes()42   virtual int64_t GetSingleBurstBytes() const override {
43     return refill_bytes_per_period_.load(std::memory_order_relaxed);
44   }
45 
46   virtual int64_t GetTotalBytesThrough(
47       const Env::IOPriority pri = Env::IO_TOTAL) const override {
48     MutexLock g(&request_mutex_);
49     if (pri == Env::IO_TOTAL) {
50       return total_bytes_through_[Env::IO_LOW] +
51              total_bytes_through_[Env::IO_HIGH];
52     }
53     return total_bytes_through_[pri];
54   }
55 
56   virtual int64_t GetTotalRequests(
57       const Env::IOPriority pri = Env::IO_TOTAL) const override {
58     MutexLock g(&request_mutex_);
59     if (pri == Env::IO_TOTAL) {
60       return total_requests_[Env::IO_LOW] + total_requests_[Env::IO_HIGH];
61     }
62     return total_requests_[pri];
63   }
64 
GetBytesPerSecond()65   virtual int64_t GetBytesPerSecond() const override {
66     return rate_bytes_per_sec_;
67   }
68 
69  private:
70   void Refill();
71   int64_t CalculateRefillBytesPerPeriod(int64_t rate_bytes_per_sec);
72   Status Tune();
73 
NowMicrosMonotonic(Env * env)74   uint64_t NowMicrosMonotonic(Env* env) {
75     return env->NowNanos() / std::milli::den;
76   }
77 
78   // This mutex guard all internal states
79   mutable port::Mutex request_mutex_;
80 
81   const int64_t kMinRefillBytesPerPeriod = 100;
82 
83   const int64_t refill_period_us_;
84 
85   int64_t rate_bytes_per_sec_;
86   // This variable can be changed dynamically.
87   std::atomic<int64_t> refill_bytes_per_period_;
88   Env* const env_;
89 
90   bool stop_;
91   port::CondVar exit_cv_;
92   int32_t requests_to_wait_;
93 
94   int64_t total_requests_[Env::IO_TOTAL];
95   int64_t total_bytes_through_[Env::IO_TOTAL];
96   int64_t available_bytes_;
97   int64_t next_refill_us_;
98 
99   int32_t fairness_;
100   Random rnd_;
101 
102   struct Req;
103   Req* leader_;
104   std::deque<Req*> queue_[Env::IO_TOTAL];
105 
106   bool auto_tuned_;
107   int64_t num_drains_;
108   int64_t prev_num_drains_;
109   const int64_t max_bytes_per_sec_;
110   std::chrono::microseconds tuned_time_;
111 };
112 
113 }  // namespace ROCKSDB_NAMESPACE
114