Merge pull request #3 from byrnedj/acquire_sergey

Fix for race in acquire with test case

Author: vinser52

cachelib/CMakeLists.txt

@@ -396,4 +396,10 @@ if (BUILD_TESTS)
     RENAME
       Makefile
   )
+  install(
+    FILES
+      ${CACHELIB_HOME}/allocator/tests/ChainedItemParentAcquireAfterMove.gdb
+    DESTINATION
+      ${TESTS_INSTALL_DIR}
+  )
 endif()

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);

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();

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;
 

cachelib/allocator/Refcount.h

@@ -379,6 +379,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();

cachelib/allocator/tests/AllocatorTypeTest.cpp

@@ -292,6 +292,10 @@ TYPED_TEST(BaseAllocatorTest, TransferChainAfterMoving) {
   this->testTransferChainAfterMoving();
 }
 
+TYPED_TEST(BaseAllocatorTest, ChainedItemParentAcquireAfterMove) {
+  this->testChainedItemParentAcquireAfterMove();
+}
+
 TYPED_TEST(BaseAllocatorTest, AddAndPopChainedItemMultithread) {
   this->testAddAndPopChainedItemMultithread();
 }

cachelib/allocator/tests/BaseAllocatorTest.h

@@ -19,6 +19,12 @@
 #include <folly/Random.h>
 #include <folly/Singleton.h>
 #include <folly/synchronization/Baton.h>
+#include <folly/synchronization/Latch.h>
+
+#include <fcntl.h>
+#include <unistd.h>
+#include <ctype.h>
+#include <semaphore.h>
 
 #include <algorithm>
 #include <chrono>
@@ -5189,6 +5195,152 @@ class BaseAllocatorTest : public AllocatorTest<AllocatorT> {
     }
   }
 
+  // helper function to check if gdb is running and attached to test
+  bool isGdbAttached() {
+    char buf[4096];
+    const int status_fd = open("/proc/self/status", O_RDONLY);
+    if (status_fd == -1) {
+      return false;
+    }
+    const ssize_t num_read = read(status_fd, buf, sizeof(buf) - 1);
+    close(status_fd);
+  
+    if (num_read <= 0) {
+      return false;
+    }
+    buf[num_read] = '\0';
+    constexpr char tracerPidString[] = "TracerPid:";
+    const auto tracer_pid_ptr = strstr(buf, tracerPidString);
+    if (!tracer_pid_ptr) {
+      return false;
+    }
+    for (const char* characterPtr = tracer_pid_ptr + sizeof(tracerPidString) - 1; 
+            characterPtr <= buf + num_read; ++characterPtr) {
+      if (isspace(*characterPtr)) {
+        continue;
+      } else {
+        return isdigit(*characterPtr) != 0 && *characterPtr != '0';
+      }
+    }
+    return false;
+  }
+
+  void gdb_sync1() { sleep(1); }
+  void gdb_sync2() { sleep(1); }
+  void gdb_sync3() { sleep(1); }
+  void gdb_sync4() { sleep(1); }
+  void testChainedItemParentAcquireAfterMove() {
+    sem_unlink("/gdb1_sem");
+    sem_t *sem = sem_open("/gdb1_sem", O_CREAT | O_EXCL, S_IRUSR | S_IWUSR, 0);
+    int ppid = getpid(); //parent pid
+    int fpid = fork();
+    if (fpid == 0) {
+        sem_wait(sem);
+        sem_close(sem);
+        sem_unlink("/gdb1_sem");
+        char cmdpid[256];
+        sprintf(cmdpid,"%d",ppid);
+        int f = execlp("gdb","gdb","--pid",cmdpid,
+                "--batch-silent","--command=ChainedItemParentAcquireAfterMove.gdb",(char*) 0);
+        ASSERT(f != -1);
+    }
+    
+    // create an allocator worth 4 slabs 
+    // first slab is for overhead, second is parent class
+    // thrid is chained item 1 and fourth is for new chained item alloc
+    // to move to.
+    typename AllocatorT::Config config;
+    config.configureChainedItems();
+
+    config.setCacheSize(4 * Slab::kSize);
+
+    using Item = typename AllocatorT::Item;
+
+    std::atomic<uint64_t> numRemovedKeys{0};
+    config.setRemoveCallback(
+        [&](const typename AllocatorT::RemoveCbData&) { ++numRemovedKeys; });
+
+    std::atomic<uint64_t> numMoves{0};
+    config.enableMovingOnSlabRelease(
+        [&](Item& oldItem, Item& newItem, Item* /* parentPtr */) {
+          assert(oldItem.getSize() == newItem.getSize());
+          assert(oldItem.isChainedItem());
+          std::memcpy(newItem.getMemory(), oldItem.getMemory(),
+                      oldItem.getSize());
+          ++numMoves;
+        }
+    );
+
+    AllocatorT alloc(config);
+    const size_t numBytes = alloc.getCacheMemoryStats().ramCacheSize;
+    const auto poolSize = numBytes;
+    //items are 100K and 200K
+    const std::set<uint32_t> allocSizes = {2*1024*1024 + 1024 , 3*1024*1024};
+    const auto pid = alloc.addPool("one", poolSize, allocSizes);
+
+    // Allocate 1 parent items and for each parent item, 1 chained
+    // allocation
+    auto allocFn = [&](std::string keyPrefix, std::vector<uint32_t> sizes) {
+      for (unsigned int loop = 0; loop < 1; ++loop) {
+        std::vector<uint8_t*> bufList;
+        std::vector<typename AllocatorT::WriteHandle> parentHandles;
+        for (unsigned int i = 0; i < 1; ++i) {
+          const auto key = keyPrefix + folly::to<std::string>(loop) + "_" +
+                           folly::to<std::string>(i);
+
+          auto itemHandle = util::allocateAccessible(alloc, pid, key, sizes[0]);
+
+          for (unsigned int j = 0; j < 1; ++j) {
+            auto childItem =
+                alloc.allocateChainedItem(itemHandle, sizes[1]);
+            ASSERT_NE(nullptr, childItem);
+
+            uint8_t* buf = reinterpret_cast<uint8_t*>(childItem->getMemory());
+            // write first 50 bytes here
+            for (uint8_t k = 0; k < 50; ++k) {
+              buf[k] = k;
+            }
+            bufList.push_back(buf);
+
+            alloc.addChainedItem(itemHandle, std::move(childItem));
+          }
+        }
+      }
+    };
+    auto sizes = std::vector<uint32_t>{2*1024*1024, 3*1024*1024 - 1000};
+    allocFn(std::string{"yolo"}, sizes);
+
+    sem_post(sem);
+    while (!isGdbAttached()) {
+      std::this_thread::sleep_for(std::chrono::milliseconds(1));
+    }
+    //need to suspend thread 1 - who is doing the eviction
+    //gdb will do this for us
+    folly::Latch latch(1);
+    std::unique_ptr<std::thread> r = std::make_unique<std::thread>([&](){
+      ClassId cid = static_cast<ClassId>(1);
+      //gdb will interupt this thread when we hit acquire
+      //we will then switch to the other thread and attempt 
+      //to evict the parent
+      gdb_sync1();
+      latch.count_down();
+      alloc.releaseSlab(pid, cid, SlabReleaseMode::kRebalance);
+      gdb_sync3();
+    });
+    r->detach();
+    
+    latch.wait();
+    //this alloc should fail because item can't be marked for eviction
+    //and we exceed evict attempts
+    auto itemHandle = util::allocateAccessible(alloc, pid, std::to_string(10), sizes[0]);
+    EXPECT_EQ(itemHandle, nullptr);
+    //complete the move
+    gdb_sync2();
+    auto itemHandle2 = util::allocateAccessible(alloc, pid, std::to_string(11), sizes[0]);
+    EXPECT_NE(itemHandle2, nullptr);
+    gdb_sync4();
+
+  }
   // Test stats on counting chained items
   // 1. Alloc an item and several chained items
   //    * Before inserting chained items, make sure count is zero

cachelib/allocator/tests/ChainedItemParentAcquireAfterMove.gdb

@@ -0,0 +1,22 @@
+break gdb_sync1
+c
+set scheduler-locking on
+break acquire
+c
+c
+c
+c
+thread 1
+set scheduler-locking on
+del 2
+break gdb_sync2
+c
+thread 4
+set scheduler-locking on
+break gdb_sync3
+c
+thread 1
+set scheduler-locking off
+break gdb_sync4
+c
+c