1 | #include <Access/QuotaContext.h> |
2 | #include <Common/Exception.h> |
3 | #include <Common/quoteString.h> |
4 | #include <ext/chrono_io.h> |
5 | #include <ext/range.h> |
6 | #include <boost/range/algorithm/fill.hpp> |
7 | |
8 | |
9 | namespace DB |
10 | { |
11 | namespace ErrorCodes |
12 | { |
13 | extern const int QUOTA_EXPIRED; |
14 | } |
15 | |
16 | struct QuotaContext::Impl |
17 | { |
18 | [[noreturn]] static void throwQuotaExceed( |
19 | const String & user_name, |
20 | const String & quota_name, |
21 | ResourceType resource_type, |
22 | ResourceAmount used, |
23 | ResourceAmount max, |
24 | std::chrono::seconds duration, |
25 | std::chrono::system_clock::time_point end_of_interval) |
26 | { |
27 | std::function<String(UInt64)> amount_to_string = [](UInt64 amount) { return std::to_string(amount); }; |
28 | if (resource_type == Quota::EXECUTION_TIME) |
29 | amount_to_string = [&](UInt64 amount) { return ext::to_string(std::chrono::nanoseconds(amount)); }; |
30 | |
31 | throw Exception( |
32 | "Quota for user " + backQuote(user_name) + " for " + ext::to_string(duration) + " has been exceeded: " |
33 | + Quota::getNameOfResourceType(resource_type) + " = " + amount_to_string(used) + "/" + amount_to_string(max) + ". " |
34 | + "Interval will end at " + ext::to_string(end_of_interval) + ". " + "Name of quota template: " + backQuote(quota_name), |
35 | ErrorCodes::QUOTA_EXPIRED); |
36 | } |
37 | |
38 | |
39 | static std::chrono::system_clock::time_point getEndOfInterval( |
40 | const Interval & interval, std::chrono::system_clock::time_point current_time, bool * counters_were_reset = nullptr) |
41 | { |
42 | auto & end_of_interval = interval.end_of_interval; |
43 | auto end_loaded = end_of_interval.load(); |
44 | auto end = std::chrono::system_clock::time_point{end_loaded}; |
45 | if (current_time < end) |
46 | { |
47 | if (counters_were_reset) |
48 | *counters_were_reset = false; |
49 | return end; |
50 | } |
51 | |
52 | const auto duration = interval.duration; |
53 | |
54 | do |
55 | { |
56 | end = end + (current_time - end + duration) / duration * duration; |
57 | if (end_of_interval.compare_exchange_strong(end_loaded, end.time_since_epoch())) |
58 | { |
59 | boost::range::fill(interval.used, 0); |
60 | break; |
61 | } |
62 | end = std::chrono::system_clock::time_point{end_loaded}; |
63 | } |
64 | while (current_time >= end); |
65 | |
66 | if (counters_were_reset) |
67 | *counters_were_reset = true; |
68 | return end; |
69 | } |
70 | |
71 | |
72 | static void used( |
73 | const String & user_name, |
74 | const Intervals & intervals, |
75 | ResourceType resource_type, |
76 | ResourceAmount amount, |
77 | std::chrono::system_clock::time_point current_time, |
78 | bool check_exceeded) |
79 | { |
80 | for (const auto & interval : intervals.intervals) |
81 | { |
82 | ResourceAmount used = (interval.used[resource_type] += amount); |
83 | ResourceAmount max = interval.max[resource_type]; |
84 | if (max == Quota::UNLIMITED) |
85 | continue; |
86 | if (used > max) |
87 | { |
88 | bool counters_were_reset = false; |
89 | auto end_of_interval = getEndOfInterval(interval, current_time, &counters_were_reset); |
90 | if (counters_were_reset) |
91 | { |
92 | used = (interval.used[resource_type] += amount); |
93 | if ((used > max) && check_exceeded) |
94 | throwQuotaExceed(user_name, intervals.quota_name, resource_type, used, max, interval.duration, end_of_interval); |
95 | } |
96 | else if (check_exceeded) |
97 | throwQuotaExceed(user_name, intervals.quota_name, resource_type, used, max, interval.duration, end_of_interval); |
98 | } |
99 | } |
100 | } |
101 | |
102 | static void checkExceeded( |
103 | const String & user_name, |
104 | const Intervals & intervals, |
105 | ResourceType resource_type, |
106 | std::chrono::system_clock::time_point current_time) |
107 | { |
108 | for (const auto & interval : intervals.intervals) |
109 | { |
110 | ResourceAmount used = interval.used[resource_type]; |
111 | ResourceAmount max = interval.max[resource_type]; |
112 | if (max == Quota::UNLIMITED) |
113 | continue; |
114 | if (used > max) |
115 | { |
116 | bool used_counters_reset = false; |
117 | std::chrono::system_clock::time_point end_of_interval = getEndOfInterval(interval, current_time, &used_counters_reset); |
118 | if (!used_counters_reset) |
119 | throwQuotaExceed(user_name, intervals.quota_name, resource_type, used, max, interval.duration, end_of_interval); |
120 | } |
121 | } |
122 | } |
123 | |
124 | static void checkExceeded( |
125 | const String & user_name, |
126 | const Intervals & intervals, |
127 | std::chrono::system_clock::time_point current_time) |
128 | { |
129 | for (auto resource_type : ext::range_with_static_cast<Quota::ResourceType>(Quota::MAX_RESOURCE_TYPE)) |
130 | checkExceeded(user_name, intervals, resource_type, current_time); |
131 | } |
132 | }; |
133 | |
134 | |
135 | QuotaContext::Interval & QuotaContext::Interval::operator =(const Interval & src) |
136 | { |
137 | randomize_interval = src.randomize_interval; |
138 | duration = src.duration; |
139 | end_of_interval.store(src.end_of_interval.load()); |
140 | for (auto resource_type : ext::range(MAX_RESOURCE_TYPE)) |
141 | { |
142 | max[resource_type] = src.max[resource_type]; |
143 | used[resource_type].store(src.used[resource_type].load()); |
144 | } |
145 | return *this; |
146 | } |
147 | |
148 | |
149 | QuotaUsageInfo QuotaContext::Intervals::getUsageInfo(std::chrono::system_clock::time_point current_time) const |
150 | { |
151 | QuotaUsageInfo info; |
152 | info.quota_id = quota_id; |
153 | info.quota_name = quota_name; |
154 | info.quota_key = quota_key; |
155 | info.intervals.reserve(intervals.size()); |
156 | for (const auto & in : intervals) |
157 | { |
158 | info.intervals.push_back({}); |
159 | auto & out = info.intervals.back(); |
160 | out.duration = in.duration; |
161 | out.randomize_interval = in.randomize_interval; |
162 | out.end_of_interval = Impl::getEndOfInterval(in, current_time); |
163 | for (auto resource_type : ext::range(MAX_RESOURCE_TYPE)) |
164 | { |
165 | out.max[resource_type] = in.max[resource_type]; |
166 | out.used[resource_type] = in.used[resource_type]; |
167 | } |
168 | } |
169 | return info; |
170 | } |
171 | |
172 | |
173 | QuotaContext::QuotaContext() |
174 | : atomic_intervals(std::make_shared<Intervals>()) /// Unlimited quota. |
175 | { |
176 | } |
177 | |
178 | |
179 | QuotaContext::QuotaContext( |
180 | const String & user_name_, |
181 | const Poco::Net::IPAddress & address_, |
182 | const String & client_key_) |
183 | : user_name(user_name_), address(address_), client_key(client_key_) |
184 | { |
185 | } |
186 | |
187 | |
188 | QuotaContext::~QuotaContext() = default; |
189 | |
190 | |
191 | void QuotaContext::used(ResourceType resource_type, ResourceAmount amount, bool check_exceeded) |
192 | { |
193 | used({resource_type, amount}, check_exceeded); |
194 | } |
195 | |
196 | |
197 | void QuotaContext::used(const std::pair<ResourceType, ResourceAmount> & resource, bool check_exceeded) |
198 | { |
199 | auto intervals_ptr = std::atomic_load(&atomic_intervals); |
200 | auto current_time = std::chrono::system_clock::now(); |
201 | Impl::used(user_name, *intervals_ptr, resource.first, resource.second, current_time, check_exceeded); |
202 | } |
203 | |
204 | |
205 | void QuotaContext::used(const std::pair<ResourceType, ResourceAmount> & resource1, const std::pair<ResourceType, ResourceAmount> & resource2, bool check_exceeded) |
206 | { |
207 | auto intervals_ptr = std::atomic_load(&atomic_intervals); |
208 | auto current_time = std::chrono::system_clock::now(); |
209 | Impl::used(user_name, *intervals_ptr, resource1.first, resource1.second, current_time, check_exceeded); |
210 | Impl::used(user_name, *intervals_ptr, resource2.first, resource2.second, current_time, check_exceeded); |
211 | } |
212 | |
213 | |
214 | void QuotaContext::used(const std::pair<ResourceType, ResourceAmount> & resource1, const std::pair<ResourceType, ResourceAmount> & resource2, const std::pair<ResourceType, ResourceAmount> & resource3, bool check_exceeded) |
215 | { |
216 | auto intervals_ptr = std::atomic_load(&atomic_intervals); |
217 | auto current_time = std::chrono::system_clock::now(); |
218 | Impl::used(user_name, *intervals_ptr, resource1.first, resource1.second, current_time, check_exceeded); |
219 | Impl::used(user_name, *intervals_ptr, resource2.first, resource2.second, current_time, check_exceeded); |
220 | Impl::used(user_name, *intervals_ptr, resource3.first, resource3.second, current_time, check_exceeded); |
221 | } |
222 | |
223 | |
224 | void QuotaContext::used(const std::vector<std::pair<ResourceType, ResourceAmount>> & resources, bool check_exceeded) |
225 | { |
226 | auto intervals_ptr = std::atomic_load(&atomic_intervals); |
227 | auto current_time = std::chrono::system_clock::now(); |
228 | for (const auto & resource : resources) |
229 | Impl::used(user_name, *intervals_ptr, resource.first, resource.second, current_time, check_exceeded); |
230 | } |
231 | |
232 | |
233 | void QuotaContext::checkExceeded() |
234 | { |
235 | auto intervals_ptr = std::atomic_load(&atomic_intervals); |
236 | Impl::checkExceeded(user_name, *intervals_ptr, std::chrono::system_clock::now()); |
237 | } |
238 | |
239 | |
240 | void QuotaContext::checkExceeded(ResourceType resource_type) |
241 | { |
242 | auto intervals_ptr = std::atomic_load(&atomic_intervals); |
243 | Impl::checkExceeded(user_name, *intervals_ptr, resource_type, std::chrono::system_clock::now()); |
244 | } |
245 | |
246 | |
247 | QuotaUsageInfo QuotaContext::getUsageInfo() const |
248 | { |
249 | auto intervals_ptr = std::atomic_load(&atomic_intervals); |
250 | return intervals_ptr->getUsageInfo(std::chrono::system_clock::now()); |
251 | } |
252 | |
253 | |
254 | QuotaUsageInfo::QuotaUsageInfo() : quota_id(UUID(UInt128(0))) |
255 | { |
256 | } |
257 | |
258 | |
259 | QuotaUsageInfo::Interval::Interval() |
260 | { |
261 | boost::range::fill(used, 0); |
262 | boost::range::fill(max, 0); |
263 | } |
264 | } |
265 | |