*: Fix invalid memory access when shutting down (#8409) (#8412)

close #8408

dbms/src/Common/MemoryTracker.cpp

@@ -39,7 +39,7 @@ MemoryTracker::~MemoryTracker()
     if (is_global_root)
         return;
 
-    if (peak)
+    if (peak && log_peak_memory_usage_in_destructor)
     {
         try
         {
@@ -305,19 +305,23 @@ void initStorageMemoryTracker(Int64 limit, Int64 larger_than_limit)
         "Storage task memory limit={}, larger_than_limit={}",
         formatReadableSizeWithBinarySuffix(limit),
         formatReadableSizeWithBinarySuffix(larger_than_limit));
+
+    // When these (static) mem tracker are reset, it means the process is shutdown and the logging system is stopped.
+    // Do not log down the usage info in dctor of them.
+    bool log_in_destructor = false;
     RUNTIME_CHECK(sub_root_of_query_storage_task_mem_trackers == nullptr);
-    sub_root_of_query_storage_task_mem_trackers = MemoryTracker::create(limit);
+    sub_root_of_query_storage_task_mem_trackers = MemoryTracker::create(limit, nullptr, log_in_destructor);
     sub_root_of_query_storage_task_mem_trackers->setBytesThatRssLargerThanLimit(larger_than_limit);
     sub_root_of_query_storage_task_mem_trackers->setAmountMetric(CurrentMetrics::MemoryTrackingQueryStorageTask);
 
     RUNTIME_CHECK(fetch_pages_mem_tracker == nullptr);
-    fetch_pages_mem_tracker = MemoryTracker::create();
-    fetch_pages_mem_tracker->setNext(sub_root_of_query_storage_task_mem_trackers.get());
+    fetch_pages_mem_tracker
+        = MemoryTracker::create(0, sub_root_of_query_storage_task_mem_trackers.get(), log_in_destructor);
     fetch_pages_mem_tracker->setAmountMetric(CurrentMetrics::MemoryTrackingFetchPages);
 
     RUNTIME_CHECK(shared_column_data_mem_tracker == nullptr);
-    shared_column_data_mem_tracker = MemoryTracker::create();
-    shared_column_data_mem_tracker->setNext(sub_root_of_query_storage_task_mem_trackers.get());
+    shared_column_data_mem_tracker
+        = MemoryTracker::create(0, sub_root_of_query_storage_task_mem_trackers.get(), log_in_destructor);
     shared_column_data_mem_tracker->setAmountMetric(CurrentMetrics::MemoryTrackingSharedColumnData);
 }
 

dbms/src/Common/MemoryTracker.h

@@ -18,6 +18,7 @@
 #include <common/types.h>
 
 #include <atomic>
+#include <boost/noncopyable.hpp>
 
 extern std::atomic<Int64> real_rss, proc_num_threads, baseline_of_query_mem_tracker;
 extern std::atomic<UInt64> proc_virt_size;
@@ -46,6 +47,7 @@ class MemoryTracker : public std::enable_shared_from_this<MemoryTracker>
     double fault_probability = 0;
 
     bool is_global_root = false;
+    bool log_peak_memory_usage_in_destructor = true;
 
     /// To test the accuracy of memory track, it throws an exception when the part exceeding the tracked amount is greater than accuracy_diff_for_test.
     std::atomic<Int64> accuracy_diff_for_test{0};
@@ -65,7 +67,6 @@ class MemoryTracker : public std::enable_shared_from_this<MemoryTracker>
     std::atomic<const char *> description = nullptr;
 
     /// Make constructors private to ensure all objects of this class is created by `MemoryTracker::create`.
-    MemoryTracker() = default;
     explicit MemoryTracker(Int64 limit_)
         : limit(limit_)
     {}
@@ -79,16 +80,15 @@ class MemoryTracker : public std::enable_shared_from_this<MemoryTracker>
 
 public:
     /// Using `std::shared_ptr` and `new` instread of `std::make_shared` is because `std::make_shared` cannot call private constructors.
-    static MemoryTrackerPtr create(Int64 limit = 0)
+    static MemoryTrackerPtr create(
+        Int64 limit = 0,
+        MemoryTracker * parent = nullptr,
+        bool log_peak_memory_usage_in_destructor = true)
     {
-        if (limit == 0)
-        {
-            return std::shared_ptr<MemoryTracker>(new MemoryTracker);
-        }
-        else
-        {
-            return std::shared_ptr<MemoryTracker>(new MemoryTracker(limit));
-        }
+        auto p = std::shared_ptr<MemoryTracker>(new MemoryTracker(limit));
+        p->setParent(parent);
+        p->log_peak_memory_usage_in_destructor = log_peak_memory_usage_in_destructor;
+        return p;
     }
 
     static MemoryTrackerPtr createGlobalRoot() { return std::shared_ptr<MemoryTracker>(new MemoryTracker(0, true)); }
@@ -128,7 +128,7 @@ class MemoryTracker : public std::enable_shared_from_this<MemoryTracker>
     void setAccuracyDiffForTest(Int64 value) { accuracy_diff_for_test.store(value, std::memory_order_relaxed); }
 
     /// next should be changed only once: from nullptr to some value.
-    void setNext(MemoryTracker * elem)
+    void setParent(MemoryTracker * elem)
     {
         MemoryTracker * old_val = nullptr;
         if (!next.compare_exchange_strong(old_val, elem, std::memory_order_seq_cst, std::memory_order_relaxed))
@@ -187,8 +187,6 @@ void free(Int64 size);
 } // namespace CurrentMemoryTracker
 
 
-#include <boost/noncopyable.hpp>
-
 struct TemporarilyDisableMemoryTracker : private boost::noncopyable
 {
     MemoryTracker * memory_tracker;

dbms/src/Common/tests/gtest_memtracker.cpp

@@ -39,8 +39,7 @@ TEST_F(MemTrackerTest, testRootAndChild)
 try
 {
     auto root_mem_tracker = MemoryTracker::create();
-    auto child_mem_tracker = MemoryTracker::create(512);
-    child_mem_tracker->setNext(root_mem_tracker.get());
+    auto child_mem_tracker = MemoryTracker::create(512, root_mem_tracker.get());
     // alloc 500
     child_mem_tracker->alloc(500);
     ASSERT_EQ(500, child_mem_tracker->get());
@@ -71,10 +70,8 @@ TEST_F(MemTrackerTest, testRootAndMultipleChild)
 try
 {
     auto root = MemoryTracker::create(512); // limit 512
-    auto child1 = MemoryTracker::create(512); // limit 512
-    auto child2 = MemoryTracker::create(512); // limit 512
-    child1->setNext(root.get());
-    child2->setNext(root.get());
+    auto child1 = MemoryTracker::create(512, root.get()); // limit 512
+    auto child2 = MemoryTracker::create(512, root.get()); // limit 512
     // alloc 500 on child1
     child1->alloc(500);
     ASSERT_EQ(500, child1->get());

dbms/src/Interpreters/ProcessList.cpp

@@ -172,7 +172,7 @@ ProcessList::EntryPtr ProcessList::insert(
             user_process_list.user_memory_tracker->setOrRaiseLimit(
                 settings.max_memory_usage_for_user.getActualBytes(total_memory));
             user_process_list.user_memory_tracker->setDescription("(for user)");
-            current_memory_tracker->setNext(user_process_list.user_memory_tracker.get());
+            current_memory_tracker->setParent(user_process_list.user_memory_tracker.get());
 
             /// Track memory usage for all simultaneously running queries.
             /// You should specify this value in configuration for default profile,
@@ -183,7 +183,7 @@ ProcessList::EntryPtr ProcessList::insert(
             total_memory_tracker->setBytesThatRssLargerThanLimit(settings.bytes_that_rss_larger_than_limit);
             total_memory_tracker->setDescription("(total)");
             total_memory_tracker->setAccuracyDiffForTest(settings.memory_tracker_accuracy_diff_for_test);
-            user_process_list.user_memory_tracker->setNext(total_memory_tracker.get());
+            user_process_list.user_memory_tracker->setParent(total_memory_tracker.get());
         }
 
         if (!total_network_throttler && settings.max_network_bandwidth_for_all_users)

dbms/src/Storages/BackgroundProcessingPool.cpp

@@ -213,8 +213,7 @@ void BackgroundProcessingPool::threadFunction(size_t thread_idx) noexcept
     }
 
     // set up the thread local memory tracker
-    auto memory_tracker = MemoryTracker::create();
-    memory_tracker->setNext(root_of_non_query_mem_trackers.get());
+    auto memory_tracker = MemoryTracker::create(0, root_of_non_query_mem_trackers.get());
     memory_tracker->setMetric(CurrentMetrics::MemoryTrackingInBackgroundProcessingPool);
     current_memory_tracker = memory_tracker.get();