1 | #pragma once |
2 | |
3 | #include <atomic> |
4 | #include <common/Types.h> |
5 | #include <Common/CurrentMetrics.h> |
6 | #include <Common/SimpleActionBlocker.h> |
7 | #include <Common/VariableContext.h> |
8 | |
9 | |
10 | /** Tracks memory consumption. |
11 | * It throws an exception if amount of consumed memory become greater than certain limit. |
12 | * The same memory tracker could be simultaneously used in different threads. |
13 | */ |
14 | class MemoryTracker |
15 | { |
16 | std::atomic<Int64> amount {0}; |
17 | std::atomic<Int64> peak {0}; |
18 | std::atomic<Int64> limit {0}; |
19 | |
20 | /// To test exception safety of calling code, memory tracker throws an exception on each memory allocation with specified probability. |
21 | double fault_probability = 0; |
22 | |
23 | /// Singly-linked list. All information will be passed to subsequent memory trackers also (it allows to implement trackers hierarchy). |
24 | /// In terms of tree nodes it is the list of parents. Lifetime of these trackers should "include" lifetime of current tracker. |
25 | std::atomic<MemoryTracker *> parent {}; |
26 | |
27 | /// You could specify custom metric to track memory usage. |
28 | CurrentMetrics::Metric metric = CurrentMetrics::end(); |
29 | |
30 | /// This description will be used as prefix into log messages (if isn't nullptr) |
31 | const char * description = nullptr; |
32 | |
33 | public: |
34 | MemoryTracker(VariableContext level_ = VariableContext::Thread) : level(level_) {} |
35 | MemoryTracker(Int64 limit_, VariableContext level_ = VariableContext::Thread) : limit(limit_), level(level_) {} |
36 | MemoryTracker(MemoryTracker * parent_, VariableContext level_ = VariableContext::Thread) : parent(parent_), level(level_) {} |
37 | |
38 | ~MemoryTracker(); |
39 | |
40 | VariableContext level; |
41 | |
42 | /** Call the following functions before calling of corresponding operations with memory allocators. |
43 | */ |
44 | void alloc(Int64 size); |
45 | |
46 | void realloc(Int64 old_size, Int64 new_size) |
47 | { |
48 | Int64 addition = new_size - old_size; |
49 | if (addition > 0) |
50 | alloc(addition); |
51 | else |
52 | free(-addition); |
53 | } |
54 | |
55 | /** This function should be called after memory deallocation. |
56 | */ |
57 | void free(Int64 size); |
58 | |
59 | Int64 get() const |
60 | { |
61 | return amount.load(std::memory_order_relaxed); |
62 | } |
63 | |
64 | Int64 getPeak() const |
65 | { |
66 | return peak.load(std::memory_order_relaxed); |
67 | } |
68 | |
69 | void setLimit(Int64 limit_) |
70 | { |
71 | limit.store(limit_, std::memory_order_relaxed); |
72 | } |
73 | |
74 | /** Set limit if it was not set. |
75 | * Otherwise, set limit to new value, if new value is greater than previous limit. |
76 | */ |
77 | void setOrRaiseLimit(Int64 value); |
78 | |
79 | void setFaultProbability(double value) |
80 | { |
81 | fault_probability = value; |
82 | } |
83 | |
84 | /// next should be changed only once: from nullptr to some value. |
85 | /// NOTE: It is not true in MergeListElement |
86 | void setParent(MemoryTracker * elem) |
87 | { |
88 | parent.store(elem, std::memory_order_relaxed); |
89 | } |
90 | |
91 | MemoryTracker * getParent() |
92 | { |
93 | return parent.load(std::memory_order_relaxed); |
94 | } |
95 | |
96 | /// The memory consumption could be shown in realtime via CurrentMetrics counter |
97 | void setMetric(CurrentMetrics::Metric metric_) |
98 | { |
99 | metric = metric_; |
100 | } |
101 | |
102 | void setDescription(const char * description_) |
103 | { |
104 | description = description_; |
105 | } |
106 | |
107 | /// Reset the accumulated data |
108 | void resetCounters(); |
109 | |
110 | /// Reset the accumulated data and the parent. |
111 | void reset(); |
112 | |
113 | /// Prints info about peak memory consumption into log. |
114 | void logPeakMemoryUsage() const; |
115 | |
116 | /// To be able to temporarily stop memory tracker |
117 | DB::SimpleActionBlocker blocker; |
118 | }; |
119 | |
120 | |
121 | /// Convenience methods, that use current thread's memory_tracker if it is available. |
122 | namespace CurrentMemoryTracker |
123 | { |
124 | void alloc(Int64 size); |
125 | void realloc(Int64 old_size, Int64 new_size); |
126 | void free(Int64 size); |
127 | } |
128 | |
129 | |
130 | DB::SimpleActionLock getCurrentMemoryTrackerActionLock(); |
131 | |