Fix erase in Iterate

Summary:
The iterator returned from erase api could skip nodes. The fix is to
initialize the returned iterator with right value.

Reviewed By: djwatson

Differential Revision: D6579707

fbshipit-source-id: a45f47a53e106d22daa9cf57be6c40c4f6a430d9

diff --git a/folly/concurrency/ConcurrentHashMap.h b/folly/concurrency/ConcurrentHashMap.h
index b4d2fc3a4a3e6c895c631b72e0786d34d3fa7109..a2cf218acfb081740f8f925855cde22e30a6066b 100644 (file)
--- a/folly/concurrency/ConcurrentHashMap.h
+++ b/folly/concurrency/ConcurrentHashMap.h
@@ -421,7 +421,9 @@ class ConcurrentHashMap {
     }
 
     ConstIterator(const ConcurrentHashMap* parent, uint64_t segment)
-        : segment_(segment), parent_(parent) {}
+        : it_(parent->ensureSegment(segment)->cbegin()),
+          segment_(segment),
+          parent_(parent) {}
 
    private:
     // cbegin iterator

diff --git a/folly/concurrency/test/ConcurrentHashMapTest.cpp b/folly/concurrency/test/ConcurrentHashMapTest.cpp
index b1b18e42d7a34d2eaf86eb2c062209dd1f2f01e9..c16985cdb14b2c92abf9eb4b9ce2745fd431ae0c 100644 (file)
--- a/folly/concurrency/test/ConcurrentHashMapTest.cpp
+++ b/folly/concurrency/test/ConcurrentHashMapTest.cpp
@@ -257,6 +257,24 @@ 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);
+  }
+  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