diff --git a/cachelib/allocator/ChainedHashTable.h b/cachelib/allocator/ChainedHashTable.h index a075497f9..522f8e210 100644 --- a/cachelib/allocator/ChainedHashTable.h +++ b/cachelib/allocator/ChainedHashTable.h @@ -1169,7 +1169,10 @@ void ChainedHashTable::Container::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. } diff --git a/cachelib/allocator/tests/AccessTypeTest.h b/cachelib/allocator/tests/AccessTypeTest.h index f21d8b131..76304fc09 100644 --- a/cachelib/allocator/tests/AccessTypeTest.h +++ b/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 @@ -585,6 +588,69 @@ void AccessTypeTest::testRemoveIf() { ASSERT_EQ(c.find(node->getKey()).get(), node.get()); } } + +template +void AccessTypeTest::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 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 visitedKeys1; + std::set 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 diff --git a/cachelib/allocator/tests/ChainedHashTest.cpp b/cachelib/allocator/tests/ChainedHashTest.cpp index 6ee15b131..d9557aeee 100644 --- a/cachelib/allocator/tests/ChainedHashTest.cpp +++ b/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;