1 | #pragma once |
2 | |
3 | #include <Access/Quota.h> |
4 | #include <Core/UUID.h> |
5 | #include <Poco/Net/IPAddress.h> |
6 | #include <ext/shared_ptr_helper.h> |
7 | #include <boost/noncopyable.hpp> |
8 | #include <atomic> |
9 | #include <chrono> |
10 | #include <memory> |
11 | |
12 | |
13 | namespace DB |
14 | { |
15 | struct QuotaUsageInfo; |
16 | |
17 | |
18 | /// Instances of `QuotaContext` are used to track resource consumption. |
19 | class QuotaContext : public boost::noncopyable |
20 | { |
21 | public: |
22 | using ResourceType = Quota::ResourceType; |
23 | using ResourceAmount = Quota::ResourceAmount; |
24 | |
25 | /// Default constructors makes an unlimited quota. |
26 | QuotaContext(); |
27 | |
28 | ~QuotaContext(); |
29 | |
30 | /// Tracks resource consumption. If the quota exceeded and `check_exceeded == true`, throws an exception. |
31 | void used(ResourceType resource_type, ResourceAmount amount, bool check_exceeded = true); |
32 | void used(const std::pair<ResourceType, ResourceAmount> & resource, bool check_exceeded = true); |
33 | void used(const std::pair<ResourceType, ResourceAmount> & resource1, const std::pair<ResourceType, ResourceAmount> & resource2, bool check_exceeded = true); |
34 | void used(const std::pair<ResourceType, ResourceAmount> & resource1, const std::pair<ResourceType, ResourceAmount> & resource2, const std::pair<ResourceType, ResourceAmount> & resource3, bool check_exceeded = true); |
35 | void used(const std::vector<std::pair<ResourceType, ResourceAmount>> & resources, bool check_exceeded = true); |
36 | |
37 | /// Checks if the quota exceeded. If so, throws an exception. |
38 | void checkExceeded(); |
39 | void checkExceeded(ResourceType resource_type); |
40 | |
41 | /// Returns the information about this quota context. |
42 | QuotaUsageInfo getUsageInfo() const; |
43 | |
44 | private: |
45 | friend class QuotaContextFactory; |
46 | friend struct ext::shared_ptr_helper<QuotaContext>; |
47 | |
48 | /// Instances of this class are created by QuotaContextFactory. |
49 | QuotaContext(const String & user_name_, const Poco::Net::IPAddress & address_, const String & client_key_); |
50 | |
51 | static constexpr size_t MAX_RESOURCE_TYPE = Quota::MAX_RESOURCE_TYPE; |
52 | |
53 | struct Interval |
54 | { |
55 | mutable std::atomic<ResourceAmount> used[MAX_RESOURCE_TYPE]; |
56 | ResourceAmount max[MAX_RESOURCE_TYPE]; |
57 | std::chrono::seconds duration; |
58 | bool randomize_interval; |
59 | mutable std::atomic<std::chrono::system_clock::duration> end_of_interval; |
60 | |
61 | Interval() {} |
62 | Interval(const Interval & src) { *this = src; } |
63 | Interval & operator =(const Interval & src); |
64 | }; |
65 | |
66 | struct Intervals |
67 | { |
68 | std::vector<Interval> intervals; |
69 | UUID quota_id; |
70 | String quota_name; |
71 | String quota_key; |
72 | |
73 | QuotaUsageInfo getUsageInfo(std::chrono::system_clock::time_point current_time) const; |
74 | }; |
75 | |
76 | struct Impl; |
77 | |
78 | const String user_name; |
79 | const Poco::Net::IPAddress address; |
80 | const String client_key; |
81 | std::shared_ptr<const Intervals> atomic_intervals; /// atomically changed by QuotaUsageManager |
82 | }; |
83 | |
84 | using QuotaContextPtr = std::shared_ptr<QuotaContext>; |
85 | |
86 | |
87 | /// The information about a quota context. |
88 | struct QuotaUsageInfo |
89 | { |
90 | using ResourceType = Quota::ResourceType; |
91 | using ResourceAmount = Quota::ResourceAmount; |
92 | static constexpr size_t MAX_RESOURCE_TYPE = Quota::MAX_RESOURCE_TYPE; |
93 | |
94 | struct Interval |
95 | { |
96 | ResourceAmount used[MAX_RESOURCE_TYPE]; |
97 | ResourceAmount max[MAX_RESOURCE_TYPE]; |
98 | std::chrono::seconds duration = std::chrono::seconds::zero(); |
99 | bool randomize_interval = false; |
100 | std::chrono::system_clock::time_point end_of_interval; |
101 | Interval(); |
102 | }; |
103 | |
104 | std::vector<Interval> intervals; |
105 | UUID quota_id; |
106 | String quota_name; |
107 | String quota_key; |
108 | QuotaUsageInfo(); |
109 | }; |
110 | } |
111 | |