7

be/src/olap/page_cache.cpp

@@ -30,20 +30,13 @@ PageBase<TAllocator>::PageBase(size_t b, bool use_cache, segment_v2::PageTypePB
           LRUCacheValueBase(),
           _size(b),
           _capacity(b) {
-    if (use_cache) {
-        SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(
-                StoragePageCache::instance()->mem_tracker(page_type));
-        _data = reinterpret_cast<char*>(TAllocator::alloc(_capacity, ALLOCATOR_ALIGNMENT_16));
-    } else {
-        _data = reinterpret_cast<char*>(TAllocator::alloc(_capacity, ALLOCATOR_ALIGNMENT_16));
-    }
+    _data = reinterpret_cast<char*>(TAllocator::alloc(_capacity, ALLOCATOR_ALIGNMENT_16));
 }
 
 template <typename TAllocator>
 PageBase<TAllocator>::~PageBase() {
     if (_data != nullptr) {
         DCHECK(_capacity != 0 && _size != 0);
-        SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(TAllocator::mem_tracker_);
         TAllocator::free(_data, _capacity);
     }
 }

be/src/runtime/thread_context.h

@@ -452,27 +452,34 @@ class SwitchThreadMemTrackerLimiter {
             const std::shared_ptr<doris::MemTrackerLimiter>& mem_tracker) {
         DCHECK(mem_tracker);
         doris::ThreadLocalHandle::create_thread_local_if_not_exits();
-        _old_mem_tracker = thread_context()->thread_mem_tracker_mgr->limiter_mem_tracker();
-        thread_context()->thread_mem_tracker_mgr->attach_limiter_tracker(mem_tracker);
+        if (mem_tracker != thread_context()->thread_mem_tracker_mgr->limiter_mem_tracker()) {
+            _old_mem_tracker = thread_context()->thread_mem_tracker_mgr->limiter_mem_tracker();
+            thread_context()->thread_mem_tracker_mgr->attach_limiter_tracker(mem_tracker);
+        }
     }
 
     explicit SwitchThreadMemTrackerLimiter(const doris::QueryThreadContext& query_thread_context) {
         doris::ThreadLocalHandle::create_thread_local_if_not_exits();
         DCHECK(thread_context()->task_id() ==
                query_thread_context.query_id); // workload group alse not change
         DCHECK(query_thread_context.query_mem_tracker);
-        _old_mem_tracker = thread_context()->thread_mem_tracker_mgr->limiter_mem_tracker();
-        thread_context()->thread_mem_tracker_mgr->attach_limiter_tracker(
-                query_thread_context.query_mem_tracker);
+        if (query_thread_context.query_mem_tracker !=
+            thread_context()->thread_mem_tracker_mgr->limiter_mem_tracker()) {
+            _old_mem_tracker = thread_context()->thread_mem_tracker_mgr->limiter_mem_tracker();
+            thread_context()->thread_mem_tracker_mgr->attach_limiter_tracker(
+                    query_thread_context.query_mem_tracker);
+        }
     }
 
     ~SwitchThreadMemTrackerLimiter() {
-        thread_context()->thread_mem_tracker_mgr->detach_limiter_tracker(_old_mem_tracker);
+        if (_old_mem_tracker != nullptr) {
+            thread_context()->thread_mem_tracker_mgr->detach_limiter_tracker(_old_mem_tracker);
+        }
         doris::ThreadLocalHandle::del_thread_local_if_count_is_zero();
     }
 
 private:
-    std::shared_ptr<doris::MemTrackerLimiter> _old_mem_tracker;
+    std::shared_ptr<doris::MemTrackerLimiter> _old_mem_tracker {nullptr};
 };
 
 class AddThreadMemTrackerConsumer {

be/src/vec/common/allocator.cpp

@@ -48,10 +48,6 @@ Allocator<clear_memory_, mmap_populate, use_mmap, MemoryAllocator>::Allocator(
         const std::shared_ptr<doris::MemTrackerLimiter>& tracker) {
     if (tracker) {
         mem_tracker_ = tracker;
-    } else {
-        // Note that when constructing an object that inherits Allocator, the memory tracker may not be attached
-        // in tls. memory tracker may be attached in tls only when the memory is allocated for the first time.
-        mem_tracker_ = doris::thread_context()->thread_mem_tracker_mgr->limiter_mem_tracker();
     }
 }
 
@@ -224,29 +220,21 @@ 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 {
-    CONSUME_THREAD_MEM_TRACKER(size);
+    if (mem_tracker_) {
+        SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(mem_tracker_);
+        CONSUME_THREAD_MEM_TRACKER(size);
+    } else {
+        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 {
-    doris::ThreadContext* thread_context = doris::thread_context(true);
-    if ((thread_context && thread_context->thread_mem_tracker()->label() != "Orphan") ||
-        mem_tracker_ == nullptr) {
-        // If thread_context exist and the label of thread_mem_tracker not equal to `Orphan`,
-        // this means that in the scope of SCOPED_ATTACH_TASK,
-        // so thread_mem_tracker should be used to release memory.
-        // If mem_tracker_ is nullptr there is a scenario where an object that inherits Allocator
-        // has never called alloc, but free memory.
-        // in phmap, the memory alloced by an object may be transferred to another object and then free.
-        // in this case, thread context must attach a memory tracker other than Orphan,
-        // otherwise memory tracking will be wrong.
+    if (mem_tracker_) {
+        SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(mem_tracker_);
         RELEASE_THREAD_MEM_TRACKER(size);
     } else {
-        // if thread_context does not exist or the label of thread_mem_tracker is equal to
-        // `Orphan`, it usually happens during object destruction. This means that
-        // the scope of SCOPED_ATTACH_TASK has been left,  so release memory using Allocator tracker.
-        SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(mem_tracker_);
         RELEASE_THREAD_MEM_TRACKER(size);
     }
 }

be/src/vec/common/allocator.h

@@ -234,6 +234,13 @@ 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'.
+    // some special cases:
+    // 1. objects that inherit Allocator will not be shared by multiple queries.
+    //  non-compliant: page cache, ORC ByteBuffer.
+    // 2. objects that inherit Allocator will only free memory allocated by themselves.
+    //  non-compliant: phmap, the memory alloced by an object may be transferred to another object and then free.
+    // 3. the memory tracker in TLS is the same during the construction of objects that inherit Allocator
+    //  and during subsequent memory allocation.
     void consume_memory(size_t size) const;
     void release_memory(size_t size) const;
     void throw_bad_alloc(const std::string& err) const;