Fix moveRegularItemWithSync and add tests

Author: vinser52

cachelib/allocator/CacheAllocator-inl.h

@@ -1308,15 +1308,15 @@ CacheAllocator<CacheTrait>::moveRegularItemWithSync(
   // make sure that no other thread removed it, and only then replaces it.
   if (!replaceInMMContainer(oldItem, *newItemHdl)) {
     accessContainer_->remove(*newItemHdl);
-    return {};
+    return acquire(&oldItem);
   }
 
   // Replacing into the MM container was successful, but someone could have
   // called insertOrReplace() or remove() before or after the
   // replaceInMMContainer() operation, which would invalidate newItemHdl.
   if (!newItemHdl->isAccessible()) {
     removeFromMMContainer(*newItemHdl);
-    return {};
+    return acquire(&oldItem);
   }
 
   // no one can add or remove chained items at this point

cachelib/allocator/CacheAllocator.h

@@ -1496,8 +1496,9 @@ class CacheAllocator : public CacheBase {
   // @param oldItem     Reference to the item being moved
   // @param newItemHdl  Reference to the handle of the new item being moved into
   //
-  // @return true  If the move was completed, and the containers were updated
-  //               successfully.
+  // @return            the handle to the oldItem if the move was completed
+  //                    and the oldItem can be recycled.
+  //                    Otherwise an empty handle is returned.
   template <typename P>
   WriteHandle moveRegularItemWithSync(Item& oldItem, WriteHandle& newItemHdl, P&& predicate);
 

cachelib/allocator/tests/AllocatorMemoryTiersTest.cpp

@@ -28,6 +28,8 @@ TEST_F(LruAllocatorMemoryTiersTest, MultiTiersFromFileValid) { this->testMultiTi
 TEST_F(LruAllocatorMemoryTiersTest, MultiTiersValidMixed) { this->testMultiTiersValidMixed(); }
 TEST_F(LruAllocatorMemoryTiersTest, MultiTiersNumaBindingsSysVValid) { this->testMultiTiersNumaBindingsSysVValid(); }
 TEST_F(LruAllocatorMemoryTiersTest, MultiTiersNumaBindingsPosixValid) { this->testMultiTiersNumaBindingsPosixValid(); }
+TEST_F(LruAllocatorMemoryTiersTest, MultiTiersRemoveDuringEviction) { this->testMultiTiersRemoveDuringEviction(); }
+TEST_F(LruAllocatorMemoryTiersTest, MultiTiersReplaceDuringEviction) { this->testMultiTiersReplaceDuringEviction(); }
 
 } // end of namespace tests
 } // end of namespace cachelib

cachelib/allocator/tests/AllocatorMemoryTiersTest.h

@@ -124,6 +124,98 @@ class AllocatorMemoryTiersTest : public AllocatorTest<AllocatorT> {
     ASSERT(handle != nullptr);
     ASSERT_NO_THROW(alloc->insertOrReplace(handle));
   }
+
+  void testMultiTiersRemoveDuringEviction() {
+    bool quit = false;
+    typename AllocatorT::Config config;
+    config.setCacheSize(4 * Slab::kSize);
+    config.enableCachePersistence("/tmp");
+    config.configureMemoryTiers({
+        MemoryTierCacheConfig::fromShm()
+            .setRatio(1).setMemBind({0}),
+        MemoryTierCacheConfig::fromShm()
+            .setRatio(1).setMemBind({0})
+    });
+
+    std::unique_ptr<AllocatorT> alloc;
+    std::unique_ptr<std::thread> t;
+
+    const auto moveCb = [&] (typename AllocatorT::Item& oldItem,
+                             typename AllocatorT::Item& newItem,
+                             typename AllocatorT::Item* /* parentPtr */) {
+      memcpy(newItem.getMemory(), oldItem.getMemory(), oldItem.getSize());
+      auto key = oldItem.getKey();
+      t = std::make_unique<std::thread>([&](){ alloc->remove(key); });
+      // sleep to make sure async thread calls remove(key)
+      std::this_thread::sleep_for(std::chrono::milliseconds(100));
+      quit = true;
+    };
+    config.enableMovingOnSlabRelease(moveCb, {} /* ChainedItemsMoveSync */,
+                                     -1 /* movingAttemptsLimit */);
+
+    alloc = std::make_unique<AllocatorT>(AllocatorT::SharedMemNew, config);
+    ASSERT(alloc != nullptr);
+    auto pool = alloc->addPool("default", alloc->getCacheMemoryStats().cacheSize);
+
+    int i = 0;
+    while(!quit) {
+      auto handle = alloc->allocate(pool, std::to_string(++i), std::string("value").size());
+      ASSERT(handle != nullptr);
+      ASSERT_NO_THROW(alloc->insertOrReplace(handle));
+    }
+
+    t->join();
+  }
+
+  void testMultiTiersReplaceDuringEviction() {
+    bool quit = false;
+    typename AllocatorT::Config config;
+    config.setCacheSize(4 * Slab::kSize);
+    config.enableCachePersistence("/tmp");
+    config.configureMemoryTiers({
+        MemoryTierCacheConfig::fromShm()
+            .setRatio(1).setMemBind({0}),
+        MemoryTierCacheConfig::fromShm()
+            .setRatio(1).setMemBind({0})
+    });
+
+    std::unique_ptr<AllocatorT> alloc;
+    PoolId pool;
+    std::unique_ptr<std::thread> t;
+
+    const auto moveCb = [&] (typename AllocatorT::Item& oldItem,
+                             typename AllocatorT::Item& newItem,
+                             typename AllocatorT::Item* /* parentPtr */) {
+      memcpy(newItem.getMemory(), oldItem.getMemory(), oldItem.getSize());
+      auto key = oldItem.getKey();
+      if(!quit) {
+        // we need to replace only once because subsequent allocate calls
+        // will cause evictions recursevly
+        quit = true;
+        t = std::make_unique<std::thread>([&](){
+              auto handle = alloc->allocate(pool, key, std::string("new value").size());
+              ASSERT_NO_THROW(alloc->insertOrReplace(handle));
+            });
+        // sleep to make sure async thread calls remove(key)
+        std::this_thread::sleep_for(std::chrono::milliseconds(100));
+      }
+    };
+    config.enableMovingOnSlabRelease(moveCb, {} /* ChainedItemsMoveSync */,
+                                     -1 /* movingAttemptsLimit */);
+
+    alloc = std::make_unique<AllocatorT>(AllocatorT::SharedMemNew, config);
+    ASSERT(alloc != nullptr);
+    pool = alloc->addPool("default", alloc->getCacheMemoryStats().cacheSize);
+
+    int i = 0;
+    while(!quit) {
+      auto handle = alloc->allocate(pool, std::to_string(++i), std::string("value").size());
+      ASSERT(handle != nullptr);
+      ASSERT_NO_THROW(alloc->insertOrReplace(handle));
+    }
+
+    t->join();
+  }
 };
 } // namespace tests
 } // namespace cachelib