Implemented async Item movement between tiers

Author: vinser52

cachelib/allocator/CacheAllocator-inl.h

@@ -1101,12 +1101,42 @@ CacheAllocator<CacheTrait>::insertOrReplace(const ItemHandle& handle) {
   return replaced;
 }
 
+/* Next two methods are used to asynchronously move Item between memory tiers.
+ *
+ * The thread, which moves Item allocate, new Item in the tier we are moving to
+ * and calls moveRegularItemOnEviction() method. This method does the following:
+ *  1. Create MoveCtx and put it to the movesMap.
+ *  2. Updates the access container with the new item from the tier we are
+ *     moving to. This Item has kNotReady flag set.
+ *  3. Copy data from the old Item to the new one.
+ *  4. Unset thekNotReady flag and Notify MoveCtx
+ * 
+ * Concurrent threads which are getting handle to the same key:
+ *  1. When a handle is created it checks if the kNotReady flag is set
+ *  2. If so, Handle implementation creates waitContext and adds it to the
+ *     MoveCtx by calling addWaitContextForMovingItem() method.
+ *  3. Wait until the moving thread will complete its job.
+ */
+template <typename CacheTrait>
+bool CacheAllocator<CacheTrait>::addWaitContextForMovingItem(
+    folly::StringPiece key, std::shared_ptr<WaitContext<ItemHandle>> waiter) {
+  auto shard = getShardForKey(key);
+  auto& movesMap = getMoveMapForShard(shard);
+  auto lock = getMoveLockForShard(shard);
+  auto it = movesMap.find(key);
+  if (it == movesMap.end()) {
+    return false;
+  }
+  auto ctx = it->second.get();
+  ctx->addWaiter(std::move(waiter));
+  return true;
+}
+
 template <typename CacheTrait>
 bool CacheAllocator<CacheTrait>::moveRegularItemOnEviction(
-    Item& oldItem,
-    ItemHandle& newItemHdl) {
-  // TODO: should we introduce new latency tracker. E.g. evictRegularLatency_ ???
-  // util::LatencyTracker tracker{stats_.evictRegularLatency_};
+    Item& oldItem, ItemHandle& newItemHdl) {
+  // TODO: should we introduce new latency tracker. E.g. evictRegularLatency_
+  // ??? util::LatencyTracker tracker{stats_.evictRegularLatency_};
 
   if (!oldItem.isAccessible() || oldItem.isExpired()) {
     return false;
@@ -1122,22 +1152,36 @@ bool CacheAllocator<CacheTrait>::moveRegularItemOnEviction(
     newItemHdl->markNvmClean();
   }
 
-  if(config_.moveCb) {
-    // Execute the move callback. We cannot make any guarantees about the
-    // consistency of the old item beyond this point, because the callback can
-    // do more than a simple memcpy() e.g. update external references. If there
-    // are any remaining handles to the old item, it is the caller's
-    // responsibility to invalidate them. The move can only fail after this
-    // statement if the old item has been removed or replaced, in which case it
-    // should be fine for it to be left in an inconsistent state.
-    config_.moveCb(oldItem, *newItemHdl, nullptr);
-  } else {
-    std::memcpy(newItemHdl->getWritableMemory(), oldItem.getMemory(), oldItem.getSize());
+  folly::StringPiece key(oldItem.getKey());
+  auto shard = getShardForKey(key);
+  auto& movesMap = getMoveMapForShard(shard);
+  MoveCtx* ctx(nullptr);
+  {
+    auto lock = getMoveLockForShard(shard);
+    auto res = movesMap.try_emplace(key, std::make_unique<MoveCtx>());
+    if (!res.second) {
+      return false;
+    }
+    ctx = res.first->second.get();
   }
 
-  // TODO: Possible data race. We copied Item's memory to the newItemHdl
-  // but have not updated accessContainer yet. Concurrent threads might get handle
-  // to the old Item. 
+  auto resHdl = ItemHandle{};
+  auto guard = folly::makeGuard([key, this, ctx, shard, &resHdl]() {
+    auto& movesMap = getMoveMapForShard(shard);
+    resHdl->unmarkNotReady();
+    auto lock = getMoveLockForShard(shard);
+    ctx->setItemHandle(std::move(resHdl));
+    movesMap.erase(key);
+  });
+
+  // TODO: Possibly we can use markMoving() instead. But today
+  // moveOnSlabRelease logic assume that we mark as moving old Item
+  // and than do copy and replace old Item with the new one in access
+  // container. Furthermore, Item can be marked as Moving only
+  // if it is linked to MM container. In our case we mark the new Item
+  // and update access container before the new Item is ready (content is
+  // copied).
+  newItemHdl->markNotReady();
 
   // Inside the access container's lock, this checks if the old item is
   // accessible and its refcount is zero. If the item is not accessible,
@@ -1147,10 +1191,25 @@ bool CacheAllocator<CacheTrait>::moveRegularItemOnEviction(
   // this item through an API such as remove(ItemHandle). In this case,
   // it is unsafe to replace the old item with a new one, so we should
   // also abort.
-  if (!accessContainer_->replaceIf(oldItem, *newItemHdl, itemEvictionPredicate)) {
+  if (!accessContainer_->replaceIf(oldItem, *newItemHdl,
+                                   itemEvictionPredicate)) {
     return false;
   }
 
+  if (config_.moveCb) {
+    // Execute the move callback. We cannot make any guarantees about the
+    // consistency of the old item beyond this point, because the callback can
+    // do more than a simple memcpy() e.g. update external references. If there
+    // are any remaining handles to the old item, it is the caller's
+    // responsibility to invalidate them. The move can only fail after this
+    // statement if the old item has been removed or replaced, in which case it
+    // should be fine for it to be left in an inconsistent state.
+    config_.moveCb(oldItem, *newItemHdl, nullptr);
+  } else {
+    std::memcpy(newItemHdl->getWritableMemory(), oldItem.getMemory(),
+                oldItem.getSize());
+  }
+
   // Inside the MM container's lock, this checks if the old item exists to
   // make sure that no other thread removed it, and only then replaces it.
   if (!replaceInMMContainer(oldItem, *newItemHdl)) {
@@ -1185,6 +1244,7 @@ bool CacheAllocator<CacheTrait>::moveRegularItemOnEviction(
     XDCHECK(newItemHdl->hasChainedItem());
   }
   newItemHdl.unmarkNascent();
+  resHdl = newItemHdl.clone(); // guard will assign it to ctx under lock
   return true;
 }
 

cachelib/allocator/CacheAllocator.h

@@ -21,6 +21,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 <functional>
 #include <memory>
@@ -1754,6 +1756,84 @@ class CacheAllocator : public CacheBase {
     return 0; // TODO
   }
 
+  bool addWaitContextForMovingItem(
+      folly::StringPiece key, std::shared_ptr<WaitContext<ItemHandle>> waiter);
+
+  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(ItemHandle _it) { it = std::move(_it); }
+
+    // enqueue a waiter into the waiter list
+    // @param  waiter       WaitContext
+    void addWaiter(std::shared_ptr<WaitContext<ItemHandle>> waiter) {
+      XDCHECK(waiter);
+      waiters.push_back(std::move(waiter));
+    }
+
+   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(ItemHandle{});
+          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;
+        }
+      }
+    }
+
+    ItemHandle it; // will be set when Context is being filled
+    std::vector<std::shared_ptr<WaitContext<ItemHandle>>> 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
@@ -1871,6 +1951,18 @@ class CacheAllocator : public CacheBase {
   // indicates if the shutdown of cache is in progress or not
   std::atomic<bool> shutDownInProgress_{false};
 
+  static constexpr size_t kShards = 8192; // TODO: need to define right value
+
+  // a map of all pending moves
+  struct {
+    alignas(folly::hardware_destructive_interference_size) MoveMap movesMap_;
+  } movesMap_[kShards];
+
+  // a map of move locks for each shard
+  struct {
+    alignas(folly::hardware_destructive_interference_size) std::mutex moveLock_;
+  } moveLock_[kShards];
+
   // END private members
 
   // Make this friend to give access to acquire and release

cachelib/allocator/CacheItem-inl.h

@@ -264,6 +264,21 @@ bool CacheItem<CacheTrait>::isNvmEvicted() const noexcept {
   return ref_.isNvmEvicted();
 }
 
+template <typename CacheTrait>
+void CacheItem<CacheTrait>::markNotReady() noexcept {
+  ref_.markNotReady();
+}
+
+template <typename CacheTrait>
+void CacheItem<CacheTrait>::unmarkNotReady() noexcept {
+  ref_.unmarkNotReady();
+}
+
+template <typename CacheTrait>
+bool CacheItem<CacheTrait>::isNotReady() const noexcept {
+  return ref_.isNotReady();
+}
+
 template <typename CacheTrait>
 void CacheItem<CacheTrait>::markIsChainedItem() noexcept {
   XDCHECK(!hasChainedItem());

cachelib/allocator/CacheItem.h

@@ -240,6 +240,14 @@ class CACHELIB_PACKED_ATTR CacheItem {
   void unmarkNvmEvicted() noexcept;
   bool isNvmEvicted() const noexcept;
 
+  /**
+   * Marks that the item is migrating between memory tiers and
+   * not ready for access now. Accessing thread should wait.
+   */
+  void markNotReady() noexcept;
+  void unmarkNotReady() noexcept;
+  bool isNotReady() const noexcept;
+
   /**
    * Function to set the timestamp for when to expire an item
    * Employs a best-effort approach to update the expiryTime. Item's expiry

cachelib/allocator/Handle.h

@@ -464,7 +464,14 @@ struct HandleImpl {
 
   // Handle which has the item already
   FOLLY_ALWAYS_INLINE HandleImpl(Item* it, CacheT& alloc) noexcept
-      : alloc_(&alloc), it_(it) {}
+      : alloc_(&alloc), it_(it) {
+    if (it_ && it_->isNotReady()) {
+      waitContext_ = std::make_shared<ItemWaitContext>(alloc);
+      if (!alloc_->addWaitContextForMovingItem(it->getKey(), waitContext_)) {
+        waitContext_.reset();
+      }
+    }
+  }
 
   // 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/Refcount.h

@@ -116,6 +116,10 @@ class FOLLY_PACK_ATTR RefcountWithFlags {
     // unevictable in the past.
     kUnevictable_NOOP,
 
+    // Item is accecible but content is not ready yet. Used by eviction
+    // when Item is moved between memory tiers.
+    kNotReady,
+
     // Unused. This is just to indciate the maximum number of flags
     kFlagMax,
   };
@@ -329,6 +333,14 @@ class FOLLY_PACK_ATTR RefcountWithFlags {
   void unmarkNvmEvicted() noexcept { return unSetFlag<kNvmEvicted>(); }
   bool isNvmEvicted() const noexcept { return isFlagSet<kNvmEvicted>(); }
 
+  /**
+   * Marks that the item is migrating between memory tiers and
+   * not ready for access now. Accessing thread should wait.
+   */
+  void markNotReady() noexcept { return setFlag<kNotReady>(); }
+  void unmarkNotReady() noexcept { return unSetFlag<kNotReady>(); }
+  bool isNotReady() const noexcept { return isFlagSet<kNotReady>(); }
+
   // Whether or not an item is completely drained of access
   // Refcount is 0 and the item is not linked, accessible, nor moving
   bool isDrained() const noexcept { return getRefWithAccessAndAdmin() == 0; }

cachelib/allocator/tests/ItemHandleTest.cpp

@@ -39,6 +39,10 @@ struct TestItem {
   using ChainedItem = int;
 
   void reset() {}
+
+  folly::StringPiece getKey() const { return folly::StringPiece(); }
+
+  bool isNotReady() const { return false; }
 };
 
 struct TestNvmCache;
@@ -79,6 +83,12 @@ struct TestAllocator {
 
   void adjustHandleCountForThread_private(int i) { tlRef_.tlStats() += i; }
 
+  bool addWaitContextForMovingItem(
+      folly::StringPiece key,
+      std::shared_ptr<WaitContext<TestItemHandle>> waiter) {
+    return false;
+  }
+
   util::FastStats<int> tlRef_;
 };
 } // namespace