CS Patch Part 2 for mulit-tier cachelib:

- transparent item movement
- multi-tier combined locking with exclusive bit (#38)
with refactored incRef to support returning the result of
markMoving (fail if already moving or exclusvie bit is set) option.
- add tests (updated for numa bindings - post combined locking)
for transparent item movement

Author: byrnedj

cachelib/allocator/CacheAllocator-inl.h

@@ -123,6 +123,7 @@ ShmSegmentOpts CacheAllocator<CacheTrait>::createShmCacheOpts(TierId tid) {
 
 template <typename CacheTrait>
 size_t CacheAllocator<CacheTrait>::memoryTierSize(TierId tid) const {
+  auto& memoryTierConfigs = config_.memoryTierConfigs;
   auto partitions = std::accumulate(memoryTierConfigs.begin(), memoryTierConfigs.end(), 0UL,
   [](const size_t i, const MemoryTierCacheConfig& config){
     return i + config.getRatio();
@@ -1231,7 +1232,7 @@ CacheAllocator<CacheTrait>::insertOrReplace(const WriteHandle& handle) {
  * Concurrent threads which are getting handle to the same key:
  *  1. When a handle is created it checks if the moving flag is set
  *  2. If so, Handle implementation creates waitContext and adds it to the
- *     MoveCtx by calling handleWithWaitContextForMovingItem() method.
+ *     MoveCtx by calling tryGetHandleWithWaitContextForMovingItem() method.
  *  3. Wait until the moving thread will complete its job.
  */
 template <typename CacheTrait>
@@ -1399,9 +1400,10 @@ CacheAllocator<CacheTrait>::getNextCandidate(TierId tid,
   Item* toRecycle = nullptr;
   Item* candidate = nullptr;
   auto& mmContainer = getMMContainer(tid, pid, cid);
+  bool lastTier = tid+1 >= getNumTiers();
 
   mmContainer.withEvictionIterator([this, pid, cid, &candidate, &toRecycle,
-                                    &searchTries, &mmContainer,
+                                    &searchTries, &mmContainer, &lastTier,
                                     &token](auto&& itr) {
     if (!itr) {
       ++searchTries;
@@ -1421,16 +1423,21 @@ CacheAllocator<CacheTrait>::getNextCandidate(TierId tid,
               ? &toRecycle_->asChainedItem().getParentItem(compressor_)
               : toRecycle_;
 
-      auto putToken = createPutToken(*candidate_);
+      typename NvmCacheT::PutToken putToken;
+      if (lastTier) {
+        // if it's last tier, the item will be evicted
+        // need to create put token before marking it exclusive
+        putToken = createPutToken(*candidate_);
+      }
 
-      if (shouldWriteToNvmCache(*candidate_) && !putToken.isValid()) {
+      if (lastTier && shouldWriteToNvmCache(*candidate_) && !putToken.isValid()) {
         stats_.evictFailConcurrentFill.inc();
         ++itr;
         continue;
       }
 
-      auto markedForEviction = candidate_->markForEviction();
-      if (!markedForEviction) {
+      auto marked = lastTier ? candidate_->markForEviction() : candidate_->markMoving();
+      if (!marked) {
         if (candidate_->hasChainedItem()) {
           stats_.evictFailParentAC.inc();
         } else {
@@ -1440,8 +1447,10 @@ CacheAllocator<CacheTrait>::getNextCandidate(TierId tid,
         continue;
       }
 
+      XDCHECK(candidate_->isMoving() || candidate_->isMarkedForEviction());
       // markForEviction to make sure no other thead is evicting the item
-      // nor holding a handle to that item
+      // nor holding a handle to that item if this is last tier
+      // since we won't be moving the item to the next tier
       toRecycle = toRecycle_;
       candidate = candidate_;
       token = std::move(putToken);
@@ -1464,13 +1473,44 @@ CacheAllocator<CacheTrait>::getNextCandidate(TierId tid,
 
   XDCHECK(toRecycle);
   XDCHECK(candidate);
-  XDCHECK(candidate->isMarkedForEviction());
+  XDCHECK(candidate->isMoving() || candidate->isMarkedForEviction());
 
-  unlinkItemForEviction(*candidate);
+  auto evictedToNext = lastTier ? nullptr
+      : tryEvictToNextMemoryTier(*candidate, false);
+  if (!evictedToNext) {
+    if (!token.isValid()) {
+      token = createPutToken(*candidate);
+    }
+    // tryEvictToNextMemoryTier should only fail if allocation of the new item fails
+    // in that case, it should be still possible to mark item as exclusive.
+    //
+    // in case that we are on the last tier, we whould have already marked
+    // as exclusive since we will not be moving the item to the next tier
+    // but rather just evicting all together, no need to
+    // markExclusiveWhenMoving
+    auto ret = lastTier ? true : candidate->markForEvictionWhenMoving();
+    XDCHECK(ret);
+
+    unlinkItemForEviction(*candidate);
+    // wake up any readers that wait for the move to complete
+    // it's safe to do now, as we have the item marked exclusive and
+    // no other reader can be added to the waiters list
+    wakeUpWaiters(candidate->getKey(), {});
+
+    if (token.isValid() && shouldWriteToNvmCacheExclusive(*candidate)) {
+      nvmCache_->put(*candidate, std::move(token));
+    }
+  } else {
+    XDCHECK(!evictedToNext->isMarkedForEviction() && !evictedToNext->isMoving());
+    XDCHECK(!candidate->isMarkedForEviction() && !candidate->isMoving());
+    XDCHECK(!candidate->isAccessible());
+    XDCHECK(candidate->getKey() == evictedToNext->getKey());
 
-  if (token.isValid() && shouldWriteToNvmCacheExclusive(*candidate)) {
-    nvmCache_->put(*candidate, std::move(token));
+    wakeUpWaiters(candidate->getKey(), std::move(evictedToNext));
   }
+
+  XDCHECK(!candidate->isMarkedForEviction() && !candidate->isMoving());
+
   return {candidate, toRecycle};
 }
 
@@ -1563,6 +1603,54 @@ bool CacheAllocator<CacheTrait>::shouldWriteToNvmCacheExclusive(
   return true;
 }
 
+template <typename CacheTrait>
+typename CacheAllocator<CacheTrait>::WriteHandle
+CacheAllocator<CacheTrait>::tryEvictToNextMemoryTier(
+    TierId tid, PoolId pid, Item& item, bool fromBgThread) {
+  XDCHECK(item.isMoving());
+  XDCHECK(item.getRefCount() == 0);
+  if(item.hasChainedItem()) return WriteHandle{}; // TODO: We do not support ChainedItem yet
+  if(item.isExpired()) {
+    accessContainer_->remove(item);
+    item.unmarkMoving();
+    return acquire(&item);
+  }
+
+  TierId nextTier = tid; // TODO - calculate this based on some admission policy
+  while (++nextTier < getNumTiers()) { // try to evict down to the next memory tiers
+    // allocateInternal might trigger another eviction
+    auto newItemHdl = allocateInternalTier(nextTier, pid,
+                     item.getKey(),
+                     item.getSize(),
+                     item.getCreationTime(),
+                     item.getExpiryTime(),
+                     fromBgThread);
+
+    if (newItemHdl) {
+      
+      bool moveSuccess = moveRegularItem(item, newItemHdl);
+      if (!moveSuccess) {
+        return WriteHandle{};
+      }
+      XDCHECK_EQ(newItemHdl->getSize(), item.getSize());
+      item.unmarkMoving();
+      return newItemHdl;
+    } else {
+      return WriteHandle{};
+    }
+  }
+
+  return {};
+}
+
+template <typename CacheTrait>
+typename CacheAllocator<CacheTrait>::WriteHandle
+CacheAllocator<CacheTrait>::tryEvictToNextMemoryTier(Item& item, bool fromBgThread) {
+  auto tid = getTierId(item);
+  auto pid = allocator_[tid]->getAllocInfo(item.getMemory()).poolId;
+  return tryEvictToNextMemoryTier(tid, pid, item, fromBgThread);
+}
+
 template <typename CacheTrait>
 typename CacheAllocator<CacheTrait>::RemoveRes
 CacheAllocator<CacheTrait>::remove(typename Item::Key key) {
@@ -2088,8 +2176,7 @@ std::vector<std::string> CacheAllocator<CacheTrait>::dumpEvictionIterator(
 
   std::vector<std::string> content;
 
-  size_t i = 0;
-  while (i < numItems && tid >= 0) {
+  while (tid >= 0) {
     auto& mm = *mmContainers_[tid][pid][cid];
     mm.withEvictionIterator([&content, numItems](auto&& itr) {
       while (itr && content.size() < numItems) {
@@ -2099,7 +2186,6 @@ std::vector<std::string> CacheAllocator<CacheTrait>::dumpEvictionIterator(
     });
     --tid;
   }
-
   return content;
 }
 

cachelib/allocator/CacheAllocator.h

@@ -23,6 +23,8 @@
 #include <folly/fibers/TimedMutex.h>
 #include <folly/logging/xlog.h>
 #include <folly/synchronization/SanitizeThread.h>
+#include <folly/hash/Hash.h>
+#include <folly/container/F14Map.h>
 #include <gtest/gtest.h>
 
 #include <chrono>
@@ -1572,15 +1574,6 @@ class CacheAllocator : public CacheBase {
   //              not exist.
   FOLLY_ALWAYS_INLINE WriteHandle findFastImpl(Key key, AccessMode mode);
 
-  // Moves a regular item to a different memory tier.
-  //
-  // @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.
-  bool moveRegularItemOnEviction(Item& oldItem, WriteHandle& newItemHdl);
-
   // Moves a regular item to a different slab. This should only be used during
   // slab release after the item's exclusive bit has been set. The user supplied
   // callback is responsible for copying the contents and fixing the semantics
@@ -1778,6 +1771,27 @@ class CacheAllocator : public CacheBase {
 
   using EvictionIterator = typename MMContainer::LockedIterator;
 
+  // Try to move the item down to the next memory tier
+  //
+  // @param tid current tier ID of the item
+  // @param pid the pool ID the item belong to.
+  // @param item the item to evict
+  // @param fromBgThread whether this is called from BG thread
+  //
+  // @return valid handle to the item. This will be the last
+  //         handle to the item. On failure an empty handle.
+  WriteHandle tryEvictToNextMemoryTier(TierId tid, PoolId pid, Item& item,
+                                       bool fromBgThread);
+
+  // Try to move the item down to the next memory tier
+  //
+  // @param item the item to evict
+  // @param fromBgThread whether this is called from BG thread
+  //
+  // @return valid handle to the item. This will be the last
+  //         handle to the item. On failure an empty handle. 
+  WriteHandle tryEvictToNextMemoryTier(Item& item, bool fromBgThread);
+
   // Wakes up waiters if there are any
   //
   // @param item    wakes waiters that are waiting on that item

cachelib/allocator/MMLru-inl.h

@@ -229,6 +229,18 @@ void MMLru::Container<T, HookPtr>::withEvictionIterator(F&& fun) {
   }
 }
 
+//template <typename T, MMLru::Hook<T> T::*HookPtr>
+//template <typename F>
+//void
+//MMLru::Container<T, HookPtr>::withPromotionIterator(F&& fun) {
+//  if (config_.useCombinedLockForIterators) {
+//    lruMutex_->lock_combine([this, &fun]() { fun(Iterator{lru_.begin()}); });
+//  } else {
+//    LockHolder lck{*lruMutex_};
+//    fun(Iterator{lru_.begin()});
+//  }
+//}
+
 template <typename T, MMLru::Hook<T> T::*HookPtr>
 template <typename F>
 void MMLru::Container<T, HookPtr>::withContainerLock(F&& fun) {

cachelib/allocator/MMLru.h

@@ -235,7 +235,7 @@ class MMLru {
     std::chrono::seconds mmReconfigureIntervalSecs{};
 
     // Whether to use combined locking for withEvictionIterator.
-    bool useCombinedLockForIterators{false};
+    bool useCombinedLockForIterators{true};
   };
 
   // The container object which can be used to keep track of objects of type

cachelib/allocator/tests/AllocatorMemoryTiersTest.cpp

@@ -26,6 +26,8 @@ using LruAllocatorMemoryTiersTest = AllocatorMemoryTiersTest<LruAllocator>;
 TEST_F(LruAllocatorMemoryTiersTest, MultiTiersInvalid) { this->testMultiTiersInvalid(); }
 TEST_F(LruAllocatorMemoryTiersTest, MultiTiersValid) { this->testMultiTiersValid(); }
 TEST_F(LruAllocatorMemoryTiersTest, MultiTiersValidMixed) { this->testMultiTiersValidMixed(); }
+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

@@ -20,12 +20,42 @@
 #include "cachelib/allocator/MemoryTierCacheConfig.h"
 #include "cachelib/allocator/tests/TestBase.h"
 
+#include <folly/synchronization/Latch.h>
+
 namespace facebook {
 namespace cachelib {
 namespace tests {
 
 template <typename AllocatorT>
 class AllocatorMemoryTiersTest : public AllocatorTest<AllocatorT> {
+ private:
+  template<typename MvCallback>
+  void testMultiTiersAsyncOpDuringMove(std::unique_ptr<AllocatorT>& alloc,
+                                       PoolId& pool, bool& quit, MvCallback&& moveCb) {
+    typename AllocatorT::Config config;
+    config.setCacheSize(4 * Slab::kSize);
+    config.enableCachePersistence("/tmp");
+    config.configureMemoryTiers({
+        MemoryTierCacheConfig::fromShm()
+            .setRatio(1).setMemBind(std::string("0")),
+        MemoryTierCacheConfig::fromShm()
+            .setRatio(1).setMemBind(std::string("0"))
+    });
+
+    config.enableMovingOnSlabRelease(moveCb, {} /* ChainedItemsMoveSync */,
+                                     -1 /* movingAttemptsLimit */);
+
+    alloc = std::make_unique<AllocatorT>(AllocatorT::SharedMemNew, config);
+    ASSERT(alloc != nullptr);
+    pool = alloc->addPool("default", alloc->getCacheMemoryStats().ramCacheSize);
+
+    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));
+    }
+  }
  public:
   void testMultiTiersInvalid() {
     typename AllocatorT::Config config;
@@ -74,6 +104,70 @@ class AllocatorMemoryTiersTest : public AllocatorTest<AllocatorT> {
     ASSERT(handle != nullptr);
     ASSERT_NO_THROW(alloc->insertOrReplace(handle));
   }
+
+  void testMultiTiersRemoveDuringEviction() {
+    std::unique_ptr<AllocatorT> alloc;
+    PoolId pool;
+    std::unique_ptr<std::thread> t;
+    folly::Latch latch(1);
+    bool quit = false;
+
+    auto moveCb = [&] (typename AllocatorT::Item& oldItem,
+                       typename AllocatorT::Item& newItem,
+                       typename AllocatorT::Item* /* parentPtr */) {
+      
+      auto key = oldItem.getKey();
+      t = std::make_unique<std::thread>([&](){
+            // remove() function is blocked by wait context
+            // till item is moved to next tier. So that, we should
+            // notify latch before calling remove()
+            latch.count_down();
+            alloc->remove(key);
+          });
+      // wait till async thread is running
+      latch.wait();
+      memcpy(newItem.getMemory(), oldItem.getMemory(), oldItem.getSize());
+      quit = true;
+    };
+
+    testMultiTiersAsyncOpDuringMove(alloc, pool, quit, moveCb);
+
+    t->join();
+  }
+
+  void testMultiTiersReplaceDuringEviction() {
+    std::unique_ptr<AllocatorT> alloc;
+    PoolId pool;
+    std::unique_ptr<std::thread> t;
+    folly::Latch latch(1);
+    bool quit = false;
+
+    auto moveCb = [&] (typename AllocatorT::Item& oldItem,
+                       typename AllocatorT::Item& newItem,
+                       typename AllocatorT::Item* /* parentPtr */) {
+      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());
+              // insertOrReplace() function is blocked by wait context
+              // till item is moved to next tier. So that, we should
+              // notify latch before calling insertOrReplace()
+              latch.count_down();
+              ASSERT_NO_THROW(alloc->insertOrReplace(handle));
+            });
+        // wait till async thread is running
+        latch.wait();
+      }
+      memcpy(newItem.getMemory(), oldItem.getMemory(), oldItem.getSize());
+    };
+
+    testMultiTiersAsyncOpDuringMove(alloc, pool, quit, moveCb);
+
+    t->join();
+  }
 };
 } // namespace tests
 } // namespace cachelib