ceph/src/rocksdb/util/rate_limiter.h

   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
  17 #include "port/port.h"
  18 #include "rocksdb/env.h"
  19 #include "rocksdb/rate_limiter.h"
  20 #include "rocksdb/status.h"
  21 #include "rocksdb/system_clock.h"
  22 #include "util/mutexlock.h"
  23 #include "util/random.h"
  24
  25 namespace ROCKSDB_NAMESPACE {
  26
  27 class GenericRateLimiter : public RateLimiter {
  28  public:
  29   GenericRateLimiter(int64_t refill_bytes, int64_t refill_period_us,
  30                      int32_t fairness, RateLimiter::Mode mode,
  31                      const std::shared_ptr<SystemClock>& clock,
  32                      bool auto_tuned);
  33
  34   virtual ~GenericRateLimiter();
  35
  36   // This API allows user to dynamically change rate limiter's bytes per second.
  37   virtual void SetBytesPerSecond(int64_t bytes_per_second) override;
  38
  39   // Request for token to write bytes. If this request can not be satisfied,
  40   // the call is blocked. Caller is responsible to make sure
  41   // bytes <= GetSingleBurstBytes() and bytes >= 0. Negative bytes
  42   // passed in will be rounded up to 0.
  43   using RateLimiter::Request;
  44   virtual void Request(const int64_t bytes, const Env::IOPriority pri,
  45                        Statistics* stats) override;
  46
  47   virtual int64_t GetSingleBurstBytes() const override {
  48     return refill_bytes_per_period_.load(std::memory_order_relaxed);
  49   }
  50
  51   virtual int64_t GetTotalBytesThrough(
  52       const Env::IOPriority pri = Env::IO_TOTAL) const override {
  53     MutexLock g(&request_mutex_);
  54     if (pri == Env::IO_TOTAL) {
  55       int64_t total_bytes_through_sum = 0;
  56       for (int i = Env::IO_LOW; i < Env::IO_TOTAL; ++i) {
  57         total_bytes_through_sum += total_bytes_through_[i];
  58       }
  59       return total_bytes_through_sum;
  60     }
  61     return total_bytes_through_[pri];
  62   }
  63
  64   virtual int64_t GetTotalRequests(
  65       const Env::IOPriority pri = Env::IO_TOTAL) const override {
  66     MutexLock g(&request_mutex_);
  67     if (pri == Env::IO_TOTAL) {
  68       int64_t total_requests_sum = 0;
  69       for (int i = Env::IO_LOW; i < Env::IO_TOTAL; ++i) {
  70         total_requests_sum += total_requests_[i];
  71       }
  72       return total_requests_sum;
  73     }
  74     return total_requests_[pri];
  75   }
  76
  77   virtual Status GetTotalPendingRequests(
  78       int64_t* total_pending_requests,
  79       const Env::IOPriority pri = Env::IO_TOTAL) const override {
  80     assert(total_pending_requests != nullptr);
  81     MutexLock g(&request_mutex_);
  82     if (pri == Env::IO_TOTAL) {
  83       int64_t total_pending_requests_sum = 0;
  84       for (int i = Env::IO_LOW; i < Env::IO_TOTAL; ++i) {
  85         total_pending_requests_sum += static_cast<int64_t>(queue_[i].size());
  86       }
  87       *total_pending_requests = total_pending_requests_sum;
  88     } else {
  89       *total_pending_requests = static_cast<int64_t>(queue_[pri].size());
  90     }
  91     return Status::OK();
  92   }
  93
  94   virtual int64_t GetBytesPerSecond() const override {
  95     return rate_bytes_per_sec_.load(std::memory_order_relaxed);
  96   }
  97
  98   virtual void TEST_SetClock(std::shared_ptr<SystemClock> clock) {
  99     MutexLock g(&request_mutex_);
 100     clock_ = std::move(clock);
 101     next_refill_us_ = NowMicrosMonotonicLocked();
 102   }
 103
 104  private:
 105   void RefillBytesAndGrantRequestsLocked();
 106   std::vector<Env::IOPriority> GeneratePriorityIterationOrderLocked();
 107   int64_t CalculateRefillBytesPerPeriodLocked(int64_t rate_bytes_per_sec);
 108   Status TuneLocked();
 109   void SetBytesPerSecondLocked(int64_t bytes_per_second);
 110
 111   uint64_t NowMicrosMonotonicLocked() {
 112     return clock_->NowNanos() / std::milli::den;
 113   }
 114
 115   // This mutex guard all internal states
 116   mutable port::Mutex request_mutex_;
 117
 118   const int64_t refill_period_us_;
 119
 120   std::atomic<int64_t> rate_bytes_per_sec_;
 121   std::atomic<int64_t> refill_bytes_per_period_;
 122   std::shared_ptr<SystemClock> clock_;
 123
 124   bool stop_;
 125   port::CondVar exit_cv_;
 126   int32_t requests_to_wait_;
 127
 128   int64_t total_requests_[Env::IO_TOTAL];
 129   int64_t total_bytes_through_[Env::IO_TOTAL];
 130   int64_t available_bytes_;
 131   int64_t next_refill_us_;
 132
 133   int32_t fairness_;
 134   Random rnd_;
 135
 136   struct Req;
 137   std::deque<Req*> queue_[Env::IO_TOTAL];
 138   bool wait_until_refill_pending_;
 139
 140   bool auto_tuned_;
 141   int64_t num_drains_;
 142   const int64_t max_bytes_per_sec_;
 143   std::chrono::microseconds tuned_time_;
 144 };
 145
 146 }  // namespace ROCKSDB_NAMESPACE