-
Notifications
You must be signed in to change notification settings - Fork 8.2k
Expand file tree
/
Copy pathCurrentMemoryTracker.cpp
More file actions
144 lines (119 loc) · 4.48 KB
/
CurrentMemoryTracker.cpp
File metadata and controls
144 lines (119 loc) · 4.48 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
#include <Common/CurrentMemoryTracker.h>
#include <Common/CurrentThread.h>
#include <Common/Exception.h>
#include <Common/MemoryTracker.h>
#include <Common/MemoryTrackerBlockerInThread.h>
#ifdef MEMORY_TRACKER_DEBUG_CHECKS
thread_local bool memory_tracker_always_throw_logical_error_on_allocation = false;
#endif
namespace DB
{
namespace ErrorCodes
{
extern const int LOGICAL_ERROR;
}
}
namespace
{
MemoryTracker * getMemoryTracker()
{
if (auto * thread_memory_tracker = DB::CurrentThread::getMemoryTracker())
return thread_memory_tracker;
/// total_memory_tracker can be used before MainThreadStatus is initialized,
/// but only after its own initialization and before teardown.
if (DB::MainThreadStatus::initialized() || isTotalMemoryTrackerInitialized())
return &total_memory_tracker;
return nullptr;
}
}
using DB::current_thread;
AllocationTrace CurrentMemoryTracker::allocImpl(Int64 size, bool throw_if_memory_exceeded)
{
#ifdef MEMORY_TRACKER_DEBUG_CHECKS
if (unlikely(memory_tracker_always_throw_logical_error_on_allocation))
{
memory_tracker_always_throw_logical_error_on_allocation = false;
throw DB::Exception(DB::ErrorCodes::LOGICAL_ERROR, "Memory tracker: allocations not allowed.");
}
#endif
if (auto * memory_tracker = getMemoryTracker())
{
if (!current_thread)
{
/// total_memory_tracker only, ignore untracked_memory
return memory_tracker->allocImpl(size, throw_if_memory_exceeded);
}
/// Make sure we do memory tracker calls with the correct level in MemoryTrackerBlockerInThread.
/// E.g. suppose allocImpl is called twice: first for 2 MB with blocker set to
/// VariableContext::User, then for 3 MB with no blocker. This should increase the
/// Global memory tracker by 5 MB and the User memory tracker by 3 MB. So we can't group
/// these two calls into one memory_tracker->allocImpl call.
VariableContext blocker_level = MemoryTrackerBlockerInThread::getLevel();
if (blocker_level != current_thread->untracked_memory_blocker_level)
{
current_thread->flushUntrackedMemory();
}
current_thread->untracked_memory_blocker_level = blocker_level;
Int64 previous_untracked_memory = current_thread->untracked_memory;
current_thread->untracked_memory += size;
if (current_thread->untracked_memory > current_thread->untracked_memory_limit)
{
Int64 current_untracked_memory = current_thread->untracked_memory;
current_thread->untracked_memory = 0;
try
{
return memory_tracker->allocImpl(current_untracked_memory, throw_if_memory_exceeded);
}
catch (...)
{
current_thread->untracked_memory += previous_untracked_memory;
throw;
}
}
return AllocationTrace(memory_tracker->getSampleProbability(size));
}
return AllocationTrace(0);
}
void CurrentMemoryTracker::check()
{
if (auto * memory_tracker = getMemoryTracker())
std::ignore = memory_tracker->allocImpl(0, true);
}
AllocationTrace CurrentMemoryTracker::alloc(Int64 size)
{
return allocImpl(size, /*throw_if_memory_exceeded=*/ true);
}
AllocationTrace CurrentMemoryTracker::allocNoThrow(Int64 size)
{
return allocImpl(size, /*throw_if_memory_exceeded=*/ false);
}
AllocationTrace CurrentMemoryTracker::free(Int64 size)
{
if (auto * memory_tracker = getMemoryTracker())
{
if (!current_thread)
{
return memory_tracker->free(size);
}
VariableContext blocker_level = MemoryTrackerBlockerInThread::getLevel();
if (blocker_level != current_thread->untracked_memory_blocker_level)
{
current_thread->flushUntrackedMemory();
}
current_thread->untracked_memory_blocker_level = blocker_level;
current_thread->untracked_memory -= size;
if (current_thread->untracked_memory < -current_thread->untracked_memory_limit)
{
Int64 untracked_memory = current_thread->untracked_memory;
current_thread->untracked_memory = 0;
return memory_tracker->free(-untracked_memory);
}
return AllocationTrace(memory_tracker->getSampleProbability(size));
}
return AllocationTrace(0);
}
void CurrentMemoryTracker::injectFault()
{
if (auto * memory_tracker = getMemoryTracker())
memory_tracker->injectFault();
}