Block readers when item is moving

This simplifies moving code a lot.
We don't need to worry about any failure other
than allocation failure since item is protected
from any readers (and by extension, removes).

Author: igchor

cachelib/allocator/CacheAllocator-inl.h

@@ -1389,7 +1389,7 @@ class CacheAllocator : public CacheBase {
   double slabsApproxFreePercentage(TierId tid) const;
 
   // wrapper around Item's refcount and active handle tracking
-  FOLLY_ALWAYS_INLINE bool incRef(Item& it);
+  FOLLY_ALWAYS_INLINE RefcountWithFlags::incResult incRef(Item& it, bool failIfMoving);
   FOLLY_ALWAYS_INLINE RefcountWithFlags::Value decRef(Item& it);
 
   // drops the refcount and if needed, frees the allocation back to the memory
@@ -1622,8 +1622,7 @@ class CacheAllocator : public CacheBase {
   //
   // @return true  If the move was completed, and the containers were updated
   //               successfully.
-  template <typename P>
-  bool moveRegularItemWithSync(Item& oldItem, WriteHandle& newItemHdl, P&& predicate);
+  void moveRegularItemWithSync(Item& oldItem, WriteHandle& newItemHdl);
 
   // Moves a regular item to a different slab. This should only be used during
   // slab release after the item's exclusive bit has been set. The user supplied
@@ -1792,19 +1791,19 @@ class CacheAllocator : public CacheBase {
   //
   // @return valid handle to the item. This will be the last
   //         handle to the item. On failure an empty handle.
-  bool tryEvictToNextMemoryTier(TierId tid, PoolId pid, Item& item, bool fromBgThread);
+  WriteHandle tryEvictToNextMemoryTier(TierId tid, PoolId pid, Item& item, bool fromBgThread);
 
-  bool tryPromoteToNextMemoryTier(TierId tid, PoolId pid, Item& item, bool fromBgThread);
+  WriteHandle tryPromoteToNextMemoryTier(TierId tid, PoolId pid, Item& item, bool fromBgThread);
 
-  bool tryPromoteToNextMemoryTier(Item& item, bool fromBgThread);
+  WriteHandle tryPromoteToNextMemoryTier(Item& item, bool fromBgThread);
 
   // 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. 
-  bool tryEvictToNextMemoryTier(Item& item, bool fromBgThread);
+  WriteHandle tryEvictToNextMemoryTier(Item& item, bool fromBgThread);
 
   size_t memoryTierSize(TierId tid) const;
 
@@ -2235,8 +2234,9 @@ class CacheAllocator : public CacheBase {
     return memoryTierConfigs.size();
   }
 
-  bool addWaitContextForMovingItem(
-      folly::StringPiece key, std::shared_ptr<WaitContext<ReadHandle>> waiter);
+  WriteHandle handleWithWaitContextForMovingItem(Item& item);
+
+  size_t wakeUpWaiters(folly::StringPiece key, WriteHandle&& handle);
 
   class MoveCtx {
    public:
@@ -2261,6 +2261,8 @@ class CacheAllocator : public CacheBase {
       waiters.push_back(std::move(waiter));
     }
 
+    size_t numWaiters() const { return waiters.size(); }
+
    private:
     // notify all pending waiters that are waiting for the fetch.
     void wakeUpWaiters() {

cachelib/allocator/CacheAllocator.h

@@ -237,8 +237,8 @@ bool CacheItem<CacheTrait>::markExclusiveWhenMoving() {
 }
 
 template <typename CacheTrait>
-bool CacheItem<CacheTrait>::markMoving() {
-  return ref_.markMoving();
+bool CacheItem<CacheTrait>::markMoving(bool failIfRefNotZero) {
+  return ref_.markMoving(failIfRefNotZero);
 }
 
 template <typename CacheTrait>
@@ -286,21 +286,6 @@ bool CacheItem<CacheTrait>::isNvmEvicted() const noexcept {
   return ref_.isNvmEvicted();
 }
 
-template <typename CacheTrait>
-void CacheItem<CacheTrait>::markIncomplete() noexcept {
-  ref_.markIncomplete();
-}
-
-template <typename CacheTrait>
-void CacheItem<CacheTrait>::unmarkIncomplete() noexcept {
-  ref_.unmarkIncomplete();
-}
-
-template <typename CacheTrait>
-bool CacheItem<CacheTrait>::isIncomplete() const noexcept {
-  return ref_.isIncomplete();
-}
-
 template <typename CacheTrait>
 void CacheItem<CacheTrait>::markIsChainedItem() noexcept {
   XDCHECK(!hasChainedItem());

cachelib/allocator/CacheItem-inl.h

@@ -242,14 +242,6 @@ class CACHELIB_PACKED_ATTR CacheItem {
   void unmarkNvmEvicted() noexcept;
   bool isNvmEvicted() const noexcept;
 
-  /**
-   * Marks that the item is migrating between memory tiers and
-   * not ready for access now. Accessing thread should wait.
-   */
-  void markIncomplete() noexcept;
-  void unmarkIncomplete() noexcept;
-  bool isIncomplete() const noexcept;
-
   /**
    * Function to set the timestamp for when to expire an item
    *
@@ -318,9 +310,9 @@ class CACHELIB_PACKED_ATTR CacheItem {
   //
   // @return true on success, failure if item is marked as exclusive
   // @throw exception::RefcountOverflow on ref count overflow
-  FOLLY_ALWAYS_INLINE bool incRef() {
+  FOLLY_ALWAYS_INLINE RefcountWithFlags::incResult incRef(bool failIfMoving) {
     try {
-      return ref_.incRef();
+      return ref_.incRef(failIfMoving);
     } catch (exception::RefcountOverflow& e) {
       throw exception::RefcountOverflow(
           folly::sformat("{} item: {}", e.what(), toString()));
@@ -387,7 +379,7 @@ class CACHELIB_PACKED_ATTR CacheItem {
    * Unmarking moving will also return the refcount at the moment of
    * unmarking.
    */
-  bool markMoving();
+  bool markMoving(bool failIfRefNotZero);
   RefcountWithFlags::Value unmarkMoving() noexcept;
   bool isMoving() const noexcept;
   bool isOnlyMoving() const noexcept;

cachelib/allocator/CacheItem.h

@@ -488,13 +488,6 @@ struct ReadHandleImpl {
   // Handle which has the item already
   FOLLY_ALWAYS_INLINE ReadHandleImpl(Item* it, CacheT& alloc) noexcept
       : alloc_(&alloc), it_(it) {
-    if (it_ && it_->isIncomplete()) {
-      waitContext_ = std::make_shared<ItemWaitContext>(alloc);
-      if (!alloc_->addWaitContextForMovingItem(it->getKey(), waitContext_)) {
-        waitContext_->discard();
-        waitContext_.reset();
-      }
-    }
   }
 
   // handle that has a wait context allocated. Used for async handles

cachelib/allocator/Handle.h

@@ -116,10 +116,6 @@ class FOLLY_PACK_ATTR RefcountWithFlags {
     // unevictable in the past.
     kUnevictable_NOOP,
 
-    // Item is accecible but content is not ready yet. Used by eviction
-    // when Item is moved between memory tiers.
-    kIncomplete,
-
     // Unused. This is just to indciate the maximum number of flags
     kFlagMax,
   };
@@ -135,12 +131,18 @@ class FOLLY_PACK_ATTR RefcountWithFlags {
   RefcountWithFlags(RefcountWithFlags&&) = delete;
   RefcountWithFlags& operator=(RefcountWithFlags&&) = delete;
 
+  enum incResult {
+    incOk,
+    incFailedMoving,
+    incFailedExclusive
+  };
+
   // Bumps up the reference count only if the new count will be strictly less
   // than or equal to the maxCount and the item is not exclusive
   // @return true if refcount is bumped. false otherwise (if item is exclusive)
   // @throw  exception::RefcountOverflow if new count would be greater than
   // maxCount
-  FOLLY_ALWAYS_INLINE bool incRef() {
+  FOLLY_ALWAYS_INLINE incResult incRef(bool failIfMoving) {
     Value* const refPtr = &refCount_;
     unsigned int nCASFailures = 0;
     constexpr bool isWeak = false;
@@ -155,12 +157,15 @@ class FOLLY_PACK_ATTR RefcountWithFlags {
         throw exception::RefcountOverflow("Refcount maxed out.");
       }
       if (alreadyExclusive && (oldVal & kAccessRefMask) == 0) {
-        return false;
+        return incFailedExclusive;
+      }
+      else if (alreadyExclusive && failIfMoving) {
+        return incFailedMoving;
       }
 
       if (__atomic_compare_exchange_n(refPtr, &oldVal, newCount, isWeak,
                                       __ATOMIC_ACQ_REL, __ATOMIC_RELAXED)) {
-        return true;
+        return incOk;
       }
 
       if ((++nCASFailures % 4) == 0) {
@@ -327,8 +332,8 @@ class FOLLY_PACK_ATTR RefcountWithFlags {
    *
    * Unmarking moving does not depend on `isInMMContainer`
    */
-  bool markMoving() {
-    auto predicate = [](const Value curValue) {
+  bool markMoving(bool failIfRefNotZero) {
+    auto predicate = [failIfRefNotZero](const Value curValue) {
       Value conditionBitMask = getAdminRef<kLinked>();
       const bool flagSet = curValue & conditionBitMask;
       const bool alreadyExclusive = curValue & getAdminRef<kExclusive>();
@@ -337,6 +342,9 @@ class FOLLY_PACK_ATTR RefcountWithFlags {
       if (!flagSet || alreadyExclusive) {
         return false;
       }
+      if (failIfRefNotZero && (curValue & kAccessRefMask) != 0) {
+        return false;
+      }
       if (UNLIKELY((curValue & kAccessRefMask) == (kAccessRefMask))) {
         throw exception::RefcountOverflow("Refcount maxed out.");
       }
@@ -424,14 +432,6 @@ class FOLLY_PACK_ATTR RefcountWithFlags {
   void unmarkNvmEvicted() noexcept { return unSetFlag<kNvmEvicted>(); }
   bool isNvmEvicted() const noexcept { return isFlagSet<kNvmEvicted>(); }
 
-  /**
-   * Marks that the item is migrating between memory tiers and
-   * not ready for access now. Accessing thread should wait.
-   */
-  void markIncomplete() noexcept { return setFlag<kIncomplete>(); }
-  void unmarkIncomplete() noexcept { return unSetFlag<kIncomplete>(); }
-  bool isIncomplete() const noexcept { return isFlagSet<kIncomplete>(); }
-
   // Whether or not an item is completely drained of access
   // Refcount is 0 and the item is not linked, accessible, nor exclusive
   bool isDrained() const noexcept { return getRefWithAccessAndAdmin() == 0; }

cachelib/allocator/Refcount.h

@@ -41,8 +41,6 @@ struct TestItem {
   void reset() {}
 
   folly::StringPiece getKey() const { return folly::StringPiece(); }
-
-  bool isIncomplete() const { return false; }
 };
 
 struct TestNvmCache;

cachelib/allocator/tests/ItemHandleTest.cpp

@@ -51,7 +51,7 @@ void RefCountTest::testMultiThreaded() {
         nLocalRef--;
         ref.markAccessible();
       } else {
-        ref.incRef();
+        ref.incRef(true);
         nLocalRef++;
         ref.unmarkAccessible();
       }
@@ -100,12 +100,12 @@ void RefCountTest::testBasic() {
   ASSERT_FALSE(ref.template isFlagSet<RefcountWithFlags::Flags::kMMFlag1>());
 
   for (uint32_t i = 0; i < RefcountWithFlags::kAccessRefMask; i++) {
-    ASSERT_TRUE(ref.incRef());
+    ASSERT_TRUE(ref.incRef(true));
   }
 
   // Incrementing past the max will fail
   auto rawRef = ref.getRaw();
-  ASSERT_THROW(ref.incRef(), std::overflow_error);
+  ASSERT_THROW(ref.incRef(true), std::overflow_error);
   ASSERT_EQ(rawRef, ref.getRaw());
 
   // Bumping up access ref shouldn't affect admin ref and flags