Repository: tikv/tikv (language: rust)

Issue:
Scaling: optimize the clearing peer mechanism to mitigate the impacts to tail latency.
## Development Task

This task is a sub-task of https://github.com/tikv/tikv/issues/18042, which is used to mitigate the latency impacts during scaling operations.
It mainly consists of two parts:
- Do not use ingest operations when clearing stale ranges during the balancing regions process.
- Do not clear data of offline stores during the scale-in process, as this data will be automatically cleared when the corresponding node goes offline.



Discussion / hints:

Fix patch:
diff --git a/Cargo.lock b/Cargo.lock
index 2474f8093b7..061a2e2559f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3743,7 +3743,7 @@ dependencies = [
 [[package]]
 name = "kvproto"
 version = "0.0.2"
-source = "git+https://github.com/pingcap/kvproto.git#3b77f2c65c635567751d27afd6c58c7958645785"
+source = "git+https://github.com/pingcap/kvproto.git#258e68e813aa1af3a9817ce87630e2ec638851cd"
 dependencies = [
  "futures 0.3.15",
  "grpcio",
diff --git a/components/raftstore/src/store/snap.rs b/components/raftstore/src/store/snap.rs
index 90a880f2539..19959540bff 100644
--- a/components/raftstore/src/store/snap.rs
+++ b/components/raftstore/src/store/snap.rs
@@ -1473,17 +1473,18 @@ struct SnapManagerCore {
     max_per_file_size: Arc<AtomicU64>,
     enable_multi_snapshot_files: Arc<AtomicBool>,
     stats: Arc<Mutex<Vec<SnapshotStat>>>,
+    max_total_size: Arc<AtomicU64>,
     // Minimal column family size & kv counts for applying by ingest.
     min_ingest_cf_size: u64,
     min_ingest_cf_kvs: u64,
+    // Marker to represent the relative store is marked with Offline.
+    offlined: Arc<AtomicBool>,
 }
 
 /// `SnapManagerCore` trace all current processing snapshots.
 #[derive(Clone)]
 pub struct SnapManager {
     core: SnapManagerCore,
-    max_total_size: Arc<AtomicU64>,
-
     // only used to receive snapshot from v2
     tablet_snap_manager: Option<TabletSnapManager>,
 }
@@ -1767,11 +1768,13 @@ impl SnapManager {
     }
 
     pub fn max_total_snap_size(&self) -> u64 {
-        self.max_total_size.load(Ordering::Acquire)
+        self.core.max_total_size.load(Ordering::Acquire)
     }
 
     pub fn set_max_total_snap_size(&self, max_total_size: u64) {
-        self.max_total_size.store(max_total_size, Ordering::Release);
+        self.core
+            .max_total_size
+            .store(max_total_size, Ordering::Release);
     }
 
     pub fn set_max_per_file_size(&mut self, max_per_file_size: u64) {
@@ -1945,6 +1948,14 @@ impl SnapManager {
             .recv_concurrency_limiter
             .set_reserved_capacity(num_pending_applies)
     }
+
+    pub fn set_offline(&mut self, state: bool) {
+        self.core.offlined.store(state, Ordering::Release);
+    }
+
+    pub fn is_offlined(&self) -> bool {
+        self.core.offlined.load(Ordering::Acquire)
+    }
 }
 
 impl SnapManagerCore {
@@ -2237,11 +2248,12 @@ impl SnapManagerBuilder {
                 enable_multi_snapshot_files: Arc::new(AtomicBool::new(
                     self.enable_multi_snapshot_files,
                 )),
+                max_total_size: Arc::new(AtomicU64::new(max_total_size)),
                 stats: Default::default(),
                 min_ingest_cf_size: self.min_ingest_snapshot_size,
                 min_ingest_cf_kvs: self.min_ingest_snapshot_kvs,
+                offlined: Arc::new(AtomicBool::new(false)),
             },
-            max_total_size: Arc::new(AtomicU64::new(max_total_size)),
             tablet_snap_manager,
         };
         snapshot.set_max_per_file_size(self.max_per_file_size); // set actual max_per_file_size
@@ -2738,9 +2750,11 @@ pub mod tests {
             encryption_key_manager: None,
             max_per_file_size: Arc::new(AtomicU64::new(max_per_file_size)),
             enable_multi_snapshot_files: Arc::new(AtomicBool::new(true)),
+            max_total_size: Arc::new(AtomicU64::new(u64::MAX)),
             stats: Default::default(),
             min_ingest_cf_size: 0,
             min_ingest_cf_kvs: 0,
+            offlined: Arc::new(AtomicBool::new(false)),
         }
     }
 
diff --git a/components/raftstore/src/store/worker/pd.rs b/components/raftstore/src/store/worker/pd.rs
index 5909d825888..bcc33b45281 100644
--- a/components/raftstore/src/store/worker/pd.rs
+++ b/components/raftstore/src/store/worker/pd.rs
@@ -1274,6 +1274,7 @@ where
 
         let scheduler = self.scheduler.clone();
         let router = self.router.clone();
+        let mut snap_mgr = self.snap_mgr.clone();
         let resp = self
             .pd_client
             .store_heartbeat(stats, store_report, dr_autosync_status);
@@ -1346,6 +1347,36 @@ where
                             warn!("fail to schedule control grpc task"; "err" => ?e);
                         }
                     }
+                    // NodeState for this store.
+                    {
+                        let state = (|| {
+                            #[cfg(feature = "failpoints")]
+                            fail_point!("manually_set_store_offline", |_| {
+                                metapb::NodeState::Removing
+                            });
+                            resp.get_state()
+                        })();
+                        match state {
+                            metapb::NodeState::Removing | metapb::NodeState::Removed => {
+                                let is_offlined = snap_mgr.is_offlined();
+                                if !is_offlined {
+                                    snap_mgr.set_offline(true);
+                                    info!("store is offlined by pd");
+                                }
+                            }
+                            // As for NodeState::Preparing | Serving, if `is_offlined == true`,
+                            // it means the store has been re-added into the cluster and
+                            // the offline operation is terminated. Therefore, the state
+                            // `is_offlined` should be reset with `false`.
+                            _ => {
+                                let is_offlined = snap_mgr.is_offlined();
+                                if is_offlined {
+                                    snap_mgr.set_offline(false);
+                                    info!("store is remarked with serving state by pd");
+                                }
+                            }
+                        }
+                    }
                 }
                 Err(e) => {
                     error!("store heartbeat failed"; "err" => ?e);
diff --git a/components/raftstore/src/store/worker/region.rs b/components/raftstore/src/store/worker/region.rs
index 3e7a1e7848f..6fbfbc385ae 100644
--- a/components/raftstore/src/store/worker/region.rs
+++ b/components/raftstore/src/store/worker/region.rs
@@ -487,7 +487,7 @@ where
     /// with it.
     fn clean_overlap_ranges(&mut self, start_key: Vec<u8>, end_key: Vec<u8>) -> Result<()> {
         let (start_key, end_key) = self.clean_overlap_ranges_roughly(start_key, end_key);
-        self.delete_all_in_range(&[Range::new(&start_key, &end_key)])
+        self.delete_all_in_range(&[Range::new(&start_key, &end_key)], false)
     }
 
     /// Inserts a new pending range, and it will be cleaned up with some delay.
@@ -510,6 +510,10 @@ where
 
     /// Cleans up stale ranges.
     fn clean_stale_ranges(&mut self) {
+        if self.mgr.is_offlined() {
+            info!("no need to clear stale range as store is marked offlined");
+            return;
+        }
         STALE_PEER_PENDING_DELETE_RANGE_GAUGE.set(self.pending_delete_ranges.len() as f64);
         if self.ingest_maybe_stall() {
             return;
@@ -538,7 +542,7 @@ where
                 Range::new(start, end)
             })
             .collect();
-
+        // Clear sst files belonging to the given range directly.
         self.engine
             .delete_ranges_cfs(
                 &WriteOptions::default(),
@@ -549,10 +553,13 @@ where
                 error!("failed to delete files in range"; "err" => %e);
             })
             .unwrap();
-        if let Err(e) = self.delete_all_in_range(&ranges) {
+        // Remove all overlapped ranges directly without ingesting.
+        if let Err(e) = self.delete_all_in_range(&ranges, true) {
             error!("failed to cleanup stale range"; "err" => %e);
             return;
         }
+        // Clear related blob files belonging to the given range directly after clearing
+        // all stale keys in sst files.
         self.engine
             .delete_ranges_cfs(
                 &WriteOptions::default(),
@@ -588,7 +595,7 @@ where
         false
     }
 
-    fn delete_all_in_range(&self, ranges: &[Range<'_>]) -> Result<()> {
+    fn delete_all_in_range(&self, ranges: &[Range<'_>], forcely_by_key: bool) -> Result<()> {
         let wopts = WriteOptions::default();
         for cf in self.engine.cf_names() {
             // CF_LOCK usually contains fewer keys than other CFs, so we delete them by key.
@@ -602,7 +609,18 @@ where
                     DeleteStrategy::DeleteByRange,
                     &CLEAR_OVERLAP_REGION_DURATION.by_range,
                 )
+            // Meanwhile, if `forcely_by_key` is specified and
+            // `cfg.use_delete_range` is not enabled, it should remove the
+            // range by delete keys rather than ingestion.
+            } else if forcely_by_key {
+                (
+                    DeleteStrategy::DeleteByKey,
+                    &CLEAR_OVERLAP_REGION_DURATION.by_key,
+                )
             } else {
+                // Deprecated method for clearing stale ranges due to potential latency
+                // jitters. It's still applicable when clearing overlapped regions
+                // before applying the newly added snapshot.
                 (
                     DeleteStrategy::DeleteByWriter {
                         sst_path: self.mgr.get_temp_path_for_ingest(),


Test patch:
diff --git a/tests/failpoints/cases/test_stale_peer.rs b/tests/failpoints/cases/test_stale_peer.rs
index ee1171ae114..bd2a108751b 100644
--- a/tests/failpoints/cases/test_stale_peer.rs
+++ b/tests/failpoints/cases/test_stale_peer.rs
@@ -1,12 +1,14 @@
 // Copyright 2018 TiKV Project Authors. Licensed under Apache-2.0.
 
 use std::{
+    fs,
+    path::Path,
     sync::{atomic::AtomicBool, mpsc, Arc},
     thread,
     time::Duration,
 };
 
-use engine_traits::{RaftEngineDebug, RaftEngineReadOnly};
+use engine_traits::{RaftEngine, RaftEngineDebug, RaftEngineReadOnly};
 use futures::executor::block_on;
 use kvproto::raft_serverpb::{PeerState, RaftLocalState, RaftMessage};
 use pd_client::PdClient;
@@ -45,6 +47,65 @@ fn test_one_node_leader_missing() {
     fail::remove(check_stale_state);
 }
 
+#[test]
+fn test_clean_stale_peer() {
+    let validate_data_files = |dir: &std::path::Path| -> bool {
+        let mut data_file = 0_usize;
+        for entry in fs::read_dir(dir).unwrap() {
+            let path = entry.unwrap().path();
+            if path.is_file() && path.to_str().unwrap().ends_with(".log") {
+                data_file += 1;
+            }
+        }
+        data_file > 0
+    };
+    let mut cluster = new_node_cluster(0, 3);
+    cluster.cfg.raft_store.snap_mgr_gc_tick_interval = ReadableDuration::millis(100);
+    cluster.cfg.raft_store.pd_heartbeat_tick_interval = ReadableDuration::millis(100);
+    // Disable raft log gc in this test case.
+    cluster.cfg.raft_store.raft_log_gc_tick_interval = ReadableDuration::secs(60);
+
+    let pd_client = Arc::clone(&cluster.pd_client);
+    // Disable default max peer count check.
+    pd_client.disable_default_operator();
+
+    let region_id = cluster.run_conf_change();
+
+    // Manually set store 3 is offlined.
+    fail::cfg("manually_set_store_offline", "return").unwrap();
+    pd_client.must_add_peer(region_id, new_peer(2, 2));
+    pd_client.must_add_peer(region_id, new_peer(3, 3));
+
+    (0..10).for_each(|_| cluster.must_put(b"k1", b"v1"));
+    // Send heartbeats to pd to ensure that heartbeat responses have been received
+    // by node 3.
+    cluster.must_send_store_heartbeat(3);
+    must_get_equal(&cluster.get_engine(3), b"k1", b"v1");
+    pd_client.must_remove_peer(region_id, new_peer(3, 3));
+    // Since the node 3 is marked with NodeState::Removing, its `offlined == true`
+    // in SnapManager, and no need to clear the corresponding data.
+    sleep_ms(500);
+    let engine = cluster.get_engine(3);
+    let engine_path = Path::new(engine.get_engine_path());
+    assert!(validate_data_files(engine_path));
+    fail::remove("manually_set_store_offline");
+    // After the state is cleared, the store should be marked with
+    // NodeState::Serving. And after adding a new peer to store 3,
+    // all previous data should be synced to it.
+    cluster.must_send_store_heartbeat(3);
+    cluster.must_put(b"k10", b"v10");
+    let engine = cluster.get_engine(3);
+    let (key, value) = (b"k10", b"v10");
+    must_get_none(&engine, key);
+    pd_client.must_add_peer(region_id, new_peer(3, 4));
+    must_get_equal(&engine, key, value);
+    must_get_equal(&cluster.get_engine(3), b"k1", b"v1");
+    // Remove the peer 4.
+    pd_client.must_remove_peer(region_id, new_peer(3, 4));
+    sleep_ms(500);
+    must_get_none(&engine, key);
+}
+
 #[test_case(test_raftstore::new_node_cluster)]
 #[test_case(test_raftstore_v2::new_node_cluster)]
 fn test_node_update_localreader_after_removed() {