Increment item refCount under AC lock

to avoid races with remove/replace with predicate.

Also, refact tryEvictRegularItem and tryEvictChainedItem into
a single function.

Author: igchor

cachelib/allocator/CacheAllocator-inl.h

@@ -1466,51 +1466,58 @@ CacheAllocator<CacheTrait>::findEviction(TierId tid, PoolId pid, ClassId cid) {
       ++itr;
       continue;
     }
+
+    if (candidate->isChainedItem()) {
+      candidate = &candidate->asChainedItem().getParentItem(compressor_);
+    }
+
+    {
+      auto toReleaseHandle = accessContainer_->find(*candidate);
+      if (!toReleaseHandle) {
+        ++itr;
+        continue;
+      }
+      toReleaseHandle.release();
+    }
+    // incRef(*candidate);
 
-    incRef(*candidate);
     itr.destroy();
 
     // for chained items, the ownership of the parent can change. We try to
     // evict what we think as parent and see if the eviction of parent
     // recycles the child we intend to.
 
-    WriteHandle toReleaseHandle = tryEvictToNextMemoryTier(tid, pid, candidate);
     bool movedToNextTier = false;
-    if(toReleaseHandle) {
-      movedToNextTier = true;
-    } else {
-      toReleaseHandle = candidate->isChainedItem()
-                               ? tryEvictChainedItem(mmContainer, *candidate)
-                               : tryEvictRegularItem(mmContainer, *candidate);
-                               // XXX: fix chained item
+    bool evicted = false;
+    {
+      auto h = tryEvictToNextMemoryTier(tid, pid, candidate);
+      if (h) {
+        movedToNextTier = true;
+        evicted = true;
+        h.release();
+        decRef(*candidate);
+      }
     }
-
-    if (toReleaseHandle) {
-      if (toReleaseHandle->hasChainedItem()) {
+    if(!movedToNextTier) {
+      evicted = tryEvictItem(mmContainer, *candidate);
+    }
+    if (evicted) {
+      if (candidate->hasChainedItem()) {
         (*stats_.chainedItemEvictions)[pid][cid].inc();
       } else {
         (*stats_.regularItemEvictions)[pid][cid].inc();
       }
 
-      // we must be the last handle and for chained items, this will be
-      // the parent.
-      XDCHECK(toReleaseHandle.get() == candidate || candidate->isChainedItem());
-      XDCHECK_EQ(2u, toReleaseHandle->getRefCount());
-
-      // We manually release the item here because we don't want to
-      // invoke the Item Handle's destructor which will be decrementing
-      // an already zero refcount, which will throw exception
-      auto& itemToRelease = *toReleaseHandle.release();
+      XDCHECK_EQ(1u, candidate->getRefCount());
 
       // Decrementing the refcount because we want to recycle the item
-      decRef(itemToRelease);
-      const auto ref = decRef(itemToRelease);
+      const auto ref = decRef(*candidate);
       XDCHECK_EQ(0u, ref);
 
       // check if by releasing the item we intend to, we actually
       // recycle the candidate.
       if (ReleaseRes::kRecycled ==
-          releaseBackToAllocator(itemToRelease, RemoveContext::kEviction,
+          releaseBackToAllocator(*candidate, RemoveContext::kEviction,
                                  /* isNascent */ movedToNextTier, candidate)) {
         return candidate;
       }
@@ -1611,9 +1618,8 @@ CacheAllocator<CacheTrait>::tryEvictToNextMemoryTier(Item* item) {
 }
 
 template <typename CacheTrait>
-typename CacheAllocator<CacheTrait>::WriteHandle
-CacheAllocator<CacheTrait>::tryEvictRegularItem(MMContainer& mmContainer,
-                                                Item& item) {
+bool CacheAllocator<CacheTrait>::tryEvictItem(MMContainer& mmContainer,
+                                                Item &item) {
   // we should flush this to nvmcache if it is not already present in nvmcache
   // and the item is not expired.
   const bool evictToNvmCache = shouldWriteToNvmCache(item);
@@ -1623,21 +1629,25 @@ CacheAllocator<CacheTrait>::tryEvictRegularItem(MMContainer& mmContainer,
   // record the in-flight eviciton. If not, we move on to next item to avoid
   // stalling eviction.
   if (evictToNvmCache && !token.isValid()) {
-    stats_.evictFailConcurrentFill.inc();
+    if (item.isChainedItem())
+      stats_.evictFailConcurrentFill.inc();
+    else
+      stats_.evictFailConcurrentFill.inc();
     decRef(item);
-    return WriteHandle{};
+    return false;
   }
 
-  // If there are other accessors, we should abort. Acquire a handle here since
-  // if we remove the item from both access containers and mm containers
-  // below, we will need a handle to ensure proper cleanup in case we end up
-  // not evicting this item
+  // If there are other accessors, we should abort. We should be the only
+  // handle owner.
   auto evictHandle = accessContainer_->removeIf(item, &itemEvictionPredicate);
+  decRef(item);
 
   if (!evictHandle) {
-    stats_.evictFailAC.inc();
-    decRef(item);
-    return evictHandle;
+    if (item.isChainedItem())
+      stats_.evictFailParentAC.inc();
+    else
+      stats_.evictFailAC.inc();
+    return false;
   }
 
   auto removed = mmContainer.remove(item);
@@ -1652,81 +1662,24 @@ CacheAllocator<CacheTrait>::tryEvictRegularItem(MMContainer& mmContainer,
   // for eviction. It is safe to destroy the handle here since the moving bit
   // is set. Iterator was already advance by the remove call above.
   if (evictHandle->isMoving()) {
-    stats_.evictFailMove.inc();
-    decRef(item);
-    return WriteHandle{};
+    if (item.isChainedItem())
+      stats_.evictFailParentMove.inc();
+    else
+      stats_.evictFailMove.inc();
+    return false;
   }
 
   // Ensure that there are no accessors after removing from the access
-  // container
-  XDCHECK(evictHandle->getRefCount() == 2);
+  // container. handle and evictHandle should be the only two handles.
+  XDCHECK(evictHandle->getRefCount() == 1);
 
   if (evictToNvmCache && shouldWriteToNvmCacheExclusive(item)) {
     XDCHECK(token.isValid());
     nvmCache_->put(evictHandle, std::move(token));
   }
-  return evictHandle;
-}
-
-template <typename CacheTrait>
-typename CacheAllocator<CacheTrait>::WriteHandle
-CacheAllocator<CacheTrait>::tryEvictChainedItem(MMContainer& mmContainer, Item& item) {
-  XDCHECK(item.isChainedItem());
-
-  ChainedItem* candidate = &item.asChainedItem();
-
-  // The parent could change at any point through transferChain. However, if
-  // that happens, we would realize that the releaseBackToAllocator return
-  // kNotRecycled and we would try another chained item, leading to transient
-  // failure.
-  auto& parent = candidate->getParentItem(compressor_);
-
-  const bool evictToNvmCache = shouldWriteToNvmCache(parent);
 
-  auto token = evictToNvmCache ? nvmCache_->createPutToken(parent.getKey())
-                               : typename NvmCacheT::PutToken{};
-
-  // if token is invalid, return. iterator is already advanced.
-  if (evictToNvmCache && !token.isValid()) {
-    stats_.evictFailConcurrentFill.inc();
-    return ItemHandle{};
-  }
-
-  // check if the parent exists in the hashtable and refcount is drained.
-  auto parentHandle =
-      accessContainer_->removeIf(parent, &itemEvictionPredicate);
-  if (!parentHandle) {
-    stats_.evictFailParentAC.inc();
-    return parentHandle;
-  }
-
-  // Ensure we have the correct parent and we're the only user of the
-  // parent, then free it from access container. Otherwise, we abort
-  auto removed = mmContainer.remove(item);
-  XDCHECK(removed);
-  XDCHECK_EQ(reinterpret_cast<uintptr_t>(&parent),
-             reinterpret_cast<uintptr_t>(parentHandle.get()));
-  XDCHECK_EQ(2u, parent.getRefCount()); // XXX?
-  XDCHECK(!parent.isInMMContainer());
-  XDCHECK(!parent.isAccessible());
-
-  // TODO: add multi-tier support (similar as for unchained items)
-
-  // We need to make sure the parent is not marked as moving
-  // and we're the only holder of the parent item. Safe to destroy the handle
-  // here since moving bit is set.
-  if (parentHandle->isMoving()) {
-    stats_.evictFailParentMove.inc();
-    return ItemHandle{};
-  }
-
-  if (evictToNvmCache && shouldWriteToNvmCacheExclusive(*parentHandle)) {
-    XDCHECK(token.isValid());
-    XDCHECK(parentHandle->hasChainedItem());
-    nvmCache_->put(parentHandle, std::move(token));
-  }
-
-  return parentHandle;
+  evictHandle.release();
+  return true;
 }
 
 template <typename CacheTrait>

cachelib/allocator/CacheAllocator.h

@@ -1566,17 +1566,8 @@ class CacheAllocator : public CacheBase {
   // @param  mmContainer  the container to look for evictions.
   // @param  item         item to evict
   //
-  // @return  valid handle to regular item on success. This will be the last
-  //          handle to the item. On failure an empty handle.
-  WriteHandle tryEvictRegularItem(MMContainer& mmContainer, Item& item);
-
-  // Try to evict a chained item.
-  //
-  // @param  item         item to evict
-  //
-  // @return  valid handle to the parent item on success. This will be the last
-  //          handle to the item
-  WriteHandle tryEvictChainedItem(MMContainer& mmContainer, Item& item);
+  // @return whther eviction succeeded
+  bool tryEvictItem(MMContainer& mmContainer, Item& item);
 
   // Try to move the item down to the next memory tier
   //

cachelib/allocator/ChainedHashTable-inl.h

@@ -460,6 +460,21 @@ typename T::Handle ChainedHashTable::Container<T, HookPtr, LockT>::find(
   return handleMaker_(ht_.findInBucket(key, bucket));
 }
 
+template <typename T,
+          typename ChainedHashTable::Hook<T> T::*HookPtr,
+          typename LockT>
+typename T::Handle ChainedHashTable::Container<T, HookPtr, LockT>::find(
+    T& node) const {
+  const auto bucket = ht_.getBucket(node.getKey());
+  auto l = locks_.lockShared(bucket);
+
+  if (!node.isAccessible()) {
+    return {};
+  }
+
+  return handleMaker_(&node);
+}
+
 template <typename T,
           typename ChainedHashTable::Hook<T> T::*HookPtr,
           typename LockT>

cachelib/allocator/ChainedHashTable.h

@@ -490,6 +490,17 @@ class ChainedHashTable {
     //        creating this item handle.
     Handle find(Key key) const;
 
+    // returns a handle to specified node.
+    //
+    // @param  node       requested node
+    //
+    // @return handle to the node if find was sucessfull. returns a
+    // null handle if the node was not in the container.
+    //
+    // @throw std::overflow_error is the maximum item refcount is execeeded by
+    //        creating this item handle.
+    Handle find(T& node) const;
+
     // for saving the state of the hash table
     //
     // precondition:  serialization must happen without any reader or writer