Fix bug in iteration logic that may return null handle

Summary:
Reported by user: https://fb.workplace.com/groups/cachelibusers/posts/2528004787386648

A couple of months ago we changed our eviction synchronization logic significantly. The outcome is that on lookup, we can return a miss if an item is concurrently marked as for eviction. We have missed updating the iteration logic to take this into account.

This diff fixes that and only returns valid handles.

Reviewed By: SLieve

Differential Revision: D58193077

fbshipit-source-id: 83331a700ca7aeb8006401375219892a139faf13

cachelib/allocator/ChainedHashTable.h

@@ -1169,7 +1169,10 @@ void ChainedHashTable::Container<T, HookPtr, LockT>::getBucketElems(
   ht_.forEachBucketElem(bucket, [this, &handles](T* e) {
     try {
       XDCHECK(e);
-      handles.emplace_back(handleMaker_(e));
+      auto h = handleMaker_(e);
+      if (h) {
+        handles.emplace_back(std::move(h));
+      }
     } catch (const std::exception&) {
       // if we are not able to acquire a handle, skip over them.
     }

cachelib/allocator/tests/AccessTypeTest.h

@@ -155,6 +155,9 @@ struct AccessTypeTest : public SlabAllocatorTestBase {
 
   // try to iterate while calling saveState()
   void testIteratorWithSerialization();
+
+  // test iterator when handler maker returns null
+  void testIteratorMayContainNull();
 };
 
 template <typename AccessType>
@@ -585,6 +588,69 @@ void AccessTypeTest<AccessType>::testRemoveIf() {
     ASSERT_EQ(c.find(node->getKey()).get(), node.get());
   }
 }
+
+template <typename AccessType>
+void AccessTypeTest<AccessType>::testIteratorMayContainNull() {
+  std::string keyToReturnNull{};
+  Container c{Config{}, PtrCompressor{},
+              [&keyToReturnNull](Node* node) -> Node::Handle {
+                if (!node) {
+                  return nullptr;
+                }
+                if (node->getKey() == keyToReturnNull) {
+                  return nullptr;
+                }
+                node->incRef();
+                return typename Node::Handle{node};
+              }};
+
+  // empty container should not have anything when you iterate.
+  ASSERT_EQ(0, iterateAndGetKeys(c).size());
+
+  auto nodes = createSimpleContainer(c);
+
+  std::set<std::string> existingKeys;
+  for (const auto& node : nodes) {
+    existingKeys.insert(node->getKey().str());
+  }
+
+  // all nodes must be present in the absence of any writes.
+  ASSERT_EQ(existingKeys, iterateAndGetKeys(c));
+
+  // Mark one key for the handle maker to return null
+  keyToReturnNull = *existingKeys.begin();
+
+  std::set<std::string> visitedKeys1;
+  std::set<std::string> visitedKeys2;
+  // start two iterators and visit them with interleaving.
+  auto iter1 = c.begin();
+  while (iter1 != c.end() && visitedKeys1.size() < nodes.size() / 2) {
+    visitedKeys1.insert(iter1->getKey().str());
+    ++iter1;
+  }
+
+  auto iter2 = c.begin();
+  while (iter2 != c.end() && visitedKeys2.size() < nodes.size() / 2) {
+    visitedKeys2.insert(iter2->getKey().str());
+    ++iter2;
+  }
+
+  while (iter1 != c.end() || iter2 != c.end()) {
+    auto rand = folly::Random::rand32();
+    if (iter1 != c.end() && rand % 2) {
+      visitedKeys1.insert(iter1->getKey().str());
+      ++iter1;
+    }
+
+    if (iter2 != c.end() && rand % 3) {
+      visitedKeys2.insert(iter2->getKey().str());
+      ++iter2;
+    }
+  }
+  ASSERT_EQ(visitedKeys1, visitedKeys2);
+  existingKeys.erase(keyToReturnNull);
+  ASSERT_EQ(existingKeys, visitedKeys1);
+}
 } // namespace tests
 } // namespace cachelib
 } // namespace facebook

cachelib/allocator/tests/ChainedHashTest.cpp

@@ -53,6 +53,10 @@ TEST_F(ChainedHashTest, HandleIteration) {
 
 TEST_F(ChainedHashTest, RemoveIf) { testRemoveIf(); }
 
+TEST_F(ChainedHashTest, testIteratorMayContainNull) {
+  testIteratorMayContainNull();
+}
+
 TEST_F(ChainedHashTest, Chaining) {
   using HashConfig = ChainedHashTable::Config;
   const unsigned int bucketsPower = 5;