From fabd0464a7fd7aaa18cad57517e169030d546f54 Mon Sep 17 00:00:00 2001
From: Xinyi Zou <zouxinyi02@gmail.com>
Date: Wed, 10 Apr 2024 22:44:40 +0800
Subject: [PATCH] [fix](memory) Fix memory tracker destructor deadlock (#33497)

---
 be/src/common/daemon.cpp                      |  1 +
 be/src/runtime/memory/mem_tracker_limiter.cpp | 40 ++++++++++---------
 be/src/runtime/memory/mem_tracker_limiter.h   |  2 +-
 3 files changed, 23 insertions(+), 20 deletions(-)

diff --git a/be/src/common/daemon.cpp b/be/src/common/daemon.cpp
index 4c6a0a012c2274..3a217413db8feb 100644
--- a/be/src/common/daemon.cpp
+++ b/be/src/common/daemon.cpp
@@ -211,6 +211,7 @@ void Daemon::memory_maintenance_thread() {
         // Update and print memory stat when the memory changes by 256M.
         if (abs(last_print_proc_mem - PerfCounters::get_vm_rss()) > 268435456) {
             last_print_proc_mem = PerfCounters::get_vm_rss();
+            doris::MemTrackerLimiter::clean_tracker_limiter_group();
             doris::MemTrackerLimiter::enable_print_log_process_usage();
 
             // Refresh mem tracker each type counter.
diff --git a/be/src/runtime/memory/mem_tracker_limiter.cpp b/be/src/runtime/memory/mem_tracker_limiter.cpp
index 08d8ea4d5e2f05..2218bed6959ddc 100644
--- a/be/src/runtime/memory/mem_tracker_limiter.cpp
+++ b/be/src/runtime/memory/mem_tracker_limiter.cpp
@@ -89,12 +89,9 @@ std::shared_ptr<MemTrackerLimiter> MemTrackerLimiter::create_shared(MemTrackerLi
     DCHECK(ExecEnv::tracking_memory());
     std::lock_guard<std::mutex> l(
             ExecEnv::GetInstance()->mem_tracker_limiter_pool[tracker->group_num()].group_lock);
-    tracker->tracker_limiter_group_it =
-            ExecEnv::GetInstance()->mem_tracker_limiter_pool[tracker->group_num()].trackers.insert(
-                    ExecEnv::GetInstance()
-                            ->mem_tracker_limiter_pool[tracker->group_num()]
-                            .trackers.end(),
-                    tracker);
+    ExecEnv::GetInstance()->mem_tracker_limiter_pool[tracker->group_num()].trackers.insert(
+            ExecEnv::GetInstance()->mem_tracker_limiter_pool[tracker->group_num()].trackers.end(),
+            tracker);
 #endif
     return tracker;
 }
@@ -128,19 +125,6 @@ MemTrackerLimiter::~MemTrackerLimiter() {
         }
         _consumption->set(0);
     }
-#ifndef BE_TEST
-    if (ExecEnv::tracking_memory()) {
-        std::lock_guard<std::mutex> l(
-                ExecEnv::GetInstance()->mem_tracker_limiter_pool[_group_num].group_lock);
-        if (tracker_limiter_group_it !=
-            ExecEnv::GetInstance()->mem_tracker_limiter_pool[_group_num].trackers.end()) {
-            ExecEnv::GetInstance()->mem_tracker_limiter_pool[_group_num].trackers.erase(
-                    tracker_limiter_group_it);
-            tracker_limiter_group_it =
-                    ExecEnv::GetInstance()->mem_tracker_limiter_pool[_group_num].trackers.end();
-        }
-    }
-#endif
     g_memtrackerlimiter_cnt << -1;
 }
 
@@ -173,6 +157,24 @@ void MemTrackerLimiter::refresh_global_counter() {
     }
 }
 
+void MemTrackerLimiter::clean_tracker_limiter_group() {
+#ifndef BE_TEST
+    if (ExecEnv::tracking_memory()) {
+        for (auto& group : ExecEnv::GetInstance()->mem_tracker_limiter_pool) {
+            std::lock_guard<std::mutex> l(group.group_lock);
+            auto it = group.trackers.begin();
+            while (it != group.trackers.end()) {
+                if ((*it).expired()) {
+                    it = group.trackers.erase(it);
+                } else {
+                    ++it;
+                }
+            }
+        }
+    }
+#endif
+}
+
 void MemTrackerLimiter::make_process_snapshots(std::vector<MemTracker::Snapshot>* snapshots) {
     MemTrackerLimiter::refresh_global_counter();
     int64_t all_tracker_mem_sum = 0;
diff --git a/be/src/runtime/memory/mem_tracker_limiter.h b/be/src/runtime/memory/mem_tracker_limiter.h
index d22a79c73f7a22..fcb319a8d770f9 100644
--- a/be/src/runtime/memory/mem_tracker_limiter.h
+++ b/be/src/runtime/memory/mem_tracker_limiter.h
@@ -160,6 +160,7 @@ class MemTrackerLimiter final : public MemTracker {
     }
 
     static void refresh_global_counter();
+    static void clean_tracker_limiter_group();
 
     Snapshot make_snapshot() const override;
     // Returns a list of all the valid tracker snapshots.
@@ -241,7 +242,6 @@ class MemTrackerLimiter final : public MemTracker {
     }
 
     // Iterator into mem_tracker_limiter_pool for this object. Stored to have O(1) remove.
-    std::list<std::weak_ptr<MemTrackerLimiter>>::iterator tracker_limiter_group_it;
     std::list<std::weak_ptr<MemTrackerLimiter>>::iterator tg_tracker_limiter_group_it;
 
 private: