From fa6dd8fc6f8d1bd7f9cf9e95789c3ffc1435ea46 Mon Sep 17 00:00:00 2001
From: Xinyi Zou <zouxinyi02@gmail.com>
Date: Tue, 6 Aug 2024 17:03:04 +0800
Subject: [PATCH] 1

---
 be/src/runtime/thread_context.h |  6 +++---
 be/src/vec/common/allocator.cpp | 13 +++++++++++--
 be/src/vec/common/allocator.h   |  8 +++++++-
 3 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/be/src/runtime/thread_context.h b/be/src/runtime/thread_context.h
index b009affa53fc7c7..11be7f7ddbe5343 100644
--- a/be/src/runtime/thread_context.h
+++ b/be/src/runtime/thread_context.h
@@ -188,9 +188,9 @@ class ThreadContext {
         ss << std::this_thread::get_id();
         return ss.str();
     }
-    // After thread_mem_tracker_mgr is initialized, the current thread Hook starts to
-    // consume/release mem_tracker.
-    // Note that the use of shared_ptr will cause a crash. The guess is that there is an
+    // Note that if set global Memory Hook, After thread_mem_tracker_mgr is initialized,
+    // the current thread Hook starts to consume/release mem_tracker.
+    // the use of shared_ptr will cause a crash. The guess is that there is an
     // intermediate state during the copy construction of shared_ptr. Shared_ptr is not equal
     // to nullptr, but the object it points to is not initialized. At this time, when the memory
     // is released somewhere, the hook is triggered to cause the crash.
diff --git a/be/src/vec/common/allocator.cpp b/be/src/vec/common/allocator.cpp
index ed25b6a8f7787a4..2bd0fd190b8d085 100644
--- a/be/src/vec/common/allocator.cpp
+++ b/be/src/vec/common/allocator.cpp
@@ -209,14 +209,23 @@ void Allocator<clear_memory_, mmap_populate, use_mmap, MemoryAllocator>::memory_
 
 template <bool clear_memory_, bool mmap_populate, bool use_mmap, typename MemoryAllocator>
 void Allocator<clear_memory_, mmap_populate, use_mmap, MemoryAllocator>::consume_memory(
-        size_t size) const {
+        size_t size) {
+    if (_tracker == nullptr) { // first time called
+        _tracker = doris::thread_context()->thread_mem_tracker_mgr->limiter_mem_tracker();
+    }
     CONSUME_THREAD_MEM_TRACKER(size);
 }
 
 template <bool clear_memory_, bool mmap_populate, bool use_mmap, typename MemoryAllocator>
 void Allocator<clear_memory_, mmap_populate, use_mmap, MemoryAllocator>::release_memory(
         size_t size) const {
-    RELEASE_THREAD_MEM_TRACKER(size);
+    doris::ThreadContext* thread_context = doris::thread_context(true);
+    DCHECK(_tracker);
+    if (thread_context && thread_context->thread_mem_tracker()->label() == _tracker->label()) {
+        RELEASE_THREAD_MEM_TRACKER(size);
+    } else {
+        _tracker->release(size);
+    }
 }
 
 template <bool clear_memory_, bool mmap_populate, bool use_mmap, typename MemoryAllocator>
diff --git a/be/src/vec/common/allocator.h b/be/src/vec/common/allocator.h
index 2bcce9a9c689413..9135196401a010c 100644
--- a/be/src/vec/common/allocator.h
+++ b/be/src/vec/common/allocator.h
@@ -209,6 +209,10 @@ class RecordSizeMemoryAllocator {
     static std::mutex _mutex;
 };
 
+namespace doris {
+class MemTrackerLimiter;
+}
+
 /** Responsible for allocating / freeing memory. Used, for example, in PODArray, Arena.
   * Also used in hash tables.
   * The interface is different from std::allocator
@@ -228,7 +232,7 @@ class Allocator {
     // alloc will continue to execute, so the consume memtracker is forced.
     void memory_check(size_t size) const;
     // Increases consumption of this tracker by 'bytes'.
-    void consume_memory(size_t size) const;
+    void consume_memory(size_t size);
     void release_memory(size_t size) const;
     void throw_bad_alloc(const std::string& err) const;
 #ifndef NDEBUG
@@ -412,6 +416,8 @@ class Allocator {
                                       | (mmap_populate ? MAP_POPULATE : 0)
 #endif
             ;
+
+    std::shared_ptr<doris::MemTrackerLimiter> _tracker {nullptr};
 };
 
 /** Allocator with optimization to place small memory ranges in automatic memory.