memory tracked new/delete concept

contrib/jemalloc-cmake/CMakeLists.txt

@@ -15,7 +15,6 @@ ${JEMALLOC_SOURCE_DIR}/src/extent_mmap.c
 ${JEMALLOC_SOURCE_DIR}/src/hash.c
 ${JEMALLOC_SOURCE_DIR}/src/hook.c
 ${JEMALLOC_SOURCE_DIR}/src/jemalloc.c
-${JEMALLOC_SOURCE_DIR}/src/jemalloc_cpp.cpp
 ${JEMALLOC_SOURCE_DIR}/src/large.c
 ${JEMALLOC_SOURCE_DIR}/src/log.c
 ${JEMALLOC_SOURCE_DIR}/src/malloc_io.c

dbms/src/Common/CurrentThread.cpp

@@ -46,6 +46,12 @@ MemoryTracker * CurrentThread::getMemoryTracker()
     return &current_thread->memory_tracker;
 }
 
+Int64 & CurrentThread::getUntrackedMemory()
+{
+    /// It assumes that (current_thread != nullptr) is already checked with getMemoryTracker()
+    return current_thread->untracked_memory;
+}
+
 void CurrentThread::updateProgressIn(const Progress & value)
 {
     if (unlikely(!current_thread))

dbms/src/Common/CurrentThread.h

@@ -47,6 +47,7 @@ public:
 
     static ProfileEvents::Counters & getProfileEvents();
     static MemoryTracker * getMemoryTracker();
+    static Int64 & getUntrackedMemory();
 
     /// Update read and write rows (bytes) statistics (used in system.query_thread_log)
     static void updateProgressIn(const Progress & value);

dbms/src/Common/MemoryTracker.cpp

@@ -17,6 +17,7 @@ namespace DB
 
 
 static constexpr size_t log_peak_memory_usage_every = 1ULL << 30;
+static constexpr Int64 untracked_memory_limit = 4 * 1024 * 1024;
 
 
 MemoryTracker::~MemoryTracker()
@@ -191,19 +192,38 @@ namespace CurrentMemoryTracker
     void alloc(Int64 size)
     {
         if (auto memory_tracker = DB::CurrentThread::getMemoryTracker())
-            memory_tracker->alloc(size);
+        {
+            Int64 & untracked = DB::CurrentThread::getUntrackedMemory();
+            untracked += size;
+            if (untracked > untracked_memory_limit)
+            {
+                memory_tracker->alloc(untracked);
+                untracked = 0;
+            }
+        }
     }
 
     void realloc(Int64 old_size, Int64 new_size)
     {
-        if (auto memory_tracker = DB::CurrentThread::getMemoryTracker())
-            memory_tracker->alloc(new_size - old_size);
+        Int64 addition = new_size - old_size;
+        if (addition > 0)
+            alloc(addition);
+        else
+            free(-addition);
     }
 
     void free(Int64 size)
     {
         if (auto memory_tracker = DB::CurrentThread::getMemoryTracker())
-            memory_tracker->free(size);
+        {
+            Int64 & untracked = DB::CurrentThread::getUntrackedMemory();
+            untracked -= size;
+            if (untracked < -untracked_memory_limit)
+            {
+                memory_tracker->free(-untracked);
+                untracked = 0;
+            }
+        }
     }
 }
 

dbms/src/Common/MemoryTracker.h

@@ -45,7 +45,11 @@ public:
 
     void realloc(Int64 old_size, Int64 new_size)
     {
-        alloc(new_size - old_size);
+        Int64 addition = new_size - old_size;
+        if (addition > 0)
+            alloc(addition);
+        else
+            free(-addition);
     }
 
     /** This function should be called after memory deallocation.

dbms/src/Common/ThreadStatus.cpp

@@ -50,6 +50,11 @@ ThreadStatus::ThreadStatus()
 
 ThreadStatus::~ThreadStatus()
 {
+    if (untracked_memory > 0)
+        memory_tracker.alloc(untracked_memory);
+    else
+        memory_tracker.free(-untracked_memory);
+
     if (deleter)
         deleter();
     current_thread = nullptr;

dbms/src/Common/ThreadStatus.h

@@ -92,6 +92,8 @@ public:
     /// TODO: merge them into common entity
     ProfileEvents::Counters performance_counters{VariableContext::Thread};
     MemoryTracker memory_tracker{VariableContext::Thread};
+    /// Small amount of untracked memory (per thread atomic-less counter)
+    Int64 untracked_memory = 0;
 
     /// Statistics of read and write rows/bytes
     Progress progress_in;

dbms/src/Common/new_delete.cpp

@@ -0,0 +1,34 @@
+#include <new>
+
+#include <common/likely.h>
+#include <Common/MemoryTracker.h>
+
+/// Replace default new/delete with memory tracking versions.
+/// @sa https://en.cppreference.com/w/cpp/memory/new/operator_new
+///     https://en.cppreference.com/w/cpp/memory/new/operator_delete
+#if 1
+
+void * operator new (std::size_t size)
+{
+    CurrentMemoryTracker::alloc(size);
+
+    auto * ptr = malloc(size);
+    if (likely(ptr != nullptr))
+        return ptr;
+
+    CurrentMemoryTracker::free(size);
+
+    /// @note no std::get_new_handler logic implemented
+    std::__throw_bad_alloc();
+}
+
+/// Called instead of 'delete(void * ptr)' if a user-defined replacement is provided
+void operator delete (void * ptr, std::size_t size) noexcept
+{
+    CurrentMemoryTracker::free(size);
+
+    if (likely(ptr != nullptr))
+        free(ptr);
+}
+
+#endif