Back out "Shorten critical section in findEviction"

Summary:
Original commit changeset: 4d44d7518253

Original Phabricator Diff: D45071460

Reviewed By: haowu14

Differential Revision: D46566152

fbshipit-source-id: d62121506527e71a88dbd48f7ec83b57971552a7

Author: Tushar Pankaj

cachelib/allocator/CacheAllocator-inl.h

@@ -1240,19 +1240,6 @@ bool CacheAllocator<CacheTrait>::moveChainedItem(ChainedItem& oldItem,
   return true;
 }
 
-template <typename CacheTrait>
-void CacheAllocator<CacheTrait>::unlinkItemForEviction(Item& it) {
-  XDCHECK(it.isMarkedForEviction());
-  XDCHECK_EQ(0u, it.getRefCount());
-  accessContainer_->remove(it);
-  removeFromMMContainer(it);
-
-  // Since we managed to mark the item for eviction we must be the only
-  // owner of the item.
-  const auto ref = it.unmarkForEviction();
-  XDCHECK_EQ(0u, ref);
-}
-
 template <typename CacheTrait>
 typename CacheAllocator<CacheTrait>::Item*
 CacheAllocator<CacheTrait>::findEviction(PoolId pid, ClassId cid) {
@@ -1261,109 +1248,76 @@ 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;
-  while (config_.evictionSearchTries == 0 ||
-         config_.evictionSearchTries > searchTries) {
-    Item* toRecycle = nullptr;
-    Item* candidate = nullptr;
-    typename NvmCacheT::PutToken token;
-
-    mmContainer.withEvictionIterator([this, pid, cid, &candidate, &toRecycle,
-                                      &searchTries, &mmContainer,
-                                      &token](auto&& itr) {
-      if (!itr) {
-        ++searchTries;
-        (*stats_.evictionAttempts)[pid][cid].inc();
-        return;
-      }
-
-      while ((config_.evictionSearchTries == 0 ||
-              config_.evictionSearchTries > searchTries) &&
-             itr) {
-        ++searchTries;
-        (*stats_.evictionAttempts)[pid][cid].inc();
-
-        auto* toRecycle_ = itr.get();
-        auto* candidate_ =
-            toRecycle_->isChainedItem()
-                ? &toRecycle_->asChainedItem().getParentItem(compressor_)
-                : toRecycle_;
-
-        const bool evictToNvmCache = shouldWriteToNvmCache(*candidate_);
-        auto putToken = evictToNvmCache
-                            ? nvmCache_->createPutToken(candidate_->getKey())
-                            : typename NvmCacheT::PutToken{};
-
-        if (evictToNvmCache && !putToken.isValid()) {
-          stats_.evictFailConcurrentFill.inc();
-          ++itr;
-          continue;
-        }
-
-        auto markedForEviction = candidate_->markForEviction();
-        if (!markedForEviction) {
-          if (candidate_->hasChainedItem()) {
-            stats_.evictFailParentAC.inc();
-          } else {
-            stats_.evictFailAC.inc();
-          }
-          ++itr;
-          continue;
-        }
-
-        XDCHECK(candidate_->isMarkedForEviction());
-        // markForEviction to make sure no other thead is evicting the item
-        // nor holding a handle to that item
-        toRecycle = toRecycle_;
-        candidate = candidate_;
-        token = std::move(putToken);
-
-        // Check if parent changed for chained items - if yes, we cannot
-        // remove the child from the mmContainer as we will not be evicting
-        // it. We could abort right here, but we need to cleanup in case
-        // unmarkForEviction() returns 0 - so just go through normal path.
-        if (!toRecycle_->isChainedItem() ||
-            &toRecycle->asChainedItem().getParentItem(compressor_) ==
-                candidate) {
-          mmContainer.remove(itr);
-        }
-        return;
-      }
-    });
-
-    if (!toRecycle) {
+  auto itr = mmContainer.getEvictionIterator();
+  while ((config_.evictionSearchTries == 0 ||
+          config_.evictionSearchTries > searchTries) &&
+         itr) {
+    ++searchTries;
+    (*stats_.evictionAttempts)[pid][cid].inc();
+
+    Item* toRecycle = itr.get();
+
+    Item* candidate =
+        toRecycle->isChainedItem()
+            ? &toRecycle->asChainedItem().getParentItem(compressor_)
+            : toRecycle;
+
+    // make sure no other thead is evicting the item
+    if (candidate->getRefCount() != 0 || !candidate->markMoving()) {
+      ++itr;
       continue;
     }
 
-    XDCHECK(toRecycle);
-    XDCHECK(candidate);
-
-    unlinkItemForEviction(*candidate);
+    // 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.
+    bool evictionSuccessful = false;
+    {
+      auto toReleaseHandle =
+          itr->isChainedItem()
+              ? advanceIteratorAndTryEvictChainedItem(itr)
+              : advanceIteratorAndTryEvictRegularItem(mmContainer, itr);
+      evictionSuccessful = toReleaseHandle != nullptr;
+      // destroy toReleaseHandle. The item won't be released to allocator
+      // since we marked for eviction.
+    }
+
+    const auto ref = candidate->unmarkMoving();
+    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.
+      if (candidate->hasChainedItem()) {
+        (*stats_.chainedItemEvictions)[pid][cid].inc();
+      } else {
+        (*stats_.regularItemEvictions)[pid][cid].inc();
+      }
 
-    if (token.isValid() && shouldWriteToNvmCacheExclusive(*candidate)) {
-      nvmCache_->put(*candidate, std::move(token));
-    }
+      if (auto eventTracker = getEventTracker()) {
+        eventTracker->record(AllocatorApiEvent::DRAM_EVICT, candidate->getKey(),
+                             AllocatorApiResult::EVICTED, candidate->getSize(),
+                             candidate->getConfiguredTTL().count());
+      }
 
-    // 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.
-    if (candidate->hasChainedItem()) {
-      (*stats_.chainedItemEvictions)[pid][cid].inc();
+      // check if by releasing the item we intend to, we actually
+      // recycle the candidate.
+      if (ReleaseRes::kRecycled ==
+          releaseBackToAllocator(*candidate, RemoveContext::kEviction,
+                                 /* isNascent */ false, toRecycle)) {
+        return toRecycle;
+      }
     } else {
-      (*stats_.regularItemEvictions)[pid][cid].inc();
+      XDCHECK(!evictionSuccessful);
     }
 
-    if (auto eventTracker = getEventTracker()) {
-      eventTracker->record(AllocatorApiEvent::DRAM_EVICT, candidate->getKey(),
-                           AllocatorApiResult::EVICTED, candidate->getSize(),
-                           candidate->getConfiguredTTL().count());
-    }
-
-    // check if by releasing the item we intend to, we actually
-    // recycle the candidate.
-    auto ret = releaseBackToAllocator(*candidate, RemoveContext::kEviction,
-                                      /* isNascent */ false, toRecycle);
-    if (ret == ReleaseRes::kRecycled) {
-      return toRecycle;
+    // If we destroyed the itr to possibly evict and failed, we restart
+    // from the beginning again
+    if (!itr) {
+      itr.resetToBegin();
     }
   }
   return nullptr;
@@ -1507,7 +1461,7 @@ bool CacheAllocator<CacheTrait>::pushToNvmCacheFromRamForTesting(
 
   if (handle && nvmCache_ && shouldWriteToNvmCache(*handle) &&
       shouldWriteToNvmCacheExclusive(*handle)) {
-    nvmCache_->put(*handle, nvmCache_->createPutToken(handle->getKey()));
+    nvmCache_->put(handle, nvmCache_->createPutToken(handle->getKey()));
     return true;
   }
   return false;
@@ -2714,6 +2668,126 @@ void CacheAllocator<CacheTrait>::evictForSlabRelease(
   }
 }
 
+template <typename CacheTrait>
+typename CacheAllocator<CacheTrait>::WriteHandle
+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()) {
+    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) {
@@ -2746,7 +2820,7 @@ CacheAllocator<CacheTrait>::evictNormalItemForSlabRelease(Item& item) {
   // now that we are the only handle and we actually removed something from
   // the RAM cache, we enqueue it to nvmcache.
   if (evictToNvmCache && shouldWriteToNvmCacheExclusive(item)) {
-    nvmCache_->put(*handle, std::move(token));
+    nvmCache_->put(handle, std::move(token));
   }
 
   return handle;
@@ -2847,7 +2921,7 @@ CacheAllocator<CacheTrait>::evictChainedItemForSlabRelease(ChainedItem& child) {
   // the RAM cache, we enqueue it to nvmcache.
   if (evictToNvmCache && shouldWriteToNvmCacheExclusive(*parentHandle)) {
     DCHECK(parentHandle->hasChainedItem());
-    nvmCache_->put(*parentHandle, std::move(token));
+    nvmCache_->put(parentHandle, std::move(token));
   }
 
   return parentHandle;

cachelib/allocator/CacheAllocator.h

@@ -1664,12 +1664,6 @@ class CacheAllocator : public CacheBase {
                        bool removeFromNvm = true,
                        bool recordApiEvent = true);
 
-  // Must be called by the thread which called markForEviction and
-  // succeeded. After this call, the item is unlinked from Access and
-  // MM Containers. The item is no longer marked as exclusive and it's
-  // ref count is 0 - it's available for recycling.
-  void unlinkItemForEviction(Item& it);
-
   // Implementation to find a suitable eviction from the container. The
   // two parameters together identify a single container.
   //
@@ -1680,6 +1674,25 @@ class CacheAllocator : public CacheBase {
 
   using EvictionIterator = typename MMContainer::LockedIterator;
 
+  // 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