[BugFix] Fix race condition in merge_isomorphic_profiles (backport StarRocks#59809) (StarRocks#59823)

Signed-off-by: stdpain <[email protected]>
Co-authored-by: stdpain <[email protected]>

Author: mergify[bot]

be/src/util/runtime_profile.cpp

@@ -360,6 +360,16 @@ RuntimeProfile* RuntimeProfile::get_child(const std::string& name) {
     return it->second;
 }
 
+RuntimeProfile* RuntimeProfile::get_child_unlock(const std::string& name) {
+    auto it = _child_map.find(name);
+
+    if (it == _child_map.end()) {
+        return nullptr;
+    }
+
+    return it->second;
+}
+
 RuntimeProfile* RuntimeProfile::get_child(const size_t index) {
     std::lock_guard<std::mutex> l(_children_lock);
     if (index >= _children.size()) {
@@ -857,6 +867,13 @@ RuntimeProfile::EventSequence* RuntimeProfile::add_event_sequence(const std::str
     _event_sequence_map[name] = timer;
     return timer;
 }
+template <class Visitor>
+void RuntimeProfile::foreach_children(Visitor&& callback) {
+    std::lock_guard guard(_children_lock);
+    for (auto& child : _children) {
+        callback(child.first, child.second);
+    }
+}
 
 RuntimeProfile* RuntimeProfile::merge_isomorphic_profiles(ObjectPool* obj_pool, std::vector<RuntimeProfile*>& profiles,
                                                           bool require_identical) {
@@ -1027,19 +1044,28 @@ RuntimeProfile* RuntimeProfile::merge_isomorphic_profiles(ObjectPool* obj_pool,
         size_t max_child_size = 0;
         RuntimeProfile* profile_with_full_child = nullptr;
         for (auto* profile : profiles) {
-            if (profile->_children.size() > max_child_size) {
+            if (profile->num_children() > max_child_size) {
                 max_child_size = profile->_children.size();
                 profile_with_full_child = profile;
             }
         }
         if (profile_with_full_child != nullptr) {
             bool identical = true;
-            for (size_t i = 0; i < max_child_size; i++) {
-                auto& prototype_kv = profile_with_full_child->_children[i];
-                const std::string& child_name = prototype_kv.first->name();
+
+            profile_with_full_child->foreach_children([&obj_pool, &profiles, &identical, require_identical,
+                                                       merged_profile,
+                                                       profile_with_full_child](RuntimeProfile* child, bool indent) {
+                const std::string& child_name = child->name();
                 std::vector<RuntimeProfile*> sub_profiles;
                 for (auto* profile : profiles) {
-                    auto* child = profile->get_child(child_name);
+                    RuntimeProfile* child = nullptr;
+                    // We've already acquired the profile's lock, don't acquire again
+                    if (profile == profile_with_full_child) {
+                        child = profile->get_child_unlock(child_name);
+                    } else {
+                        child = profile->get_child(child_name);
+                    }
+
                     if (child == nullptr) {
                         identical = false;
                         if (require_identical) {
@@ -1051,8 +1077,9 @@ RuntimeProfile* RuntimeProfile::merge_isomorphic_profiles(ObjectPool* obj_pool,
                     sub_profiles.push_back(child);
                 }
                 auto* merged_child = merge_isomorphic_profiles(obj_pool, sub_profiles, require_identical);
-                merged_profile->add_child(merged_child, prototype_kv.second, nullptr);
-            }
+                merged_profile->add_child(merged_child, indent, nullptr);
+            });
+
             if (require_identical && !identical) {
                 merged_profile->add_info_string("NotIdentical");
             }

be/src/util/runtime_profile.h

@@ -40,6 +40,7 @@
 #include <atomic>
 #include <functional>
 #include <iostream>
+#include <mutex>
 #include <optional>
 #include <thread>
 #include <unordered_set>
@@ -438,6 +439,9 @@ class RuntimeProfile {
         return add_child_counter(name, type, strategy, ROOT_COUNTER);
     }
 
+    template <class Visitor>
+    void foreach_children(Visitor&& callback);
+
     // Add a derived counter with 'name'/'type'. The counter is owned by the
     // RuntimeProfile object.
     // If parent_name is a non-empty string, the counter is added as a child of
@@ -506,7 +510,10 @@ class RuntimeProfile {
     // Divides all counters by n
     void divide(int n);
 
-    size_t num_children() const { return _child_map.size(); }
+    size_t num_children() const {
+        std::lock_guard guard(_children_lock);
+        return _child_map.size();
+    }
 
     // Get child of given name
     RuntimeProfile* get_child(const std::string& name);
@@ -579,6 +586,8 @@ class RuntimeProfile {
     Counter* add_counter_unlock(const std::string& name, TUnit::type type, const TCounterStrategy& strategy,
                                 const std::string& parent_name);
 
+    RuntimeProfile* get_child_unlock(const std::string& name);
+
     RuntimeProfile* _parent;
 
     // Pool for allocated counters. Usually owned by the creator of this

be/test/util/runtime_profile_test.cpp

@@ -19,7 +19,12 @@
 
 #include <gtest/gtest.h>
 
+#include <thread>
+#include <tuple>
+
 #include "common/logging.h"
+#include "common/object_pool.h"
+#include "fmt/format.h"
 #include "gen_cpp/RuntimeProfile_types.h"
 
 namespace starrocks {
@@ -533,4 +538,31 @@ TEST(TestRuntimeProfile, testUpdateWithOldAndNewProfile) {
     ASSERT_EQ(2, child_profile->get_version());
 }
 
+TEST(TestRuntimeProfile, testRaceMergeProfiles) {
+    ObjectPool pool;
+    std::vector<RuntimeProfile*> profiles;
+    const size_t init_profile_size = 10;
+    for (size_t i = 0; i < init_profile_size; ++i) {
+        profiles.push_back(pool.add(new RuntimeProfile("test-name")));
+    }
+
+    const size_t loop_test_size = 10;
+
+    // race between merge_isomorphic_profiles and create_child
+    std::thread thr1([&profiles]() {
+        for (size_t i = 0; i < loop_test_size; ++i) {
+            ObjectPool pool;
+            std::ignore = RuntimeProfile::merge_isomorphic_profiles(&pool, profiles);
+        }
+    });
+    std::thread thr2([&profiles]() {
+        for (size_t i = 0; i < loop_test_size; ++i) {
+            profiles[0]->create_child(fmt::format("MERGE{}", i), true, true);
+        }
+    });
+
+    thr1.join();
+    thr2.join();
+}
+
 } // namespace starrocks