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

@@ -1186,14 +1186,13 @@ bool CacheAllocator<CacheTrait>::addWaitContextForMovingItem(
 }
 
 template <typename CacheTrait>
-template <typename ItemPtr>
 typename CacheAllocator<CacheTrait>::ItemHandle
 CacheAllocator<CacheTrait>::moveRegularItemOnEviction(
-    ItemPtr& oldItemPtr, ItemHandle& newItemHdl) {
+    Item& oldItem, ItemHandle& newItemHdl) {
+  XDCHECK(oldItem.isMoving());
   // TODO: should we introduce new latency tracker. E.g. evictRegularLatency_
   // ??? util::LatencyTracker tracker{stats_.evictRegularLatency_};
 
-  Item& oldItem = *oldItemPtr;
   if (!oldItem.isAccessible() || oldItem.isExpired()) {
     return {};
   }
@@ -1249,7 +1248,7 @@ CacheAllocator<CacheTrait>::moveRegularItemOnEviction(
   // it is unsafe to replace the old item with a new one, so we should
   // also abort.
   if (!accessContainer_->replaceIf(oldItem, *newItemHdl,
-                                   itemEvictionPredicate)) {
+                                   itemMovingPredicate)) {
     return {};
   }
 
@@ -1269,7 +1268,7 @@ CacheAllocator<CacheTrait>::moveRegularItemOnEviction(
 
   // Inside the MM container's lock, this checks if the old item exists to
   // make sure that no other thread removed it, and only then replaces it.
-  if (!replaceInMMContainer(oldItemPtr, *newItemHdl)) {
+  if (!replaceInMMContainer(oldItem, *newItemHdl)) {
     accessContainer_->remove(*newItemHdl);
     return {};
   }
@@ -1459,20 +1458,15 @@ CacheAllocator<CacheTrait>::findEviction(TierId tid, PoolId pid, ClassId cid) {
          itr) {
     ++searchTries;
 
-    Item* candidate = itr.get();
-
-    // make sure no other thead is evicting the item
-    if (candidate->getRefCount() != 0) {
-      ++itr;
-      continue;
-    }
+    Item* toRecycle = itr.get();
 
-    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;
     }
@@ -1482,52 +1476,57 @@ CacheAllocator<CacheTrait>::findEviction(TierId tid, 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 movedToNextTier = false;
-    bool evicted = false;
-    {
-      auto h = tryEvictToNextMemoryTier(tid, pid, candidate);
-      if (h) {
-        movedToNextTier = true;
-        evicted = true;
-        h.release();
-        toReleaseHandle.release();
-
-        decRef(*candidate);
-      }
-    }
-    if(!movedToNextTier) {
-      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 {
         (*stats_.regularItemEvictions)[pid][cid].inc();
       }
+    } 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());
 
-      XDCHECK_EQ(1u, candidate->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 */ movedToNextTier, candidate)) {
-        return candidate;
+                                 /* isNascent */ false, toRecycle)) {
+        return toRecycle;
       }
     }
 
-    // If we destroyed the itr to possibly evict and failed, we restart
-    // from the beginning again
-    if (!itr) {
-      itr.resetToBegin();
-      for (int i = 0; i < searchTries; i++)
-        ++itr;
-    }
+    itr.resetToBegin();
   }
   return nullptr;
 }
@@ -1581,24 +1580,23 @@ bool CacheAllocator<CacheTrait>::shouldWriteToNvmCacheExclusive(
 }
 
 template <typename CacheTrait>
-template <typename ItemPtr>
 typename CacheAllocator<CacheTrait>::WriteHandle
 CacheAllocator<CacheTrait>::tryEvictToNextMemoryTier(
-    TierId tid, PoolId pid, ItemPtr& item) {
-  if(item->isChainedItem()) return {}; // TODO: We do not support ChainedItem yet
-  if(item->isExpired()) return acquire(item);
+    TierId tid, PoolId pid, Item& item) {
+  if(item.isChainedItem()) return {}; // TODO: We do not support ChainedItem yet
+  if(item.isExpired()) return acquire(&item);
 
   TierId nextTier = tid; // TODO - calculate this based on some admission policy
   while (++nextTier < numTiers_) { // try to evict down to the next memory tiers
     // allocateInternal might trigger another eviction
     auto newItemHdl = allocateInternalTier(nextTier, pid,
-                     item->getKey(),
-                     item->getSize(),
-                     item->getCreationTime(),
-                     item->getExpiryTime());
+                     item.getKey(),
+                     item.getSize(),
+                     item.getCreationTime(),
+                     item.getExpiryTime());
 
     if (newItemHdl) {
-      XDCHECK_EQ(newItemHdl->getSize(), item->getSize());
+      XDCHECK_EQ(newItemHdl->getSize(), item.getSize());
 
       return moveRegularItemOnEviction(item, newItemHdl);
     }
@@ -1609,83 +1607,10 @@ CacheAllocator<CacheTrait>::tryEvictToNextMemoryTier(
 
 template <typename CacheTrait>
 typename CacheAllocator<CacheTrait>::WriteHandle
-CacheAllocator<CacheTrait>::tryEvictToNextMemoryTier(Item* item) {
-  auto tid = getTierId(*item);
-  auto pid = allocator_[tid]->getAllocInfo(item->getMemory()).poolId;
-  return tryEvictToNextMemoryTier(tid, pid, item);
-}
-
-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.release();
-    decRef(item);
-    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.release();
-  decRef(item);
-
-  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;
+CacheAllocator<CacheTrait>::tryEvictToNextMemoryTier(Item& item) {
+    auto tid = getTierId(item);
+    auto pid = allocator_[tid]->getAllocInfo(item.getMemory()).poolId;
+    return tryEvictToNextMemoryTier(tid, pid, item);
 }
 
 template <typename CacheTrait>
@@ -2885,7 +2810,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.
@@ -2942,14 +2867,15 @@ void CacheAllocator<CacheTrait>::evictForSlabRelease(
 
 template <typename CacheTrait>
 typename CacheAllocator<CacheTrait>::ItemHandle
-CacheAllocator<CacheTrait>::evictNormalItemForSlabRelease(Item& item) {
+CacheAllocator<CacheTrait>::evictNormalItem(Item& item,
+                                            bool skipIfTokenInvalid) {
   XDCHECK(item.isMoving());
 
   if (item.isOnlyMoving()) {
     return ItemHandle{};
   }
 
-  auto evictHandle = tryEvictToNextMemoryTier(&item);
+  auto evictHandle = tryEvictToNextMemoryTier(item);
   if(evictHandle) return evictHandle;
 
   auto predicate = [](const Item& it) { return it.getRefCount() == 0; };
@@ -2958,6 +2884,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 ItemHandle{};
+  }
+
   // 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

@@ -1401,8 +1401,7 @@ class CacheAllocator : public CacheBase {
   //
   // @return true  If the move was completed, and the containers were updated
   //               successfully.
-  template <typename ItemPtr>
-  ItemHandle moveRegularItemOnEviction(ItemPtr& oldItem, ItemHandle& newItemHdl);
+  ItemHandle moveRegularItemOnEviction(Item& oldItem, ItemHandle& newItemHdl);
 
   // Moves a regular item to a different slab. This should only be used during
   // slab release after the item's moving bit has been set. The user supplied
@@ -1561,14 +1560,6 @@ class CacheAllocator : public CacheBase {
   // @return An evicted item or nullptr  if there is no suitable candidate.
   Item* findEviction(TierId tid, PoolId pid, ClassId cid);
 
-  // Try to evict a regular item.
-  //
-  // @param  mmContainer  the container to look for evictions.
-  // @param  item         item to evict
-  //
-  // @return whther eviction succeeded
-  bool tryEvictItem(MMContainer& mmContainer, WriteHandle&& handle);
-
   // Try to move the item down to the next memory tier
   //
   // @param tid current tier ID of the item
@@ -1577,16 +1568,15 @@ class CacheAllocator : public CacheBase {
   //
   // @return valid handle to the item. This will be the last
   //         handle to the item. On failure an empty handle.
-  template <typename ItemPtr>
-  WriteHandle tryEvictToNextMemoryTier(TierId tid, PoolId pid, ItemPtr& item);
+  WriteHandle tryEvictToNextMemoryTier(TierId tid, PoolId pid, Item& item);
 
   // Try to move the item down to the next memory tier
   //
   // @param item the item to evict
   //
   // @return valid handle to the item. This will be the last
   //         handle to the item. On failure an empty handle. 
-  WriteHandle tryEvictToNextMemoryTier(Item* item);
+  WriteHandle tryEvictToNextMemoryTier(Item& item);
 
   // Deserializer CacheAllocatorMetadata and verify the version
   //
@@ -1713,7 +1703,7 @@ class CacheAllocator : public CacheBase {
   //
   // @return last handle for corresponding to item on success. empty handle on
   // failure. caller can retry if needed.
-  ItemHandle evictNormalItemForSlabRelease(Item& item);
+  ItemHandle 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
@@ -1834,10 +1824,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

@@ -229,8 +229,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

@@ -378,7 +378,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;