stats: keep a set of ref-counted parent histograms in ThreadLocalStore so that two with the same name map to the same histogram object.

.bazelversion

@@ -1 +1 @@
-3.4.1
+2.2.0

source/common/stats/allocator_impl.h

@@ -58,29 +58,10 @@ class AllocatorImpl : public Allocator {
   friend class TextReadoutImpl;
   friend class NotifyingAllocatorImpl;
 
-  struct HeapStatHash {
-    using is_transparent = void; // NOLINT(readability-identifier-naming)
-    size_t operator()(const Metric* a) const { return a->statName().hash(); }
-    size_t operator()(StatName a) const { return a.hash(); }
-  };
-
-  struct HeapStatCompare {
-    using is_transparent = void; // NOLINT(readability-identifier-naming)
-    bool operator()(const Metric* a, const Metric* b) const {
-      return a->statName() == b->statName();
-    }
-    bool operator()(const Metric* a, StatName b) const { return a->statName() == b; }
-  };
-
   void removeCounterFromSetLockHeld(Counter* counter) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
   void removeGaugeFromSetLockHeld(Gauge* gauge) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
   void removeTextReadoutFromSetLockHeld(Counter* counter) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
 
-  // An unordered set of HeapStatData pointers which keys off the key()
-  // field in each object. This necessitates a custom comparator and hasher, which key off of the
-  // StatNamePtr's own StatNamePtrHash and StatNamePtrCompare operators.
-  template <class StatType>
-  using StatSet = absl::flat_hash_set<StatType*, HeapStatHash, HeapStatCompare>;
   StatSet<Counter> counters_ ABSL_GUARDED_BY(mutex_);
   StatSet<Gauge> gauges_ ABSL_GUARDED_BY(mutex_);
   StatSet<TextReadout> text_readouts_ ABSL_GUARDED_BY(mutex_);

source/common/stats/metric_impl.h

@@ -32,10 +32,35 @@ class MetricHelper {
   void iterateTagStatNames(const Metric::TagStatNameIterFn& fn) const;
   void clear(SymbolTable& symbol_table) { stat_names_.clear(symbol_table); }
 
+  // Hasher for metrics.
+  struct Hash {
+    using is_transparent = void; // NOLINT(readability-identifier-naming)
+    size_t operator()(const Metric* a) const { return a->statName().hash(); }
+    size_t operator()(StatName a) const { return a.hash(); }
+  };
+
+  // Comparator for metrics.
+  struct Compare {
+    using is_transparent = void; // NOLINT(readability-identifier-naming)
+    bool operator()(const Metric* a, const Metric* b) const {
+      return a->statName() == b->statName();
+    }
+    bool operator()(const Metric* a, StatName b) const { return a->statName() == b; }
+  };
+
 private:
   StatNameList stat_names_;
 };
 
+// An unordered set of stat pointers. which keys off Metric::statName().
+// This necessitates a custom comparator and hasher, using the StatNamePtr's
+// own StatNamePtrHash and StatNamePtrCompare operators.
+//
+// This is used by AllocatorImpl for counters, gauges, and text-readouts, and
+// is also used by thread_local_store.h for histograms.
+template <class StatType>
+using StatSet = absl::flat_hash_set<StatType*, MetricHelper::Hash, MetricHelper::Compare>;
+
 /**
  * Partial implementation of the Metric interface on behalf of Counters, Gauges,
  * and Histograms. It leaves symbolTable() unimplemented so that implementations

source/common/stats/thread_local_store.cc

@@ -42,6 +42,7 @@ ThreadLocalStoreImpl::~ThreadLocalStoreImpl() {
   ASSERT(shutting_down_ || !threading_ever_initialized_);
   default_scope_.reset();
   ASSERT(scopes_.empty());
+  ASSERT(histograms_to_cleanup_.empty());
 }
 
 void ThreadLocalStoreImpl::setHistogramSettings(HistogramSettingsConstPtr&& histogram_settings) {
@@ -161,11 +162,6 @@ std::vector<TextReadoutSharedPtr> ThreadLocalStoreImpl::textReadouts() const {
 std::vector<ParentHistogramSharedPtr> ThreadLocalStoreImpl::histograms() const {
   std::vector<ParentHistogramSharedPtr> ret;
   Thread::LockGuard lock(lock_);
-  // TODO(ramaraochavali): As histograms don't share storage, there is a chance of duplicate names
-  // here. We need to create global storage for histograms similar to how we have a central storage
-  // in shared memory for counters/gauges. In the interim, no de-dup is done here. This may result
-  // in histograms with duplicate names, but until shared storage is implemented it's ultimately
-  // less confusing for users who have such configs.
   for (ScopeImpl* scope : scopes_) {
     for (const auto& name_histogram_pair : scope->central_cache_->histograms_) {
       const ParentHistogramSharedPtr& parent_hist = name_histogram_pair.second;
@@ -189,6 +185,10 @@ void ThreadLocalStoreImpl::initializeThreading(Event::Dispatcher& main_thread_di
 void ThreadLocalStoreImpl::shutdownThreading() {
   // This will block both future cache fills as well as cache flushes.
   shutting_down_ = true;
+  Thread::LockGuard lock(hist_mutex_);
+  for (ParentHistogramImpl* histogram : histogram_set_) {
+    histogram->setShuttingDown(true);
+  }
 }
 
 void ThreadLocalStoreImpl::mergeHistograms(PostMergeCb merge_complete_cb) {
@@ -197,12 +197,9 @@ void ThreadLocalStoreImpl::mergeHistograms(PostMergeCb merge_complete_cb) {
     merge_in_progress_ = true;
     tls_->runOnAllThreads(
         [this]() -> void {
-          for (const auto& scope : tls_->getTyped<TlsCache>().scope_cache_) {
-            const TlsCacheEntry& tls_cache_entry = scope.second;
-            for (const auto& name_histogram_pair : tls_cache_entry.histograms_) {
-              const TlsHistogramSharedPtr& tls_hist = name_histogram_pair.second;
-              tls_hist->beginMerge();
-            }
+          for (const auto& id_hist : tls_->getTyped<TlsCache>().histogram_cache_) {
+            const TlsHistogramSharedPtr& tls_hist = id_hist.second;
+            tls_hist->beginMerge();
           }
         },
         [this, merge_complete_cb]() -> void { mergeInternal(merge_complete_cb); });
@@ -257,19 +254,49 @@ void ThreadLocalStoreImpl::releaseScopeCrossThread(ScopeImpl* scope) {
   if (!shutting_down_ && main_thread_dispatcher_) {
     const uint64_t scope_id = scope->scope_id_;
     lock.release();
+
+    // TODO(jmarantz): consider batching all the scope IDs that should be
+    // cleared from TLS caches to reduce bursts of runOnAllThreads on a large
+    // config update. See the pattern below used for histograms.
     main_thread_dispatcher_->post([this, central_cache, scope_id]() {
       sync_.syncPoint(MainDispatcherCleanupSync);
       clearScopeFromCaches(scope_id, central_cache);
     });
   }
 }
 
+void ThreadLocalStoreImpl::releaseHistogramCrossThread(uint64_t histogram_id) {
+  // This can happen from any thread. We post() back to the main thread which will initiate the
+  // cache flush operation.
+  if (!shutting_down_ && main_thread_dispatcher_) {
+    // It's possible that many different histograms will be deleted at the same
+    // time, before the main thread gets a chance to run
+    // clearHistogramsFromCaches. If a new histogram is deleted before that
+    // post runs, we add it to our list of histograms to clear, and there's no
+    // need to issue another post.
+    bool need_post = false;
+    {
+      Thread::LockGuard lock(hist_mutex_);
+      need_post = histograms_to_cleanup_.empty();
+      histograms_to_cleanup_.push_back(histogram_id);
+    }
+    if (need_post) {
+      main_thread_dispatcher_->post([this]() { clearHistogramsFromCaches(); });
+    }
+  }
+}
+
 ThreadLocalStoreImpl::TlsCacheEntry&
 ThreadLocalStoreImpl::TlsCache::insertScope(uint64_t scope_id) {
   return scope_cache_[scope_id];
 }
 
 void ThreadLocalStoreImpl::TlsCache::eraseScope(uint64_t scope_id) { scope_cache_.erase(scope_id); }
+void ThreadLocalStoreImpl::TlsCache::eraseHistograms(const std::vector<uint64_t>& histograms) {
+  for (uint64_t histogram_id : histograms) {
+    histogram_cache_.erase(histogram_id);
+  }
+}
 
 void ThreadLocalStoreImpl::clearScopeFromCaches(uint64_t scope_id,
                                                 CentralCacheEntrySharedPtr central_cache) {
@@ -283,6 +310,27 @@ void ThreadLocalStoreImpl::clearScopeFromCaches(uint64_t scope_id,
   }
 }
 
+void ThreadLocalStoreImpl::clearHistogramsFromCaches() {
+  // If we are shutting down we no longer perform cache flushes as workers may be shutting down
+  // at the same time.
+  if (!shutting_down_) {
+    // Capture all the pending histograms in a local, clearing the list held in
+    // this. Once this occurs, if a new histogram is deleted, a new post will be
+    // required.
+    auto histograms = new std::vector<uint64_t>();
+    {
+      Thread::LockGuard lock(hist_mutex_);
+      histograms->swap(histograms_to_cleanup_);
+    }
+
+    // Perform a cache flush on all threads. The local copy of the
+    // histograms-list is deleted explicitly after all the threads complete.
+    tls_->runOnAllThreads(
+        [this, histograms]() { tls_->getTyped<TlsCache>().eraseHistograms(*histograms); },
+        [histograms]() { delete histograms; });
+  }
+}
+
 ThreadLocalStoreImpl::ScopeImpl::ScopeImpl(ThreadLocalStoreImpl& parent, const std::string& prefix)
     : scope_id_(parent.next_scope_id_++), parent_(parent),
       prefix_(Utility::sanitizeStatsName(prefix), parent.alloc_.symbolTable()),
@@ -566,9 +614,23 @@ Histogram& ThreadLocalStoreImpl::ScopeImpl::histogramFromStatNameWithTags(
                                      [&buckets, this](absl::string_view stat_name) {
                                        buckets = &parent_.histogram_settings_->buckets(stat_name);
                                      });
-    RefcountPtr<ParentHistogramImpl> stat(new ParentHistogramImpl(
-        final_stat_name, unit, parent_, *this, tag_helper.tagExtractedName(),
-        tag_helper.statNameTags(), *buckets));
+
+    RefcountPtr<ParentHistogramImpl> stat;
+    {
+      Thread::LockGuard lock(parent_.hist_mutex_);
+      auto iter = parent_.histogram_set_.find(final_stat_name);
+      if (iter != parent_.histogram_set_.end()) {
+        stat = RefcountPtr<ParentHistogramImpl>(*iter);
+      } else {
+        stat = new ParentHistogramImpl(final_stat_name, unit, parent_,
+                                       tag_helper.tagExtractedName(), tag_helper.statNameTags(),
+                                       *buckets, parent_.next_histogram_id_++);
+        if (!parent_.shutting_down_) {
+          parent_.histogram_set_.insert(stat.get());
+        }
+      }
+    }
+
     central_ref = &central_cache_->histograms_[stat->statName()];
     *central_ref = stat;
   }
@@ -639,34 +701,34 @@ TextReadoutOptConstRef ThreadLocalStoreImpl::ScopeImpl::findTextReadout(StatName
   return findStatLockHeld<TextReadout>(name, central_cache_->text_readouts_);
 }
 
-Histogram& ThreadLocalStoreImpl::ScopeImpl::tlsHistogram(StatName name,
-                                                         ParentHistogramImpl& parent) {
+Histogram& ThreadLocalStoreImpl::tlsHistogram(ParentHistogramImpl& parent, uint64_t id) {
   // tlsHistogram() is generally not called for a histogram that is rejected by
   // the matcher, so no further rejection-checking is needed at this level.
   // TlsHistogram inherits its reject/accept status from ParentHistogram.
 
   // See comments in counterFromStatName() which explains the logic here.
 
-  StatNameHashMap<TlsHistogramSharedPtr>* tls_cache = nullptr;
-  if (!parent_.shutting_down_ && parent_.tls_) {
-    tls_cache = &parent_.tls_->getTyped<TlsCache>().scope_cache_[this->scope_id_].histograms_;
-    auto iter = tls_cache->find(name);
-    if (iter != tls_cache->end()) {
-      return *iter->second;
+  TlsHistogramSharedPtr* tls_histogram = nullptr;
+  if (!shutting_down_ && tls_ != nullptr) {
+    TlsCache& tls_cache = tls_->getTyped<TlsCache>();
+    tls_histogram = &tls_cache.histogram_cache_[id];
+    if (*tls_histogram != nullptr) {
+      return **tls_histogram;
     }
   }
 
-  StatNameTagHelper tag_helper(parent_, name, absl::nullopt);
+  StatNameTagHelper tag_helper(*this, parent.statName(), absl::nullopt);
 
   TlsHistogramSharedPtr hist_tls_ptr(
-      new ThreadLocalHistogramImpl(name, parent.unit(), tag_helper.tagExtractedName(),
+      new ThreadLocalHistogramImpl(parent.statName(), parent.unit(), tag_helper.tagExtractedName(),
                                    tag_helper.statNameTags(), symbolTable()));
 
   parent.addTlsHistogram(hist_tls_ptr);
 
-  if (tls_cache) {
-    tls_cache->insert(std::make_pair(hist_tls_ptr->statName(), hist_tls_ptr));
+  if (tls_histogram != nullptr) {
+    *tls_histogram = hist_tls_ptr;
   }
+
   return *hist_tls_ptr;
 }
 
@@ -682,7 +744,7 @@ ThreadLocalHistogramImpl::ThreadLocalHistogramImpl(StatName name, Histogram::Uni
 }
 
 ThreadLocalHistogramImpl::~ThreadLocalHistogramImpl() {
-  MetricImpl::clear(symbolTable());
+  MetricImpl::clear(symbol_table_);
   hist_free(histograms_[0]);
   hist_free(histograms_[1]);
 }
@@ -699,26 +761,83 @@ void ThreadLocalHistogramImpl::merge(histogram_t* target) {
   hist_clear(*other_histogram);
 }
 
-ParentHistogramImpl::ParentHistogramImpl(StatName name, Histogram::Unit unit, Store& parent,
-                                         TlsScope& tls_scope, StatName tag_extracted_name,
+ParentHistogramImpl::ParentHistogramImpl(StatName name, Histogram::Unit unit,
+                                         ThreadLocalStoreImpl& parent, StatName tag_extracted_name,
                                          const StatNameTagVector& stat_name_tags,
-                                         ConstSupportedBuckets& supported_buckets)
+                                         ConstSupportedBuckets& supported_buckets, uint64_t id)
     : MetricImpl(name, tag_extracted_name, stat_name_tags, parent.symbolTable()), unit_(unit),
-      parent_(parent), tls_scope_(tls_scope), interval_histogram_(hist_alloc()),
-      cumulative_histogram_(hist_alloc()),
+      parent_(parent), interval_histogram_(hist_alloc()), cumulative_histogram_(hist_alloc()),
       interval_statistics_(interval_histogram_, supported_buckets),
-      cumulative_statistics_(cumulative_histogram_, supported_buckets), merged_(false) {}
+      cumulative_statistics_(cumulative_histogram_, supported_buckets), merged_(false), id_(id) {}
 
 ParentHistogramImpl::~ParentHistogramImpl() {
+  parent_.releaseHistogramCrossThread(id_);
+  ASSERT(ref_count_ == 0);
   MetricImpl::clear(parent_.symbolTable());
   hist_free(interval_histogram_);
   hist_free(cumulative_histogram_);
 }
 
+void ParentHistogramImpl::incRefCount() { ++ref_count_; }
+
+bool ParentHistogramImpl::decRefCount() {
+  bool ret;
+  if (shutting_down_) {
+    ret = --ref_count_ == 0;
+  } else {
+    ret = parent_.decHistogramRefCount(*this, ref_count_);
+  }
+  return ret;
+}
+
+bool ThreadLocalStoreImpl::decHistogramRefCount(ParentHistogramImpl& hist,
+                                                std::atomic<uint32_t>& ref_count) {
+  // We must hold the store's histogram lock when decrementing the
+  // refcount. Otherwise another thread may simultaneously try to allocate the
+  // same name'd stat after we decrement it, and we'll wind up with a
+  // dtor/update race. To avoid this we must hold the lock until the stat is
+  // removed from the map.
+  Thread::LockGuard lock(hist_mutex_);
+  ASSERT(ref_count >= 1);
+  if (--ref_count == 0) {
+    if (!shutting_down_) {
+      const size_t count = histogram_set_.erase(hist.statName());
+      ASSERT(shutting_down_ || count == 1);
+    }
+    return true;
+  }
+  return false;
+}
+
+uint32_t ThreadLocalStoreImpl::numHistograms() const {
+  Thread::LockGuard lock(hist_mutex_);
+  return histogram_set_.size();
+}
+
+uint32_t ThreadLocalStoreImpl::numTlsHistogramsForTesting() const {
+  std::atomic<uint32_t> num_tls_histograms = 0;
+  absl::Mutex mutex;
+  bool done = false;
+  tls_->runOnAllThreads(
+      [this, &num_tls_histograms]() {
+        TlsCache& tls_cache = tls_->getTyped<TlsCache>();
+        num_tls_histograms += tls_cache.histogram_cache_.size();
+      },
+      [&mutex, &done]() {
+        absl::MutexLock lock(&mutex);
+        done = true;
+      });
+  absl::MutexLock lock(&mutex);
+  mutex.Await(absl::Condition(&done));
+  return num_tls_histograms;
+}
+
+SymbolTable& ParentHistogramImpl::symbolTable() { return parent_.symbolTable(); }
+
 Histogram::Unit ParentHistogramImpl::unit() const { return unit_; }
 
 void ParentHistogramImpl::recordValue(uint64_t value) {
-  Histogram& tls_histogram = tls_scope_.tlsHistogram(statName(), *this);
+  Histogram& tls_histogram = parent_.tlsHistogram(*this, id_);
   tls_histogram.recordValue(value);
   parent_.deliverHistogramToSinks(*this, value);
 }