[BugFix] fix thread local scope not work in pk table (StarRocks#48858)

Signed-off-by: luohaha <[email protected]>

Author: luohaha

be/src/storage/tablet_updates.cpp

@@ -912,13 +912,9 @@ DEFINE_FAIL_POINT(tablet_apply_load_compaction_state_failed);
 DEFINE_FAIL_POINT(tablet_apply_load_segments_failed);
 
 void TabletUpdates::do_apply() {
-    if (config::enable_pk_strict_memcheck) {
-        SCOPED_THREAD_LOCAL_CHECK_MEM_LIMIT_SETTER(true);
-        SCOPED_THREAD_LOCAL_SINGLETON_CHECK_MEM_TRACKER_SETTER(
-                StorageEngine::instance()->update_manager()->mem_tracker());
-    } else {
-        SCOPED_THREAD_LOCAL_CHECK_MEM_LIMIT_SETTER(false);
-    }
+    SCOPED_THREAD_LOCAL_CHECK_MEM_LIMIT_SETTER(config::enable_pk_strict_memcheck);
+    SCOPED_THREAD_LOCAL_SINGLETON_CHECK_MEM_TRACKER_SETTER(
+            config::enable_pk_strict_memcheck ? StorageEngine::instance()->update_manager()->mem_tracker() : nullptr);
     // only 1 thread at max is running this method
     bool first = true;
     while (!_apply_stopped) {

be/src/storage/update_manager.cpp

@@ -494,10 +494,8 @@ Status UpdateManager::set_cached_del_vec(const TabletSegmentId& tsid, const DelV
 }
 
 Status UpdateManager::on_rowset_finished(Tablet* tablet, Rowset* rowset) {
-    if (config::enable_pk_strict_memcheck) {
-        SCOPED_THREAD_LOCAL_CHECK_MEM_LIMIT_SETTER(true);
-        SCOPED_THREAD_LOCAL_SINGLETON_CHECK_MEM_TRACKER_SETTER(mem_tracker());
-    }
+    SCOPED_THREAD_LOCAL_MEM_SETTER(GlobalEnv::GetInstance()->process_mem_tracker(), config::enable_pk_strict_memcheck);
+    SCOPED_THREAD_LOCAL_SINGLETON_CHECK_MEM_TRACKER_SETTER(config::enable_pk_strict_memcheck ? mem_tracker() : nullptr);
     if (!rowset->has_data_files() || tablet->tablet_state() == TABLET_NOTREADY) {
         // if rowset is empty or tablet is in schemachange, we can skip preparing updatestates and pre-loading primary index
         return Status::OK();