Implemented async Item movement between tiers

Author: vinser52

cachelib/allocator/CacheAllocator-inl.h

@@ -1104,6 +1104,22 @@ CacheAllocator<CacheTrait>::insertOrReplace(const ItemHandle& handle) {
   return replaced;
 }
 
+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,
@@ -1125,22 +1141,35 @@ 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,
@@ -1154,6 +1183,19 @@ bool CacheAllocator<CacheTrait>::moveRegularItemOnEviction(
     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)) {
@@ -1188,6 +1230,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>
@@ -1745,6 +1747,79 @@ 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
@@ -1862,6 +1937,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,11 @@ 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