feat(server): flush slots traverse better yield (#4821)

Signed-off-by: adi_holden <adi@dragonflydb.io>

Author: adiholden

src/server/cluster/cluster_family_test.cc

@@ -723,7 +723,7 @@ TEST_F(ClusterFamilyTest, FlushSlots) {
   ExpectConditionWithinTimeout([&]() {
     return RunPrivileged({"dflycluster", "flushslots", "0", "0"}) == "OK";
   });
-
+  util::ThisFiber::SleepFor(1ms);
   EXPECT_THAT(RunPrivileged({"dflycluster", "getslotinfo", "slots", "0", "1"}),
               RespArray(ElementsAre(
                   RespArray(ElementsAre(IntArg(0), "key_count", IntArg(0), "total_reads", _,
@@ -732,7 +732,7 @@ TEST_F(ClusterFamilyTest, FlushSlots) {
                                         "total_writes", _, "memory_bytes", _)))));
 
   EXPECT_EQ(RunPrivileged({"dflycluster", "flushslots", "0", "1"}), "OK");
-
+  util::ThisFiber::SleepFor(1ms);
   EXPECT_THAT(
       RunPrivileged({"dflycluster", "getslotinfo", "slots", "0", "1"}),
       RespArray(ElementsAre(RespArray(ElementsAre(IntArg(0), "key_count", IntArg(0), "total_reads",

src/server/db_slice.cc

@@ -787,63 +787,56 @@ void DbSlice::Del(Context cntx, Iterator it) {
 }
 
 void DbSlice::FlushSlotsFb(const cluster::SlotSet& slot_ids) {
+  VLOG(1) << "Start FlushSlotsFb";
   // Slot deletion can take time as it traverses all the database, hence it runs in fiber.
   // We want to flush all the data of a slot that was added till the time the call to FlushSlotsFb
   // was made. Therefore we delete slots entries with version < next_version
   uint64_t next_version = 0;
+  uint64_t del_count = 0;
 
   std::string tmp;
-  auto del_entry_cb = [&](PrimeTable::iterator it) {
-    std::string_view key = it->first.GetSlice(&tmp);
-    SlotId sid = KeySlot(key);
-    if (slot_ids.Contains(sid) && it.GetVersion() < next_version) {
-      PerformDeletion(Iterator::FromPrime(it), db_arr_[0].get());
+  auto iterate_bucket = [&](PrimeTable::bucket_iterator it) {
+    it.AdvanceIfNotOccupied();
+    while (!it.is_done()) {
+      std::string_view key = it->first.GetSlice(&tmp);
+      SlotId sid = KeySlot(key);
+      if (slot_ids.Contains(sid) && it.GetVersion() < next_version) {
+        PerformDeletion(Iterator::FromPrime(it), db_arr_[0].get());
+        ++del_count;
+      }
+      ++it;
     }
-    return true;
   };
 
   auto on_change = [&](DbIndex db_index, const ChangeReq& req) {
     FiberAtomicGuard fg;
     PrimeTable* table = GetTables(db_index).first;
 
-    auto iterate_bucket = [&](DbIndex db_index, PrimeTable::bucket_iterator it) {
-      it.AdvanceIfNotOccupied();
-      while (!it.is_done()) {
-        del_entry_cb(it);
-        ++it;
-      }
-    };
-
     if (const PrimeTable::bucket_iterator* bit = req.update()) {
       if (!bit->is_done() && bit->GetVersion() < next_version) {
-        iterate_bucket(db_index, *bit);
+        iterate_bucket(*bit);
       }
     } else {
       string_view key = get<string_view>(req.change);
-      table->CVCUponInsert(
-          next_version, key,
-          [db_index, next_version, iterate_bucket](PrimeTable::bucket_iterator it) {
-            DCHECK_LT(it.GetVersion(), next_version);
-            iterate_bucket(db_index, it);
-          });
+      table->CVCUponInsert(next_version, key,
+                           [next_version, iterate_bucket](PrimeTable::bucket_iterator it) {
+                             DCHECK_LT(it.GetVersion(), next_version);
+                             iterate_bucket(it);
+                           });
     }
   };
   next_version = RegisterOnChange(std::move(on_change));
 
   ServerState& etl = *ServerState::tlocal();
   PrimeTable* pt = &db_arr_[0]->prime;
   PrimeTable::Cursor cursor;
-  uint64_t i = 0;
+
   do {
-    PrimeTable::Cursor next = pt->Traverse(cursor, del_entry_cb);
-    ++i;
+    PrimeTable::Cursor next = pt->TraverseBuckets(cursor, iterate_bucket);
     cursor = next;
-    if (i % 100 == 0) {
-      ThisFiber::Yield();
-    }
-
+    ThisFiber::Yield();
   } while (cursor && etl.gstate() != GlobalState::SHUTTING_DOWN);
-
+  VLOG(1) << "FlushSlotsFb del count is: " << del_count;
   UnregisterOnChange(next_version);
 
   etl.DecommitMemory(ServerState::kDataHeap);