From ae8914d5c80fbdb0156b58ba07607d7a2cdcc175 Mon Sep 17 00:00:00 2001
From: Xinyi Zou <zouxinyi02@gmail.com>
Date: Mon, 5 Jun 2023 17:09:13 +0800
Subject: [PATCH] [fix](memory) Fix meta_tool start failed #20045

---
 .../runtime/memory/thread_mem_tracker_mgr.cpp |  2 ++
 .../runtime/memory/thread_mem_tracker_mgr.h   | 21 ++++++++++---------
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/be/src/runtime/memory/thread_mem_tracker_mgr.cpp b/be/src/runtime/memory/thread_mem_tracker_mgr.cpp
index f0ffaebaad..d53a3731d8 100644
--- a/be/src/runtime/memory/thread_mem_tracker_mgr.cpp
+++ b/be/src/runtime/memory/thread_mem_tracker_mgr.cpp
@@ -28,6 +28,7 @@ void ThreadMemTrackerMgr::attach_limiter_tracker(
         const std::shared_ptr<MemTrackerLimiter>& mem_tracker,
         const TUniqueId& fragment_instance_id) {
     DCHECK(mem_tracker);
+    CHECK(init());
     flush_untracked_mem();
     _fragment_instance_id = fragment_instance_id;
     _limiter_tracker = mem_tracker;
@@ -37,6 +38,7 @@ void ThreadMemTrackerMgr::attach_limiter_tracker(
 
 void ThreadMemTrackerMgr::detach_limiter_tracker(
         const std::shared_ptr<MemTrackerLimiter>& old_mem_tracker) {
+    CHECK(init());
     flush_untracked_mem();
     _fragment_instance_id = TUniqueId();
     _limiter_tracker = old_mem_tracker;
diff --git a/be/src/runtime/memory/thread_mem_tracker_mgr.h b/be/src/runtime/memory/thread_mem_tracker_mgr.h
index b613883792..dd7b3b8b32 100644
--- a/be/src/runtime/memory/thread_mem_tracker_mgr.h
+++ b/be/src/runtime/memory/thread_mem_tracker_mgr.h
@@ -46,7 +46,7 @@ public:
         if (_init) flush_untracked_mem();
     }
 
-    void init();
+    bool init();
 
     // After attach, the current thread Memory Hook starts to consume/release task mem_tracker
     void attach_limiter_tracker(const std::shared_ptr<MemTrackerLimiter>& mem_tracker,
@@ -77,16 +77,16 @@ public:
     // must increase the control to avoid entering infinite recursion, otherwise it may cause crash or stuck,
     // Returns whether the memory exceeds limit, and will consume mem trcker no matter whether the limit is exceeded.
     void consume(int64_t size);
-    bool flush_untracked_mem();
+    void flush_untracked_mem();
 
     bool is_attach_query() { return _fragment_instance_id != TUniqueId(); }
 
     std::shared_ptr<MemTrackerLimiter> limiter_mem_tracker() {
-        init();
+        CHECK(init());
         return _limiter_tracker;
     }
     MemTrackerLimiter* limiter_mem_tracker_raw() {
-        init();
+        CHECK(init());
         return _limiter_tracker_raw;
     }
 
@@ -130,16 +130,18 @@ private:
     TUniqueId _fragment_instance_id = TUniqueId();
 };
 
-inline void ThreadMemTrackerMgr::init() {
+inline bool ThreadMemTrackerMgr::init() {
     // 1. Initialize in the thread context when the thread starts
     // 2. ExecEnv not initialized when thread start, initialized in limiter_mem_tracker().
-    if (!_init) {
-        DCHECK(_limiter_tracker == nullptr);
+    if (_init) return true;
+    if (ExecEnv::GetInstance()->orphan_mem_tracker() != nullptr) {
         _limiter_tracker = ExecEnv::GetInstance()->orphan_mem_tracker();
         _limiter_tracker_raw = ExecEnv::GetInstance()->orphan_mem_tracker_raw();
         _wait_gc = true;
         _init = true;
+        return true;
     }
+    return false;
 }
 
 inline bool ThreadMemTrackerMgr::push_consumer_tracker(MemTracker* tracker) {
@@ -179,11 +181,11 @@ inline void ThreadMemTrackerMgr::consume(int64_t size) {
     }
 }
 
-inline bool ThreadMemTrackerMgr::flush_untracked_mem() {
+inline void ThreadMemTrackerMgr::flush_untracked_mem() {
     // Temporary memory may be allocated during the consumption of the mem tracker, which will lead to entering
     // the Memory Hook again, so suspend consumption to avoid falling into an infinite loop.
+    if (!init()) return;
     _stop_consume = true;
-    init();
     DCHECK(_limiter_tracker_raw);
 
     old_untracked_mem = _untracked_mem;
@@ -194,7 +196,6 @@ inline bool ThreadMemTrackerMgr::flush_untracked_mem() {
     }
     _untracked_mem -= old_untracked_mem;
     _stop_consume = false;
-    return true;
 }
 
 } // namespace doris