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
13namespace DB
14{
15struct QuotaUsageInfo;
16
17
18/// Instances of `QuotaContext` are used to track resource consumption.
19class QuotaContext : public boost::noncopyable
20{
21public:
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
44private:
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
84using QuotaContextPtr = std::shared_ptr<QuotaContext>;
85
86
87/// The information about a quota context.
88struct 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