ConcurrentHashMap: Fix cloning of non-copyable items

Summary:
Fix an algorithm bug that occurs when the Key or Value types are not nothrow copy constructible.
- Nodes (the entities protectable by iterators) may be cloned as a result of rehashing.
- Key-value item objects are unique (because they are not copyable).
- Multiple nodes may point to the same key-value item.
- In the incorrect algorithm, each node has a bool member that indicates if it owns the key-value item or not (i.e., if it is responsible for deleting it).
- When a node is cloned, its bool member is cleared.
- The problem is that if the key is erased and the clone (or  descendant clone) may be reclaimed (because it is unprotected by iterators) then the key value item will be reclaimed, even if a cloned-from node is still protected by iterators.
- The failure happens when an iterator holder to a cloned-from node tries to access the already-reclaimed key-value, expecting it (correctly) to be still protected.

The fix:
- Eliminate the bool owned member in the node structure.
- Add a link counter to the key-value item structure.
- Increment (atomically) the link counter when the (latest) node to the key-value item is cloned.
- Decrement (atomically) the link counter when a node pointing to the item is reclaimed.
- Reclaim the item only when its last link is released.
- Note that atomic increments  and decrements are used only if the node pointing to an item is cloned.

Added a test to detect the incorrect behavior.
This change fixes only the regular (non-SIMD) version.
Currently, the SIMD version fails the added test, therefore (for now) the test is enabled only for the regular version.

Reviewed By: davidtgoldblatt

Differential Revision: D30371519

fbshipit-source-id: 29e143712afb3ef794f9f6f4d99919d5b4688416

folly/concurrency/detail/ConcurrentHashMap-detail.h

@@ -49,6 +49,8 @@ template <
     typename KeyType,
     typename ValueType,
     typename Allocator,
+    template <typename>
+    class Atom,
     typename Enabled = void>
 class ValueHolder {
  public:
@@ -71,44 +73,83 @@ class ValueHolder {
 // If the ValueType is not copy constructible, we can instead add
 // an extra indirection.  Adds more allocations / deallocations and
 // pulls in an extra cacheline.
-template <typename KeyType, typename ValueType, typename Allocator>
+template <
+    typename KeyType,
+    typename ValueType,
+    typename Allocator,
+    template <typename>
+    class Atom>
 class ValueHolder<
     KeyType,
     ValueType,
     Allocator,
+    Atom,
     std::enable_if_t<
         !std::is_nothrow_copy_constructible<ValueType>::value ||
         !std::is_nothrow_copy_constructible<KeyType>::value>> {
- public:
   typedef std::pair<const KeyType, ValueType> value_type;
 
+  struct CountedItem {
+    value_type kv_;
+    Atom<uint32_t> numlinks_{1}; // Number of incoming links
+
+    template <typename Arg, typename... Args>
+    CountedItem(std::piecewise_construct_t, Arg&& k, Args&&... args)
+        : kv_(std::piecewise_construct,
+              std::forward_as_tuple(std::forward<Arg>(k)),
+              std::forward_as_tuple(std::forward<Args>(args)...)) {}
+
+    value_type& getItem() { return kv_; }
+
+    void acquireLink() {
+      uint32_t count = numlinks_.fetch_add(1, std::memory_order_release);
+      DCHECK_GE(count, 1);
+    }
+
+    bool releaseLink() {
+      uint32_t count = numlinks_.load(std::memory_order_acquire);
+      DCHECK_GE(count, 1);
+      if (count > 1) {
+        count = numlinks_.fetch_sub(1, std::memory_order_acq_rel);
+      }
+      return count == 1;
+    }
+  }; // CountedItem
+  // Back to ValueHolder specialization
+
+  CountedItem* item_; // Link to unique key-value item.
+
+ public:
   explicit ValueHolder(const ValueHolder& other) {
-    other.owned_ = false;
+    DCHECK(other.item_);
     item_ = other.item_;
+    item_->acquireLink();
   }
 
+  ValueHolder& operator=(const ValueHolder&) = delete;
+
   template <typename Arg, typename... Args>
   ValueHolder(std::piecewise_construct_t, Arg&& k, Args&&... args) {
-    item_ = (value_type*)Allocator().allocate(sizeof(value_type));
-    new (item_) value_type(
+    item_ = (CountedItem*)Allocator().allocate(sizeof(CountedItem));
+    new (item_) CountedItem(
         std::piecewise_construct,
-        std::forward_as_tuple(std::forward<Arg>(k)),
-        std::forward_as_tuple(std::forward<Args>(args)...));
+        std::forward<Arg>(k),
+        std::forward<Args>(args)...);
   }
 
   ~ValueHolder() {
-    if (owned_) {
-      item_->~value_type();
-      Allocator().deallocate((uint8_t*)item_, sizeof(value_type));
+    DCHECK(item_);
+    if (item_->releaseLink()) {
+      item_->~CountedItem();
+      Allocator().deallocate((uint8_t*)item_, sizeof(CountedItem));
     }
   }
 
-  value_type& getItem() { return *item_; }
-
- private:
-  value_type* item_;
-  mutable bool owned_{true};
-};
+  value_type& getItem() {
+    DCHECK(item_);
+    return item_->getItem();
+  }
+}; // ValueHolder specialization
 
 // hazptr deleter that can use an allocator.
 template <typename Allocator>
@@ -186,7 +227,7 @@ class NodeT : public hazptr_obj_base_linked<
     this->acquire_link_safe(); // defined in hazptr_obj_base_linked
   }
 
-  ValueHolder<KeyType, ValueType, Allocator> item_;
+  ValueHolder<KeyType, ValueType, Allocator, Atom> item_;
 };
 
 template <

folly/concurrency/test/ConcurrentHashMapTest.cpp

@@ -1017,6 +1017,32 @@ TYPED_TEST_P(ConcurrentHashMapTest, InsertOrAssignIterator) {
   EXPECT_EQ(itr2->second, 2);
 }
 
+TYPED_TEST_P(ConcurrentHashMapTest, EraseClonedNonCopyable) {
+  // Using a non-copyable value type to use the node structure with an
+  // extra level of indirection to key-value items.
+  using Value = std::unique_ptr<int>;
+  // [TODO] Fix the SIMD version to pass this test, then change the
+  // map type to CHM.
+  ConcurrentHashMap<int, Value> map;
+  int cloned = 32; // The item that will end up being cloned.
+  for (int i = 0; i < cloned; i++) {
+    map.try_emplace(256 * i, std::make_unique<int>(0));
+  }
+  auto [iter, _] = map.try_emplace(256 * cloned, std::make_unique<int>(0));
+  // Add more items to cause rehash that clones the item.
+  int num = 10000;
+  for (int i = cloned + 1; i < num; i++) {
+    map.try_emplace(256 * i, std::make_unique<int>(0));
+  }
+  // Erase items to invoke hazard pointer asynchronous reclamation.
+  for (int i = 0; i < num; i++) {
+    map.erase(256 * i);
+  }
+  // The cloned node and the associated key-value item should still be
+  // protected by iter from being reclaimed.
+  EXPECT_EQ(iter->first, 256 * cloned);
+}
+
 REGISTER_TYPED_TEST_CASE_P(
     ConcurrentHashMapTest,
     MapTest,
@@ -1057,7 +1083,8 @@ REGISTER_TYPED_TEST_CASE_P(
     IteratorLoop,
     HeterogeneousLookup,
     HeterogeneousInsert,
-    InsertOrAssignIterator);
+    InsertOrAssignIterator,
+    EraseClonedNonCopyable);
 
 using folly::detail::concurrenthashmap::bucket::BucketTable;
 

folly/support/gdb.py

@@ -387,7 +387,7 @@ class ConcurrentHashMapPrinter:
     def children(self):
         "Returns an iterator of tuple(name, value)"
         return (
-            (f'[{str(item["first"])}]', item["second"])
+            (f'[{str(item["kv_"]["first"])}]', item["kv_"]["second"])
             for idx, item in enumerate(ConcurrentHashMapIterator(self.val))
         )