From 807dbda2ac9376065c0bde41131e672a063a386a Mon Sep 17 00:00:00 2001
From: Daniel Byrne <byrnedj12@gmail.com>
Date: Wed, 25 Oct 2023 05:13:12 -0700
Subject: [PATCH] fix race for acquire parent handle on chained items

---
 cachelib/allocator/CacheAllocator-inl.h | 20 +++++------------
 cachelib/allocator/CacheItem-inl.h      |  5 +++++
 cachelib/allocator/CacheItem.h          |  1 +
 cachelib/allocator/Refcount.h           | 29 +++++++++++++++++++++++++
 4 files changed, 41 insertions(+), 14 deletions(-)

diff --git a/cachelib/allocator/CacheAllocator-inl.h b/cachelib/allocator/CacheAllocator-inl.h
index f9dfd24a30..4e36cce579 100644
--- a/cachelib/allocator/CacheAllocator-inl.h
+++ b/cachelib/allocator/CacheAllocator-inl.h
@@ -2639,20 +2639,12 @@ bool CacheAllocator<CacheTrait>::moveForSlabRelease(
   const auto allocInfo = allocator_->getAllocInfo(oldItem.getMemory());
   if (chainedItem) {
     newItemHdl.reset();
-    auto ref = parentItem->unmarkMoving();
-    if (UNLIKELY(ref == 0)) {
-      wakeUpWaiters(*parentItem, {});
-      const auto res =
-          releaseBackToAllocator(*parentItem, RemoveContext::kNormal, false);
-      XDCHECK(res == ReleaseRes::kReleased);
-      return true;
-    } else {
-      XDCHECK_NE(ref, 0);
-      auto parentHdl = acquire(parentItem);
-      if (parentHdl) {
-        wakeUpWaiters(*parentItem, std::move(parentHdl));
-      }
-    }
+    auto ref = parentItem->unmarkMovingAndIncRef();
+    XDCHECK_NE(ref, 0);
+    auto parentHdl = acquire(parentItem);
+    XDCHECK(parentHdl);
+    parentHdl->decRef();
+    wakeUpWaiters(*parentItem, std::move(parentHdl));
   } else {
     auto ref = unmarkMovingAndWakeUpWaiters(oldItem, std::move(newItemHdl));
     XDCHECK(ref == 0);
diff --git a/cachelib/allocator/CacheItem-inl.h b/cachelib/allocator/CacheItem-inl.h
index bf77b43aa5..24500e634b 100644
--- a/cachelib/allocator/CacheItem-inl.h
+++ b/cachelib/allocator/CacheItem-inl.h
@@ -247,6 +247,11 @@ RefcountWithFlags::Value CacheItem<CacheTrait>::unmarkMoving() noexcept {
   return ref_.unmarkMoving();
 }
 
+template <typename CacheTrait>
+RefcountWithFlags::Value CacheItem<CacheTrait>::unmarkMovingAndIncRef() noexcept {
+  return ref_.unmarkMovingAndIncRef();
+}
+
 template <typename CacheTrait>
 bool CacheItem<CacheTrait>::isMoving() const noexcept {
   return ref_.isMoving();
diff --git a/cachelib/allocator/CacheItem.h b/cachelib/allocator/CacheItem.h
index 4c32ece794..d89d7746dd 100644
--- a/cachelib/allocator/CacheItem.h
+++ b/cachelib/allocator/CacheItem.h
@@ -380,6 +380,7 @@ class CACHELIB_PACKED_ATTR CacheItem {
    */
   bool markMoving();
   RefcountWithFlags::Value unmarkMoving() noexcept;
+  RefcountWithFlags::Value unmarkMovingAndIncRef() noexcept;
   bool isMoving() const noexcept;
   bool isOnlyMoving() const noexcept;
 
diff --git a/cachelib/allocator/Refcount.h b/cachelib/allocator/Refcount.h
index 1288dd1ed3..06a1d4c4b4 100644
--- a/cachelib/allocator/Refcount.h
+++ b/cachelib/allocator/Refcount.h
@@ -384,6 +384,35 @@ class FOLLY_PACK_ATTR RefcountWithFlags {
 
     return retValue & kRefMask;
   }
+ 
+  /*
+   * this is used when we immediately call acquire after unmarking
+   * moving - this is currently done in the case of moving a
+   * chained item when the parent is unmarked moving and we
+   * need to wake up the waiters with the parent handle BUT
+   * we don't want the parent item to be marked moving/exclusive 
+   * between unmarking moving and acquire - so we do not
+   * modify the refcount (moving state = exclusive bit set and refcount == 1)
+   */
+  Value unmarkMovingAndIncRef() noexcept {
+    XDCHECK(isMoving());
+    auto predicate = [](const Value curValue) {
+      XDCHECK((curValue & kAccessRefMask) != 0);
+      return true;
+    };
+
+    Value retValue;
+    auto newValue = [&retValue](const Value curValue) {
+      retValue =
+          (curValue) & ~getAdminRef<kExclusive>();
+      return retValue;
+    };
+
+    auto updated = atomicUpdateValue(predicate, newValue);
+    XDCHECK(updated);
+
+    return retValue & kRefMask;
+  }
 
   bool isMoving() const noexcept {
     auto raw = getRaw();