raftstore-v2: clean up import sst file only if flushed epoch is stale.

components/engine_traits/src/flush.rs

@@ -20,6 +20,7 @@ use std::{
     },
 };
 
+use kvproto::import_sstpb::SstMeta;
 use slog_global::info;
 use tikv_util::set_panic_mark;
 
@@ -61,7 +62,8 @@ struct FlushProgress {
 /// if the flushed index greater than it .
 #[derive(Debug, Clone)]
 pub struct SstApplyState {
-    sst_map: Arc<RwLock<HashMap<Vec<u8>, u64>>>,
+    // Map from cf to SstApplyEntry.
+    sst_map: Arc<RwLock<HashMap<usize, Vec<SstApplyEntry>>>>,
 }
 
 impl Default for SstApplyState {
@@ -72,25 +74,51 @@ impl Default for SstApplyState {
     }
 }
 
+#[derive(Debug)]
+pub struct SstApplyEntry {
+    pub applied_index: u64,
+    pub sst: SstMeta,
+}
+
+impl SstApplyEntry {
+    pub fn new(applied_index: u64, sst: SstMeta) -> Self {
+        Self { applied_index, sst }
+    }
+}
+
 impl SstApplyState {
     #[inline]
-    pub fn registe_ssts(&self, uuids: Vec<Vec<u8>>, sst_applied_index: u64) {
+    pub fn register_ssts(&self, applied_index: u64, ssts: Vec<SstMeta>) {
         let mut map = self.sst_map.write().unwrap();
-        for uuid in uuids {
-            map.insert(uuid, sst_applied_index);
+        for sst in ssts {
+            let cf_index = data_cf_offset(sst.get_cf_name());
+            let entry = SstApplyEntry::new(applied_index, sst);
+            map.entry(cf_index).or_default().push(entry);
         }
     }
 
-    /// Query the sst applied index.
     #[inline]
-    pub fn sst_applied_index(&self, uuid: &Vec<u8>) -> Option<u64> {
-        self.sst_map.read().unwrap().get(uuid).copied()
+    pub fn stale_ssts(&self, cf: &str, flushed_index: u64) -> Vec<SstMeta> {
+        let map = self.sst_map.read().unwrap();
+        let cf_index = data_cf_offset(cf);
+        if let Some(ssts) = map.get(&cf_index) {
+            return ssts
+                .iter()
+                .filter(|entry| entry.applied_index <= flushed_index)
+                .map(|entry| entry.sst.clone())
+                .collect();
+        }
+        vec![]
     }
 
-    pub fn delete_ssts(&self, uuids: Vec<Vec<u8>>) {
-        let mut map = self.sst_map.write().unwrap();
-        for uuid in uuids {
-            map.remove(&uuid);
+    pub fn delete_ssts(&self, ssts: &Vec<SstMeta>) {
+        let mut map: std::sync::RwLockWriteGuard<'_, HashMap<usize, Vec<SstApplyEntry>>> =
+            self.sst_map.write().unwrap();
+        for sst in ssts {
+            let cf_index = data_cf_offset(sst.get_cf_name());
+            if let Some(metas) = map.get_mut(&cf_index) {
+                metas.drain_filter(|entry| entry.sst.get_uuid() == sst.get_uuid());
+            }
         }
     }
 }
@@ -270,3 +298,25 @@ impl<R: RaftEngine> StateStorage for R {
         self.consume(&mut batch, true).unwrap();
     }
 }
+
+#[cfg(test)]
+mod test {
+    use std::vec;
+
+    use kvproto::import_sstpb::SstMeta;
+
+    use super::SstApplyState;
+
+    #[test]
+    pub fn test_sst_apply_state() {
+        let stat = SstApplyState::default();
+        let mut sst = SstMeta::default();
+        sst.set_cf_name("write".to_owned());
+        sst.set_uuid(vec![1, 2, 3, 4]);
+        stat.register_ssts(10, vec![sst.clone()]);
+        assert!(stat.stale_ssts("default", 10).is_empty());
+        let sst = stat.stale_ssts("write", 10);
+        assert_eq!(sst[0].get_uuid(), vec![1, 2, 3, 4]);
+        stat.delete_ssts(&sst);
+    }
+}

components/engine_traits/src/lib.rs

@@ -254,6 +254,7 @@
 #![feature(linked_list_cursors)]
 #![feature(let_chains)]
 #![feature(str_split_as_str)]
+#![feature(drain_filter)]
 
 #[macro_use(fail_point)]
 extern crate fail;

components/raftstore-v2/src/fsm/peer.rs

@@ -291,9 +291,12 @@ impl<'a, EK: KvEngine, ER: RaftEngine, T: Transport> PeerFsmDelegate<'a, EK, ER,
                     tablet_index,
                     flushed_index,
                 } => {
-                    self.fsm
-                        .peer_mut()
-                        .on_data_flushed(cf, tablet_index, flushed_index);
+                    self.fsm.peer_mut().on_data_flushed(
+                        self.store_ctx,
+                        cf,
+                        tablet_index,
+                        flushed_index,
+                    );
                 }
                 PeerMsg::PeerUnreachable { to_peer_id } => {
                     self.fsm.peer_mut().on_peer_unreachable(to_peer_id)

components/raftstore-v2/src/operation/command/write/ingest.rs

@@ -8,7 +8,7 @@ use raftstore::{
     store::{check_sst_for_ingestion, metrics::PEER_WRITE_CMD_COUNTER, util},
     Result,
 };
-use slog::error;
+use slog::{error, info};
 use tikv_util::{box_try, slog_panic};
 
 use crate::{
@@ -67,38 +67,27 @@ impl<EK: KvEngine, ER: RaftEngine> Peer<EK, ER> {
         ctx: &mut StoreContext<EK, ER, T>,
         ssts: Box<[SstMeta]>,
     ) {
-        let mut stale_ssts = Vec::from(ssts);
-        let epoch = self.region().get_region_epoch();
-        stale_ssts.retain(|sst| {
-            fail::fail_point!("on_cleanup_import_sst", |_| true);
-            util::is_epoch_stale(sst.get_region_epoch(), epoch)
-        });
-
-        // some sst needs to be kept if the log didn't flush the disk.
-        let flushed_indexes = self.storage().apply_trace().flushed_indexes();
-        stale_ssts.retain(|sst| {
-            let off = data_cf_offset(sst.get_cf_name());
-            let uuid = sst.get_uuid().to_vec();
-            let sst_index = self.sst_apply_state().sst_applied_index(&uuid);
-            if let Some(index) = sst_index {
-                return flushed_indexes.as_ref()[off] >= index;
-            }
-            true
-        });
+        let mut stale_ssts: Vec<SstMeta> = Vec::from(ssts);
+        let flushed_epoch = self.storage().flushed_epoch();
+        stale_ssts.retain(|sst| util::is_epoch_stale(sst.get_region_epoch(), flushed_epoch));
 
         fail::fail_point!("on_cleanup_import_sst_schedule");
         if stale_ssts.is_empty() {
             return;
         }
-        let uuids = stale_ssts
-            .iter()
-            .map(|sst| sst.get_uuid().to_vec())
-            .collect();
-        self.sst_apply_state().delete_ssts(uuids);
+        info!(
+            self.logger,
+            "clean up import sst file";
+            "flushed_epoch" => ?flushed_epoch,
+            "stale_ssts" => ?stale_ssts);
+
+        self.sst_apply_state().delete_ssts(&stale_ssts);
         let _ = ctx
             .schedulers
             .tablet
-            .schedule(tablet::Task::CleanupImportSst(stale_ssts.into()));
+            .schedule(tablet::Task::CleanupImportSst(
+                stale_ssts.into_boxed_slice(),
+            ));
     }
 }
 
@@ -143,12 +132,9 @@ impl<EK: KvEngine, R: ApplyResReporter> Apply<EK, R> {
             if let Err(e) = self.sst_importer().ingest(&infos, self.tablet()) {
                 slog_panic!(self.logger, "ingest fail"; "ssts" => ?ssts, "error" => ?e);
             }
+            let metas: Vec<SstMeta> = infos.iter().map(|info| info.meta.clone()).collect();
+            self.sst_apply_state().register_ssts(index, metas);
         }
-        let uuids = infos
-            .iter()
-            .map(|info| info.meta.get_uuid().to_vec())
-            .collect::<Vec<_>>();
-        self.set_sst_applied_index(uuids, index);
 
         self.metrics.size_diff_hint += size;
         self.metrics.written_bytes += size as u64;

components/raftstore-v2/src/operation/ready/apply_trace.rs

@@ -57,6 +57,7 @@ use crate::{
     },
     raft::{Peer, Storage},
     router::PeerMsg,
+    worker::tablet,
     Result, StoreRouter,
 };
 
@@ -519,6 +520,13 @@ impl<EK: KvEngine, ER: RaftEngine> Storage<EK, ER> {
         }
         let region_id = self.region().get_id();
         let raft_engine = self.entry_storage().raft_engine();
+        let epoch = raft_engine
+            .get_region_state(region_id, trace.admin.flushed)
+            .unwrap()
+            .unwrap()
+            .get_region()
+            .get_region_epoch()
+            .clone();
         let tablet_index = self.tablet_index();
         let lb = write_task
             .extra_write
@@ -529,11 +537,19 @@ impl<EK: KvEngine, ER: RaftEngine> Storage<EK, ER> {
             .unwrap();
         trace.try_persist = false;
         trace.persisted_applied = trace.admin.flushed;
+
+        self.set_flushed_epoch(epoch);
     }
 }
 
 impl<EK: KvEngine, ER: RaftEngine> Peer<EK, ER> {
-    pub fn on_data_flushed(&mut self, cf: &str, tablet_index: u64, index: u64) {
+    pub fn on_data_flushed<T>(
+        &mut self,
+        ctx: &mut StoreContext<EK, ER, T>,
+        cf: &str,
+        tablet_index: u64,
+        index: u64,
+    ) {
         trace!(self.logger, "data flushed"; "cf" => cf, "tablet_index" => tablet_index, "index" => index, "trace" => ?self.storage().apply_trace());
         if tablet_index < self.storage().tablet_index() {
             // Stale tablet.
@@ -543,6 +559,24 @@ impl<EK: KvEngine, ER: RaftEngine> Peer<EK, ER> {
         let apply_trace = self.storage_mut().apply_trace_mut();
         apply_trace.on_flush(cf, index);
         apply_trace.maybe_advance_admin_flushed(apply_index);
+        let stale_ssts = self.sst_apply_state().stale_ssts(cf, index);
+        if stale_ssts.is_empty() {
+            return;
+        }
+        info!(
+            self.logger,
+            "delete stale ssts after flush";
+            "stale_ssts" => ?stale_ssts,
+            "apply_index" => apply_index,
+            "cf" => cf,
+            "flushed_index" => index,
+        );
+        let _ = ctx
+            .schedulers
+            .tablet
+            .schedule(tablet::Task::CleanupImportSst(
+                stale_ssts.into_boxed_slice(),
+            ));
     }
 
     pub fn on_data_modified(&mut self, modification: DataTrace) {

components/raftstore-v2/src/operation/ready/mod.rs

@@ -1200,6 +1200,7 @@ impl<EK: KvEngine, ER: RaftEngine> Storage<EK, ER> {
                     slog_panic!(self.logger(), "failed to clean up region"; "error" => ?e);
                 });
             self.init_apply_trace(write_task);
+            self.set_flushed_epoch(self.region_state().get_region().get_region_epoch().clone());
             self.set_ever_persisted();
         }
         if self.apply_trace().should_persist() {

components/raftstore-v2/src/operation/ready/snapshot.rs

@@ -698,6 +698,7 @@ impl<EK: KvEngine, ER: RaftEngine> Storage<EK, ER> {
         for cf in ALL_CFS {
             lb.put_flushed_index(region_id, cf, last_index, last_index)
                 .unwrap();
+            self.set_flushed_epoch(self.region_state().get_region().get_region_epoch().clone());
         }
 
         let (path, clean_split) = match self.split_init_mut() {

components/raftstore-v2/src/raft/apply.rs

@@ -298,8 +298,8 @@ impl<EK: KvEngine, R> Apply<EK, R> {
     }
 
     #[inline]
-    pub fn set_sst_applied_index(&mut self, uuid: Vec<Vec<u8>>, apply_index: u64) {
-        self.sst_apply_state.registe_ssts(uuid, apply_index);
+    pub fn sst_apply_state(&self) -> &SstApplyState {
+        &self.sst_apply_state
     }
 
     #[inline]

components/raftstore-v2/src/raft/storage.rs

@@ -9,6 +9,7 @@ use collections::HashMap;
 use engine_traits::{KvEngine, RaftEngine};
 use kvproto::{
     metapb,
+    metapb::RegionEpoch,
     raft_serverpb::{RaftApplyState, RaftLocalState, RegionLocalState},
 };
 use raft::{
@@ -46,6 +47,8 @@ pub struct Storage<EK: KvEngine, ER> {
     split_init: Option<Box<SplitInit>>,
     /// The flushed index of all CFs.
     apply_trace: ApplyTrace,
+    // The flushed epoch means that the epoch has persisted into the raft engine.
+    flushed_epoch: RegionEpoch,
 }
 
 impl<EK: KvEngine, ER> Debug for Storage<EK, ER> {
@@ -129,6 +132,18 @@ impl<EK: KvEngine, ER> Storage<EK, ER> {
     pub fn has_dirty_data(&self) -> bool {
         self.has_dirty_data
     }
+
+    #[inline]
+    pub fn set_flushed_epoch(&mut self, epoch: RegionEpoch) {
+        if util::is_epoch_stale(&self.flushed_epoch, &epoch) {
+            self.flushed_epoch = epoch;
+        }
+    }
+
+    #[inline]
+    pub fn flushed_epoch(&self) -> &RegionEpoch {
+        &self.flushed_epoch
+    }
 }
 
 impl<EK: KvEngine, ER: RaftEngine> Storage<EK, ER> {
@@ -151,6 +166,7 @@ impl<EK: KvEngine, ER: RaftEngine> Storage<EK, ER> {
             }
         };
         let region = region_state.get_region();
+        let epoch = region.get_region_epoch().clone();
         let logger = logger.new(o!("region_id" => region.id, "peer_id" => peer.get_id()));
         let has_dirty_data =
             match engine.get_dirty_mark(region.get_id(), region_state.get_tablet_index()) {
@@ -183,6 +199,7 @@ impl<EK: KvEngine, ER: RaftEngine> Storage<EK, ER> {
             gen_snap_task: RefCell::new(Box::new(None)),
             split_init: None,
             apply_trace,
+            flushed_epoch: epoch,
         })
     }