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
28namespace ailego {
29
30/*! RateLimiter abstract class
31 */
32struct 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 */
73class 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