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

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;

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;

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