Synchronize with SlabRelease before eviction

attempts. Mark the item as exclusive before trying
to remove from AC or MMContainer.

Author: igchor

cachelib/allocator/CacheAllocator-inl.h

@@ -1221,55 +1221,44 @@ CacheAllocator<CacheTrait>::findEviction(PoolId pid, ClassId cid) {
   // Keep searching for a candidate until we were able to evict it
   // or until the search limit has been exhausted
   unsigned int searchTries = 0;
+  auto itr = mmContainer.getEvictionIterator();
   while ((config_.evictionSearchTries == 0 ||
-          config_.evictionSearchTries > searchTries)) {
+          config_.evictionSearchTries > searchTries) &&
+         itr) {
     ++searchTries;
     (*stats_.evictionAttempts)[pid][cid].inc();
 
-    Item* toRecycle = nullptr;
-    Item* candidate = nullptr;
-
-    mmContainer.withEvictionIterator(
-        [this, &candidate, &toRecycle, &searchTries](auto&& itr) {
-          while ((config_.evictionSearchTries == 0 ||
-                  config_.evictionSearchTries > searchTries) &&
-                 itr) {
-            ++searchTries;
-
-            auto* toRecycle_ = itr.get();
-            auto* candidate_ =
-                toRecycle_->isChainedItem()
-                    ? &toRecycle_->asChainedItem().getParentItem(compressor_)
-                    : toRecycle_;
-
-            // make sure no other thead is evicting the item
-            if (candidate_->getRefCount() == 0 && candidate_->markExclusive()) {
-              toRecycle = toRecycle_;
-              candidate = candidate_;
-              return;
-            }
-
-            ++itr;
-          }
-        });
-
-    if (!toRecycle)
-      continue;
+    Item* toRecycle = itr.get();
+
+    Item* candidate =
+        toRecycle->isChainedItem()
+            ? &toRecycle->asChainedItem().getParentItem(compressor_)
+            : toRecycle;
 
-    XDCHECK(toRecycle);
-    XDCHECK(candidate);
+    // make sure no other thead is evicting the item
+    if (candidate->getRefCount() != 0 || !candidate->markExclusive()) {
+      ++itr;
+      continue;
+    }
 
     // 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.
     {
-      auto toReleaseHandle = evictNormalItem(*candidate);
+      auto toReleaseHandle =
+          itr->isChainedItem()
+              ? advanceIteratorAndTryEvictChainedItem(itr)
+              : advanceIteratorAndTryEvictRegularItem(mmContainer, itr);
       // destroy toReleseHandle. The item won't be release to allocator
       // since we marked it as exclusive.
     }
-    const auto ref = candidate->unmarkExclusive();
 
+    const auto ref = candidate->unmarkExclusive();
     if (ref == 0u) {
+      // Invalidate iterator since later on we may use this mmContainer
+      // again, which cannot be done unless we drop this iterator
+      itr.destroy();
+
       // recycle the item. it's safe to do so, even if toReleaseHandle was
       // NULL. If `ref` == 0 then it means that we are the last holder of
       // that item.
@@ -1299,6 +1288,8 @@ CacheAllocator<CacheTrait>::findEviction(PoolId pid, ClassId cid) {
         stats_.evictFailAC.inc();
       }
     }
+
+    itr.resetToBegin();
   }
   return nullptr;
 }
@@ -2589,7 +2580,7 @@ void CacheAllocator<CacheTrait>::evictForSlabRelease(
     auto owningHandle =
         item.isChainedItem()
             ? evictChainedItemForSlabRelease(item.asChainedItem())
-            : evictNormalItem(item);
+            : evictNormalItemForSlabRelease(item);
 
     // we managed to evict the corresponding owner of the item and have the
     // last handle for the owner.
@@ -2644,7 +2635,128 @@ void CacheAllocator<CacheTrait>::evictForSlabRelease(
 
 template <typename CacheTrait>
 typename CacheAllocator<CacheTrait>::WriteHandle
-CacheAllocator<CacheTrait>::evictNormalItem(Item& item) {
+CacheAllocator<CacheTrait>::advanceIteratorAndTryEvictRegularItem(
+    MMContainer& mmContainer, EvictionIterator& itr) {
+  // we should flush this to nvmcache if it is not already present in nvmcache
+  // and the item is not expired.
+  Item& item = *itr;
+  const bool evictToNvmCache = shouldWriteToNvmCache(item);
+
+  auto token = evictToNvmCache ? nvmCache_->createPutToken(item.getKey())
+                               : typename NvmCacheT::PutToken{};
+
+  // record the in-flight eviciton. If not, we move on to next item to avoid
+  // stalling eviction.
+  if (evictToNvmCache && !token.isValid()) {
+    ++itr;
+    stats_.evictFailConcurrentFill.inc();
+    return WriteHandle{};
+  }
+
+  // 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
+  auto evictHandle = accessContainer_->removeIf(item, &itemExclusivePredicate);
+  if (!evictHandle) {
+    ++itr;
+    stats_.evictFailAC.inc();
+    return evictHandle;
+  }
+
+  mmContainer.remove(itr);
+  XDCHECK_EQ(reinterpret_cast<uintptr_t>(evictHandle.get()),
+             reinterpret_cast<uintptr_t>(&item));
+  XDCHECK(!evictHandle->isInMMContainer());
+  XDCHECK(!evictHandle->isAccessible());
+
+  // Invalidate iterator since later on if we are not evicting this
+  // item, we may need to rely on the handle we created above to ensure
+  // proper cleanup if the item's raw refcount has dropped to 0.
+  // And since this item may be a parent item that has some child items
+  // in this very same mmContainer, we need to make sure we drop this
+  // exclusive iterator so we can gain access to it when we're cleaning
+  // up the child items
+  itr.destroy();
+
+  // Ensure that there are no accessors after removing from the access
+  // container
+  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>::advanceIteratorAndTryEvictChainedItem(
+    EvictionIterator& itr) {
+  XDCHECK(itr->isChainedItem());
+
+  ChainedItem* candidate = &itr->asChainedItem();
+  ++itr;
+
+  // 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()) {
+    ++itr;
+    stats_.evictFailConcurrentFill.inc();
+    return WriteHandle{};
+  }
+
+  // check if the parent exists in the hashtable and refcount is drained.
+  auto parentHandle =
+      accessContainer_->removeIf(parent, &itemExclusivePredicate);
+  if (!parentHandle) {
+    stats_.evictFailParentAC.inc();
+    return parentHandle;
+  }
+
+  // Invalidate iterator since later on we may use the mmContainer
+  // associated with this iterator which cannot be done unless we
+  // drop this iterator
+  //
+  // This must be done once we know the parent is not nullptr.
+  // Since we can very well be the last holder of this parent item,
+  // which may have a chained item that is linked in this MM container.
+  itr.destroy();
+
+  // Ensure we have the correct parent and we're the only user of the
+  // parent, then free it from access container. Otherwise, we abort
+  XDCHECK_EQ(reinterpret_cast<uintptr_t>(&parent),
+             reinterpret_cast<uintptr_t>(parentHandle.get()));
+  XDCHECK_EQ(1u, parent.getRefCount());
+
+  removeFromMMContainer(*parentHandle);
+
+  XDCHECK(!parent.isInMMContainer());
+  XDCHECK(!parent.isAccessible());
+
+  if (evictToNvmCache && shouldWriteToNvmCacheExclusive(*parentHandle)) {
+    XDCHECK(token.isValid());
+    XDCHECK(parentHandle->hasChainedItem());
+    nvmCache_->put(parentHandle, std::move(token));
+  }
+
+  return parentHandle;
+}
+
+template <typename CacheTrait>
+typename CacheAllocator<CacheTrait>::WriteHandle
+CacheAllocator<CacheTrait>::evictNormalItemForSlabRelease(Item& item) {
   XDCHECK(item.isExclusive());
 
   if (item.isOnlyExclusive()) {
@@ -2841,7 +2953,6 @@ bool CacheAllocator<CacheTrait>::markExclusiveForSlabRelease(
     itemFreed = true;
 
     if (shutDownInProgress_) {
-      XDCHECK(!static_cast<Item*>(alloc)->isExclusive());
       allocator_->abortSlabRelease(ctx);
       throw exception::SlabReleaseAborted(
           folly::sformat("Slab Release aborted while still trying to mark"

cachelib/allocator/CacheAllocator.h

@@ -1661,6 +1661,25 @@ class CacheAllocator : public CacheBase {
 
   using EvictionIterator = typename MMContainer::Iterator;
 
+  // Advance the current iterator and try to evict a regular item
+  //
+  // @param  mmContainer  the container to look for evictions.
+  // @param  itr          iterator holding the item
+  //
+  // @return  valid handle to regular item on success. This will be the last
+  //          handle to the item. On failure an empty handle.
+  WriteHandle advanceIteratorAndTryEvictRegularItem(MMContainer& mmContainer,
+                                                    EvictionIterator& itr);
+
+  // Advance the current iterator and try to evict a chained item
+  // Iterator may also be reset during the course of this function
+  //
+  // @param  itr          iterator holding the item
+  //
+  // @return  valid handle to the parent item on success. This will be the last
+  //          handle to the item
+  WriteHandle advanceIteratorAndTryEvictChainedItem(EvictionIterator& itr);
+
   // Deserializer CacheAllocatorMetadata and verify the version
   //
   // @param  deserializer   Deserializer object
@@ -1778,7 +1797,7 @@ class CacheAllocator : public CacheBase {
   //
   // @return last handle for corresponding to item on success. empty handle on
   // failure. caller can retry if needed.
-  WriteHandle evictNormalItem(Item& item);
+  WriteHandle evictNormalItemForSlabRelease(Item& item);
 
   // Helper function to evict a child item for slab release
   // As a side effect, the parent item is also evicted

cachelib/allocator/CacheItem.h

@@ -353,7 +353,9 @@ class CACHELIB_PACKED_ATTR CacheItem {
    * User can also query if an item "isOnlyExclusive". This returns true only
    * if the refcount is 0 and only the exclusive bit is set.
    *
-   * Unmarking exclusive does not depend on `isInMMContainer`
+   * Unmarking exclusive does not depend on `isInMMContainer`.
+   * Unmarking exclusive will also return the refcount at the moment of
+   * unmarking.
    */
   bool markExclusive() noexcept;
   RefcountWithFlags::Value unmarkExclusive() noexcept;

cachelib/allocator/MM2Q-inl.h

@@ -250,7 +250,18 @@ MM2Q::Container<T, HookPtr>::getEvictionIterator() const noexcept {
   // arbitrary amount of time outside a lambda-friendly piece of code (eg. they
   // can return the iterator from functions, pass it to functions, etc)
   //
-  // to get advantage of combining, use withEvictionIterator
+  // it would be theoretically possible to refactor this interface into
+  // something like the following to allow combining
+  //
+  //    mm2q.withEvictionIterator([&](auto iterator) {
+  //      // user code
+  //    });
+  //
+  // at the time of writing it is unclear if the gains from combining are
+  // reasonable justification for the codemod required to achieve combinability
+  // as we don't expect this critical section to be the hotspot in user code.
+  // This is however subject to change at some time in the future as and when
+  // this assertion becomes false.
   LockHolder l(*lruMutex_);
   return Iterator{std::move(l), lru_.rbegin()};
 }

cachelib/allocator/Refcount.h

@@ -256,7 +256,9 @@ class FOLLY_PACK_ATTR RefcountWithFlags {
    * User can also query if an item "isOnlyExclusive". This returns true only
    * if the refcount is 0 and only the exclusive bit is set.
    *
-   * Unmarking exclusive does not depend on `isInMMContainer`
+   * Unmarking exclusive does not depend on `isInMMContainer`.
+   * Unmarking exclusive will also return the refcount at the moment of
+   * unmarking.
    */
   bool markExclusive() noexcept {
     Value bitMask = getAdminRef<kExclusive>();
@@ -298,7 +300,8 @@ class FOLLY_PACK_ATTR RefcountWithFlags {
   bool isOnlyExclusive() const noexcept {
     // An item is only exclusive when its refcount is zero and only the
     // exclusive bit among all the control bits is set. This indicates an item
-    // is already on its way out of cache and does not need to be moved.
+    // is exclusive to the current thread. No other thread is allowed to
+    // do anything with it.
     auto ref = getRefWithAccessAndAdmin();
     bool anyOtherBitSet = ref & ~getAdminRef<kExclusive>();
     if (anyOtherBitSet) {