fix: Run defrag on dbs > 0 as well

Signed-off-by: Vladislav Oleshko <vlad@dragonflydb.io>
dragonflydb · Aug 24, 2023 · 15fae99 · 15fae99
1 parent 5aa63a6
commit 15fae99
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 6 deletions.
diff --git a/src/server/dragonfly_test.cc b/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");

diff --git a/src/server/engine_shard_set.cc b/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
@@ -94,6 +93,11 @@ void EngineShard::DefragTaskState::UpdateScanState(uint64_t cursor_val) {
   underutilized_found = false;
 }
 
+void EngineShard::DefragTaskState::ResetScanState() {
+  dbid = cursor = 0u;
+  underutilized_found = false;
+}
+
 // This function checks 3 things:
 // 1. Don't try memory fragmentation if we don't use "enough" memory (control by
 // mem_defrag_threshold flag)
@@ -137,8 +141,8 @@ 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);
+  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 +162,15 @@ bool EngineShard::DoDefrag() {
   } while (traverses_count < kMaxTraverses && cur);
 
   defrag_state_.UpdateScanState(cur.value());
+
+  // Once we're done with a db, jump to the next
+  if (defrag_state_.cursor == kCursorDoneState) {
+    defrag_state_.dbid++;
+    // Skip null dbs that still take up slots in the array
+    while (defrag_state_.dbid < slice.db_array_size() && !slice.IsDbValid(defrag_state_.dbid))
+      defrag_state_.dbid++;
+  }
+
   if (reallocations > 0) {
     VLOG(1) << "shard " << slice.shard_id() << ": successfully defrag  " << reallocations
             << " times, did it in " << traverses_count << " cursor is at the "
@@ -168,10 +181,15 @@ 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;
+
+  bool items_left = slice.IsDbValid(defrag_state_.dbid);
+  if (!items_left)
+    defrag_state_.dbid = defrag_state_.cursor = 0u;
+  return items_left;
 }
 
 // the memory defragmentation task is as follow:

diff --git a/src/server/engine_shard_set.h b/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);