1

xinyiZzz · May 23, 2024 · 9adf326 · 9adf326
1 parent e4f41db
commit 9adf326
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 11 deletions.
diff --git a/be/src/runtime/memory/thread_mem_tracker_mgr.cpp b/be/src/runtime/memory/thread_mem_tracker_mgr.cpp
@@ -44,20 +44,28 @@ class AsyncCancelQueryTask : public Runnable {
 void ThreadMemTrackerMgr::attach_limiter_tracker(
         const std::shared_ptr<MemTrackerLimiter>& mem_tracker) {
     DCHECK(mem_tracker);
-    DCHECK(_reserved_mem == 0);
     CHECK(init());
-    flush_untracked_mem();
+    if (_reserved_mem == 0) {
+        flush_untracked_mem();
+    }
     _limiter_tracker = mem_tracker;
     _limiter_tracker_raw = mem_tracker.get();
+    _attach_level++;
 }
 
 void ThreadMemTrackerMgr::detach_limiter_tracker(
         const std::shared_ptr<MemTrackerLimiter>& old_mem_tracker) {
     CHECK(init());
-    release_reserved();
-    flush_untracked_mem();
+    // If reserve memory is called in nested attach, release reserved in outermost detach.
+    if (_attach_level == 1) {
+        release_reserved();
+    }
+    if (_reserved_mem == 0) {
+        flush_untracked_mem();
+    }
     _limiter_tracker = old_mem_tracker;
     _limiter_tracker_raw = old_mem_tracker.get();
+    _attach_level--;
 }
 
 void ThreadMemTrackerMgr::cancel_query(const std::string& exceed_msg) {

diff --git a/be/src/runtime/memory/thread_mem_tracker_mgr.h b/be/src/runtime/memory/thread_mem_tracker_mgr.h
@@ -149,6 +149,9 @@ class ThreadMemTrackerMgr {
     bool _stop_consume = false;
     TUniqueId _query_id = TUniqueId();
     bool _is_query_cancelled = false;
+    // SCOPED_ATTACH_TASK cannot be nested, but SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER can continue to be called,
+    // so `attach_limiter_tracker` may be nested, _attach_level is the number of nesting levels.
+    size_t _attach_level = 0;
 };
 
 inline bool ThreadMemTrackerMgr::init() {
@@ -281,21 +284,21 @@ inline bool ThreadMemTrackerMgr::try_reserve(int64_t size) {
         tracker->consume(size);
     }
     _reserved_mem += size;
-    DCHECK(_untracked_mem == 0);
     return true;
 }
 
 inline void ThreadMemTrackerMgr::release_reserved() {
-    flush_untracked_mem();
     if (_reserved_mem > 0) {
-        doris::GlobalMemoryArbitrator::release_process_reserved_memory(_reserved_mem);
-        _limiter_tracker_raw->consume(-_reserved_mem);
+        doris::GlobalMemoryArbitrator::release_process_reserved_memory(_reserved_mem +
+                                                                       _untracked_mem);
+        _limiter_tracker_raw->release(_reserved_mem);
         if (_count_scope_mem) {
             _scope_mem -= _reserved_mem;
         }
         for (auto* tracker : _consumer_tracker_stack) {
-            tracker->consume(-_reserved_mem);
+            tracker->release(_reserved_mem);
         }
+        _untracked_mem = 0;
         _reserved_mem = 0;
     }
 }

diff --git a/be/src/util/mem_info.cpp b/be/src/util/mem_info.cpp
@@ -71,8 +71,8 @@ std::atomic<int64_t> MemInfo::refresh_interval_memory_growth = 0;
 
 static std::unordered_map<std::string, int64_t> _mem_info_bytes;
 std::atomic<int64_t> MemInfo::_s_sys_mem_available = -1;
-int64_t MemInfo::_s_sys_mem_available_low_water_mark = -1;
-int64_t MemInfo::_s_sys_mem_available_warning_water_mark = -1;
+int64_t MemInfo::_s_sys_mem_available_low_water_mark = std::numeric_limits<int64_t>::min();
+int64_t MemInfo::_s_sys_mem_available_warning_water_mark = std::numeric_limits<int64_t>::min();
 std::atomic<int64_t> MemInfo::_s_process_minor_gc_size = -1;
 std::atomic<int64_t> MemInfo::_s_process_full_gc_size = -1;
 std::mutex MemInfo::je_purge_dirty_pages_lock;