Shorten critical section in findEviction

Remove the item from mmContainer and drop the lock before attempting
eviction.

Author: igchor

cachelib/allocator/CacheAllocator-inl.h

@@ -1221,13 +1221,15 @@ CacheAllocator<CacheTrait>::findEviction(PoolId pid, ClassId cid) {
     ++searchTries;
 
     Item* candidate = itr.get();
+    mmContainer.remove(itr);
+    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.
-    auto toReleaseHandle =
-        itr->isChainedItem()
-            ? advanceIteratorAndTryEvictChainedItem(itr)
-            : advanceIteratorAndTryEvictRegularItem(mmContainer, itr);
+    auto toReleaseHandle = candidate->isChainedItem()
+                               ? tryEvictChainedItem(*candidate)
+                               : tryEvictRegularItem(mmContainer, *candidate);
 
     if (toReleaseHandle) {
       if (toReleaseHandle->hasChainedItem()) {
@@ -1236,10 +1238,6 @@ CacheAllocator<CacheTrait>::findEviction(PoolId pid, ClassId cid) {
         (*stats_.regularItemEvictions)[pid][cid].inc();
       }
 
-      // Invalidate iterator since later on we may use this mmContainer
-      // again, which cannot be done unless we drop this iterator
-      itr.destroy();
-
       // we must be the last handle and for chained items, this will be
       // the parent.
       XDCHECK(toReleaseHandle.get() == candidate || candidate->isChainedItem());
@@ -1263,11 +1261,9 @@ CacheAllocator<CacheTrait>::findEviction(PoolId pid, ClassId cid) {
       }
     }
 
-    // If we destroyed the itr to possibly evict and failed, we restart
-    // from the beginning again
-    if (!itr) {
-      itr.resetToBegin();
-    }
+    // Insert item back to the mmContainer if eviction failed.
+    mmContainer.add(*candidate);
+    itr.resetToBegin();
   }
   return nullptr;
 }
@@ -1322,19 +1318,17 @@ bool CacheAllocator<CacheTrait>::shouldWriteToNvmCacheExclusive(
 
 template <typename CacheTrait>
 typename CacheAllocator<CacheTrait>::WriteHandle
-CacheAllocator<CacheTrait>::advanceIteratorAndTryEvictRegularItem(
-    MMContainer& mmContainer, EvictionIterator& itr) {
+CacheAllocator<CacheTrait>::tryEvictRegularItem(MMContainer& mmContainer,
+                                                Item& item) {
   // 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{};
   }
@@ -1346,12 +1340,10 @@ CacheAllocator<CacheTrait>::advanceIteratorAndTryEvictRegularItem(
   auto evictHandle = accessContainer_->removeIf(item, &itemEvictionPredicate);
 
   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());
@@ -1366,15 +1358,6 @@ CacheAllocator<CacheTrait>::advanceIteratorAndTryEvictRegularItem(
     return WriteHandle{};
   }
 
-  // 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);
@@ -1388,12 +1371,10 @@ CacheAllocator<CacheTrait>::advanceIteratorAndTryEvictRegularItem(
 
 template <typename CacheTrait>
 typename CacheAllocator<CacheTrait>::WriteHandle
-CacheAllocator<CacheTrait>::advanceIteratorAndTryEvictChainedItem(
-    EvictionIterator& itr) {
-  XDCHECK(itr->isChainedItem());
+CacheAllocator<CacheTrait>::tryEvictChainedItem(Item& item) {
+  XDCHECK(item.isChainedItem());
 
-  ChainedItem* candidate = &itr->asChainedItem();
-  ++itr;
+  ChainedItem* candidate = &item.asChainedItem();
 
   // The parent could change at any point through transferChain. However, if
   // that happens, we would realize that the releaseBackToAllocator return
@@ -1420,23 +1401,11 @@ CacheAllocator<CacheTrait>::advanceIteratorAndTryEvictChainedItem(
     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());
 

cachelib/allocator/CacheAllocator.h

@@ -1580,24 +1580,22 @@ class CacheAllocator : public CacheBase {
 
   using EvictionIterator = typename MMContainer::Iterator;
 
-  // Advance the current iterator and try to evict a regular item
+  // Try to evict a regular item.
   //
   // @param  mmContainer  the container to look for evictions.
-  // @param  itr          iterator holding the item
+  // @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 advanceIteratorAndTryEvictRegularItem(MMContainer& mmContainer,
-                                                    EvictionIterator& itr);
+  WriteHandle tryEvictRegularItem(MMContainer& mmContainer, Item& item);
 
-  // Advance the current iterator and try to evict a chained item
-  // Iterator may also be reset during the course of this function
+  // Try to evict a chained item.
   //
-  // @param  itr          iterator holding the 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 advanceIteratorAndTryEvictChainedItem(EvictionIterator& itr);
+  WriteHandle tryEvictChainedItem(Item& item);
 
   // Deserializer CacheAllocatorMetadata and verify the version
   //