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
9namespace DB
10{
11namespace ErrorCodes
12{
13 extern const int QUOTA_EXPIRED;
14}
15
16struct 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
135QuotaContext::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
149QuotaUsageInfo 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
173QuotaContext::QuotaContext()
174 : atomic_intervals(std::make_shared<Intervals>()) /// Unlimited quota.
175{
176}
177
178
179QuotaContext::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
188QuotaContext::~QuotaContext() = default;
189
190
191void QuotaContext::used(ResourceType resource_type, ResourceAmount amount, bool check_exceeded)
192{
193 used({resource_type, amount}, check_exceeded);
194}
195
196
197void 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
205void 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
214void 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
224void 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
233void 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
240void 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
247QuotaUsageInfo 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
254QuotaUsageInfo::QuotaUsageInfo() : quota_id(UUID(UInt128(0)))
255{
256}
257
258
259QuotaUsageInfo::Interval::Interval()
260{
261 boost::range::fill(used, 0);
262 boost::range::fill(max, 0);
263}
264}
265