Fix erase in Iterate (2)

Summary:
Previously D6579707.

Correctly advance to next item if we erase the current element.  Corner cases were slightly off if we were at the end of a hash chain.

Reviewed By: yfeldblum

Differential Revision: D6603518

fbshipit-source-id: acb959e5bcd5da1c3df642b75985d464fdd3b23d

diff --git a/folly/concurrency/ConcurrentHashMap.h b/folly/concurrency/ConcurrentHashMap.h
index b4d2fc3a4a3e6c895c631b72e0786d34d3fa7109..9a3328094312c32f07dda191b162c1b9946b8510 100644 (file)
--- a/folly/concurrency/ConcurrentHashMap.h
+++ b/folly/concurrency/ConcurrentHashMap.h
@@ -326,7 +326,6 @@ class ConcurrentHashMap {
   ConstIterator erase(ConstIterator& pos) {
     auto segment = pickSegment(pos->first);
     ConstIterator res(this, segment);
-    res.next();
     ensureSegment(segment)->erase(res.it_, pos.it_);
     res.next(); // May point to segment end, and need to advance.
     return res;

diff --git a/folly/concurrency/detail/ConcurrentHashMap-detail.h b/folly/concurrency/detail/ConcurrentHashMap-detail.h
index 60688726f5bf355ad159864d97b1cd3fa6ee18c5..51c1375b0fac326a9cd1d1a389e4967676a86754 100644 (file)
--- a/folly/concurrency/detail/ConcurrentHashMap-detail.h
+++ b/folly/concurrency/detail/ConcurrentHashMap-detail.h
@@ -563,6 +563,7 @@ class FOLLY_ALIGNED(64) ConcurrentHashMapSegment {
             iter->hazptrs_[0].reset(buckets);
             iter->setNode(
                 node->next_.load(std::memory_order_acquire), buckets, idx);
+            iter->next();
           }
           size_--;
           break;

diff --git a/folly/concurrency/test/ConcurrentHashMapTest.cpp b/folly/concurrency/test/ConcurrentHashMapTest.cpp
index b1b18e42d7a34d2eaf86eb2c062209dd1f2f01e9..9ac3e0a183d04b14766d12270e7226c576e0580e 100644 (file)
--- a/folly/concurrency/test/ConcurrentHashMapTest.cpp
+++ b/folly/concurrency/test/ConcurrentHashMapTest.cpp
@@ -257,6 +257,25 @@ TEST(ConcurrentHashMap, EraseTest) {
   foomap.erase(f1);
 }
 
+TEST(ConcurrentHashMap, EraseInIterateTest) {
+  ConcurrentHashMap<uint64_t, uint64_t> foomap(3);
+  for (uint64_t k = 0; k < 10; ++k) {
+    foomap.insert(k, k);
+  }
+  EXPECT_EQ(10, foomap.size());
+  for (auto it = foomap.cbegin(); it != foomap.cend();) {
+    if (it->second > 3) {
+      it = foomap.erase(it);
+    } else {
+      ++it;
+    }
+  }
+  EXPECT_EQ(4, foomap.size());
+  for (auto it = foomap.cbegin(); it != foomap.cend(); ++it) {
+    EXPECT_GE(3, it->second);
+  }
+}
+
 // TODO: hazptrs must support DeterministicSchedule
 
 #define Atom std::atomic // DeterministicAtomic