[BugFix] destruct workgroup's MemTracker prematurely (backport #55134) (

#55157)

Co-authored-by: satanson <[email protected]>

be/src/exec/pipeline/query_context.cpp

@@ -193,10 +193,20 @@ Status QueryContext::init_query_once(workgroup::WorkGroup* wg, bool enable_group
 void QueryContext::release_workgroup_token_once() {
     auto* old = _wg_running_query_token_atomic_ptr.load();
     if (old != nullptr && _wg_running_query_token_atomic_ptr.compare_exchange_strong(old, nullptr)) {
+        // The release_workgroup_token_once function is called by FragmentContext::cancel
+        // to detach the QueryContext from the workgroup.
+        // When the workgroup undergoes a configuration change, the old version of the workgroup is released,
+        // and a new version is created. The old workgroup will only be physically destroyed once no
+        // QueryContext is attached to it.
+        // However, the MemTracker of the old workgroup outlives the workgroup itself because
+        // it is accessed during the destruction of the QueryContext through its MemTracker
+        // (the workgroup's MemTracker serves as the parent of the QueryContext's MemTracker).
+        // To prevent the MemTracker from being released prematurely, it must be explicitly retained
+        // to ensure it remains valid until it is no longer needed.
+        _wg_mem_tracker = _wg_running_query_token_ptr->get_wg()->grab_mem_tracker();
         _wg_running_query_token_ptr.reset();
     }
 }
-
 void QueryContext::set_query_trace(std::shared_ptr<starrocks::debug::QueryTrace> query_trace) {
     std::call_once(_query_trace_init_flag, [this, &query_trace]() { _query_trace = std::move(query_trace); });
 }

be/src/exec/pipeline/query_context.h

@@ -258,8 +258,6 @@ class QueryContext : public std::enable_shared_from_this<QueryContext> {
     int64_t _big_query_profile_threshold_ns = 0;
     int64_t _runtime_profile_report_interval_ns = std::numeric_limits<int64_t>::max();
     TPipelineProfileLevel::type _profile_level;
-    std::shared_ptr<MemTracker> _mem_tracker;
-    std::shared_ptr<MemTracker> _connector_scan_mem_tracker;
     ObjectPool _object_pool;
     DescriptorTbl* _desc_tbl = nullptr;
     std::once_flag _query_trace_init_flag;
@@ -295,8 +293,14 @@ class QueryContext : public std::enable_shared_from_this<QueryContext> {
     std::shared_ptr<QueryStatisticsRecvr> _sub_plan_query_statistics_recvr; // For receive
 
     int64_t _scan_limit = 0;
+    // _wg_mem_tracker is used to grab mem_tracker in workgroup to prevent it from
+    // being released prematurely in FragmentContext::cancel, otherwise accessing
+    // workgroup's mem_tracker in QueryContext's dtor shall cause segmentation fault.
+    std::shared_ptr<MemTracker> _wg_mem_tracker = nullptr;
     workgroup::RunningQueryTokenPtr _wg_running_query_token_ptr;
     std::atomic<workgroup::RunningQueryToken*> _wg_running_query_token_atomic_ptr = nullptr;
+    std::shared_ptr<MemTracker> _mem_tracker;
+    std::shared_ptr<MemTracker> _connector_scan_mem_tracker;
 
     // STREAM MV
     std::shared_ptr<StreamEpochManager> _stream_epoch_manager;

be/src/exec/workgroup/work_group.h

@@ -101,6 +101,7 @@ using WorkGroupScanSchedEntity = WorkGroupSchedEntity<ScanTaskQueue>;
 struct RunningQueryToken {
     explicit RunningQueryToken(WorkGroupPtr wg) : wg(std::move(wg)) {}
     ~RunningQueryToken();
+    WorkGroupPtr get_wg() { return wg; }
 
 private:
     WorkGroupPtr wg;
@@ -127,6 +128,7 @@ class WorkGroup : public std::enable_shared_from_this<WorkGroup> {
     void copy_metrics(const WorkGroup& rhs);
 
     MemTracker* mem_tracker() { return _mem_tracker.get(); }
+    std::shared_ptr<MemTracker> grab_mem_tracker() { return _mem_tracker; }
     const MemTracker* mem_tracker() const { return _mem_tracker.get(); }
     MemTracker* connector_scan_mem_tracker() { return _connector_scan_mem_tracker.get(); }