blob: 3c99fd123535f092cac7e412fa7e5871c0702130 (
plain) (
blame)
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
|
#pragma once
#include <Common/CurrentThread.h>
#include <Common/MemoryTracker.h>
namespace DB
{
struct MemoryTrackerSwitcher
{
explicit MemoryTrackerSwitcher(MemoryTracker * new_tracker)
{
/// current_thread is not initialized for the main thread, so simply do not switch anything
if (!current_thread)
return;
auto * thread_tracker = CurrentThread::getMemoryTracker();
prev_untracked_memory = current_thread->untracked_memory;
prev_memory_tracker_parent = thread_tracker->getParent();
current_thread->untracked_memory = 0;
thread_tracker->setParent(new_tracker);
}
~MemoryTrackerSwitcher()
{
/// current_thread is not initialized for the main thread, so simply do not switch anything
if (!current_thread)
return;
CurrentThread::flushUntrackedMemory();
auto * thread_tracker = CurrentThread::getMemoryTracker();
current_thread->untracked_memory = prev_untracked_memory;
thread_tracker->setParent(prev_memory_tracker_parent);
}
private:
MemoryTracker * prev_memory_tracker_parent = nullptr;
Int64 prev_untracked_memory = 0;
};
}
|
