Increase Hazptr link counts from 16bit -> 32bit

Summary:
PROBLEM
For linked objects, the ref count and link count are stored in 16 bit integers. It's easily possible to overflow either counts.

SOLUTION
Increase the counts to be stored using 32 bits. While this theoretically increases the footprint of the `hazptr_obj_base_linked` struct from 28 bytes to 32 bytes, the increase is small; and in many cases struct padding would anyways cause the `hazptr_obj_base_linked` to be 32 bytes.

Note: This is a mitigation for facebook/folly#2097 . Theoretically its possible to increase contention more to cause the issue to recur, but in practice that's hard to do. See facebook/folly#2107 for a potential fix to the underlying issue.

Reviewed By: ot

Differential Revision: D51829892

fbshipit-source-id: f3ef8f7cf245dd7ff0e1ba6f6ee7bb15ead532ef

third-party/folly/src/folly/concurrency/test/ConcurrentHashMapTest.cpp

@@ -17,8 +17,11 @@
 #include <folly/concurrency/ConcurrentHashMap.h>
 
 #include <atomic>
+#include <latch>
+#include <limits>
 #include <memory>
 #include <thread>
+#include <vector>
 
 #include <folly/Traits.h>
 #include <folly/container/test/TrackingTypes.h>
@@ -1131,6 +1134,52 @@ TYPED_TEST_P(ConcurrentHashMapTest, ConcurrentInsertClear) {
   }
 }
 
+TYPED_TEST_P(ConcurrentHashMapTest, StressTestReclamation) {
+  // Create a map where we keep reclaiming a lot of objects that are linked to
+  // one node.
+
+  // Ensure all entries are mapped to a single segment.
+  auto constant_hash = [](unsigned long) -> uint64_t { return 0; };
+  CHM<unsigned long, unsigned long, decltype(constant_hash)> map;
+  static constexpr unsigned long key_prev =
+      0; // A key that the test key has a link to - to guard against immediate
+         // reclamation.
+  static constexpr unsigned long key_test =
+      1; // A key that keeps being reclaimed repeatedly.
+  static constexpr unsigned long key_link_explosion =
+      2; // A key that is linked to the test key.
+
+  EXPECT_TRUE(map.insert(std::make_pair(key_prev, 0)).second);
+  EXPECT_TRUE(map.insert(std::make_pair(key_test, 0)).second);
+  EXPECT_TRUE(map.insert(std::make_pair(key_link_explosion, 0)).second);
+
+  std::vector<std::thread> threads;
+  // Test with (2^16)+ threads, enough to overflow a 16 bit integer.
+  // It should be uncommon to have more than 2^32 concurrent accesses.
+  static constexpr uint64_t num_threads = std::numeric_limits<uint16_t>::max();
+  static constexpr uint64_t iters = 100;
+  std::latch start{num_threads};
+  for (uint64_t t = 0; t < num_threads; t++) {
+    threads.push_back(lib::thread([t, &map, &start]() {
+      start.arrive_and_wait();
+      static constexpr uint64_t progress_report_pct =
+          (iters / 20); // Every 5% we log progress
+      for (uint64_t i = 0; i < iters; i++) {
+        if (t == 0 && (i % progress_report_pct) == 0) {
+          // To a casual observer - to know that the test is progressing, even
+          // if slowly
+          LOG(INFO) << "Progress: " << (i * 100 / iters);
+        }
+
+        map.insert_or_assign(key_test, i * num_threads);
+      }
+    }));
+  }
+  for (auto& t : threads) {
+    join;
+  }
+}
+
 REGISTER_TYPED_TEST_SUITE_P(
     ConcurrentHashMapTest,
     MapTest,
@@ -1174,7 +1223,8 @@ REGISTER_TYPED_TEST_SUITE_P(
     HeterogeneousInsert,
     InsertOrAssignIterator,
     EraseClonedNonCopyable,
-    ConcurrentInsertClear);
+    ConcurrentInsertClear,
+    StressTestReclamation);
 
 using folly::detail::concurrenthashmap::bucket::BucketTable;
 

third-party/folly/src/folly/synchronization/HazptrObjLinked.h

@@ -86,11 +86,11 @@ class hazptr_root {
  */
 template <template <typename> class Atom>
 class hazptr_obj_linked : public hazptr_obj<Atom> {
-  using Count = uint32_t;
+  using Count = uint64_t;
 
-  static constexpr Count kRef = 1u;
-  static constexpr Count kLink = 1u << 16;
-  static constexpr Count kRefMask = kLink - 1u;
+  static constexpr Count kRef = Count{1};
+  static constexpr Count kLink = Count{1} << 32;
+  static constexpr Count kRefMask = kLink - Count{1};
   static constexpr Count kLinkMask = ~kRefMask;
 
   Atom<Count> count_{0};