folly ConcurrentHashMap: add erase_if_equal function

Summary: Erase if and only if key k is equal to expected

Reviewed By: magedm

Differential Revision: D13542801

fbshipit-source-id: dd9e3b91a7e3104b18315043d5f81b6194a407eb

folly/concurrency/ConcurrentHashMap.h

@@ -348,6 +348,16 @@ class ConcurrentHashMap {
     return res;
   }
 
+  // Erase if and only if key k is equal to expected
+  size_type erase_if_equal(const key_type& k, const ValueType& expected) {
+    auto segment = pickSegment(k);
+    auto seg = segments_[segment].load(std::memory_order_acquire);
+    if (!seg) {
+      return 0;
+    }
+    return seg->erase_if_equal(k, expected);
+  }
+
   // NOT noexcept, initializes new shard segments vs.
   void clear() {
     for (uint64_t i = 0; i < NumShards; i++) {

folly/concurrency/detail/ConcurrentHashMap-detail.h

@@ -553,10 +553,18 @@ class alignas(64) ConcurrentHashMapSegment {
 
   // Listed separately because we need a prev pointer.
   size_type erase(const key_type& key) {
-    return erase_internal(key, nullptr);
+    return erase_internal(key, nullptr, [](const ValueType&) { return true; });
   }
 
-  size_type erase_internal(const key_type& key, Iterator* iter) {
+  size_type erase_if_equal(const key_type& key, const ValueType& expected) {
+    return erase_internal(key, nullptr, [&expected](const ValueType& v) {
+      return v == expected;
+    });
+  }
+
+  template <typename MatchFunc>
+  size_type
+  erase_internal(const key_type& key, Iterator* iter, MatchFunc match) {
     Node* node{nullptr};
     auto h = HashFn()(key);
     {
@@ -570,6 +578,9 @@ class alignas(64) ConcurrentHashMapSegment {
       Node* prev = nullptr;
       while (node) {
         if (KeyEqual()(key, node->getItem().first)) {
+          if (!match(node->getItem().second)) {
+            return 0;
+          }
           auto next = node->next_.load(std::memory_order_relaxed);
           if (next) {
             next->acquire_link(); // defined in hazptr_obj_base_linked
@@ -613,7 +624,7 @@ class alignas(64) ConcurrentHashMapSegment {
   // This is a small departure from standard stl containers: erase may
   // throw if hash or key_eq functions throw.
   void erase(Iterator& res, Iterator& pos) {
-    erase_internal(pos->first, &res);
+    erase_internal(pos->first, &res, [](const ValueType&) { return true; });
     // Invalidate the iterator.
     pos = cend();
   }

folly/concurrency/test/ConcurrentHashMapTest.cpp

@@ -50,6 +50,9 @@ TEST(ConcurrentHashMap, MapTest) {
   EXPECT_TRUE(foomap.insert_or_assign(2, 0).second);
   EXPECT_TRUE(foomap.assign_if_equal(2, 0, 3));
   EXPECT_TRUE(foomap.insert(3, 0).second);
+  EXPECT_FALSE(foomap.erase_if_equal(3, 1));
+  EXPECT_TRUE(foomap.erase_if_equal(3, 0));
+  EXPECT_TRUE(foomap.insert(3, 0).second);
   EXPECT_NE(foomap.find(1), foomap.cend());
   EXPECT_NE(foomap.find(2), foomap.cend());
   EXPECT_EQ(foomap.find(2)->second, 3);
@@ -268,6 +271,16 @@ TEST(ConcurrentHashMap, EraseTest) {
   foomap.erase(f1);
 }
 
+TEST(ConcurrentHashMap, EraseIfEqualTest) {
+  ConcurrentHashMap<uint64_t, uint64_t> foomap(3);
+  foomap.insert(1, 0);
+  EXPECT_FALSE(foomap.erase_if_equal(1, 1));
+  auto f1 = foomap.find(1);
+  EXPECT_EQ(0, f1->second);
+  EXPECT_TRUE(foomap.erase_if_equal(1, 0));
+  EXPECT_EQ(foomap.find(1), foomap.cend());
+}
+
 TEST(ConcurrentHashMap, CopyIterator) {
   ConcurrentHashMap<int, int> map;
   map.insert(0, 0);
@@ -388,7 +401,11 @@ TEST(ConcurrentHashMap, EraseStressTest) {
         unsigned long k = folly::hash::jenkins_rev_mix32((i + offset));
         auto res = m.insert(k, k).second;
         if (res) {
+          if (i % 2 == 0) {
             res = m.erase(k);
+          } else {
+            res = m.erase_if_equal(k, k);
+          }
           if (!res) {
             printf("Faulre to erase thread %i val %li\n", t, k);
             exit(0);
@@ -450,7 +467,11 @@ TEST(ConcurrentHashMap, IterateStressTest) {
         unsigned long k = folly::hash::jenkins_rev_mix32((i + offset));
         auto res = m.insert(k, k).second;
         if (res) {
+          if (i % 2 == 0) {
             res = m.erase(k);
+          } else {
+            res = m.erase_if_equal(k, k);
+          }
           if (!res) {
             printf("Faulre to erase thread %i val %li\n", t, k);
             exit(0);