simplify cache local/lifetime coordination in SingletonRelaxedCounter

Summary:
Rather than giving up on constraining the cache thread_local local and cache thread_local lifetime to be in the same DSO, now constrain it and assert on the constraint. This allows the cache thread_local local and the cache thread_local lifetime to be 1:1 with each other, rather than M:N with each other without the constraint.

Note that there is now no cross-DSO merging of these thread_local variables - they are now explicitly per-DSO.

Differential Revision: D59144274

fbshipit-source-id: 73b3a426059321bb9abfcb92fdbdd3547656adc8

folly/concurrency/BUCK

@@ -165,6 +165,7 @@ cpp_library(
         "//folly:utility",
         "//folly/detail:static_singleton_manager",
         "//folly/detail:thread_local_globals",
+        "//folly/lang:safe_assert",
         "//folly/synchronization:atomic_ref",
     ],
 )

folly/concurrency/SingletonRelaxedCounter.h

@@ -27,6 +27,7 @@
 #include <folly/Utility.h>
 #include <folly/detail/StaticSingletonManager.h>
 #include <folly/detail/thread_local_globals.h>
+#include <folly/lang/SafeAssert.h>
 #include <folly/synchronization/AtomicRef.h>
 
 namespace folly {
@@ -67,9 +68,7 @@ class SingletonRelaxedCounterBase {
   //  a global fallback counter.
   struct Global {
     struct Tracking {
-      using CounterSet = std::unordered_set<CounterAndCache*>;
-      std::unordered_map<CounterAndCache*, size_t> locals; // for summing
-      std::unordered_map<LocalLifetime*, CounterSet> lifetimes;
+      std::unordered_map<LocalLifetime*, CounterAndCache*> lifetimes;
     };
 
     Counter fallback; // used instead of local during thread destruction
@@ -89,9 +88,10 @@ class SingletonRelaxedCounterBase {
   //  LocalLifetime
   //
   //  Manages local().cache, global().tracking, and moving outstanding counts
-  //  from local().counter to global().counter during thread destruction.
+  //  from local().counter to global().counter during thread destruction and dso
+  //  unload.
   //
-  //  The counter-set is within Global to reduce per-thread overhead for threads
+  //  The index map is within Global to reduce per-thread overhead for threads
   //  which do not participate in counter mutations, rather than being a member
   //  field of LocalLifetime. This comes at the cost of the slow path always
   //  acquiring a unique lock on the global mutex.
@@ -101,20 +101,14 @@ class SingletonRelaxedCounterBase {
     }
     FOLLY_NOINLINE void destroy_(GetGlobal& get_global) {
       auto& global = get_global();
-      auto global_fallback = atomic_ref(global.fallback);
       auto const tracking = global.tracking.wlock();
-      auto& lifetimes = tracking->lifetimes[this];
-      for (auto ctr : lifetimes) {
-        auto const it = tracking->locals.find(ctr);
-        if (!--it->second) {
-          tracking->locals.erase(it);
-          auto ctr_counter = atomic_ref(ctr->counter);
-          auto const current = ctr_counter.load(std::memory_order_relaxed);
-          global_fallback.fetch_add(current, std::memory_order_relaxed);
-          ctr_counter.store(Signed(0), std::memory_order_relaxed);
-          ctr->cache = nullptr;
-        }
-      }
+      auto& entry = tracking->lifetimes[this];
+      FOLLY_SAFE_CHECK(entry);
+      auto entry_counter = atomic_ref(entry->counter);
+      auto const current = entry_counter.load(std::memory_order_relaxed);
+      atomic_ref(global.fallback).fetch_add(current, std::memory_order_relaxed);
+      entry_counter.store(Signed(0), std::memory_order_relaxed);
+      entry->cache = nullptr;
       tracking->lifetimes.erase(this);
     }
 
@@ -124,8 +118,9 @@ class SingletonRelaxedCounterBase {
     FOLLY_NOINLINE void track_(Global& global, CounterAndCache& state) {
       state.cache = &state.counter;
       auto const tracking = global.tracking.wlock();
-      auto const inserted = tracking->lifetimes[this].insert(&state);
-      tracking->locals[&state] += inserted.second;
+      auto& entry = tracking->lifetimes[this];
+      FOLLY_SAFE_CHECK(!entry || &state == entry);
+      entry = &state;
     }
   };
 
@@ -136,8 +131,9 @@ class SingletonRelaxedCounterBase {
     auto& global = get_global();
     auto count = atomic_ref(global.fallback).load(std::memory_order_relaxed);
     auto const tracking = global.tracking.rlock();
-    for (auto const& kvp : tracking->locals) {
-      count += atomic_ref(kvp.first->counter).load(std::memory_order_relaxed);
+    for (auto const& [_, entry] : tracking->lifetimes) {
+      FOLLY_SAFE_CHECK(entry);
+      count += atomic_ref(entry->counter).load(std::memory_order_relaxed);
     }
     return std::is_unsigned<Int>::value
         ? to_unsigned(std::max(Signed(0), count))
@@ -239,7 +235,15 @@ class SingletonRelaxedCounter
     ~MonoLocalLifetime() noexcept(false) { LocalLifetime::destroy(global); }
   };
 
-  FOLLY_NOINLINE static void mutate_slow(Signed v) noexcept {
+  //  It is an invariant that in a single call to mutate which calls mutate_slow
+  //  the thread_local local and the thread_local lifetime that are used are in
+  //  the same DSO as each other.
+  //
+  //  The following functions are all [[gnu::visibility("hidden")]] in order to
+  //  ensure this invariant.
+
+  FOLLY_ERASE_NOINLINE static void mutate_slow(Signed v) noexcept {
+    static constexpr Arg arg{global, local, lifetime};
     mutate_slow(v, arg);
   }
 
@@ -251,25 +255,19 @@ class SingletonRelaxedCounter
 
   static constexpr GetGlobal& global = folly::detail::createGlobal<Global, Tag>;
 
-  FOLLY_EXPORT FOLLY_ALWAYS_INLINE static CounterAndCache& local() {
+  FOLLY_ERASE static CounterAndCache& local() {
     //  this is a member function local instead of a class member because of
     //  https://gcc.gnu.org/bugzilla/show_bug.cgi?id=66944
     static thread_local CounterAndCache instance;
     return instance;
   }
 
-  FOLLY_EXPORT FOLLY_ALWAYS_INLINE static LocalLifetime& lifetime() {
+  FOLLY_ERASE static LocalLifetime& lifetime() {
     static thread_local MonoLocalLifetime lifetime;
     return lifetime;
   }
-
-  static constexpr Arg arg{global, local, lifetime};
 };
 
-template <typename Int, typename Tag>
-constexpr typename SingletonRelaxedCounter<Int, Tag>::Arg
-    SingletonRelaxedCounter<Int, Tag>::arg;
-
 template <typename Counted>
 class SingletonRelaxedCountableAccess;