fix: Remove index docs on eviction and cluster migration (#6062)

When we evict keys or during slot migration we should also remove index docs.

Signed-off-by: mkaruza <[email protected]>

Author: mkaruza

src/server/db_slice.cc

@@ -230,7 +230,7 @@ unsigned PrimeEvictionPolicy::Evict(const PrimeTable::HotBuckets& eb, PrimeTable
     if (auto journal = db_slice_->shard_owner()->journal(); journal) {
       RecordExpiryBlocking(cntx_.db_index, key);
     }
-    db_slice_->PerformDeletion(DbSlice::Iterator(last_slot_it, StringOrView::FromView(key)), table);
+    db_slice_->Del(cntx_, DbSlice::Iterator(last_slot_it, StringOrView::FromView(key)));
 
     ++evicted_;
   }
@@ -802,18 +802,25 @@ void DbSlice::ActivateDb(DbIndex db_ind) {
   CreateDb(db_ind);
 }
 
-void DbSlice::Del(Context cntx, Iterator it) {
+void DbSlice::Del(Context cntx, Iterator it, DbTable* db_table) {
   CHECK(IsValid(it));
 
-  auto& db = db_arr_[cntx.db_index];
+  ExpIterator exp_it;
+  DbTable* table = db_table ? db_table : db_arr_[cntx.db_index].get();
   auto obj_type = it->second.ObjType();
 
   if (doc_del_cb_ && (obj_type == OBJ_JSON || obj_type == OBJ_HASH)) {
     string tmp;
     string_view key = it->first.GetSlice(&tmp);
     doc_del_cb_(key, cntx, it->second);
   }
-  PerformDeletion(it, db.get());
+
+  if (it->second.HasExpire()) {
+    exp_it = ExpIterator::FromPrime(table->expire.Find(it->first));
+    DCHECK(!exp_it.is_done());
+  }
+
+  PerformDeletionAtomic(it, exp_it, table);
 }
 
 void DbSlice::DelMutable(Context cntx, ItAndUpdater it_updater) {
@@ -833,14 +840,19 @@ void DbSlice::FlushSlotsFb(const cluster::SlotSet& slot_ids) {
   boost::intrusive_ptr<DbTable> table = db_arr_.front();
   size_t memory_before = table->table_memory() + table->stats.obj_memory_usage;
 
+  DbContext db_cntx;
+  db_cntx.time_now_ms = GetCurrentTimeMs();
+  db_cntx.db_index = table->index;
+
   std::string tmp;
   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), table.get());
+        // We use copy of table smart pointer and pass it as table because FLLUSHALL can drop table.
+        Del(db_cntx, Iterator::FromPrime(it), table.get());
         ++del_count;
       }
       ++it;
@@ -1413,7 +1425,7 @@ int32_t DbSlice::GetNextSegmentForEviction(int32_t segment_id, DbIndex db_ind) c
          db_arr_[db_ind]->prime.GetSegmentCount();
 }
 
-pair<uint64_t, size_t> DbSlice::FreeMemWithEvictionStepAtomic(DbIndex db_ind,
+pair<uint64_t, size_t> DbSlice::FreeMemWithEvictionStepAtomic(DbIndex db_ind, const Context& cntx,
                                                               size_t starting_segment_id,
                                                               size_t increase_goal_bytes) {
   // Disable flush journal changes to prevent preemtion
@@ -1476,7 +1488,8 @@ pair<uint64_t, size_t> DbSlice::FreeMemWithEvictionStepAtomic(DbIndex db_ind,
 
         evicted_bytes += evict_it->first.MallocUsed() + evict_it->second.MallocUsed();
         ++evicted_items;
-        PerformDeletion(Iterator(evict_it, StringOrView::FromView(key)), db_table.get());
+
+        Del(cntx, Iterator(evict_it, StringOrView::FromView(key)));
 
         // returns when whichever condition is met first
         if ((evicted_items == max_eviction_per_hb) || (evicted_bytes >= increase_goal_bytes))
@@ -1824,16 +1837,6 @@ void DbSlice::PerformDeletionAtomic(const Iterator& del_it, const ExpIterator& e
   }
 }
 
-void DbSlice::PerformDeletion(Iterator del_it, DbTable* table) {
-  ExpIterator exp_it;
-  if (del_it->second.HasExpire()) {
-    exp_it = ExpIterator::FromPrime(table->expire.Find(del_it->first));
-    DCHECK(!exp_it.is_done());
-  }
-
-  PerformDeletionAtomic(del_it, exp_it, table);
-}
-
 void DbSlice::OnCbFinishBlocking() {
   if (IsCacheMode()) {
     // move fetched items to local variable

src/server/db_slice.h

@@ -322,11 +322,10 @@ class DbSlice {
   // Creates a database with index `db_ind`. If such database exists does nothing.
   void ActivateDb(DbIndex db_ind);
 
-  // Delete a key referred by its iterator.
-  void PerformDeletion(Iterator del_it, DbTable* table);
-
   // Deletes the iterator. The iterator must be valid.
-  void Del(Context cntx, Iterator it);
+  // Context argument is used only for document removal and it just needs
+  // timestamp field. Last argument, db_table, is optional and is used only in FlushSlotsCb.
+  void Del(Context cntx, Iterator it, DbTable* db_table = nullptr);
 
   // Deletes a key after FindMutable(). Runs post_updater before deletion
   // to update memory accounting while the key is still valid.
@@ -456,7 +455,7 @@ class DbSlice {
   // Evicts items with dynamically allocated data from the primary table.
   // Does not shrink tables.
   // Returns number of (elements,bytes) freed due to evictions.
-  std::pair<uint64_t, size_t> FreeMemWithEvictionStepAtomic(DbIndex db_indx,
+  std::pair<uint64_t, size_t> FreeMemWithEvictionStepAtomic(DbIndex db_indx, const Context& cntx,
                                                             size_t starting_segment_id,
                                                             size_t increase_goal_bytes);
 

src/server/engine_shard.cc

@@ -688,7 +688,7 @@ void EngineShard::RetireExpiredAndEvict() {
     if (eviction_goal) {
       uint32_t starting_segment_id = rand() % pt->GetSegmentCount();
       auto [evicted_items, evicted_bytes] =
-          db_slice.FreeMemWithEvictionStepAtomic(i, starting_segment_id, eviction_goal);
+          db_slice.FreeMemWithEvictionStepAtomic(i, db_cntx, starting_segment_id, eviction_goal);
 
       VLOG(2) << "Heartbeat eviction: Expected to evict " << eviction_goal
               << " bytes. Actually evicted " << evicted_items << " items, " << evicted_bytes

tests/dragonfly/cluster_test.py

@@ -3483,3 +3483,68 @@ async def test_SearchRequestDistribution(df_factory: DflyInstanceFactory):
     await asyncio.sleep(3)
     for node in nodes:
         await wait_for_ft_index_creation(node.client, "idx")
+
+
+async def verify_keys_match_number_of_index_docs(client, expected_num_keys):
+    # Get number of docs in index
+    index_info = await client.execute_command(f"FT.INFO idx")
+    index_info_num_docs = index_info[9]
+
+    # Get number of keys in database
+    keyspace_info = await client.info("keyspace")
+    keyspace_keys = keyspace_info["db0"]["keys"]
+
+    assert index_info_num_docs == keyspace_keys
+    assert index_info_num_docs == expected_num_keys
+    assert keyspace_keys == expected_num_keys
+
+
+@dfly_args({"proactor_threads": 2, "cluster_mode": "yes"})
+async def test_remove_docs_on_cluster_migration(df_factory):
+    instances = [
+        df_factory.create(port=next(next_port), admin_port=next(next_port)) for i in range(2)
+    ]
+
+    df_factory.start_all(instances)
+
+    nodes = [(await create_node_info(instance)) for instance in instances]
+    nodes[0].slots = [(0, 16383)]
+    nodes[1].slots = []
+
+    await push_config(json.dumps(generate_config(nodes)), [node.admin_client for node in nodes])
+
+    # Create index on both nodes
+    await nodes[0].client.execute_command(
+        "FT.CREATE", "idx", "ON", "HASH", "PREFIX", "1", "doc:", "SCHEMA", "v", "TEXT"
+    )
+
+    # Populate node 0
+    keys = 100
+    for i in range(0, keys):
+        random_string = "".join(random.choices(string.ascii_letters + string.digits, k=1_000))
+        await nodes[0].client.execute_command("HSET", f"doc:{i}", "v", random_string)
+
+    # Verify on node 0 that keys are added and index is populated
+    await verify_keys_match_number_of_index_docs(nodes[0].client, keys)
+
+    nodes[0].migrations.append(
+        MigrationInfo("127.0.0.1", instances[1].port, [(0, 16383)], nodes[1].id)
+    )
+    logging.debug("Start migration")
+    await push_config(json.dumps(generate_config(nodes)), [node.admin_client for node in nodes])
+
+    await wait_for_status(nodes[0].admin_client, nodes[1].id, "FINISHED")
+
+    nodes[0].migrations = []
+    nodes[0].slots = []
+    nodes[1].slots = [(0, 16383)]
+    logging.debug("finalize migration")
+    await push_config(json.dumps(generate_config(nodes)), [node.admin_client for node in nodes])
+
+    await asyncio.sleep(1)
+
+    # Verify on node 1 that keys are moved and index is populated
+    await verify_keys_match_number_of_index_docs(nodes[1].client, keys)
+
+    # Verify that node 0 doesn't have any keys and no index docs
+    await verify_keys_match_number_of_index_docs(nodes[0].client, 0)

tests/dragonfly/memory_test.py

@@ -323,3 +323,45 @@ async def poll():
     df.stop()
     found = df.find_in_logs("Commands squashing current reply size is overlimit")
     assert len(found) > 0
+
+
+@pytest.mark.asyncio
+async def test_remove_docs_on_eviction(df_factory):
+    max_memory = 256 * 1024**2  # 256MB
+    df_server = df_factory.create(
+        proactor_threads=1,
+        cache_mode="yes",
+        maxmemory=max_memory,
+        vmodule="engine_shard=2",
+        eviction_memory_budget_threshold=0.99,
+        enable_heartbeat_rss_eviction="no",
+    )
+    df_server.start()
+    client = df_server.client()
+
+    await client.execute_command(
+        "FT.CREATE", "idx", "ON", "HASH", "PREFIX", "1", "doc:", "SCHEMA", "v", "TEXT"
+    )
+
+    i = 0
+    while True:
+        random_string = "".join(random.choices(string.ascii_letters + string.digits, k=1_000))
+        await client.execute_command("HSET", f"doc:{i}", "v", random_string)
+        stats_info = await client.info("stats")
+        # Done when see at least 50 evictions
+        if stats_info["evicted_keys"] > 50:
+            break
+        i = i + 1
+
+    # Give some time to eviction stabilize
+    await asyncio.sleep(1)
+
+    # Get number of docs in index
+    index_info = await client.execute_command(f"FT.INFO idx")
+    index_info_num_docs = index_info[9]
+
+    # Get number of keys in database
+    keyspace_info = await client.info("keyspace")
+    keyspace_keys = keyspace_info["db0"]["keys"]
+
+    assert index_info_num_docs == keyspace_keys