From 5cb6f92fa281945cba5d7cdcf078fcae6bdb5309 Mon Sep 17 00:00:00 2001
From: UndefinedSy <shunyu.sean@gmail.com>
Date: Tue, 20 Aug 2024 10:39:56 +0800
Subject: [PATCH] fix. stats manager may cause crash when stats with new
 labels.

---
 src/common/stats/StatsManager.cpp          | 68 ++++++++++++++--------
 src/common/stats/test/StatsManagerTest.cpp | 30 ++++++++++
 2 files changed, 75 insertions(+), 23 deletions(-)

diff --git a/src/common/stats/StatsManager.cpp b/src/common/stats/StatsManager.cpp
index 14fd7d62d0c..2e3d641859b 100644
--- a/src/common/stats/StatsManager.cpp
+++ b/src/common/stats/StatsManager.cpp
@@ -177,18 +177,24 @@ CounterId StatsManager::counterWithLabels(const CounterId& id,
   }
   newIndex.back() = '}';
 
-  auto it = sm.nameMap_.find(newIndex);
-  // Get the counter if it already exists
-  if (it != sm.nameMap_.end()) {
-    return it->second.id_;
-  }
-
-  // Register a new counter if it doesn't exist
-  auto it2 = sm.nameMap_.find(index);
-  DCHECK(it2 != sm.nameMap_.end());
-  auto& methods = it2->second.methods_;
-
-  return registerStats(newIndex, methods);
+  std::vector<StatsMethod> methods;
+  {
+    folly::RWSpinLock::ReadHolder rh(sm.nameMapLock_);
+    // Check if the counter already exists
+    auto it = sm.nameMap_.find(newIndex);
+    if (it != sm.nameMap_.end()) {
+      // fast path if the counter already exists
+      return it->second.id_;
+    }
+    
+    // Register a new counter if it doesn't exist
+    // Fetch the counter indexed by the parent index and get the methods
+    auto it2 = sm.nameMap_.find(index);
+    DCHECK(it2 != sm.nameMap_.end());
+    methods = it2->second.methods_;
+  }  
+
+  return registerStats(newIndex, *methods);
 }
 
 // static
@@ -205,19 +211,35 @@ CounterId StatsManager::histoWithLabels(const CounterId& id, const std::vector<L
     newIndex.append(k).append("=").append(v).append(",");
   }
   newIndex.back() = '}';
-  auto it = sm.nameMap_.find(newIndex);
-  // Get the counter if it already exists
-  if (it != sm.nameMap_.end()) {
-    return it->second.id_;
-  }
 
-  auto it2 = sm.nameMap_.find(index);
-  DCHECK(it2 != sm.nameMap_.end());
-  auto& methods = it2->second.methods_;
-  auto& percentiles = it2->second.percentiles_;
+  std::vector<StatsMethod> methods;
+  std::vector<std::pair<std::string, double>> percentiles;
+  StatsManager::VT bucketSize, min, max;
+  {
+    folly::RWSpinLock::ReadHolder rh(sm.nameMapLock_);
+    // Check if the counter already exists
+    auto it = sm.nameMap_.find(newIndex);
+    if (it != sm.nameMap_.end()) {
+      // fast path if the counter already exists
+      return it->second.id_;
+    }
+    // Register a new counter if it doesn't exist
+    // Fetch the counter indexed by the parent index and get the methods
+    auto it2 = sm.nameMap_.find(index);
+    DCHECK(it2 != sm.nameMap_.end());
+    methods = it2->second.methods_;
+    percentiles = it2->second.percentiles_;
+    bucketSize = it2->second.bucketSize_;
+    min = it2->second.min_;
+    max = it2->second.max_;
+  }
 
-  return registerHisto(
-      newIndex, it2->second.bucketSize_, it2->second.min_, it2->second.max_, methods, percentiles);
+  return registerHisto(/* counterName */  newIndex,
+                       /* bucketSize */   bucketSize,
+                       /* min */          min,
+                       /* max */          max,
+                       /* methods */      *methods,
+                       /* percentiles */  *percentiles);
 }
 
 // static
diff --git a/src/common/stats/test/StatsManagerTest.cpp b/src/common/stats/test/StatsManagerTest.cpp
index 4f4fe51ac22..92874b3272d 100644
--- a/src/common/stats/test/StatsManagerTest.cpp
+++ b/src/common/stats/test/StatsManagerTest.cpp
@@ -163,6 +163,36 @@ TEST(StatsManager, ReadAllTest) {
   EXPECT_FALSE(counterExists(stats2, "stat04{space=test}.p95.5", val));
 }
 
+TEST(StatsManager, ConcurrentCountWithLabelTest) {
+  auto statId = StatsManager::registerStats("stat06", "RATE, sum");
+  std::vector<std::thread> threads;
+  for (int i = 0; i < 257; i++) {
+    std::string spaceName = folly::stringPrintf("space%d", i);
+    threads.emplace_back([&statId, &spaceName]() {
+      auto counterId = StatsManager::counterWithLabels(statId, {{"space", spaceName}});
+    });
+  }
+
+  for (auto& t : threads) {
+    t.join();
+  }
+}
+
+TEST(StatsManager, ConcurrentHistoWithLabelTest) {
+  auto statId = StatsManager::registerHisto("stat07", 1, 1, 100, "sum, p95, p99");
+  std::vector<std::thread> threads;
+  for (int i = 0; i < 257; i++) {
+    std::string spaceName = folly::stringPrintf("space%d", i);
+    threads.emplace_back([&statId, &spaceName]() {
+      auto histoId = StatsManager::histoWithLabels(statId, {{"space", spaceName}});
+    });
+  }
+
+  for (auto& t : threads) {
+    t.join();
+  }
+}
+
 }  // namespace stats
 }  // namespace nebula