ConcurrentHashMap: Fix a bug in replacing the value of an existing key

Summary: Add missing protection of the new node when replacing an existing node.

Differential Revision: D29271517

fbshipit-source-id: 77812f27c37d4950a6e485db674813fab0cf8772

folly/concurrency/detail/ConcurrentHashMap-detail.h

@@ -672,6 +672,7 @@ class alignas(64) BucketTable {
           }
           prev->store(cur, std::memory_order_release);
           it.setNode(cur, buckets, bcount, idx);
+          haznode.reset_protection(cur);
           g.unlock();
           // Release not under lock.
           node->release();

folly/concurrency/test/ConcurrentHashMapTest.cpp

@@ -540,6 +540,29 @@ TYPED_TEST_P(ConcurrentHashMapTest, TryEmplaceEraseStressTest) {
   }
 }
 
+TYPED_TEST_P(ConcurrentHashMapTest, InsertOrAssignStressTest) {
+  DeterministicSchedule sched(DeterministicSchedule::uniform(FLAGS_seed));
+  std::atomic<int> iterations{10000};
+  std::vector<std::thread> threads;
+  unsigned int num_threads = 32;
+  threads.reserve(num_threads);
+  folly::ConcurrentHashMap<int, int> map;
+  for (uint32_t t = 0; t < num_threads; t++) {
+    threads.push_back(lib::thread([&]() {
+      int i = 0;
+      while (--iterations >= 0) {
+        auto res = map.insert_or_assign(0, ++i);
+        ASSERT_TRUE(res.second);
+        auto v = res.first->second;
+        ASSERT_EQ(v, i);
+      }
+    }));
+  }
+  for (auto& t : threads) {
+    join;
+  }
+}
+
 TYPED_TEST_P(ConcurrentHashMapTest, IterateStressTest) {
   DeterministicSchedule sched(DeterministicSchedule::uniform(FLAGS_seed));
 
@@ -1015,6 +1038,7 @@ REGISTER_TYPED_TEST_CASE_P(
     EraseTest,
     ForEachLoop,
     TryEmplaceEraseStressTest,
+    InsertOrAssignStressTest,
     IterateStressTest,
     RefcountTest,
     UpdateStressTest,