wakeUpWaiters in SlabRelease code (based on upstream)

Author: igchor

cachelib/allocator/CacheAllocator-inl.h

@@ -1237,7 +1237,7 @@ CacheAllocator<CacheTrait>::handleWithWaitContextForMovingItem(Item& item) {
 }
 
 template <typename CacheTrait>
-size_t CacheAllocator<CacheTrait>::wakeUpWaiters(folly::StringPiece key,
+size_t CacheAllocator<CacheTrait>::wakeUpWaitersLocked(folly::StringPiece key,
   WriteHandle&& handle) {
   std::unique_ptr<MoveCtx> ctx;
   auto shard = getShardForKey(key);
@@ -1606,7 +1606,7 @@ CacheAllocator<CacheTrait>::findEviction(TierId tid, PoolId pid, ClassId cid) {
       // wake up any readers that wait for the move to complete
       // it's safe to do now, as we have the item marked exclusive and
       // no other reader can be added to the waiters list
-      wakeUpWaiters(candidate->getKey(), WriteHandle{});
+      wakeUpWaiters(*candidate, {});
 
       if (token.isValid() && shouldWriteToNvmCacheExclusive(*candidate)) {
         nvmCache_->put(*candidate, std::move(token));
@@ -1617,7 +1617,7 @@ CacheAllocator<CacheTrait>::findEviction(TierId tid, PoolId pid, ClassId cid) {
       XDCHECK(!candidate->isAccessible());
       XDCHECK(candidate->getKey() == evictedToNext->getKey());
 
-      wakeUpWaiters(candidate->getKey(), std::move(evictedToNext));
+      wakeUpWaiters(*candidate, std::move(evictedToNext));
     }
 
     XDCHECK(!candidate->isMarkedForEviction() && !candidate->isMoving());
@@ -1706,7 +1706,7 @@ CacheAllocator<CacheTrait>::tryEvictToNextMemoryTier(
   if(item.isExpired()) {
     accessContainer_->remove(item);
     item.unmarkMoving();
-    return acquire(&item); // TODO: wakeUpWaiters with null handle?
+    return acquire(&item);
   }
 
   TierId nextTier = tid; // TODO - calculate this based on some admission policy
@@ -2881,6 +2881,14 @@ void CacheAllocator<CacheTrait>::throttleWith(util::Throttler& t,
   }
 }
 
+template <typename CacheTrait>
+typename RefcountWithFlags::Value CacheAllocator<CacheTrait>::unmarkMovingAndWakeUpWaiters(Item &item, WriteHandle handle)
+{
+  auto ret = item.unmarkMoving();
+  wakeUpWaiters(item, std::move(handle));
+  return ret;
+}
+
 template <typename CacheTrait>
 bool CacheAllocator<CacheTrait>::moveForSlabRelease(
     const SlabReleaseContext& ctx, Item& oldItem, util::Throttler& throttler) {
@@ -2899,7 +2907,8 @@ bool CacheAllocator<CacheTrait>::moveForSlabRelease(
 
     // Nothing to move and the key is likely also bogus for chained items.
     if (oldItem.isOnlyMoving()) {
-      oldItem.unmarkMoving();
+      auto ret = unmarkMovingAndWakeUpWaiters(oldItem, {});
+      XDCHECK(ret == 0);
       const auto res =
           releaseBackToAllocator(oldItem, RemoveContext::kNormal, false);
       XDCHECK(res == ReleaseRes::kReleased);
@@ -2947,8 +2956,9 @@ bool CacheAllocator<CacheTrait>::moveForSlabRelease(
     });
   }
   auto tid = getTierId(oldItem);
-  auto ref = oldItem.unmarkMoving();
-  XDCHECK_EQ(ref, 0);
+  auto ref = unmarkMovingAndWakeUpWaiters(oldItem, std::move(newItemHdl));
+  XDCHECK(ref == 0);
+
   const auto allocInfo = allocator_[tid]->getAllocInfo(oldItem.getMemory());
   allocator_[tid]->free(&oldItem);
 
@@ -3066,9 +3076,26 @@ bool CacheAllocator<CacheTrait>::tryMovingForSlabRelease(
     }
   }
 
-  return oldItem.isChainedItem()
-             ? moveChainedItem(oldItem.asChainedItem(), newItemHdl)
-             : moveRegularItem(oldItem, newItemHdl);
+  // TODO: we can unify move*Item and move*ItemWithSync by always
+  // using the moving bit to block readers.
+  if (getNumTiers() == 1) {
+    return oldItem.isChainedItem()
+              ? moveChainedItem(oldItem.asChainedItem(), newItemHdl)
+              : moveRegularItem(oldItem, newItemHdl);
+  } else {
+    // TODO: add support for chained items
+    moveRegularItemWithSync(oldItem, newItemHdl);
+    return true;
+  }
+}
+
+template <typename CacheTrait>
+void CacheAllocator<CacheTrait>::wakeUpWaiters(Item& item, WriteHandle handle)
+{
+  // readers do not block on 'moving' items in case there is only one tier
+  if (getNumTiers() > 1) {
+    wakeUpWaitersLocked(item.getKey(), std::move(handle));
+  }
 }
 
 template <typename CacheTrait>
@@ -3081,7 +3108,7 @@ void CacheAllocator<CacheTrait>::evictForSlabRelease(
     stats_.numEvictionAttempts.inc();
 
     if (shutDownInProgress_) {
-      item.unmarkMoving();
+      auto ref = unmarkMovingAndWakeUpWaiters(item, {});
       allocator_[getTierId(item)]->abortSlabRelease(ctx);
       throw exception::SlabReleaseAborted(
           folly::sformat("Slab Release aborted while trying to evict"
@@ -3105,7 +3132,8 @@ void CacheAllocator<CacheTrait>::evictForSlabRelease(
     // nothing needs to be done. We simply need to call unmarkMoving and free
     // the item.
     if (item.isOnlyMoving()) {
-      item.unmarkMoving();
+      auto ref = unmarkMovingAndWakeUpWaiters(item, {});
+      XDCHECK(ref == 0);
       const auto res =
           releaseBackToAllocator(item, RemoveContext::kNormal, false);
       XDCHECK(ReleaseRes::kReleased == res);
@@ -3162,12 +3190,10 @@ void CacheAllocator<CacheTrait>::evictForSlabRelease(
         // unmark the child so it will be freed
         item.unmarkMoving();
         unlinkItemForEviction(*evicted);
-        if (getNumTiers() > 1) {
-            // wake up any readers that wait for the move to complete
-            // it's safe to do now, as we have the item marked exclusive and
-            // no other reader can be added to the waiters list
-            wakeUpWaiters(evicted->getKey(), WriteHandle{});
-        }
+        // wake up any readers that wait for the move to complete
+        // it's safe to do now, as we have the item marked exclusive and
+        // no other reader can be added to the waiters list
+        wakeUpWaiters(*evicted, {});
       } else {
         continue;
       }
@@ -3177,12 +3203,7 @@ void CacheAllocator<CacheTrait>::evictForSlabRelease(
       token = createPutToken(*evicted);
       if (evicted->markForEvictionWhenMoving()) {
         unlinkItemForEviction(*evicted);
-        if (getNumTiers() > 1) {
-      	  // wake up any readers that wait for the move to complete
-      	  // it's safe to do now, as we have the item marked exclusive and
-      	  // no other reader can be added to the waiters list
-      	  wakeUpWaiters(evicted->getKey(), WriteHandle{});
-        }
+        wakeUpWaiters(*evicted, {});
       } else {
         continue;
       }

cachelib/allocator/CacheAllocator.h

@@ -1725,6 +1725,18 @@ class CacheAllocator : public CacheBase {
   //         handle to the item. On failure an empty handle.
   WriteHandle tryEvictToNextMemoryTier(TierId tid, PoolId pid, Item& item); 
 
+  // Wakes up waiters if there are any
+  //
+  // @param item    wakes waiters that are waiting on that item
+  // @param handle  handle to pass to the waiters
+  void wakeUpWaiters(Item& item, WriteHandle handle);
+
+  // Unmarks item as moving and wakes up any waiters waiting on that item
+  //
+  // @param item    wakes waiters that are waiting on that item
+  // @param handle  handle to pass to the waiters
+  typename RefcountWithFlags::Value unmarkMovingAndWakeUpWaiters(Item &item, WriteHandle handle);
+
   // Try to move the item down to the next memory tier
   //
   // @param item the item to evict
@@ -2042,7 +2054,7 @@ class CacheAllocator : public CacheBase {
 
   WriteHandle handleWithWaitContextForMovingItem(Item& item);
 
-  size_t wakeUpWaiters(folly::StringPiece key, WriteHandle&& handle);
+  size_t wakeUpWaitersLocked(folly::StringPiece key, WriteHandle&& handle);
 
   class MoveCtx {
    public: