1 struct TokenBucket
2 {
3 uint64_t fill_interval_;
4 uint64_t gen_interval_;
5 uint64_t reset_time_;
6 // reset_time_ is equivalent to last_consume_time + token_num and
7 // the calculated token num is auto renormalized on limit change.
8
9 TokenBucket()
10 : fill_interval_(0)
11 , gen_interval_(0)
12 , reset_time_(0) {}
13
14 // eg, for 10req/30s, capacity=10, period=30
15 void set_limit(uint64_t capacity, uint64_t period)
16 {
17 fill_interval_ = period;
18 gen_interval_ = period / capacity;
19 }
20
21 bool consume(unsigned num)
22 {
23 uint64_t curr_time = get_monotonic_time();
24 uint64_t new_reset_time = std::max(reset_time_, curr_time - fill_interval_) + gen_interval_ * num;
25 if (curr_time >= new_reset_time)
26 {
27 reset_time_ = new_reset_time; // reset_time_ is monotonic increasing
28 return true;
29 }
30 return false;
31 }
32 };
