fix race in moveRegularItemWith sync where insertOrReplace can cause move to fail

byrnedj · byrnedj · commit 08bf0b4e8510 · 2023-03-03T13:39:04.000-08:00
- updated slab release logic for move failure, but there is still an issue
   with slab movement. currently investigating.
diff --git a/cachelib/allocator/CacheAllocator-inl.h b/cachelib/allocator/CacheAllocator-inl.h
@@ -1294,8 +1294,21 @@ size_t CacheAllocator<CacheTrait>::wakeUpWaitersLocked(folly::StringPiece key,
 }
 
 template <typename CacheTrait>
-void CacheAllocator<CacheTrait>::moveRegularItemWithSync(
+bool CacheAllocator<CacheTrait>::moveRegularItemWithSync(
     Item& oldItem, WriteHandle& newItemHdl) {
+  //on function exit - the new item handle is no longer moving
+  //and other threads may access it - but in case where
+  //we failed to replace in access container we can give the
+  //new item back to the allocator
+  auto guard = folly::makeGuard([&]() {
+    auto ref = newItemHdl->unmarkMoving();
+    if (UNLIKELY(ref == 0)) {
+      const auto res =
+          releaseBackToAllocator(*newItemHdl, RemoveContext::kNormal, false);
+      XDCHECK(res == ReleaseRes::kReleased);
+    }
+  });
+
   XDCHECK(oldItem.isMoving());
   XDCHECK(!oldItem.isExpired());
   // TODO: should we introduce new latency tracker. E.g. evictRegularLatency_
@@ -1326,6 +1339,22 @@ void CacheAllocator<CacheTrait>::moveRegularItemWithSync(
 
   auto replaced = accessContainer_->replaceIf(oldItem, *newItemHdl,
                                    predicate);
+  // another thread may have called insertOrReplace which could have
+  // marked this item as unaccessible causing the replaceIf
+  // in the access container to fail - in this case we want
+  // to abort the move since the item is no longer valid
+  if (!replaced) {
+      return false;
+  }
+  // what if another thread calls insertOrReplace now when
+  // the item is moving and already replaced in the hash table?
+  // 1. it succeeds in updating the hash table - so there is
+  //    no guarentee that isAccessible() is true
+  // 2. it will then try to remove from MM container
+  //     - this operation will wait for newItemHdl to
+  //       be unmarkedMoving via the waitContext
+  // 3. replaced handle is returned and eventually drops
+  //    ref to 0 and the item is recycled back to allocator.
 
   if (config_.moveCb) {
     // Execute the move callback. We cannot make any guarantees about the
@@ -1367,14 +1396,7 @@ void CacheAllocator<CacheTrait>::moveRegularItemWithSync(
     XDCHECK(newItemHdl->hasChainedItem());
   }
   newItemHdl.unmarkNascent();
-  auto ref = newItemHdl->unmarkMoving();
-  //remove because there is a chance the new item was not
-  //added to the access container
-  if (UNLIKELY(ref == 0)) {
-    const auto res =
-        releaseBackToAllocator(*newItemHdl, RemoveContext::kNormal, false);
-    XDCHECK(res == ReleaseRes::kReleased);
-  }
+  return true;
 }
 
 template <typename CacheTrait>
@@ -1529,7 +1551,6 @@ template <typename CacheTrait>
 void CacheAllocator<CacheTrait>::unlinkItemForEviction(Item& it) {
   XDCHECK(it.isMarkedForEviction());
   XDCHECK(it.getRefCount() == 0);
-
   accessContainer_->remove(it);
   removeFromMMContainer(it);
 
@@ -1624,28 +1645,43 @@ CacheAllocator<CacheTrait>::findEviction(TierId tid, PoolId pid, ClassId cid) {
     auto evictedToNext = lastTier ? nullptr
         : tryEvictToNextMemoryTier(*candidate, false);
     if (!evictedToNext) {
-      if (!token.isValid()) {
+      //if insertOrReplace was called during move
+      //then candidate will not be accessible (failed replace during tryEvict)
+      // - therefore this was why we failed to
+      //   evict to the next tier and insertOrReplace
+      //   will remove from NVM cache
+      //however, if candidate is accessible
+      //that means the allocation in the next
+      //tier failed - so we will continue to
+      //evict the item to NVM cache
+      bool failedToReplace = !candidate->isAccessible();
+      if (!token.isValid() && !failedToReplace) {
         token = createPutToken(*candidate);
       }
-      // tryEvictToNextMemoryTier should only fail if allocation of the new item fails
-      // in that case, it should be still possible to mark item as exclusive.
+      // tryEvictToNextMemoryTier can fail if:
+      //    a) allocation of the new item fails in that case,
+      //       it should be still possible to mark item for eviction.
+      //    b) another thread calls insertOrReplace and the item
+      //       is no longer accessible
       //
       // in case that we are on the last tier, we whould have already marked
       // as exclusive since we will not be moving the item to the next tier
       // but rather just evicting all together, no need to
-      // markExclusiveWhenMoving
+      // markForEvictionWhenMoving
       auto ret = lastTier ? true : candidate->markForEvictionWhenMoving();
       XDCHECK(ret);
 
       unlinkItemForEviction(*candidate);
+      
+      if (token.isValid() && shouldWriteToNvmCacheExclusive(*candidate)
+              && !failedToReplace) {
+        nvmCache_->put(*candidate, std::move(token));
+      }
       // 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, {});
 
-      if (token.isValid() && shouldWriteToNvmCacheExclusive(*candidate)) {
-        nvmCache_->put(*candidate, std::move(token));
-      }
     } else {
       XDCHECK(!evictedToNext->isMarkedForEviction() && !evictedToNext->isMoving());
       XDCHECK(!candidate->isMarkedForEviction() && !candidate->isMoving());
@@ -1756,7 +1792,10 @@ CacheAllocator<CacheTrait>::tryEvictToNextMemoryTier(
 
     if (newItemHdl) {
       XDCHECK_EQ(newItemHdl->getSize(), item.getSize());
-      moveRegularItemWithSync(item, newItemHdl);
+      if (!moveRegularItemWithSync(item, newItemHdl)) {
+          return WriteHandle{};
+      }
+      XDCHECK_EQ(newItemHdl->getKey(),item.getKey());
       item.unmarkMoving();
       return newItemHdl;
     } else {
@@ -1795,7 +1834,9 @@ CacheAllocator<CacheTrait>::tryPromoteToNextMemoryTier(
 
     if (newItemHdl) {
       XDCHECK_EQ(newItemHdl->getSize(), item.getSize());
-      moveRegularItemWithSync(item, newItemHdl);
+      if (!moveRegularItemWithSync(item, newItemHdl)) {
+          return WriteHandle{};
+      }
       item.unmarkMoving();
       return newItemHdl;
     } else {
@@ -3148,9 +3189,23 @@ bool CacheAllocator<CacheTrait>::tryMovingForSlabRelease(
       // TODO: add support for chained items
       return false;
     } else {
-      moveRegularItemWithSync(oldItem, newItemHdl);
-      removeFromMMContainer(oldItem);
-      return true;
+      //move can fail if another thread calls insertOrReplace
+      //in this case oldItem is no longer valid (not accessible, 
+      //it gets removed from MMContainer and evictForSlabRelease
+      //will send it back to the allocator
+      bool ret = moveRegularItemWithSync(oldItem, newItemHdl);
+      if (!ret) {
+          //we failed to move - newItemHdl was released back to allocator
+          //by the moveRegularItemWithSync but oldItem is not accessible
+          //and no longer valid - we need to clean it up here
+          XDCHECK(!oldItem.isAccessible());
+          oldItem.markForEvictionWhenMoving();
+          unlinkItemForEviction(oldItem);
+          wakeUpWaiters(oldItem, {});
+      } else {
+        removeFromMMContainer(oldItem);
+      }
+      return ret;
     }
   }
 }
diff --git a/cachelib/allocator/CacheAllocator.h b/cachelib/allocator/CacheAllocator.h
@@ -1617,7 +1617,7 @@ class CacheAllocator : public CacheBase {
   //
   // @return true  If the move was completed, and the containers were updated
   //               successfully.
-  void moveRegularItemWithSync(Item& oldItem, WriteHandle& newItemHdl);
+  bool 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
diff --git a/cachelib/allocator/CacheItem.h b/cachelib/allocator/CacheItem.h
@@ -46,6 +46,9 @@ class BaseAllocatorTest;
 template <typename AllocatorT>
 class AllocatorHitStatsTest;
 
+template <typename AllocatorT>
+class AllocatorMemoryTiersTest;
+
 template <typename AllocatorT>
 class MapTest;
 
@@ -473,6 +476,8 @@ class CACHELIB_PACKED_ATTR CacheItem {
   FRIEND_TEST(ItemTest, NonStringKey);
   template <typename AllocatorT>
   friend class facebook::cachelib::tests::AllocatorHitStatsTest;
+  template <typename AllocatorT>
+  friend class facebook::cachelib::tests::AllocatorMemoryTiersTest;
 };
 
 // A chained item has a hook pointing to the next chained item. The hook is
diff --git a/cachelib/allocator/tests/AllocatorMemoryTiersTest.cpp b/cachelib/allocator/tests/AllocatorMemoryTiersTest.cpp
@@ -21,14 +21,15 @@ namespace cachelib {
 namespace tests {
 
 using LruAllocatorMemoryTiersTest = AllocatorMemoryTiersTest<LruAllocator>;
-
+//using LruTestAllocatorMemoryTiersTest = AllocatorMemoryTiersTest<LruTestAllocator>;
 // TODO(MEMORY_TIER): add more tests with different eviction policies
 TEST_F(LruAllocatorMemoryTiersTest, MultiTiersInvalid) { this->testMultiTiersInvalid(); }
 TEST_F(LruAllocatorMemoryTiersTest, MultiTiersValid) { this->testMultiTiersValid(); }
 TEST_F(LruAllocatorMemoryTiersTest, MultiTiersValidMixed) { this->testMultiTiersValidMixed(); }
 TEST_F(LruAllocatorMemoryTiersTest, MultiTiersBackgroundMovers ) { this->testMultiTiersBackgroundMovers(); }
 TEST_F(LruAllocatorMemoryTiersTest, MultiTiersRemoveDuringEviction) { this->testMultiTiersRemoveDuringEviction(); }
 TEST_F(LruAllocatorMemoryTiersTest, MultiTiersReplaceDuringEviction) { this->testMultiTiersReplaceDuringEviction(); }
+TEST_F(LruAllocatorMemoryTiersTest, MultiTiersReplaceDuringEvictionWithReader) { this->testMultiTiersReplaceDuringEvictionWithReader(); }
 
 } // end of namespace tests
 } // end of namespace cachelib
diff --git a/cachelib/allocator/tests/AllocatorMemoryTiersTest.h b/cachelib/allocator/tests/AllocatorMemoryTiersTest.h
@@ -22,6 +22,10 @@
 #include "cachelib/allocator/FreeThresholdStrategy.h"
 #include "cachelib/allocator/PromotionStrategy.h"
 
+#include <fcntl.h>
+#include <unistd.h>
+#include <ctype.h>
+#include <semaphore.h>
 #include <folly/synchronization/Latch.h>
 
 namespace facebook {
@@ -58,6 +62,7 @@ class AllocatorMemoryTiersTest : public AllocatorTest<AllocatorT> {
       ASSERT_NO_THROW(alloc->insertOrReplace(handle));
     }
   }
+  
  public:
   void testMultiTiersInvalid() {
     typename AllocatorT::Config config;
@@ -201,7 +206,7 @@ class AllocatorMemoryTiersTest : public AllocatorTest<AllocatorT> {
 
     t->join();
   }
-
+  
   void testMultiTiersReplaceDuringEviction() {
     std::unique_ptr<AllocatorT> alloc;
     PoolId pool;
@@ -234,6 +239,127 @@ class AllocatorMemoryTiersTest : public AllocatorTest<AllocatorT> {
     testMultiTiersAsyncOpDuringMove(alloc, pool, quit, moveCb);
 
     t->join();
+
+  }
+
+
+  void gdb_sync1() {}
+  void gdb_sync2() {}
+  void gdb_sync3() {}
+  using ReadHandle = typename AllocatorT::ReadHandle;
+  void testMultiTiersReplaceDuringEvictionWithReader() {
+    sem_unlink ("/gdb1_sem");
+    sem_t *sem = sem_open ("/gdb1_sem", O_CREAT | O_EXCL, S_IRUSR | S_IWUSR, 0);
+    int gdbfd = open("/tmp/gdb1.gdb",O_CREAT | O_TRUNC | O_RDWR, S_IRUSR | S_IWUSR);
+    char gdbcmds[] = 
+                     "set attached=1\n"
+                     "break gdb_sync1\n"
+                     "break gdb_sync2\n"
+                     "break moveRegularItemWithSync\n"
+                     "c\n"
+                     "set scheduler-locking on\n"
+                     "thread 1\n"
+                     "c\n"
+                     "thread 4\n"
+                     "c\n"
+                     "thread 5\n"
+                     "break nativeFutexWaitImpl thread 5\n"
+                     "c\n"
+                     "thread 4\n"
+                     "break nativeFutexWaitImpl thread 4\n"
+                     "c\n"
+                     "thread 1\n"
+                     "break releaseBackToAllocator\n"
+                     "c\n"
+                     "c\n"
+                     "thread 5\n"
+                     "c\n"
+                     "thread 4\n"
+                     "c\n"
+                     "thread 1\n"
+                     "break gdb_sync3\n"
+                     "c\n"
+                     "quit\n";
+    int ret = write(gdbfd,gdbcmds,strlen(gdbcmds));
+    int ppid = getpid(); //parent pid
+    //int pid = 0;
+    int pid = fork();
+    if (pid == 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=/tmp/gdb1.gdb",(char*) 0);
+        ASSERT(f != -1);
+    }
+    sem_post(sem);
+    //wait for gdb to run
+    int attached = 0;
+    while (attached == 0);
+    
+    std::unique_ptr<AllocatorT> alloc;
+    PoolId pool;
+    bool quit = false;
+    
+    typename AllocatorT::Config config;
+    config.setCacheSize(4 * Slab::kSize);
+    config.enableCachePersistence("/tmp");
+    config.configureMemoryTiers({
+        MemoryTierCacheConfig::fromShm()
+            .setRatio(1).setMemBind(std::string("0")),
+        MemoryTierCacheConfig::fromShm()
+            .setRatio(1).setMemBind(std::string("0"))
+    });
+
+    alloc = std::make_unique<AllocatorT>(AllocatorT::SharedMemNew, config);
+    ASSERT(alloc != nullptr);
+    pool = alloc->addPool("default", alloc->getCacheMemoryStats().ramCacheSize);
+
+    int i = 0;
+    typename AllocatorT::Item* evicted;
+    std::unique_ptr<std::thread> t;
+    std::unique_ptr<std::thread> r;
+    while(!quit) {
+      auto handle = alloc->allocate(pool, std::to_string(++i), std::string("value").size());
+      ASSERT(handle != nullptr);
+      if (i == 1) {
+          evicted = static_cast<typename AllocatorT::Item*>(handle.get());
+          folly::Latch latch_t(1);
+          t = std::make_unique<std::thread>([&](){
+                auto handleNew = alloc->allocate(pool, std::to_string(1), std::string("new value").size());
+                ASSERT(handleNew != nullptr);
+                latch_t.count_down();
+                //first breakpoint will be this one because 
+                //thread 1 still has more items to fill up the
+                //cache before an evict is evicted
+                gdb_sync1();
+                ASSERT(evicted->isMoving());
+                //need to suspend thread 1 - who is doing the eviction
+                //gdb will do this for us
+                folly::Latch latch(1);
+                r = std::make_unique<std::thread>([&](){
+                    ASSERT(evicted->isMoving());
+                    latch.count_down();
+                    auto handleEvict = alloc->find(std::to_string(1));
+                    //does find block until done moving?? yes
+                    while (evicted->isMarkedForEviction()); //move will fail
+                    XDCHECK(handleEvict == nullptr) << handleEvict->toString();
+                    ASSERT(handleEvict == nullptr);
+                });
+                latch.wait();
+                gdb_sync2();
+                alloc->insertOrReplace(handleNew);
+                ASSERT(!evicted->isAccessible()); //move failed
+                quit = true;
+              });
+          latch_t.wait();
+      }
+      ASSERT_NO_THROW(alloc->insertOrReplace(handle));
+    }
+    t->join();
+    r->join();
+    gdb_sync3();
   }
 };
 } // namespace tests
