Fix issue with Envoy not reference counting across scopes under not-hot restart

source/common/stats/stats_impl.cc

@@ -150,10 +150,21 @@ bool TagExtractorImpl::extractTag(const std::string& stat_name, std::vector<Tag>
 }
 
 RawStatData* HeapRawStatDataAllocator::alloc(const std::string& name) {
-  // This must be zero-initialized
   RawStatData* data = static_cast<RawStatData*>(::calloc(RawStatData::size(), 1));
   data->initialize(name);
-  return data;
+
+  std::unique_lock<std::mutex> lock(mutex_);
+  auto ret = stats_.insert(data);
+  RawStatData* existing_data = *ret.first;
+  lock.unlock();
+
+  if (!ret.second) {
+    ::free(data);
+    ++existing_data->ref_count_;
+    return existing_data;
+  } else {
+    return data;
+  }
 }
 
 TagProducerImpl::TagProducerImpl(const envoy::config::metrics::v2::StatsConfig& config)
@@ -256,18 +267,26 @@ TagProducerImpl::addDefaultExtractors(const envoy::config::metrics::v2::StatsCon
 }
 
 void HeapRawStatDataAllocator::free(RawStatData& data) {
-  // This allocator does not ever have concurrent access to the raw data.
-  ASSERT(data.ref_count_ == 1);
+  ASSERT(data.ref_count_ > 0);
+  if (--data.ref_count_ > 0) {
+    return;
+  }
+
+  std::unique_lock<std::mutex> lock(mutex_);
+  size_t key_removed = stats_.erase(&data);
+  lock.unlock();
+
+  ASSERT(key_removed == 1);
   ::free(&data);
 }
 
 void RawStatData::initialize(absl::string_view key) {
   ASSERT(!initialized());
-  if (key.size() > maxNameLength()) {
+  if (key.size() > Stats::RawStatData::maxNameLength()) {
     ENVOY_LOG_MISC(
         warn,
         "Statistic '{}' is too long with {} characters, it will be truncated to {} characters", key,
-        key.size(), maxNameLength());
+        key.size(), Stats::RawStatData::maxNameLength());
   }
   ref_count_ = 1;
 

source/common/stats/stats_impl.h

@@ -376,8 +376,22 @@ class HistogramImpl : public Histogram, public MetricImpl {
 class HeapRawStatDataAllocator : public RawStatDataAllocator {
 public:
   // RawStatDataAllocator
+  ~HeapRawStatDataAllocator() { ASSERT(stats_.empty()); }
   RawStatData* alloc(const std::string& name) override;
   void free(RawStatData& data) override;
+
+private:
+  struct RawStatDataHash_ {
+    size_t operator()(const RawStatData* a) const { return HashUtil::xxHash64(a->key()); }
+  };
+  struct RawStatDataCompare_ {
+    bool operator()(const RawStatData* a, const RawStatData* b) const {
+      return (a->key() == b->key());
+    }
+  };
+  typedef std::unordered_set<RawStatData*, RawStatDataHash_, RawStatDataCompare_> StringRawDataSet;
+  StringRawDataSet stats_;
+  std::mutex mutex_;
 };
 
 /**

test/common/stats/stats_impl_test.cc

@@ -478,5 +478,21 @@ TEST(RawStatDataTest, Truncate) {
   alloc.free(*stat);
 }
 
+TEST(RawStatDataTest, HeapAlloc) {
+  HeapRawStatDataAllocator alloc;
+  RawStatData* stat_1 = alloc.alloc("ref_name");
+  ASSERT_NE(stat_1, nullptr);
+  RawStatData* stat_2 = alloc.alloc("ref_name");
+  ASSERT_NE(stat_2, nullptr);
+  RawStatData* stat_3 = alloc.alloc("not_ref_name");
+  ASSERT_NE(stat_3, nullptr);
+  EXPECT_EQ(stat_1, stat_2);
+  EXPECT_NE(stat_1, stat_3);
+  EXPECT_NE(stat_2, stat_3);
+  alloc.free(*stat_1);
+  alloc.free(*stat_2);
+  alloc.free(*stat_3);
+}
+
 } // namespace Stats
 } // namespace Envoy