Coverage Report

Created: 2026-08-28 17:26

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
common/cpp/token_bucket_rate_limiter.cpp
Line
Count
Source
1
// Licensed to the Apache Software Foundation (ASF) under one
2
// or more contributor license agreements.  See the NOTICE file
3
// distributed with this work for additional information
4
// regarding copyright ownership.  The ASF licenses this file
5
// to you under the Apache License, Version 2.0 (the
6
// "License"); you may not use this file except in compliance
7
// with the License.  You may obtain a copy of the License at
8
//
9
//   http://www.apache.org/licenses/LICENSE-2.0
10
//
11
// Unless required by applicable law or agreed to in writing,
12
// software distributed under the License is distributed on an
13
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14
// KIND, either express or implied.  See the License for the
15
// specific language governing permissions and limitations
16
// under the License.
17
18
#include "token_bucket_rate_limiter.h"
19
20
#include <bthread/bthread.h>
21
#include <glog/logging.h> // IWYU pragma: export
22
23
#include <atomic>
24
#include <chrono>
25
#include <mutex>
26
#include <thread>
27
28
#if defined(__APPLE__)
29
#include <ctime>
30
#endif
31
32
namespace doris {
33
// Just 10^6.
34
static constexpr auto NS = 1000000000UL;
35
36
bvar::Adder<int64_t> s3_get_rate_limit_sleep_ns("s3_get_rate_limit_sleep_ns");
37
bvar::Adder<int64_t> s3_get_rate_limit_sleep_count("s3_get_rate_limit_sleep_count");
38
bvar::Adder<int64_t> s3_get_rate_limit_rejected_count("s3_get_rate_limit_rejected_count");
39
bvar::Adder<int64_t> s3_put_rate_limit_sleep_ns("s3_put_rate_limit_sleep_ns");
40
bvar::Adder<int64_t> s3_put_rate_limit_sleep_count("s3_put_rate_limit_sleep_count");
41
bvar::Adder<int64_t> s3_put_rate_limit_rejected_count("s3_put_rate_limit_rejected_count");
42
43
static std::atomic<int64_t> s3_get_rate_limit_sleep_log_count {0};
44
static std::atomic<int64_t> s3_get_rate_limit_rejected_log_count {0};
45
static std::atomic<int64_t> s3_put_rate_limit_sleep_log_count {0};
46
static std::atomic<int64_t> s3_put_rate_limit_rejected_log_count {0};
47
48
class TokenBucketRateLimiter::SimpleSpinLock {
49
public:
50
1.72k
    SimpleSpinLock() = default;
51
    ~SimpleSpinLock() = default;
52
53
608
    void lock() {
54
608
        int spin_count = 0;
55
608
        static constexpr int MAX_SPIN_COUNT = 50;
56
644
        while (_flag.test_and_set(std::memory_order_acq_rel)) {
57
36
            spin_count++;
58
36
            if (spin_count >= MAX_SPIN_COUNT) {
59
0
                LOG(WARNING) << "Warning: Excessive spinning detected while acquiring lock. Spin "
60
0
                                "count: "
61
0
                             << spin_count;
62
0
                spin_count = 0;
63
0
            }
64
            // Spin until we acquire the lock
65
36
        }
66
608
    }
67
68
609
    void unlock() { _flag.clear(std::memory_order_release); }
69
70
private:
71
    std::atomic_flag _flag = ATOMIC_FLAG_INIT;
72
};
73
74
TokenBucketRateLimiter::TokenBucketRateLimiter(size_t max_speed, size_t max_burst, size_t limit)
75
1.72k
        : _max_speed(max_speed),
76
1.72k
          _max_burst(max_burst),
77
1.72k
          _limit(limit),
78
1.72k
          _mutex(std::make_unique<TokenBucketRateLimiter::SimpleSpinLock>()),
79
1.72k
          _remain_tokens(max_burst) {}
80
81
1.72k
TokenBucketRateLimiter::~TokenBucketRateLimiter() = default;
82
83
1.63k
TokenBucketRateLimiterHolder::~TokenBucketRateLimiterHolder() = default;
84
85
605
std::pair<size_t, double> TokenBucketRateLimiter::_update_remain_token(long now, size_t amount) {
86
    // Values obtained under lock to be checked after release
87
605
    size_t count_value;
88
605
    double tokens_value;
89
605
    {
90
605
        std::lock_guard<SimpleSpinLock> lock(*_mutex);
91
605
        now = (now < _prev_ns_count) ? _prev_ns_count : now;
92
605
        if (_max_speed) {
93
591
            double delta_seconds =
94
591
                    _prev_ns_count ? static_cast<double>(now - _prev_ns_count) / NS : 0;
95
591
            _remain_tokens = std::min<double>(_remain_tokens + _max_speed * delta_seconds - amount,
96
591
                                              _max_burst);
97
591
        }
98
605
        _count += amount;
99
605
        count_value = _count;
100
605
        if (_limit && count_value > _limit) {
101
            // Keep rejection side-effect free. Roll back before releasing the lock so
102
            // concurrent callers cannot observe debt from a request that will not run.
103
22
            _count -= amount;
104
22
            if (_max_speed) {
105
15
                _remain_tokens = std::min<double>(_remain_tokens + amount, _max_burst);
106
15
            }
107
22
        }
108
605
        tokens_value = _remain_tokens;
109
605
        _prev_ns_count = now;
110
605
    }
111
605
    return {count_value, tokens_value};
112
605
}
113
114
583
int64_t TokenBucketRateLimiter::add(size_t amount) {
115
583
    int64_t sleep_time_ns = reserve(amount);
116
583
    if (sleep_time_ns > 0) {
117
288
        bthread_usleep(sleep_time_ns / 1000);
118
288
    }
119
583
    return sleep_time_ns;
120
583
}
121
122
604
int64_t TokenBucketRateLimiter::reserve(size_t amount) {
123
    // Values obtained under lock to be checked after release
124
604
    auto duration = std::chrono::steady_clock::now().time_since_epoch();
125
604
    auto time_nano_count = std::chrono::duration_cast<std::chrono::nanoseconds>(duration).count();
126
604
    auto [count_value, tokens_value] = _update_remain_token(time_nano_count, amount);
127
128
604
    if (_limit && count_value > _limit) {
129
        // CK would throw exception
130
22
        return -1;
131
22
    }
132
133
    // Wait unless there is positive amount of remain_tokens - throttling
134
582
    int64_t sleep_time_ns = 0;
135
582
    if (_max_speed && tokens_value < 0) {
136
297
        sleep_time_ns = static_cast<int64_t>(-tokens_value / _max_speed * NS);
137
297
    }
138
139
582
    return sleep_time_ns;
140
604
}
141
142
3
void TokenBucketRateLimiter::refund(size_t amount) {
143
3
    std::lock_guard<SimpleSpinLock> lock(*_mutex);
144
3
    if (_max_speed) {
145
3
        _remain_tokens = std::min<double>(_remain_tokens + amount, _max_burst);
146
3
    }
147
3
    _count = (_count >= amount) ? _count - amount : 0;
148
3
}
149
150
TokenBucketRateLimiterHolder::TokenBucketRateLimiterHolder(size_t max_speed, size_t max_burst,
151
                                                           size_t limit,
152
                                                           std::function<void(int64_t)> metric_func)
153
1.63k
        : rate_limiter(std::make_shared<TokenBucketRateLimiter>(max_speed, max_burst, limit)),
154
1.63k
          _enabled(max_speed > 0 || limit > 0),
155
1.63k
          metric_func(std::move(metric_func)) {}
156
157
44
int64_t TokenBucketRateLimiterHolder::add(size_t amount) {
158
44
    return add_with_config(amount).sleep_duration;
159
44
}
160
161
574
TokenBucketRateLimiterResult TokenBucketRateLimiterHolder::add_with_config(size_t amount) {
162
574
    return _consume_with_config(amount, true);
163
574
}
164
165
21
TokenBucketRateLimiterResult TokenBucketRateLimiterHolder::reserve_with_config(size_t amount) {
166
21
    return _consume_with_config(amount, false);
167
21
}
168
169
TokenBucketRateLimiterResult TokenBucketRateLimiterHolder::_consume_with_config(size_t amount,
170
595
                                                                                bool wait) {
171
    // Snapshot the current limiter and consume outside the read lock. The waiting
172
    // path may sleep for a long time when throttled, and holding the read lock across
173
    // the sleep would block reset() (dynamic config update) for the whole duration.
174
595
    std::shared_ptr<TokenBucketRateLimiter> limiter;
175
595
    {
176
595
        std::shared_lock read {rate_limiter_rw_lock};
177
595
        limiter = rate_limiter;
178
595
    }
179
595
    TokenBucketRateLimiterResult result = {.sleep_duration = wait ? limiter->add(amount)
180
595
                                                                  : limiter->reserve(amount),
181
595
                                           .max_speed = limiter->get_max_speed(),
182
595
                                           .max_burst = limiter->get_max_burst(),
183
595
                                           .limit = limiter->get_limit()};
184
595
    metric_func(result.sleep_duration);
185
595
    return result;
186
595
}
187
188
7
std::shared_ptr<TokenBucketRateLimiter> TokenBucketRateLimiterHolder::charge(size_t amount) {
189
7
    std::shared_ptr<TokenBucketRateLimiter> limiter;
190
7
    {
191
7
        std::shared_lock read {rate_limiter_rw_lock};
192
7
        limiter = rate_limiter;
193
7
    }
194
7
    int64_t sleep_duration = limiter->add(amount);
195
7
    metric_func(sleep_duration);
196
7
    return sleep_duration < 0 ? nullptr : limiter;
197
7
}
198
199
85
int TokenBucketRateLimiterHolder::reset(size_t max_speed, size_t max_burst, size_t limit) {
200
85
    auto new_rate_limiter = std::make_shared<TokenBucketRateLimiter>(max_speed, max_burst, limit);
201
85
    {
202
85
        std::unique_lock write {rate_limiter_rw_lock};
203
85
        rate_limiter = std::move(new_rate_limiter);
204
85
        _enabled.store(max_speed > 0 || limit > 0, std::memory_order_release);
205
85
    }
206
85
    return 0;
207
85
}
208
209
229
size_t TokenBucketRateLimiterHolder::get_max_speed() const {
210
229
    std::shared_lock read {rate_limiter_rw_lock};
211
229
    return rate_limiter->get_max_speed();
212
229
}
213
214
68
size_t TokenBucketRateLimiterHolder::get_max_burst() const {
215
68
    std::shared_lock read {rate_limiter_rw_lock};
216
68
    return rate_limiter->get_max_burst();
217
68
}
218
219
60
size_t TokenBucketRateLimiterHolder::get_limit() const {
220
60
    std::shared_lock read {rate_limiter_rw_lock};
221
60
    return rate_limiter->get_limit();
222
60
}
223
224
53
std::string to_string(S3RateLimitType type) {
225
53
    switch (type) {
226
39
    case S3RateLimitType::GET:
227
39
        return "get";
228
14
    case S3RateLimitType::PUT:
229
14
        return "put";
230
0
    default:
231
0
        return std::to_string(static_cast<size_t>(type));
232
53
    }
233
53
}
234
235
0
S3RateLimitType string_to_s3_rate_limit_type(std::string_view value) {
236
0
    if (value == "get") {
237
0
        return S3RateLimitType::GET;
238
0
    } else if (value == "put") {
239
0
        return S3RateLimitType::PUT;
240
0
    }
241
0
    return S3RateLimitType::UNKNOWN;
242
0
}
243
244
2
std::function<void(int64_t)> s3_rate_limiter_metric_func(S3RateLimitType type) {
245
2
    switch (type) {
246
1
    case S3RateLimitType::GET:
247
1
        return metric_func_factory(s3_get_rate_limit_sleep_ns, s3_get_rate_limit_sleep_count,
248
1
                                   &s3_get_rate_limit_rejected_count);
249
1
    case S3RateLimitType::PUT:
250
1
        return metric_func_factory(s3_put_rate_limit_sleep_ns, s3_put_rate_limit_sleep_count,
251
1
                                   &s3_put_rate_limit_rejected_count);
252
0
    default:
253
0
        return [](int64_t) {};
254
2
    }
255
2
}
256
257
int64_t apply_s3_rate_limit(S3RateLimitType type, S3RateLimiterHolder* rate_limiter,
258
7
                            int64_t log_interval) {
259
7
    auto result = rate_limiter->add_with_config(1);
260
7
    auto sleep_duration = result.sleep_duration;
261
7
    if (log_interval <= 0 || sleep_duration == 0) {
262
4
        return sleep_duration;
263
4
    }
264
265
3
    auto is_get = type == S3RateLimitType::GET;
266
3
    auto* sleep_log_count =
267
3
            is_get ? &s3_get_rate_limit_sleep_log_count : &s3_put_rate_limit_sleep_log_count;
268
3
    auto* rejected_log_count =
269
3
            is_get ? &s3_get_rate_limit_rejected_log_count : &s3_put_rate_limit_rejected_log_count;
270
271
3
    if (sleep_duration > 0) {
272
0
        int64_t count = sleep_log_count->fetch_add(1, std::memory_order_relaxed) + 1;
273
0
        if (count == 1 || count % log_interval == 0) {
274
0
            LOG(INFO) << "S3 " << to_string(type) << " request is throttled by local rate limiter"
275
0
                      << ", sleep_ms=" << sleep_duration / 1000000 << ", sleep_count=" << count
276
0
                      << ", token_per_second=" << result.max_speed
277
0
                      << ", bucket_tokens=" << result.max_burst << ", token_limit=" << result.limit;
278
0
        }
279
3
    } else {
280
3
        int64_t count = rejected_log_count->fetch_add(1, std::memory_order_relaxed) + 1;
281
3
        if (count == 1 || count % log_interval == 0) {
282
            LOG(WARNING) << "S3 " << to_string(type) << " request is rejected by local rate limiter"
283
1
                         << ", rejected_count=" << count
284
1
                         << ", token_per_second=" << result.max_speed
285
1
                         << ", bucket_tokens=" << result.max_burst
286
1
                         << ", token_limit=" << result.limit;
287
1
        }
288
3
    }
289
3
    return sleep_duration;
290
7
}
291
} // namespace doris