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

source/common/stats/BUILD

@@ -24,6 +24,7 @@ envoy_cc_library(
         "//source/common/common:hash_lib",
         "//source/common/common:non_copyable",
         "//source/common/common:perf_annotation_lib",
+        "//source/common/common:thread_annotations",
         "//source/common/common:utility_lib",
         "//source/common/config:well_known_names",
         "//source/common/protobuf",

source/common/stats/stats_impl.cc

@@ -150,10 +150,26 @@ 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;
+
+  // Because the RawStatData object is initialized with and contains a truncated
+  // version of the std::string name, storing the stats in a map would require
+  // storing the name twice. Performing a lookup on the set is similarly
+  // expensive to performing a map lookup, since both require copying a truncated version of the
+  // string before doing the hash lookup.
+  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 +272,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

@@ -19,6 +19,7 @@
 #include "common/common/assert.h"
 #include "common/common/hash.h"
 #include "common/common/non_copyable.h"
+#include "common/common/thread_annotations.h"
 #include "common/common/utility.h"
 #include "common/protobuf/protobuf.h"
 
@@ -406,14 +407,34 @@ class HistogramImpl : public Histogram, public MetricImpl {
 };
 
 /**
- * Implementation of RawStatDataAllocator that just allocates a new structure in memory and returns
- * it.
+ * Implementation of RawStatDataAllocator that uses an unordered set to store
+ * RawStatData pointers.
  */
 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;
+
+  // An unordered set of RawStatData pointers which keys off the key()
+  // field in each object. This necessitates a custom comparator and hasher.
+  StringRawDataSet stats_ GUARDED_BY(mutex_);
+  // A mutex is needed here to protect the stats_ object from both alloc() and free() operations.
+  // Although alloc() operations are called under existing locking, free() operations are made from
+  // the destructors of the individual stat objects, which are not protected by locks.
+  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