mirror of
https://github.com/ClickHouse/ClickHouse.git
synced 2024-12-16 11:22:12 +00:00
82 lines
2.1 KiB
C++
82 lines
2.1 KiB
C++
|
#include <Common/MemoryTracker.h>
|
||
|
#include <Common/CurrentThread.h>
|
||
|
|
||
|
#include <Common/CurrentMemoryTracker.h>
|
||
|
|
||
|
namespace
|
||
|
{
|
||
|
|
||
|
MemoryTracker * getMemoryTracker()
|
||
|
{
|
||
|
if (auto * thread_memory_tracker = DB::CurrentThread::getMemoryTracker())
|
||
|
return thread_memory_tracker;
|
||
|
|
||
|
/// Once the main thread is initialized,
|
||
|
/// total_memory_tracker is initialized too.
|
||
|
/// And can be used, since MainThreadStatus is required for profiling.
|
||
|
if (DB::MainThreadStatus::get())
|
||
|
return &total_memory_tracker;
|
||
|
|
||
|
return nullptr;
|
||
|
}
|
||
|
|
||
|
}
|
||
|
|
||
|
namespace CurrentMemoryTracker
|
||
|
{
|
||
|
|
||
|
using DB::current_thread;
|
||
|
|
||
|
void alloc(Int64 size)
|
||
|
{
|
||
|
if (auto * memory_tracker = getMemoryTracker())
|
||
|
{
|
||
|
if (current_thread)
|
||
|
{
|
||
|
current_thread->untracked_memory += size;
|
||
|
if (current_thread->untracked_memory > current_thread->untracked_memory_limit)
|
||
|
{
|
||
|
/// Zero untracked before track. If tracker throws out-of-limit we would be able to alloc up to untracked_memory_limit bytes
|
||
|
/// more. It could be useful to enlarge Exception message in rethrow logic.
|
||
|
Int64 tmp = current_thread->untracked_memory;
|
||
|
current_thread->untracked_memory = 0;
|
||
|
memory_tracker->alloc(tmp);
|
||
|
}
|
||
|
}
|
||
|
/// total_memory_tracker only, ignore untracked_memory
|
||
|
else
|
||
|
{
|
||
|
memory_tracker->alloc(size);
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
|
||
|
void realloc(Int64 old_size, Int64 new_size)
|
||
|
{
|
||
|
Int64 addition = new_size - old_size;
|
||
|
addition > 0 ? alloc(addition) : free(-addition);
|
||
|
}
|
||
|
|
||
|
void free(Int64 size)
|
||
|
{
|
||
|
if (auto * memory_tracker = getMemoryTracker())
|
||
|
{
|
||
|
if (current_thread)
|
||
|
{
|
||
|
current_thread->untracked_memory -= size;
|
||
|
if (current_thread->untracked_memory < -current_thread->untracked_memory_limit)
|
||
|
{
|
||
|
memory_tracker->free(-current_thread->untracked_memory);
|
||
|
current_thread->untracked_memory = 0;
|
||
|
}
|
||
|
}
|
||
|
/// total_memory_tracker only, ignore untracked_memory
|
||
|
else
|
||
|
{
|
||
|
memory_tracker->free(size);
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
|
||
|
}
|