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

@@ -22,6 +22,8 @@
 #include <folly/ScopeGuard.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>
@@ -1319,7 +1321,7 @@ class CacheAllocator : public CacheBase {
 
  private:
   // wrapper around Item's refcount and active handle tracking
-  FOLLY_ALWAYS_INLINE bool incRef(Item& it);
+  FOLLY_ALWAYS_INLINE RefcountWithFlags::incResult incRef(Item& it, bool failIfMoving);
   FOLLY_ALWAYS_INLINE RefcountWithFlags::Value decRef(Item& it);
 
   // drops the refcount and if needed, frees the allocation back to the memory
@@ -1550,7 +1552,7 @@ class CacheAllocator : public CacheBase {
   //
   // @return true  If the move was completed, and the containers were updated
   //               successfully.
-  bool moveRegularItemOnEviction(Item& oldItem, WriteHandle& newItemHdl);
+  void moveRegularItemWithSync(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
@@ -1637,6 +1639,10 @@ class CacheAllocator : public CacheBase {
   //         false  if the item is not in MMContainer
   bool removeFromMMContainer(Item& item);
 
+  using EvictionIterator = typename MMContainer::LockedIterator;
+
+  WriteHandle acquire(EvictionIterator& it) { return acquire(it.get()); }
+
   // Replaces an item in the MMContainer with another item, at the same
   // position.
   //
@@ -1647,6 +1653,8 @@ class CacheAllocator : public CacheBase {
   //               destination item did not exist in the container, or if the
   //               source item already existed.
   bool replaceInMMContainer(Item& oldItem, Item& newItem);
+  bool replaceInMMContainer(Item* oldItem, Item& newItem);
+  bool replaceInMMContainer(EvictionIterator& oldItemIt, Item& newItem);
 
   // Replaces an item in the MMContainer with another item, at the same
   // position. Or, if the two chained items belong to two different MM
@@ -1705,7 +1713,35 @@ class CacheAllocator : public CacheBase {
   // @return An evicted item or nullptr  if there is no suitable candidate.
   Item* findEviction(TierId tid, PoolId pid, ClassId cid);
 
-  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
+  //
+  // @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); 
+
+  // Wakes up waiters if there are any
+  //
+  // @param item    wakes waiters that are waiting on that item
+  // @param handle  handle to pass to the waiters
+  void wakeUpWaiters(Item& item, WriteHandle handle);
+
+  // Unmarks item as moving and wakes up any waiters waiting on that item
+  //
+  // @param item    wakes waiters that are waiting on that item
+  // @param handle  handle to pass to the waiters
+  typename RefcountWithFlags::Value unmarkMovingAndWakeUpWaiters(Item &item, WriteHandle handle);
+
+  // Try to move the item down to the next memory tier
+  //
+  // @param item the item to evict
+  //
+  // @return valid handle to the item. This will be the last
+  //         handle to the item. On failure an empty handle. 
+  WriteHandle tryEvictToNextMemoryTier(Item& item);
 
   // Deserializer CacheAllocatorMetadata and verify the version
   //
@@ -1823,6 +1859,12 @@ class CacheAllocator : public CacheBase {
 
   typename NvmCacheT::PutToken createPutToken(Item& item);
 
+  // Helper function to remove a item if predicates is true.
+  //
+  // @return last handle to the item on success. empty handle on failure.
+  template <typename Fn>
+  WriteHandle removeIf(Item& item, Fn&& predicate);
+
   // Helper function to remove a item if expired.
   //
   // @return true if it item expire and removed successfully.
@@ -2008,6 +2050,87 @@ class CacheAllocator : public CacheBase {
 
   size_t memoryTierSize(TierId tid) const;
 
+  WriteHandle handleWithWaitContextForMovingItem(Item& item);
+
+  size_t wakeUpWaitersLocked(folly::StringPiece key, WriteHandle&& handle);
+
+  class MoveCtx {
+   public:
+    MoveCtx() {}
+
+    ~MoveCtx() {
+      // prevent any further enqueue to waiters
+      // Note: we don't need to hold locks since no one can enqueue
+      // after this point.
+      wakeUpWaiters();
+    }
+
+    // record the item handle. Upon destruction we will wake up the waiters
+    // and pass a clone of the handle to the callBack. By default we pass
+    // a null handle
+    void setItemHandle(WriteHandle _it) { it = std::move(_it); }
+
+    // enqueue a waiter into the waiter list
+    // @param  waiter       WaitContext
+    void addWaiter(std::shared_ptr<WaitContext<ReadHandle>> waiter) {
+      XDCHECK(waiter);
+      waiters.push_back(std::move(waiter));
+    }
+
+    size_t numWaiters() const { return waiters.size(); }
+
+   private:
+    // notify all pending waiters that are waiting for the fetch.
+    void wakeUpWaiters() {
+      bool refcountOverflowed = false;
+      for (auto& w : waiters) {
+        // If refcount overflowed earlier, then we will return miss to
+        // all subsequent waitors.
+        if (refcountOverflowed) {
+          w->set(WriteHandle{});
+          continue;
+        }
+
+        try {
+          w->set(it.clone());
+        } catch (const exception::RefcountOverflow&) {
+          // We'll return a miss to the user's pending read,
+          // so we should enqueue a delete via NvmCache.
+          // TODO: cache.remove(it);
+          refcountOverflowed = true;
+        }
+      }
+    }
+
+    WriteHandle it; // will be set when Context is being filled
+    std::vector<std::shared_ptr<WaitContext<ReadHandle>>> waiters; // list of
+                                                                   // waiters
+  };
+  using MoveMap =
+      folly::F14ValueMap<folly::StringPiece,
+                         std::unique_ptr<MoveCtx>,
+                         folly::HeterogeneousAccessHash<folly::StringPiece>>;
+
+  static size_t getShardForKey(folly::StringPiece key) {
+    return folly::Hash()(key) % kShards;
+  }
+
+  MoveMap& getMoveMapForShard(size_t shard) {
+    return movesMap_[shard].movesMap_;
+  }
+
+  MoveMap& getMoveMap(folly::StringPiece key) {
+    return getMoveMapForShard(getShardForKey(key));
+  }
+
+  std::unique_lock<std::mutex> getMoveLockForShard(size_t shard) {
+    return std::unique_lock<std::mutex>(moveLock_[shard].moveLock_);
+  }
+
+  std::unique_lock<std::mutex> getMoveLock(folly::StringPiece key) {
+    return getMoveLockForShard(getShardForKey(key));
+  }
+
   // Whether the memory allocator for this cache allocator was created on shared
   // memory. The hash table, chained item hash table etc is also created on
   // shared memory except for temporary shared memory mode when they're created
@@ -2100,6 +2223,22 @@ class CacheAllocator : public CacheBase {
   // poolResizer_, poolOptimizer_, memMonitor_, reaper_
   mutable std::mutex workersMutex_;
 
+  static constexpr size_t kShards = 8192; // TODO: need to define right value
+
+  struct MovesMapShard {
+    alignas(folly::hardware_destructive_interference_size) MoveMap movesMap_;
+  };
+
+  struct MoveLock {
+    alignas(folly::hardware_destructive_interference_size) std::mutex moveLock_;
+  };
+
+  // a map of all pending moves
+  std::vector<MovesMapShard> movesMap_;
+
+  // a map of move locks for each shard
+  std::vector<MoveLock> moveLock_;
+
   // time when the ram cache was first created
   const uint32_t cacheCreationTime_{0};
 

cachelib/allocator/CacheAllocator.h

@@ -238,8 +238,8 @@ bool CacheItem<CacheTrait>::markForEvictionWhenMoving() {
 }
 
 template <typename CacheTrait>
-bool CacheItem<CacheTrait>::markMoving() {
-  return ref_.markMoving();
+bool CacheItem<CacheTrait>::markMoving(bool failIfRefNotZero) {
+  return ref_.markMoving(failIfRefNotZero);
 }
 
 template <typename CacheTrait>

cachelib/allocator/CacheItem-inl.h

@@ -309,9 +309,9 @@ class CACHELIB_PACKED_ATTR CacheItem {
   //
   // @return true on success, failure if item is marked as exclusive
   // @throw exception::RefcountOverflow on ref count overflow
-  FOLLY_ALWAYS_INLINE bool incRef() {
+  FOLLY_ALWAYS_INLINE RefcountWithFlags::incResult incRef(bool failIfMoving) {
     try {
-      return ref_.incRef();
+      return ref_.incRef(failIfMoving);
     } catch (exception::RefcountOverflow& e) {
       throw exception::RefcountOverflow(
           folly::sformat("{} item: {}", e.what(), toString()));
@@ -378,7 +378,7 @@ class CACHELIB_PACKED_ATTR CacheItem {
    * Unmarking moving will also return the refcount at the moment of
    * unmarking.
    */
-  bool markMoving();
+  bool markMoving(bool failIfRefNotZero);
   RefcountWithFlags::Value unmarkMoving() noexcept;
   bool isMoving() const noexcept;
   bool isOnlyMoving() const noexcept;

cachelib/allocator/CacheItem.h

@@ -400,6 +400,12 @@ struct ReadHandleImpl {
       }
     }
 
+   protected:
+    friend class ReadHandleImpl;
+    // Method used only by ReadHandleImpl ctor
+    void discard() {
+      it_.store(nullptr, std::memory_order_relaxed);
+    }
    private:
     // we are waiting on Item* to be set to a value. One of the valid values is
     // nullptr. So choose something that we dont expect to indicate a ptr
@@ -479,7 +485,8 @@ struct ReadHandleImpl {
 
   // Handle which has the item already
   FOLLY_ALWAYS_INLINE ReadHandleImpl(Item* it, CacheT& alloc) noexcept
-      : alloc_(&alloc), it_(it) {}
+      : alloc_(&alloc), it_(it) {
+  }
 
   // handle that has a wait context allocated. Used for async handles
   // In this case, the it_ will be filled in asynchronously and mulitple

cachelib/allocator/Handle.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>
 void MMLru::Container<T, HookPtr>::ensureNotInsertionPoint(T& node) noexcept {
   // If we are removing the insertion point node, grow tail before we remove

cachelib/allocator/MMLru-inl.h

@@ -230,12 +230,13 @@ class MMLru {
     //     lruInsertionPointSpec = 2, we insert at a point 1/4th from tail
     uint8_t lruInsertionPointSpec{0};
 
+    // Whether to use combined locking for withEvictionIterator.
+    bool useCombinedLockForIterators{true};
+
     // Minimum interval between reconfigurations. If 0, reconfigure is never
     // called.
     std::chrono::seconds mmReconfigureIntervalSecs{};
 
-    // Whether to use combined locking for withEvictionIterator.
-    bool useCombinedLockForIterators{false};
   };
 
   // The container object which can be used to keep track of objects of type

cachelib/allocator/MMLru.h

@@ -130,30 +130,41 @@ class FOLLY_PACK_ATTR RefcountWithFlags {
   RefcountWithFlags& operator=(const RefcountWithFlags&) = delete;
   RefcountWithFlags(RefcountWithFlags&&) = delete;
   RefcountWithFlags& operator=(RefcountWithFlags&&) = delete;
-
+  enum incResult {
+     incOk,
+     incFailedMoving,
+     incFailedEviction
+   };
   // Bumps up the reference count only if the new count will be strictly less
   // than or equal to the maxCount and the item is not exclusive
   // @return true if refcount is bumped. false otherwise (if item is exclusive)
   // @throw  exception::RefcountOverflow if new count would be greater than
   // maxCount
-  FOLLY_ALWAYS_INLINE bool incRef() {
-    auto predicate = [](const Value curValue) {
-      Value bitMask = getAdminRef<kExclusive>();
-
-      const bool exlusiveBitIsSet = curValue & bitMask;
-      if (UNLIKELY((curValue & kAccessRefMask) == (kAccessRefMask))) {
-        throw exception::RefcountOverflow("Refcount maxed out.");
-      }
-
-      // Check if the item is not marked for eviction
-      return !exlusiveBitIsSet || ((curValue & kAccessRefMask) != 0);
-    };
-
-    auto newValue = [](const Value curValue) {
-      return (curValue + static_cast<Value>(1));
-    };
-
-    return atomicUpdateValue(predicate, newValue);
+  FOLLY_ALWAYS_INLINE incResult incRef(bool failIfMoving) {
+    incResult res = incOk;
+    auto predicate = [failIfMoving, &res](const Value curValue) {
+       Value bitMask = getAdminRef<kExclusive>();
+
+       const bool exlusiveBitIsSet = curValue & bitMask;
+       if (UNLIKELY((curValue & kAccessRefMask) == (kAccessRefMask))) {
+         throw exception::RefcountOverflow("Refcount maxed out.");
+       } else if (exlusiveBitIsSet && (curValue & kAccessRefMask) == 0) {
+         res = incFailedEviction;
+         return false;
+       } else if (exlusiveBitIsSet && failIfMoving) {
+         res = incFailedMoving;
+         return false;
+       }
+       res = incOk;
+       return true;
+     };
+
+     auto newValue = [](const Value curValue) {
+       return (curValue + static_cast<Value>(1));
+     };
+
+     atomicUpdateValue(predicate, newValue);
+     return res;
   }
 
   // Bumps down the reference count
@@ -309,12 +320,14 @@ class FOLLY_PACK_ATTR RefcountWithFlags {
    *
    * Unmarking moving does not depend on `isInMMContainer`
    */
-  bool markMoving() {
-    auto predicate = [](const Value curValue) {
+  bool markMoving(bool failIfRefNotZero) {
+    auto predicate = [failIfRefNotZero](const Value curValue) {
       Value conditionBitMask = getAdminRef<kLinked>();
       const bool flagSet = curValue & conditionBitMask;
       const bool alreadyExclusive = curValue & getAdminRef<kExclusive>();
-
+      if (failIfRefNotZero && (curValue & kAccessRefMask) != 0) {
+        return false;
+      }
       if (!flagSet || alreadyExclusive) {
         return false;
       }

cachelib/allocator/Refcount.h

@@ -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