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 */
14class 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
33public:
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.
122namespace CurrentMemoryTracker
123{
124 void alloc(Int64 size);
125 void realloc(Int64 old_size, Int64 new_size);
126 void free(Int64 size);
127}
128
129
130DB::SimpleActionLock getCurrentMemoryTrackerActionLock();
131