Skip to content

Commit e576f2a

Browse files
ajkrfacebook-github-bot
ajkr
authored and
facebook-github-bot
committed
Fix race conditions in GenericRateLimiter (#10374)
Summary: Made locking strict for all accesses of `GenericRateLimiter` internal state. `SetBytesPerSecond()` was the main problem since it had no locking, while the two updates it makes need to be done as one atomic operation. The test case, "ConfigOptionsTest.ConfiguringOptionsDoesNotRevertRateLimiterBandwidth", is for the issue fixed in #10378, but I forgot to include the test there. Pull Request resolved: #10374 Reviewed By: pdillinger Differential Revision: D37906367 Pulled By: ajkr fbshipit-source-id: ccde620d2a7f96d1401bdafd2bdb685cbefbafa5
1 parent 0b6bc10 commit e576f2a

File tree

5 files changed

+66
-38
lines changed

5 files changed

+66
-38
lines changed

‎HISTORY.md

+1
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99

1010
### Bug Fixes
1111
* Fix a bug where `GenericRateLimiter` could revert the bandwidth set dynamically using `SetBytesPerSecond()` when a user configures a structure enclosing it, e.g., using `GetOptionsFromString()` to configure an `Options` that references an existing `RateLimiter` object.
12+
* Fix race conditions in `GenericRateLimiter`.
1213

1314
## 7.5.0 (07/15/2022)
1415
### New Features

‎options/options_test.cc

+16
Original file line numberDiff line numberDiff line change
@@ -4958,6 +4958,22 @@ TEST_F(ConfigOptionsTest, MergeOperatorFromString) {
49584958
ASSERT_EQ(*delimiter, "&&");
49594959
}
49604960

4961+
TEST_F(ConfigOptionsTest, ConfiguringOptionsDoesNotRevertRateLimiterBandwidth) {
4962+
// Regression test for bug where rate limiter's dynamically set bandwidth
4963+
// could be silently reverted when configuring an options structure with an
4964+
// existing `rate_limiter`.
4965+
Options base_options;
4966+
base_options.rate_limiter.reset(
4967+
NewGenericRateLimiter(1 << 20 /* rate_bytes_per_sec */));
4968+
Options copy_options(base_options);
4969+
4970+
base_options.rate_limiter->SetBytesPerSecond(2 << 20);
4971+
ASSERT_EQ(2 << 20, base_options.rate_limiter->GetBytesPerSecond());
4972+
4973+
ASSERT_OK(GetOptionsFromString(base_options, "", &copy_options));
4974+
ASSERT_EQ(2 << 20, base_options.rate_limiter->GetBytesPerSecond());
4975+
}
4976+
49614977
INSTANTIATE_TEST_CASE_P(OptionsSanityCheckTest, OptionsSanityCheckTest,
49624978
::testing::Bool());
49634979
#endif // !ROCKSDB_LITE

‎util/rate_limiter.cc

+27-21
Original file line numberDiff line numberDiff line change
@@ -54,20 +54,20 @@ GenericRateLimiter::GenericRateLimiter(
5454
rate_bytes_per_sec_(auto_tuned ? rate_bytes_per_sec / 2
5555
: rate_bytes_per_sec),
5656
refill_bytes_per_period_(
57-
CalculateRefillBytesPerPeriod(rate_bytes_per_sec_)),
57+
CalculateRefillBytesPerPeriodLocked(rate_bytes_per_sec_)),
5858
clock_(clock),
5959
stop_(false),
6060
exit_cv_(&request_mutex_),
6161
requests_to_wait_(0),
6262
available_bytes_(0),
63-
next_refill_us_(NowMicrosMonotonic()),
63+
next_refill_us_(NowMicrosMonotonicLocked()),
6464
fairness_(fairness > 100 ? 100 : fairness),
6565
rnd_((uint32_t)time(nullptr)),
6666
wait_until_refill_pending_(false),
6767
auto_tuned_(auto_tuned),
6868
num_drains_(0),
6969
max_bytes_per_sec_(rate_bytes_per_sec),
70-
tuned_time_(NowMicrosMonotonic()) {
70+
tuned_time_(NowMicrosMonotonicLocked()) {
7171
for (int i = Env::IO_LOW; i < Env::IO_TOTAL; ++i) {
7272
total_requests_[i] = 0;
7373
total_bytes_through_[i] = 0;
@@ -97,10 +97,15 @@ GenericRateLimiter::~GenericRateLimiter() {
9797

9898
// This API allows user to dynamically change rate limiter's bytes per second.
9999
void GenericRateLimiter::SetBytesPerSecond(int64_t bytes_per_second) {
100+
MutexLock g(&request_mutex_);
101+
SetBytesPerSecondLocked(bytes_per_second);
102+
}
103+
104+
void GenericRateLimiter::SetBytesPerSecondLocked(int64_t bytes_per_second) {
100105
assert(bytes_per_second > 0);
101-
rate_bytes_per_sec_ = bytes_per_second;
106+
rate_bytes_per_sec_.store(bytes_per_second, std::memory_order_relaxed);
102107
refill_bytes_per_period_.store(
103-
CalculateRefillBytesPerPeriod(bytes_per_second),
108+
CalculateRefillBytesPerPeriodLocked(bytes_per_second),
104109
std::memory_order_relaxed);
105110
}
106111

@@ -115,10 +120,10 @@ void GenericRateLimiter::Request(int64_t bytes, const Env::IOPriority pri,
115120

116121
if (auto_tuned_) {
117122
static const int kRefillsPerTune = 100;
118-
std::chrono::microseconds now(NowMicrosMonotonic());
123+
std::chrono::microseconds now(NowMicrosMonotonicLocked());
119124
if (now - tuned_time_ >=
120125
kRefillsPerTune * std::chrono::microseconds(refill_period_us_)) {
121-
Status s = Tune();
126+
Status s = TuneLocked();
122127
s.PermitUncheckedError(); //**TODO: What to do on error?
123128
}
124129
}
@@ -152,7 +157,7 @@ void GenericRateLimiter::Request(int64_t bytes, const Env::IOPriority pri,
152157
// (1) Waiting for the next refill time.
153158
// (2) Refilling the bytes and granting requests.
154159
do {
155-
int64_t time_until_refill_us = next_refill_us_ - NowMicrosMonotonic();
160+
int64_t time_until_refill_us = next_refill_us_ - NowMicrosMonotonicLocked();
156161
if (time_until_refill_us > 0) {
157162
if (wait_until_refill_pending_) {
158163
// Somebody is performing (1). Trust we'll be woken up when our request
@@ -173,7 +178,7 @@ void GenericRateLimiter::Request(int64_t bytes, const Env::IOPriority pri,
173178
} else {
174179
// Whichever thread reaches here first performs duty (2) as described
175180
// above.
176-
RefillBytesAndGrantRequests();
181+
RefillBytesAndGrantRequestsLocked();
177182
if (r.granted) {
178183
// If there is any remaining requests, make sure there exists at least
179184
// one candidate is awake for future duties by signaling a front request
@@ -215,20 +220,20 @@ void GenericRateLimiter::Request(int64_t bytes, const Env::IOPriority pri,
215220
}
216221

217222
std::vector<Env::IOPriority>
218-
GenericRateLimiter::GeneratePriorityIterationOrder() {
223+
GenericRateLimiter::GeneratePriorityIterationOrderLocked() {
219224
std::vector<Env::IOPriority> pri_iteration_order(Env::IO_TOTAL /* 4 */);
220225
// We make Env::IO_USER a superior priority by always iterating its queue
221226
// first
222227
pri_iteration_order[0] = Env::IO_USER;
223228

224229
bool high_pri_iterated_after_mid_low_pri = rnd_.OneIn(fairness_);
225230
TEST_SYNC_POINT_CALLBACK(
226-
"GenericRateLimiter::GeneratePriorityIterationOrder::"
231+
"GenericRateLimiter::GeneratePriorityIterationOrderLocked::"
227232
"PostRandomOneInFairnessForHighPri",
228233
&high_pri_iterated_after_mid_low_pri);
229234
bool mid_pri_itereated_after_low_pri = rnd_.OneIn(fairness_);
230235
TEST_SYNC_POINT_CALLBACK(
231-
"GenericRateLimiter::GeneratePriorityIterationOrder::"
236+
"GenericRateLimiter::GeneratePriorityIterationOrderLocked::"
232237
"PostRandomOneInFairnessForMidPri",
233238
&mid_pri_itereated_after_low_pri);
234239

@@ -247,15 +252,16 @@ GenericRateLimiter::GeneratePriorityIterationOrder() {
247252
}
248253

249254
TEST_SYNC_POINT_CALLBACK(
250-
"GenericRateLimiter::GeneratePriorityIterationOrder::"
255+
"GenericRateLimiter::GeneratePriorityIterationOrderLocked::"
251256
"PreReturnPriIterationOrder",
252257
&pri_iteration_order);
253258
return pri_iteration_order;
254259
}
255260

256-
void GenericRateLimiter::RefillBytesAndGrantRequests() {
257-
TEST_SYNC_POINT("GenericRateLimiter::RefillBytesAndGrantRequests");
258-
next_refill_us_ = NowMicrosMonotonic() + refill_period_us_;
261+
void GenericRateLimiter::RefillBytesAndGrantRequestsLocked() {
262+
TEST_SYNC_POINT_CALLBACK(
263+
"GenericRateLimiter::RefillBytesAndGrantRequestsLocked", &request_mutex_);
264+
next_refill_us_ = NowMicrosMonotonicLocked() + refill_period_us_;
259265
// Carry over the left over quota from the last period
260266
auto refill_bytes_per_period =
261267
refill_bytes_per_period_.load(std::memory_order_relaxed);
@@ -264,7 +270,7 @@ void GenericRateLimiter::RefillBytesAndGrantRequests() {
264270
}
265271

266272
std::vector<Env::IOPriority> pri_iteration_order =
267-
GeneratePriorityIterationOrder();
273+
GeneratePriorityIterationOrderLocked();
268274

269275
for (int i = Env::IO_LOW; i < Env::IO_TOTAL; ++i) {
270276
assert(!pri_iteration_order.empty());
@@ -293,7 +299,7 @@ void GenericRateLimiter::RefillBytesAndGrantRequests() {
293299
}
294300
}
295301

296-
int64_t GenericRateLimiter::CalculateRefillBytesPerPeriod(
302+
int64_t GenericRateLimiter::CalculateRefillBytesPerPeriodLocked(
297303
int64_t rate_bytes_per_sec) {
298304
if (std::numeric_limits<int64_t>::max() / rate_bytes_per_sec <
299305
refill_period_us_) {
@@ -305,7 +311,7 @@ int64_t GenericRateLimiter::CalculateRefillBytesPerPeriod(
305311
}
306312
}
307313

308-
Status GenericRateLimiter::Tune() {
314+
Status GenericRateLimiter::TuneLocked() {
309315
const int kLowWatermarkPct = 50;
310316
const int kHighWatermarkPct = 90;
311317
const int kAdjustFactorPct = 5;
@@ -314,7 +320,7 @@ Status GenericRateLimiter::Tune() {
314320
const int kAllowedRangeFactor = 20;
315321

316322
std::chrono::microseconds prev_tuned_time = tuned_time_;
317-
tuned_time_ = std::chrono::microseconds(NowMicrosMonotonic());
323+
tuned_time_ = std::chrono::microseconds(NowMicrosMonotonicLocked());
318324

319325
int64_t elapsed_intervals = (tuned_time_ - prev_tuned_time +
320326
std::chrono::microseconds(refill_period_us_) -
@@ -349,7 +355,7 @@ Status GenericRateLimiter::Tune() {
349355
new_bytes_per_sec = prev_bytes_per_sec;
350356
}
351357
if (new_bytes_per_sec != prev_bytes_per_sec) {
352-
SetBytesPerSecond(new_bytes_per_sec);
358+
SetBytesPerSecondLocked(new_bytes_per_sec);
353359
}
354360
num_drains_ = 0;
355361
return Status::OK();

‎util/rate_limiter.h

+12-10
Original file line numberDiff line numberDiff line change
@@ -92,30 +92,32 @@ class GenericRateLimiter : public RateLimiter {
9292
}
9393

9494
virtual int64_t GetBytesPerSecond() const override {
95-
return rate_bytes_per_sec_;
95+
return rate_bytes_per_sec_.load(std::memory_order_relaxed);
9696
}
9797

9898
virtual void TEST_SetClock(std::shared_ptr<SystemClock> clock) {
9999
MutexLock g(&request_mutex_);
100100
clock_ = std::move(clock);
101-
next_refill_us_ = NowMicrosMonotonic();
101+
next_refill_us_ = NowMicrosMonotonicLocked();
102102
}
103103

104104
private:
105-
void RefillBytesAndGrantRequests();
106-
std::vector<Env::IOPriority> GeneratePriorityIterationOrder();
107-
int64_t CalculateRefillBytesPerPeriod(int64_t rate_bytes_per_sec);
108-
Status Tune();
109-
110-
uint64_t NowMicrosMonotonic() { return clock_->NowNanos() / std::milli::den; }
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+
}
111114

112115
// This mutex guard all internal states
113116
mutable port::Mutex request_mutex_;
114117

115118
const int64_t refill_period_us_;
116119

117-
int64_t rate_bytes_per_sec_;
118-
// This variable can be changed dynamically.
120+
std::atomic<int64_t> rate_bytes_per_sec_;
119121
std::atomic<int64_t> refill_bytes_per_period_;
120122
std::shared_ptr<SystemClock> clock_;
121123

‎util/rate_limiter_test.cc

+10-7
Original file line numberDiff line numberDiff line change
@@ -202,7 +202,7 @@ TEST_F(RateLimiterTest, GeneratePriorityIterationOrder) {
202202
bool mid_pri_itereated_after_low_pri_set = false;
203203
bool pri_iteration_order_verified = false;
204204
SyncPoint::GetInstance()->SetCallBack(
205-
"GenericRateLimiter::GeneratePriorityIterationOrder::"
205+
"GenericRateLimiter::GeneratePriorityIterationOrderLocked::"
206206
"PostRandomOneInFairnessForHighPri",
207207
[&](void* arg) {
208208
bool* high_pri_iterated_after_mid_low_pri = (bool*)arg;
@@ -212,7 +212,7 @@ TEST_F(RateLimiterTest, GeneratePriorityIterationOrder) {
212212
});
213213

214214
SyncPoint::GetInstance()->SetCallBack(
215-
"GenericRateLimiter::GeneratePriorityIterationOrder::"
215+
"GenericRateLimiter::GeneratePriorityIterationOrderLocked::"
216216
"PostRandomOneInFairnessForMidPri",
217217
[&](void* arg) {
218218
bool* mid_pri_itereated_after_low_pri = (bool*)arg;
@@ -222,7 +222,7 @@ TEST_F(RateLimiterTest, GeneratePriorityIterationOrder) {
222222
});
223223

224224
SyncPoint::GetInstance()->SetCallBack(
225-
"GenericRateLimiter::GeneratePriorityIterationOrder::"
225+
"GenericRateLimiter::GeneratePriorityIterationOrderLocked::"
226226
"PreReturnPriIterationOrder",
227227
[&](void* arg) {
228228
std::vector<Env::IOPriority>* pri_iteration_order =
@@ -249,13 +249,13 @@ TEST_F(RateLimiterTest, GeneratePriorityIterationOrder) {
249249
ASSERT_EQ(pri_iteration_order_verified, true);
250250
SyncPoint::GetInstance()->DisableProcessing();
251251
SyncPoint::GetInstance()->ClearCallBack(
252-
"GenericRateLimiter::GeneratePriorityIterationOrder::"
252+
"GenericRateLimiter::GeneratePriorityIterationOrderLocked::"
253253
"PreReturnPriIterationOrder");
254254
SyncPoint::GetInstance()->ClearCallBack(
255-
"GenericRateLimiter::GeneratePriorityIterationOrder::"
255+
"GenericRateLimiter::GeneratePriorityIterationOrderLocked::"
256256
"PostRandomOneInFairnessForMidPri");
257257
SyncPoint::GetInstance()->ClearCallBack(
258-
"GenericRateLimiter::GeneratePriorityIterationOrder::"
258+
"GenericRateLimiter::GeneratePriorityIterationOrderLocked::"
259259
"PostRandomOneInFairnessForHighPri");
260260
}
261261
}
@@ -387,11 +387,14 @@ TEST_F(RateLimiterTest, LimitChangeTest) {
387387
std::make_shared<GenericRateLimiter>(
388388
target, refill_period, 10, RateLimiter::Mode::kWritesOnly,
389389
SystemClock::Default(), false /* auto_tuned */);
390+
// After "GenericRateLimiter::Request:1" the mutex is held until the bytes
391+
// are refilled. This test could be improved to change the limit when lock
392+
// is released in `TimedWait()`.
390393
ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->LoadDependency(
391394
{{"GenericRateLimiter::Request",
392395
"RateLimiterTest::LimitChangeTest:changeLimitStart"},
393396
{"RateLimiterTest::LimitChangeTest:changeLimitEnd",
394-
"GenericRateLimiter::RefillBytesAndGrantRequests"}});
397+
"GenericRateLimiter::Request:1"}});
395398
Arg arg(target, Env::IO_HIGH, limiter);
396399
// The idea behind is to start a request first, then before it refills,
397400
// update limit to a different value (2X/0.5X). No starvation should

0 commit comments

Comments
 (0)