Use findInternal instead of acquire

cachelib/allocator/CacheAllocator-inl.h

@@ -862,13 +862,13 @@ CacheAllocator<CacheTrait>::releaseBackToAllocator(Item& it,
     headHandle.reset();
 
     if (head == nullptr || &head->getParentItem(compressor_) != &it) {
-      throw std::runtime_error(folly::sformat(
+      throw exception::ChainedItemInvalid(folly::sformat(
           "Mismatch parent pointer. This should not happen. Key: {}",
           it.getKey()));
     }
 
     if (!chainedItemAccessContainer_->remove(*head)) {
-      throw std::runtime_error(folly::sformat(
+      throw exception::ChainedItemInvalid(folly::sformat(
           "Chained item associated with {} cannot be removed from hash table "
           "This should not happen here.",
           it.getKey()));
@@ -2639,6 +2639,7 @@ bool CacheAllocator<CacheTrait>::moveForSlabRelease(
   const auto allocInfo = allocator_->getAllocInfo(oldItem.getMemory());
   if (chainedItem) {
     newItemHdl.reset();
+    auto parentKey = parentItem->getKey();
     auto ref = parentItem->unmarkMoving();
     if (UNLIKELY(ref == 0)) {
       wakeUpWaiters(*parentItem, {});
@@ -2648,7 +2649,7 @@ bool CacheAllocator<CacheTrait>::moveForSlabRelease(
       return true;
     } else {
       XDCHECK_NE(ref, 0);
-      auto parentHdl = acquire(parentItem);
+      auto parentHdl = findInternal(parentKey);
       if (parentHdl) {
         wakeUpWaiters(*parentItem, std::move(parentHdl));
       }

cachelib/allocator/Handle.h

@@ -26,6 +26,7 @@
 #include <mutex>
 
 #include "cachelib/allocator/nvmcache/WaitContext.h"
+#include "cachelib/common/Exceptions.h"
 
 namespace facebook {
 namespace cachelib {
@@ -70,8 +71,10 @@ struct ReadHandleImpl {
     assert(alloc_ != nullptr);
     try {
       alloc_->release(it_, isNascent());
+    } catch (const exception::ChainedItemInvalid& e) {
+      XDCHECK(false) << e.what();
     } catch (const std::exception& e) {
-      XLOGF(CRITICAL, "Failed to release {:#10x} : {}", static_cast<void*>(it_),
+      XLOGF(CRITICAL, "Failed to release {} : {}", static_cast<void*>(it_),
             e.what());
     }
     it_ = nullptr;

cachelib/allocator/tests/AllocatorTypeTest.cpp

@@ -292,6 +292,11 @@ TYPED_TEST(BaseAllocatorTest, TransferChainAfterMoving) {
   this->testTransferChainAfterMoving();
 }
 
+TYPED_TEST(BaseAllocatorTest, ChainedItemParentAcquireAfterMove) {
+  ASSERT_EXIT(this->testChainedItemParentAcquireAfterMoveLoop(),
+              testing::ExitedWithCode(0), ".*");
+}
+
 TYPED_TEST(BaseAllocatorTest, AddAndPopChainedItemMultithread) {
   this->testAddAndPopChainedItemMultithread();
 }

cachelib/allocator/tests/BaseAllocatorTest.h

@@ -19,6 +19,7 @@
 #include <folly/Random.h>
 #include <folly/Singleton.h>
 #include <folly/synchronization/Baton.h>
+#include <folly/synchronization/Latch.h>
 
 #include <algorithm>
 #include <chrono>
@@ -2469,6 +2470,80 @@ class BaseAllocatorTest : public AllocatorTest<AllocatorT> {
     }
   }
 
+  // tests case that correct parent item is acquired after move
+  void testChainedItemParentAcquireAfterMoveLoop() {
+    // create an allocator worth 250 slabs
+    // first slab is for overhead, second is parent class
+    // third is chained item 1 and rest are for new chained item alloc
+    // to move to.
+    std::unique_ptr<AllocatorT> alloc;
+    typename AllocatorT::Config config;
+    config.configureChainedItems();
+    config.setCacheSize(250 * Slab::kSize);
+
+    const std::set<uint32_t> allocSizes = {1024, 2048};
+    auto sizes = std::vector<uint32_t>{500, 1500};
+    std::atomic<uint64_t> numMoves{0};
+    std::atomic<uint64_t> numReplaces{0};
+    PoolId pid;
+
+    using Item = typename AllocatorT::Item;
+    config.enableMovingOnSlabRelease([&](Item& oldItem, Item& newItem,
+                                         Item* parentPtr) {
+      assert(oldItem.getSize() == newItem.getSize());
+      assert(oldItem.isChainedItem());
+      std::memcpy(newItem.getMemory(), oldItem.getMemory(), oldItem.getSize());
+      folly::Latch latch(1);
+      auto insertThread = std::make_unique<std::thread>([&]() {
+        ASSERT_NO_THROW({
+          auto parentReplacement =
+              alloc->allocate(pid, parentPtr->getKey(), sizes[0]);
+          Item* parentCopy = parentPtr;
+          latch.count_down();
+          while (parentCopy->isMoving())
+            ;
+          alloc->insertOrReplace(parentReplacement);
+          ++numReplaces;
+        });
+      });
+      insertThread->detach();
+      latch.wait();
+      ++numMoves;
+    });
+
+    alloc = std::make_unique<AllocatorT>(config);
+
+    const size_t numBytes = alloc->getCacheMemoryStats().ramCacheSize;
+    const auto poolSize = numBytes;
+    pid = alloc->addPool("one", poolSize, allocSizes);
+
+    auto allocFn = [&](std::string keyPrefix, std::vector<uint32_t> sizes) {
+      for (unsigned int loop = 0; loop < 20; ++loop) {
+        for (unsigned int i = 0; i < 2048; ++i) {
+          const auto key = keyPrefix + folly::to<std::string>(loop) + "_" +
+                           folly::to<std::string>(i);
+          auto itemHandle =
+              util::allocateAccessible(*alloc, pid, key, sizes[0]);
+          auto childItem = alloc->allocateChainedItem(itemHandle, sizes[1]);
+          ASSERT_NE(nullptr, childItem);
+
+          alloc->addChainedItem(itemHandle, std::move(childItem));
+        }
+      }
+    };
+    allocFn(std::string{"yolo"}, sizes);
+
+    ClassId cid = static_cast<ClassId>(1);
+    for (int i = 0; i < 20; i++) {
+      alloc->releaseSlab(pid, cid, SlabReleaseMode::kRebalance);
+    }
+    while (alloc->getSlabReleaseStats().numSlabReleaseForRebalance < 20) {
+      sleep(1);
+    }
+    // for ASSERT_EXIT
+    exit(0);
+  }
+
   // create a chain of allocations, replace the allocation and ensure that the
   // order is preserved.
   void testChainedAllocsReplaceInChain() {

cachelib/common/Exceptions.h

@@ -67,6 +67,11 @@ class SlabReleaseAborted : public std::runtime_error {
   using std::runtime_error::runtime_error;
 };
 
+class ChainedItemInvalid : public std::runtime_error {
+ public:
+  using std::runtime_error::runtime_error;
+};
+
 // An allocation error. This could be a genuine std::bad_alloc from
 // the global allocator, or it can be an internal allocation error
 // from the backing cachelib item.