Commit b555e4cd authored by Doron Roberts-Kedes's avatar Doron Roberts-Kedes Committed by Facebook Github Bot

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
parent 44f373be
...@@ -348,6 +348,16 @@ class ConcurrentHashMap { ...@@ -348,6 +348,16 @@ class ConcurrentHashMap {
return res; 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. // NOT noexcept, initializes new shard segments vs.
void clear() { void clear() {
for (uint64_t i = 0; i < NumShards; i++) { for (uint64_t i = 0; i < NumShards; i++) {
......
...@@ -553,10 +553,18 @@ class alignas(64) ConcurrentHashMapSegment { ...@@ -553,10 +553,18 @@ class alignas(64) ConcurrentHashMapSegment {
// Listed separately because we need a prev pointer. // Listed separately because we need a prev pointer.
size_type erase(const key_type& key) { 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}; Node* node{nullptr};
auto h = HashFn()(key); auto h = HashFn()(key);
{ {
...@@ -570,6 +578,9 @@ class alignas(64) ConcurrentHashMapSegment { ...@@ -570,6 +578,9 @@ class alignas(64) ConcurrentHashMapSegment {
Node* prev = nullptr; Node* prev = nullptr;
while (node) { while (node) {
if (KeyEqual()(key, node->getItem().first)) { if (KeyEqual()(key, node->getItem().first)) {
if (!match(node->getItem().second)) {
return 0;
}
auto next = node->next_.load(std::memory_order_relaxed); auto next = node->next_.load(std::memory_order_relaxed);
if (next) { if (next) {
next->acquire_link(); // defined in hazptr_obj_base_linked next->acquire_link(); // defined in hazptr_obj_base_linked
...@@ -613,7 +624,7 @@ class alignas(64) ConcurrentHashMapSegment { ...@@ -613,7 +624,7 @@ class alignas(64) ConcurrentHashMapSegment {
// This is a small departure from standard stl containers: erase may // This is a small departure from standard stl containers: erase may
// throw if hash or key_eq functions throw. // throw if hash or key_eq functions throw.
void erase(Iterator& res, Iterator& pos) { void erase(Iterator& res, Iterator& pos) {
erase_internal(pos->first, &res); erase_internal(pos->first, &res, [](const ValueType&) { return true; });
// Invalidate the iterator. // Invalidate the iterator.
pos = cend(); pos = cend();
} }
......
...@@ -50,6 +50,9 @@ TEST(ConcurrentHashMap, MapTest) { ...@@ -50,6 +50,9 @@ TEST(ConcurrentHashMap, MapTest) {
EXPECT_TRUE(foomap.insert_or_assign(2, 0).second); EXPECT_TRUE(foomap.insert_or_assign(2, 0).second);
EXPECT_TRUE(foomap.assign_if_equal(2, 0, 3)); EXPECT_TRUE(foomap.assign_if_equal(2, 0, 3));
EXPECT_TRUE(foomap.insert(3, 0).second); 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(1), foomap.cend());
EXPECT_NE(foomap.find(2), foomap.cend()); EXPECT_NE(foomap.find(2), foomap.cend());
EXPECT_EQ(foomap.find(2)->second, 3); EXPECT_EQ(foomap.find(2)->second, 3);
...@@ -268,6 +271,16 @@ TEST(ConcurrentHashMap, EraseTest) { ...@@ -268,6 +271,16 @@ TEST(ConcurrentHashMap, EraseTest) {
foomap.erase(f1); 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) { TEST(ConcurrentHashMap, CopyIterator) {
ConcurrentHashMap<int, int> map; ConcurrentHashMap<int, int> map;
map.insert(0, 0); map.insert(0, 0);
...@@ -388,7 +401,11 @@ TEST(ConcurrentHashMap, EraseStressTest) { ...@@ -388,7 +401,11 @@ TEST(ConcurrentHashMap, EraseStressTest) {
unsigned long k = folly::hash::jenkins_rev_mix32((i + offset)); unsigned long k = folly::hash::jenkins_rev_mix32((i + offset));
auto res = m.insert(k, k).second; auto res = m.insert(k, k).second;
if (res) { if (res) {
res = m.erase(k); if (i % 2 == 0) {
res = m.erase(k);
} else {
res = m.erase_if_equal(k, k);
}
if (!res) { if (!res) {
printf("Faulre to erase thread %i val %li\n", t, k); printf("Faulre to erase thread %i val %li\n", t, k);
exit(0); exit(0);
...@@ -450,7 +467,11 @@ TEST(ConcurrentHashMap, IterateStressTest) { ...@@ -450,7 +467,11 @@ TEST(ConcurrentHashMap, IterateStressTest) {
unsigned long k = folly::hash::jenkins_rev_mix32((i + offset)); unsigned long k = folly::hash::jenkins_rev_mix32((i + offset));
auto res = m.insert(k, k).second; auto res = m.insert(k, k).second;
if (res) { if (res) {
res = m.erase(k); if (i % 2 == 0) {
res = m.erase(k);
} else {
res = m.erase_if_equal(k, k);
}
if (!res) { if (!res) {
printf("Faulre to erase thread %i val %li\n", t, k); printf("Faulre to erase thread %i val %li\n", t, k);
exit(0); exit(0);
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment