fix race for acquire parent handle on chained items

facebook · Oct 27, 2023 · 807dbda · 807dbda
1 parent 468ea4f
commit 807dbda
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 14 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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();