fix: Run defrag on dbs > 0 as well

src/server/dragonfly_test.cc

@@ -601,6 +601,8 @@ TEST_F(DefragDflyEngineTest, TestDefragOption) {
 
   std::vector<std::string_view> keys(keys2delete.begin(), keys2delete.end());
 
+  Run({"SELECT", "2"});
+
   RespExpr resp = Run(
       {"DEBUG", "POPULATE", std::to_string(kNumberOfKeys), "key-name", std::to_string(kKeySize)});
   ASSERT_EQ(resp, "OK");

src/server/engine_shard_set.cc

@@ -48,7 +48,6 @@ using absl::GetFlag;
 
 namespace {
 
-constexpr DbIndex kDefaultDbIndex = 0;
 constexpr uint64_t kCursorDoneState = 0u;
 
 vector<EngineShardSet::CachedStats> cached_stats;  // initialized in EngineShardSet::Init
@@ -92,6 +91,15 @@ EngineShard::Stats& EngineShard::Stats::operator+=(const EngineShard::Stats& o)
 void EngineShard::DefragTaskState::UpdateScanState(uint64_t cursor_val) {
   cursor = cursor_val;
   underutilized_found = false;
+  // Once we're done with a db, jump to the next
+  if (cursor == kCursorDoneState) {
+    dbid++;
+  }
+}
+
+void EngineShard::DefragTaskState::ResetScanState() {
+  dbid = cursor = 0u;
+  underutilized_found = false;
 }
 
 // This function checks 3 things:
@@ -137,8 +145,19 @@ bool EngineShard::DoDefrag() {
   const float threshold = GetFlag(FLAGS_mem_defrag_page_utilization_threshold);
 
   auto& slice = db_slice();
-  DCHECK(slice.IsDbValid(kDefaultDbIndex));
-  auto [prime_table, expire_table] = slice.GetTables(kDefaultDbIndex);
+
+  // If we moved to an invalid db, skip as long as it's not the last one
+  while (!slice.IsDbValid(defrag_state_.dbid) && defrag_state_.dbid + 1 < slice.db_array_size())
+    defrag_state_.dbid++;
+
+  // If we found no valid db, we finished traversing and start from scratch next time
+  if (!slice.IsDbValid(defrag_state_.dbid)) {
+    defrag_state_.ResetScanState();
+    return false;
+  }
+
+  DCHECK(slice.IsDbValid(defrag_state_.dbid));
+  auto [prime_table, expire_table] = slice.GetTables(defrag_state_.dbid);
   PrimeTable::Cursor cur = defrag_state_.cursor;
   uint64_t reallocations = 0;
   unsigned traverses_count = 0;
@@ -158,6 +177,7 @@ bool EngineShard::DoDefrag() {
   } while (traverses_count < kMaxTraverses && cur);
 
   defrag_state_.UpdateScanState(cur.value());
+
   if (reallocations > 0) {
     VLOG(1) << "shard " << slice.shard_id() << ": successfully defrag  " << reallocations
             << " times, did it in " << traverses_count << " cursor is at the "
@@ -168,10 +188,12 @@ bool EngineShard::DoDefrag() {
             << (defrag_state_.cursor == kCursorDoneState ? "end" : "in progress")
             << " but no location for defrag were found";
   }
+
   stats_.defrag_realloc_total += reallocations;
   stats_.defrag_task_invocation_total++;
   stats_.defrag_attempt_total += attempts;
-  return defrag_state_.cursor > kCursorDoneState;
+
+  return true;
 }
 
 // the memory defragmentation task is as follow:

src/server/engine_shard_set.h

@@ -161,15 +161,17 @@ class EngineShard {
 
  private:
   struct DefragTaskState {
-    // we will add more data members later
+    size_t dbid = 0u;
     uint64_t cursor = 0u;
     bool underutilized_found = false;
 
     // check the current threshold and return true if
-    // we need to do the de-fermentation
+    // we need to do the defragmentation
     bool CheckRequired();
 
     void UpdateScanState(uint64_t cursor_val);
+
+    void ResetScanState();
   };
 
   EngineShard(util::ProactorBase* pb, bool update_db_time, mi_heap_t* heap);