Use moving bit for synchronization in findEviction

moving bit is used to give exclusive right to evict the item
to a particular thread.

Originially, there was an assumption that whoever marked the item
as moving will try to free it until he succeeds. Since we don't want
to do that in findEviction (potentially can take a long time) we need
to make sure that unmarking is safe.

This patch checks the flags after unmarking (atomically) and if ref is
zero it also recyles the item. This is needed as there might be some
concurrent thread releasing the item (and decrementing ref count). If
moving bit is set, that thread would not free the memory back to
allocator, resulting in memory leak on unmarkMoving().

Author: igchor

cachelib/allocator/CacheAllocator-inl.h

@@ -1225,20 +1225,15 @@ CacheAllocator<CacheTrait>::findEviction(PoolId pid, ClassId cid) {
          itr) {
     ++searchTries;
 
-    Item* candidate = itr.get();
+    Item* toRecycle = itr.get();
 
-    // make sure no other thead is evicting the item
-    if (candidate->getRefCount() != 0 || candidate->isMoving()) {
-      ++itr;
-      continue;
-    }
-
-    if (candidate->isChainedItem()) {
-      candidate = &candidate->asChainedItem().getParentItem(compressor_);
-    }
+    Item* candidate =
+        toRecycle->isChainedItem()
+            ? &toRecycle->asChainedItem().getParentItem(compressor_)
+            : toRecycle;
 
-    auto toReleaseHandle = accessContainer_->find(*candidate);
-    if (!toReleaseHandle || !toReleaseHandle.isReady()) {
+    // make sure no other thead is evicting the item
+    if (candidate->getRefCount() != 0 || !candidate->markMoving()) {
       ++itr;
       continue;
     }
@@ -1248,8 +1243,11 @@ CacheAllocator<CacheTrait>::findEviction(PoolId pid, ClassId cid) {
     // 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 evicted = tryEvictItem(mmContainer, std::move(toReleaseHandle));
-    if (evicted) {
+    auto toReleaseHandle =
+        evictNormalItem(*candidate, true /* skipIfTokenInvalid */);
+    auto ref = candidate->unmarkMoving();
+
+    if (toReleaseHandle || ref == 0u) {
       if (candidate->hasChainedItem()) {
         (*stats_.chainedItemEvictions)[pid][cid].inc();
       } else {
@@ -1262,17 +1260,43 @@ CacheAllocator<CacheTrait>::findEviction(PoolId pid, ClassId cid) {
             AllocatorApiResult::DRAM_EVICTED, toReleaseHandle->getSize(),
             toReleaseHandle->getConfiguredTTL().count());
       }
+    } else {
+      if (candidate->hasChainedItem()) {
+        stats_.evictFailParentAC.inc();
+      } else {
+        stats_.evictFailAC.inc();
+      }
+    }
+
+    if (toReleaseHandle) {
+      XDCHECK(toReleaseHandle.get() == candidate);
+      XDCHECK(toRecycle == candidate || toRecycle->isChainedItem());
+      XDCHECK_EQ(1u, 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();
 
       // Decrementing the refcount because we want to recycle the item
-      const auto ref = decRef(*candidate);
+      const auto ref = decRef(itemToRelease);
       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,
+                                 /* isNascent */ false, toRecycle)) {
+        return toRecycle;
+      }
+    } else if (ref == 0u) {
+      // it's safe to recycle the item here as there are no more
+      // references and the item could not been marked as moving
+      // by other thread since it's detached from MMContainer.
       if (ReleaseRes::kRecycled ==
           releaseBackToAllocator(*candidate, RemoveContext::kEviction,
-                                 /* isNascent */ false, candidate)) {
-        return candidate;
+                                 /* isNascent */ false, toRecycle)) {
+        return toRecycle;
       }
     }
 
@@ -1329,77 +1353,6 @@ bool CacheAllocator<CacheTrait>::shouldWriteToNvmCacheExclusive(
   return true;
 }
 
-template <typename CacheTrait>
-bool CacheAllocator<CacheTrait>::tryEvictItem(MMContainer& mmContainer,
-                                              WriteHandle&& handle) {
-  auto& item = *handle;
-
-  // 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);
-
-  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()) {
-    if (item.isChainedItem())
-      stats_.evictFailConcurrentFill.inc();
-    else
-      stats_.evictFailConcurrentFill.inc();
-    handle.reset();
-    return false;
-  }
-
-  // If there are other accessors, we should abort. We should be the only
-  // handle owner.
-  auto evictHandle = accessContainer_->removeIf(item, &itemEvictionPredicate);
-
-  /* We got new handle so it's safe to get rid of the old one. */
-  handle.reset();
-
-  if (!evictHandle) {
-    if (item.isChainedItem())
-      stats_.evictFailParentAC.inc();
-    else
-      stats_.evictFailAC.inc();
-    return false;
-  }
-
-  auto removed = mmContainer.remove(item);
-  XDCHECK(removed);
-  XDCHECK_EQ(reinterpret_cast<uintptr_t>(evictHandle.get()),
-             reinterpret_cast<uintptr_t>(&item));
-  XDCHECK(!evictHandle->isInMMContainer());
-  XDCHECK(!evictHandle->isAccessible());
-
-  // If the item is now marked as moving, that means its corresponding slab is
-  // being released right now. So, we look for the next item that is eligible
-  // for eviction. It is safe to destroy the handle here since the moving bit
-  // is set.
-  if (evictHandle->isMoving()) {
-    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() == 1);
-
-  if (evictToNvmCache && shouldWriteToNvmCacheExclusive(item)) {
-    XDCHECK(token.isValid());
-    nvmCache_->put(evictHandle, std::move(token));
-  }
-
-  /* Release the handle, item will be reused by the called. */
-  evictHandle.release();
-
-  return true;
-}
-
 template <typename CacheTrait>
 typename CacheAllocator<CacheTrait>::RemoveRes
 CacheAllocator<CacheTrait>::remove(typename Item::Key key) {
@@ -2589,7 +2542,7 @@ void CacheAllocator<CacheTrait>::evictForSlabRelease(
     auto owningHandle =
         item.isChainedItem()
             ? evictChainedItemForSlabRelease(item.asChainedItem())
-            : evictNormalItemForSlabRelease(item);
+            : evictNormalItem(item);
 
     // we managed to evict the corresponding owner of the item and have the
     // last handle for the owner.
@@ -2646,7 +2599,8 @@ void CacheAllocator<CacheTrait>::evictForSlabRelease(
 
 template <typename CacheTrait>
 typename CacheAllocator<CacheTrait>::WriteHandle
-CacheAllocator<CacheTrait>::evictNormalItemForSlabRelease(Item& item) {
+CacheAllocator<CacheTrait>::evictNormalItem(Item& item,
+                                            bool skipIfTokenInvalid) {
   XDCHECK(item.isMoving());
 
   if (item.isOnlyMoving()) {
@@ -2659,6 +2613,11 @@ CacheAllocator<CacheTrait>::evictNormalItemForSlabRelease(Item& item) {
   auto token = evictToNvmCache ? nvmCache_->createPutToken(item.getKey())
                                : typename NvmCacheT::PutToken{};
 
+  if (skipIfTokenInvalid && evictToNvmCache && !token.isValid()) {
+    stats_.evictFailConcurrentFill.inc();
+    return WriteHandle{};
+  }
+
   // We remove the item from both access and mm containers. It doesn't matter
   // if someone else calls remove on the item at this moment, the item cannot
   // be freed as long as we have the moving bit set.

cachelib/allocator/CacheAllocator.h

@@ -1585,15 +1585,6 @@ class CacheAllocator : public CacheBase {
 
   using EvictionIterator = typename MMContainer::Iterator;
 
-  // Try to evict an item.
-  //
-  // @param  mmContainer  the container to look for evictions.
-  // @param  handle       handle to item to evict. eviction will fail if
-  //                      this is not an owning handle.
-  //
-  // @return  whether eviction succeeded
-  bool tryEvictItem(MMContainer& mmContainer, WriteHandle&& handle);
-
   // Deserializer CacheAllocatorMetadata and verify the version
   //
   // @param  deserializer   Deserializer object
@@ -1719,7 +1710,7 @@ class CacheAllocator : public CacheBase {
   //
   // @return last handle for corresponding to item on success. empty handle on
   // failure. caller can retry if needed.
-  WriteHandle evictNormalItemForSlabRelease(Item& item);
+  WriteHandle evictNormalItem(Item& item, bool skipIfTokenInvalid = false);
 
   // Helper function to evict a child item for slab release
   // As a side effect, the parent item is also evicted
@@ -1839,10 +1830,6 @@ class CacheAllocator : public CacheBase {
     return item.getRefCount() == 0;
   }
 
-  static bool itemEvictionPredicate(const Item& item) {
-    return item.getRefCount() == 1 && !item.isMoving();
-  }
-
   static bool itemExpiryPredicate(const Item& item) {
     return item.getRefCount() == 1 && item.isExpired();
   }

cachelib/allocator/CacheItem-inl.h

@@ -219,8 +219,8 @@ bool CacheItem<CacheTrait>::markMoving() noexcept {
 }
 
 template <typename CacheTrait>
-void CacheItem<CacheTrait>::unmarkMoving() noexcept {
-  ref_.unmarkMoving();
+RefcountWithFlags::Value CacheItem<CacheTrait>::unmarkMoving() noexcept {
+  return ref_.unmarkMoving();
 }
 
 template <typename CacheTrait>

cachelib/allocator/CacheItem.h

@@ -357,7 +357,7 @@ class CACHELIB_PACKED_ATTR CacheItem {
    * Unmarking moving does not depend on `isInMMContainer`
    */
   bool markMoving() noexcept;
-  void unmarkMoving() noexcept;
+  RefcountWithFlags::Value unmarkMoving() noexcept;
   bool isMoving() const noexcept;
   bool isOnlyMoving() const noexcept;
 

cachelib/allocator/ChainedHashTable-inl.h

@@ -450,21 +450,6 @@ typename T::Handle ChainedHashTable::Container<T, HookPtr, LockT>::removeIf(
   }
 }
 
-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

@@ -496,17 +496,6 @@ 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

cachelib/allocator/Refcount.h

@@ -247,10 +247,10 @@ class FOLLY_PACK_ATTR RefcountWithFlags {
   /**
    * The following four functions are used to track whether or not
    * an item is currently in the process of being moved. This happens during a
-   * slab rebalance or resize operation.
+   * slab rebalance or resize operation or during eviction.
    *
-   * An item can only be marked moving when `isInMMContainer` returns true.
-   * This operation is atomic.
+   * An item can only be marked moving when `isInMMContainer` returns true and
+   * the item is not yet marked as moving. This operation is atomic.
    *
    * User can also query if an item "isOnlyMoving". This returns true only
    * if the refcount is 0 and only the moving bit is set.
@@ -267,7 +267,8 @@ class FOLLY_PACK_ATTR RefcountWithFlags {
     Value curValue = __atomic_load_n(refPtr, __ATOMIC_RELAXED);
     while (true) {
       const bool flagSet = curValue & conditionBitMask;
-      if (!flagSet) {
+      const bool alreadyMoving = curValue & bitMask;
+      if (!flagSet || alreadyMoving) {
         return false;
       }
 
@@ -286,9 +287,9 @@ class FOLLY_PACK_ATTR RefcountWithFlags {
       }
     }
   }
-  void unmarkMoving() noexcept {
+  Value unmarkMoving() noexcept {
     Value bitMask = ~getAdminRef<kMoving>();
-    __atomic_and_fetch(&refCount_, bitMask, __ATOMIC_ACQ_REL);
+    return __atomic_and_fetch(&refCount_, bitMask, __ATOMIC_ACQ_REL) & kRefMask;
   }
   bool isMoving() const noexcept { return getRaw() & getAdminRef<kMoving>(); }
   bool isOnlyMoving() const noexcept {