1 | /** |
2 | * Copyright 2021 Alibaba, Inc. and its affiliates. All Rights Reserved. |
3 | * |
4 | * Licensed under the Apache License, Version 2.0 (the "License"); |
5 | * you may not use this file except in compliance with the License. |
6 | * You may obtain a copy of the License at |
7 | * |
8 | * http://www.apache.org/licenses/LICENSE-2.0 |
9 | * |
10 | * Unless required by applicable law or agreed to in writing, software |
11 | * distributed under the License is distributed on an "AS IS" BASIS, |
12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
13 | * See the License for the specific language governing permissions and |
14 | * limitations under the License. |
15 | |
16 | * \author hongqing.hu |
17 | * \date Nov 2020 |
18 | * \brief Interface of Rate Limiter algorithm |
19 | */ |
20 | |
21 | #ifndef __AILEGO_ALGORITHM_RATE_LIMITER_H__ |
22 | #define __AILEGO_ALGORITHM_RATE_LIMITER_H__ |
23 | |
24 | #include <chrono> |
25 | #include <memory> |
26 | #include <mutex> |
27 | |
28 | namespace ailego { |
29 | |
30 | /*! RateLimiter abstract class |
31 | */ |
32 | struct RateLimiter { |
33 | //! Rate Limiter Shared Pointer |
34 | using Pointer = std::shared_ptr<RateLimiter>; |
35 | |
36 | //! RateLimiter Types |
37 | enum LimiterTypes { |
38 | LIMITER_BURSTY = 0, |
39 | }; |
40 | |
41 | //! Destructor |
42 | virtual ~RateLimiter(void) {} |
43 | |
44 | //! Acquire batch permits |
45 | virtual double acquire(int permits) = 0; |
46 | |
47 | //! Acquire one permit |
48 | virtual double acquire(void) = 0; |
49 | |
50 | //! Try acquire some permits within timeout ms |
51 | virtual bool try_acquire(int permits, int timeout_ms) = 0; |
52 | |
53 | //! Try acquire one permit, return at once |
54 | virtual bool try_acquire(void) = 0; |
55 | |
56 | //! Set permits per second |
57 | virtual void set_rate(double permits_per_second) = 0; |
58 | |
59 | //! Get permits per second |
60 | virtual double get_rate(void) const = 0; |
61 | |
62 | //! Create a rate limiter with type |
63 | static Pointer Create(double permits_per_second, LimiterTypes type); |
64 | |
65 | //! Create a default rate limiter |
66 | static Pointer Create(double permits_per_second) { |
67 | return RateLimiter::Create(permits_per_second, LIMITER_BURSTY); |
68 | } |
69 | }; |
70 | |
71 | /*! Bursty Rate Limiter |
72 | */ |
73 | class BurstyRateLimiter : public RateLimiter { |
74 | public: |
75 | //! Constructor |
76 | explicit BurstyRateLimiter(double max_bursty_seconds) |
77 | : max_bursty_seconds_(max_bursty_seconds) {} |
78 | |
79 | //! Destructor |
80 | virtual ~BurstyRateLimiter(void) {} |
81 | |
82 | //! Acquire batch permits |
83 | double acquire(int permits) override; |
84 | |
85 | //! Try acquire some permits with timeout |
86 | bool try_acquire(int permits, int timeout_ms) override; |
87 | |
88 | //! Set permits per second |
89 | void set_rate(double permits_per_second) override; |
90 | |
91 | //! Get permits per second |
92 | double get_rate(void) const override { |
93 | return MICROSECONDS_PER_SECOND / interval_; |
94 | } |
95 | |
96 | //! Acquire one permit |
97 | double acquire(void) override { |
98 | return this->acquire(1); |
99 | } |
100 | |
101 | //! Try acquire one permit, return at once |
102 | bool try_acquire(void) override { |
103 | return this->try_acquire(1, 0); |
104 | } |
105 | |
106 | protected: |
107 | //! Resync the information |
108 | void update_stored_permits(int64_t now_usec); |
109 | |
110 | //! Compute the wait time |
111 | int64_t compute_wait_usec(int required_permits, int64_t now_usec); |
112 | |
113 | private: |
114 | //! Disable them |
115 | BurstyRateLimiter(const BurstyRateLimiter &) = delete; |
116 | BurstyRateLimiter(BurstyRateLimiter &&) = delete; |
117 | BurstyRateLimiter &operator=(const BurstyRateLimiter &) = delete; |
118 | |
119 | //! Members |
120 | double max_bursty_seconds_{0.0}; |
121 | double max_permits_{0.0}; |
122 | double stored_permits_{0.0}; |
123 | double interval_{0.0}; |
124 | double permits_per_usec_{0.0}; |
125 | double next_free_time_{0.0}; |
126 | std::mutex mutex_{}; |
127 | |
128 | constexpr static double MICROSECONDS_PER_SECOND = 1000000.0; |
129 | }; |
130 | |
131 | } // namespace ailego |
132 | |
133 | #endif // __AILEGO_ALGORITHM_RATE_LIMITER_H__ |
134 | |