/Users/deen/code/yugabyte-db/src/yb/util/net/rate_limiter.h
Line | Count | Source (jump to first uncovered line) |
1 | | // Copyright (c) YugaByte, Inc. |
2 | | // |
3 | | // Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except |
4 | | // in compliance with the License. You may obtain a copy of the License at |
5 | | // |
6 | | // http://www.apache.org/licenses/LICENSE-2.0 |
7 | | // |
8 | | // Unless required by applicable law or agreed to in writing, software distributed under the License |
9 | | // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express |
10 | | // or implied. See the License for the specific language governing permissions and limitations |
11 | | // under the License. |
12 | | // |
13 | | |
14 | | #ifndef YB_UTIL_NET_RATE_LIMITER_H |
15 | | #define YB_UTIL_NET_RATE_LIMITER_H |
16 | | |
17 | | #include <functional> |
18 | | #include <vector> |
19 | | |
20 | | #include "yb/util/monotime.h" |
21 | | #include "yb/util/status_fwd.h" |
22 | | |
23 | | namespace yb { |
24 | | |
25 | | class RateLimiter { |
26 | | public: |
27 | | // Inactive rate limiter. When calling SendOrReceiveData, the rate limiter will update its |
28 | | // statistics and call the passed send_rcv_func, but it will not try to control the rate. |
29 | | RateLimiter(); |
30 | | |
31 | | // Active rate limiter. target_rate_updater_ will be set to max_transmission_rate_updater. |
32 | | explicit RateLimiter(const std::function<uint64_t()>& max_transmission_rate_updater); |
33 | | |
34 | | ~RateLimiter(); |
35 | | |
36 | | // This method sets the function that will provide the desired target. For example, for remote |
37 | | // bootstrap sessions, the desired target for each session will depend on the total number of |
38 | | // concurrent remote bootstrap sessions: target_aggregate_rate / number_of_sessions. |
39 | 1.00k | void SetTargetRateUpdater(std::function<uint64_t()>&& target_rate_updater) { |
40 | 1.00k | target_rate_updater_ = std::move(target_rate_updater); |
41 | 1.00k | } |
42 | | |
43 | | // This function will be in charge of sending/receiving the data by calling send_rcv_func. This |
44 | | // function might sleep before returning to keep the transmission rate as close as possible to the |
45 | | // desired target. |
46 | | CHECKED_STATUS SendOrReceiveData(std::function<Status()> send_rcv_func, |
47 | | std::function<uint64_t()> reply_size_func); |
48 | | |
49 | | // Calculates the size for the next transmission so that the transmission rate remains as close |
50 | | // as possible to the target rate. |
51 | | uint64_t GetMaxSizeForNextTransmission(); |
52 | | |
53 | | // Get the transmission rate. Calculated as the number of bytes transmitted divided by the |
54 | | // time elapsed since the object was initialized. |
55 | | uint64_t GetRate(); |
56 | | |
57 | 5.83k | bool IsInitialized() { return init_; } |
58 | | |
59 | | // If the user of RateLimiter doesn't want to use SendOrReceiveData to send or receive data, this |
60 | | // method updates the internal stats and sleeps if it determines that the current rate is higher |
61 | | // than the rate provided by target_rate_updater_. |
62 | | void UpdateDataSizeAndMaybeSleep(uint64_t data_size); |
63 | | |
64 | | void Init(); |
65 | | |
66 | | // We can only have an active rate limiter if the user has provided a function to update the rate. |
67 | 49.9k | bool active() { return target_rate_updater_ != nullptr; } |
68 | | |
69 | | void SetTargetRate(uint64_t target_rate); |
70 | | |
71 | | #if defined(OS_MACOSX) |
72 | 0 | MonoDelta total_time_slept() { return total_time_slept_; } |
73 | | |
74 | | // Only used in MacOS. Instead of using the elapsed time for the calculation, we use the time |
75 | | // we spent sleeping. Used only for testing. |
76 | | // additional_time is passed by the test to include this time in the rate calculation. |
77 | 0 | uint64_t MacOSRate(MonoDelta additional_time = MonoDelta::FromMilliseconds(0)) { |
78 | 0 | return total_time_slept_.ToMilliseconds() || additional_time.ToMilliseconds() > 0 ? |
79 | 0 | MonoTime::kMillisecondsPerSecond * total_bytes_ / |
80 | 0 | (total_time_slept_ + additional_time).ToMilliseconds() : |
81 | 0 | 0; |
82 | 0 | } |
83 | | #endif |
84 | | |
85 | 8 | uint64_t time_slot_ms() const { |
86 | 8 | return time_slot_ms_; |
87 | 8 | } |
88 | | |
89 | | private: |
90 | | void UpdateRate(); |
91 | | void UpdateTimeSlotSizeAndMaybeSleep(uint64_t data_size, MonoDelta elapsed); |
92 | | uint64_t GetSizeForNextTimeSlot(); |
93 | | |
94 | | bool init_ = false; |
95 | | |
96 | | // Time when this rate was initialized. Used to calculate the long-term rate. |
97 | | MonoTime start_time_; |
98 | | |
99 | | // Last time stats were updated. |
100 | | MonoTime end_time_; |
101 | | |
102 | | // Reset every time the rate changes |
103 | | MonoTime rate_start_time_; |
104 | | |
105 | | #if defined(OS_MACOSX) |
106 | | // Total amount of time this object has spent sleeping. |
107 | | MonoDelta total_time_slept_ = MonoDelta::FromMicroseconds(0); |
108 | | #endif |
109 | | |
110 | | // Total number of bytes sent or received by the user of this RateLimiter object. |
111 | | uint64_t total_bytes_ = 0; |
112 | | |
113 | | // Time slot in milliseconds. This is just used internally and it's never exposed to the user. |
114 | | // The transmission data size is calculated as the number of bytes that should be transmitted in |
115 | | // time_slot_ms_ milliseconds in order to achieve the desired rate. |
116 | | uint64_t time_slot_ms_ = 100; |
117 | | |
118 | | // The minimum size that we will ever use for time_slot_ms_. |
119 | | uint64_t min_time_slot_ = 10; |
120 | | |
121 | | // The maximum size that we will ever use for time_slot_ms_. |
122 | | uint64_t max_time_slot_ = 100; |
123 | | |
124 | | // Maximum transmission rate in bytes/sec. Set by calling target_rate_updater_(). |
125 | | uint64_t target_rate_ = 0; |
126 | | std::vector<uint64_t> transmissions_rates_; |
127 | | |
128 | | std::function<uint64_t()> target_rate_updater_; |
129 | | }; |
130 | | |
131 | | } // namespace yb |
132 | | |
133 | | #endif // YB_UTIL_NET_RATE_LIMITER_H |