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

source/common/stats/BUILD

@@ -17,6 +17,7 @@ envoy_cc_library(
         "//include/envoy/server:options_interface",
         "//include/envoy/stats:stats_interface",
         "//source/common/common:assert_lib",
+        "//source/common/common:block_memory_hash_set_lib",
         "//source/common/common:hash_lib",
         "//source/common/common:perf_annotation_lib",
         "//source/common/common:utility_lib",

source/common/stats/stats_impl.cc

@@ -35,6 +35,15 @@ bool regexStartsWithDot(absl::string_view regex) {
 
 } // namespace
 
+BlockMemoryHashSetOptions blockMemHashOptions(uint64_t max_stats) {
+  BlockMemoryHashSetOptions hash_set_options;
+  hash_set_options.capacity = max_stats;
+
+  // https://stackoverflow.com/questions/3980117/hash-table-why-size-should-be-prime
+  hash_set_options.num_slots = Primes::findPrimeLargerThan(hash_set_options.capacity / 2);
+  return hash_set_options;
+}
+
 size_t RawStatData::size() {
   // Normally the compiler would do this, but because name_ is a flexible-array-length
   // element, the compiler can't. RawStatData is put into an array in HotRestartImpl, so
@@ -149,13 +158,53 @@ bool TagExtractorImpl::extractTag(const std::string& stat_name, std::vector<Tag>
   return false;
 }
 
+RawStatData* BlockRawStatDataAllocator::alloc(const std::string& name) {
+  std::unique_lock<Thread::BasicLockable> lock(stat_lock_);
+
+  absl::string_view key = name;
+  if (key.size() > Stats::RawStatData::maxNameLength()) {
+    key.remove_suffix(key.size() - Stats::RawStatData::maxNameLength());
+  }
+  auto value_created = stats_set_->insert(key);
+  Stats::RawStatData* data = value_created.first;
+  if (data == nullptr) {
+    return nullptr;
+  }
+  // For new entries (value-created.second==true), BlockMemoryHashSet calls Value::initialize()
+  // automatically, but on recycled entries (value-created.second==false) we need to bump the
+  // ref-count.
+  if (!value_created.second) {
+    ++data->ref_count_;
+  }
+  return data;
+}
+
+void BlockRawStatDataAllocator::free(RawStatData& data) {
+  // We must hold the lock since the reference decrement can race with an initialize above.
+  std::unique_lock<Thread::BasicLockable> lock(stat_lock_);
+
+  ASSERT(data.ref_count_ > 0);
+  if (--data.ref_count_ > 0) {
+    return;
+  }
+  bool key_removed = stats_set_->remove(data.key());
+  ASSERT(key_removed);
+  memset(&data, 0, Stats::RawStatData::size());
+}
+
 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;
 }
 
+void HeapRawStatDataAllocator::free(RawStatData& data) {
+  // This allocator does not ever have concurrent access to the raw data.
+  ASSERT(data.ref_count_ == 1);
+  ::free(&data);
+}
+
 TagProducerImpl::TagProducerImpl(const envoy::config::metrics::v2::StatsConfig& config)
     : TagProducerImpl() {
   // To check name conflict.
@@ -255,12 +304,6 @@ TagProducerImpl::addDefaultExtractors(const envoy::config::metrics::v2::StatsCon
   return names;
 }
 
-void HeapRawStatDataAllocator::free(RawStatData& data) {
-  // This allocator does not ever have concurrent access to the raw data.
-  ASSERT(data.ref_count_ == 1);
-  ::free(&data);
-}
-
 void RawStatData::initialize(absl::string_view key) {
   ASSERT(!initialized());
   ASSERT(key.size() <= maxNameLength());

source/common/stats/stats_impl.h

@@ -17,6 +17,7 @@
 #include "envoy/stats/stats.h"
 
 #include "common/common/assert.h"
+#include "common/common/block_memory_hash_set.h"
 #include "common/common/hash.h"
 #include "common/common/utility.h"
 #include "common/protobuf/protobuf.h"
@@ -26,6 +27,10 @@
 namespace Envoy {
 namespace Stats {
 
+typedef BlockMemoryHashSet<Stats::RawStatData> RawStatDataSet;
+
+BlockMemoryHashSetOptions blockMemHashOptions(uint64_t max_stats);
+
 class TagExtractorImpl : public TagExtractor {
 public:
   /**
@@ -370,12 +375,43 @@ class HistogramImpl : public Histogram, public MetricImpl {
 };
 
 /**
- * Implementation of RawStatDataAllocator that just allocates a new structure in memory and returns
- * it.
+ * Implementation of RawStatDataAllocator which accepts a block of (potentially shared) memory.
+ */
+class BlockRawStatDataAllocator : public RawStatDataAllocator {
+public:
+  /**
+   * Creates a RawStatDataAllocator with an underlying MemoryHashMap to track stat names. Optionally
+   * works with shared memory, as in HotRestartImpl. Thread-safe but requires a lock as argument.
+   * @param options initialization parameters for BlockMemoryHashSet
+   * @param init true if the shared-memory should be initialized on construction. See the
+   * BlockMemoryHashSet constructor.
+   * @param memory the memory buffer for the set data.
+   * @param stat_lock a lock for thread safety.
+   * @return BlockRawStatDataAllocator newly constructed BlockRawStatDataAllocator.
+   */
+  BlockRawStatDataAllocator(const BlockMemoryHashSetOptions& options, bool init, uint8_t* memory,
+                            Thread::BasicLockable& stat_lock)
+      : stat_lock_(stat_lock) {
+    stats_set_ = std::make_unique<RawStatDataSet>(options, init, memory);
+  }
+
+  RawStatData* alloc(const std::string& name);
+  void free(RawStatData& data);
+  std::string version() { return stats_set_->version(); }
+
+private:
+  std::unique_ptr<RawStatDataSet> stats_set_;
+  // We must hold the stat lock when attaching to an existing memory segment
+  // because it might be actively written to while we sanityCheck it.
+  Thread::BasicLockable& stat_lock_;
+};
+
+/**
+ * Heap-based implementation of RawStatDataAllocator that just allocates a new structure in memory
+ * and returns it.
  */
 class HeapRawStatDataAllocator : public RawStatDataAllocator {
 public:
-  // RawStatDataAllocator
   RawStatData* alloc(const std::string& name) override;
   void free(RawStatData& data) override;
 };

source/exe/main_common.cc

@@ -53,7 +53,7 @@ MainCommonBase::MainCommonBase(OptionsImpl& options) : options_(options) {
     }
 #endif
     if (restarter_.get() == nullptr) {
-      restarter_.reset(new Server::HotRestartNopImpl());
+      restarter_.reset(new Server::HotRestartNopImpl(options_));
     }
 
     tls_.reset(new ThreadLocal::InstanceImpl);
@@ -69,7 +69,7 @@ MainCommonBase::MainCommonBase(OptionsImpl& options) : options_(options) {
     break;
   }
   case Server::Mode::Validate:
-    restarter_.reset(new Server::HotRestartNopImpl());
+    restarter_.reset(new Server::HotRestartNopImpl(options_));
     Logger::Registry::initialize(options_.logLevel(), options_.logFormat(), restarter_->logLock());
     break;
   }

source/server/BUILD

@@ -128,6 +128,7 @@ envoy_cc_library(
     hdrs = ["hot_restart_nop_impl.h"],
     deps = [
         "//include/envoy/server:hot_restart_interface",
+        "//source/server:hot_restart_lib",
     ],
 )
 

source/server/hot_restart_impl.cc

@@ -17,6 +17,7 @@
 #include "common/common/fmt.h"
 #include "common/common/utility.h"
 #include "common/network/utility.h"
+#include "common/stats/stats_impl.h"
 
 #include "absl/strings/string_view.h"
 
@@ -27,15 +28,6 @@ namespace Server {
 // from working. Operations code can then cope with this and do a full restart.
 const uint64_t SharedMemory::VERSION = 9;
 
-static BlockMemoryHashSetOptions blockMemHashOptions(uint64_t max_stats) {
-  BlockMemoryHashSetOptions hash_set_options;
-  hash_set_options.capacity = max_stats;
-
-  // https://stackoverflow.com/questions/3980117/hash-table-why-size-should-be-prime
-  hash_set_options.num_slots = Primes::findPrimeLargerThan(hash_set_options.capacity / 2);
-  return hash_set_options;
-}
-
 SharedMemory& SharedMemory::initialize(uint32_t stats_set_size, Options& options) {
   Api::OsSysCalls& os_sys_calls = Api::OsSysCallsSingleton::get();
 
@@ -84,8 +76,8 @@ SharedMemory& SharedMemory::initialize(uint32_t stats_set_size, Options& options
   }
 
   // Stats::RawStatData must be naturally aligned for atomics to work properly.
-  RELEASE_ASSERT((reinterpret_cast<uintptr_t>(shmem->stats_set_data_) % alignof(RawStatDataSet)) ==
-                 0);
+  RELEASE_ASSERT(
+      (reinterpret_cast<uintptr_t>(shmem->stats_set_data_) % alignof(Stats::RawStatDataSet)) == 0);
 
   // Here we catch the case where a new Envoy starts up when the current Envoy has not yet fully
   // initialized. The startup logic is quite complicated, and it's not worth trying to handle this
@@ -114,16 +106,14 @@ std::string SharedMemory::version(size_t max_num_stats, size_t max_stat_name_len
 }
 
 HotRestartImpl::HotRestartImpl(Options& options)
-    : options_(options), stats_set_options_(blockMemHashOptions(options.maxStats())),
-      shmem_(SharedMemory::initialize(RawStatDataSet::numBytes(stats_set_options_), options)),
+    : options_(options), stats_set_options_(Stats::blockMemHashOptions(options.maxStats())),
+      shmem_(
+          SharedMemory::initialize(Stats::RawStatDataSet::numBytes(stats_set_options_), options)),
       log_lock_(shmem_.log_lock_), access_log_lock_(shmem_.access_log_lock_),
       stat_lock_(shmem_.stat_lock_), init_lock_(shmem_.init_lock_) {
   {
-    // We must hold the stat lock when attaching to an existing memory segment
-    // because it might be actively written to while we sanityCheck it.
-    std::unique_lock<Thread::BasicLockable> lock(stat_lock_);
-    stats_set_.reset(new RawStatDataSet(stats_set_options_, options.restartEpoch() == 0,
-                                        shmem_.stats_set_data_));
+    stats_allocator_ = std::make_unique<Stats::BlockRawStatDataAllocator>(
+        stats_set_options_, options.restartEpoch() == 0, shmem_.stats_set_data_, stat_lock_);
   }
   my_domain_socket_ = bindDomainSocket(options.restartEpoch());
   child_address_ = createDomainSocketAddress((options.restartEpoch() + 1));
@@ -138,39 +128,6 @@ HotRestartImpl::HotRestartImpl(Options& options)
   RELEASE_ASSERT(rc != -1);
 }
 
-Stats::RawStatData* HotRestartImpl::alloc(const std::string& name) {
-  // Try to find the existing slot in shared memory, otherwise allocate a new one.
-  std::unique_lock<Thread::BasicLockable> lock(stat_lock_);
-  absl::string_view key = name;
-  if (key.size() > Stats::RawStatData::maxNameLength()) {
-    key.remove_suffix(key.size() - Stats::RawStatData::maxNameLength());
-  }
-  auto value_created = stats_set_->insert(key);
-  Stats::RawStatData* data = value_created.first;
-  if (data == nullptr) {
-    return nullptr;
-  }
-  // For new entries (value-created.second==true), BlockMemoryHashSet calls Value::initialize()
-  // automatically, but on recycled entries (value-created.second==false) we need to bump the
-  // ref-count.
-  if (!value_created.second) {
-    ++data->ref_count_;
-  }
-  return data;
-}
-
-void HotRestartImpl::free(Stats::RawStatData& data) {
-  // We must hold the lock since the reference decrement can race with an initialize above.
-  std::unique_lock<Thread::BasicLockable> lock(stat_lock_);
-  ASSERT(data.ref_count_ > 0);
-  if (--data.ref_count_ > 0) {
-    return;
-  }
-  bool key_removed = stats_set_->remove(data.key());
-  ASSERT(key_removed);
-  memset(&data, 0, Stats::RawStatData::size());
-}
-
 int HotRestartImpl::bindDomainSocket(uint64_t id) {
   Api::OsSysCalls& os_sys_calls = Api::OsSysCallsSingleton::get();
   // This actually creates the socket and binds it. We use the socket in datagram mode so we can
@@ -472,23 +429,25 @@ void HotRestartImpl::terminateParent() {
 void HotRestartImpl::shutdown() { socket_event_.reset(); }
 
 std::string HotRestartImpl::version() {
-  return versionHelper(shmem_.maxStats(), Stats::RawStatData::maxNameLength(), *stats_set_);
+  return versionHelper(shmem_.maxStats(), Stats::RawStatData::maxNameLength(),
+                       stats_allocator_->version());
 }
 
 // Called from envoy --hot-restart-version -- needs to instantiate a RawStatDataSet so it
 // can generate the version string.
 std::string HotRestartImpl::hotRestartVersion(size_t max_num_stats, size_t max_stat_name_len) {
-  const BlockMemoryHashSetOptions options = blockMemHashOptions(max_num_stats);
-  const size_t bytes = RawStatDataSet::numBytes(options);
+  const BlockMemoryHashSetOptions options = Stats::blockMemHashOptions(max_num_stats);
+  const size_t bytes = Stats::RawStatDataSet::numBytes(options);
   std::unique_ptr<uint8_t[]> mem_buffer_for_dry_run_(new uint8_t[bytes]);
-  RawStatDataSet stats_set(options, true /* init */, mem_buffer_for_dry_run_.get());
+  Stats::RawStatDataSet stats_set(options, true /* init */, mem_buffer_for_dry_run_.get());
+  auto stats_set_version = stats_set.version();
 
-  return versionHelper(max_num_stats, max_stat_name_len, stats_set);
+  return versionHelper(max_num_stats, max_stat_name_len, stats_set_version);
 }
 
 std::string HotRestartImpl::versionHelper(size_t max_num_stats, size_t max_stat_name_len,
-                                          RawStatDataSet& stats_set) {
-  return SharedMemory::version(max_num_stats, max_stat_name_len) + "." + stats_set.version();
+                                          std::string stats_set_version) {
+  return SharedMemory::version(max_num_stats, max_stat_name_len) + "." + stats_set_version;
 }
 
 } // namespace Server

source/server/hot_restart_impl.h

@@ -18,8 +18,6 @@
 namespace Envoy {
 namespace Server {
 
-typedef BlockMemoryHashSet<Stats::RawStatData> RawStatDataSet;
-
 /**
  * Shared memory segment. This structure is laid directly into shared memory and is used amongst
  * all running envoy processes.
@@ -113,9 +111,7 @@ class ProcessSharedMutex : public Thread::BasicLockable {
 /**
  * Implementation of HotRestart built for Linux.
  */
-class HotRestartImpl : public HotRestart,
-                       public Stats::RawStatDataAllocator,
-                       Logger::Loggable<Logger::Id::main> {
+class HotRestartImpl : public HotRestart, Logger::Loggable<Logger::Id::main> {
 public:
   HotRestartImpl(Options& options);
 
@@ -130,18 +126,14 @@ class HotRestartImpl : public HotRestart,
   std::string version() override;
   Thread::BasicLockable& logLock() override { return log_lock_; }
   Thread::BasicLockable& accessLogLock() override { return access_log_lock_; }
-  Stats::RawStatDataAllocator& statsAllocator() override { return *this; }
+  Stats::RawStatDataAllocator& statsAllocator() override { return *stats_allocator_; }
 
   /**
    * envoy --hot_restart_version doesn't initialize Envoy, but computes the version string
    * based on the configured options.
    */
   static std::string hotRestartVersion(uint64_t max_num_stats, uint64_t max_stat_name_len);
 
-  // RawStatDataAllocator
-  Stats::RawStatData* alloc(const std::string& name) override;
-  void free(Stats::RawStatData& data) override;
-
 private:
   enum class RpcMessageType {
     DrainListenersRequest = 1,
@@ -204,12 +196,12 @@ class HotRestartImpl : public HotRestart,
   RpcBase* receiveRpc(bool block);
   void sendMessage(sockaddr_un& address, RpcBase& rpc);
   static std::string versionHelper(uint64_t max_num_stats, uint64_t max_stat_name_len,
-                                   RawStatDataSet& stats_set);
+                                   std::string stats_set_version);
 
   Options& options_;
   BlockMemoryHashSetOptions stats_set_options_;
   SharedMemory& shmem_;
-  std::unique_ptr<RawStatDataSet> stats_set_;
+  std::unique_ptr<Stats::BlockRawStatDataAllocator> stats_allocator_;
   ProcessSharedMutex log_lock_;
   ProcessSharedMutex access_log_lock_;
   ProcessSharedMutex stat_lock_;