raftstore: Observe when receive raft message (tikv#14043)

ref tikv#13855

Introduce observers when receive raft message.

Signed-off-by: CalvinNeo <calvinneo1995@gmail.com>

Author: CalvinNeo

components/cdc/src/observer.rs

@@ -273,6 +273,7 @@ mod tests {
                 prev_lead_transferee: raft::INVALID_ID,
                 vote: raft::INVALID_ID,
                 initialized: true,
+                peer_id: raft::INVALID_ID,
             },
         );
         match rx.recv_timeout(Duration::from_millis(10)).unwrap().unwrap() {
@@ -301,6 +302,7 @@ mod tests {
                 prev_lead_transferee: 3,
                 vote: 3,
                 initialized: true,
+                peer_id: raft::INVALID_ID,
             },
         );
         match rx.recv_timeout(Duration::from_millis(10)).unwrap().unwrap() {

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

@@ -768,6 +768,7 @@ impl<EK: KvEngine, ER: RaftEngine> Peer<EK, ER> {
                     prev_lead_transferee: target,
                     vote: self.raft_group().raft.vote,
                     initialized: self.storage().is_initialized(),
+                    peer_id: self.peer().get_id(),
                 },
             );
             self.proposal_control_mut().maybe_update_term(term);

components/raftstore/src/coprocessor/dispatcher.rs

@@ -8,6 +8,7 @@ use kvproto::{
     metapb::{Region, RegionEpoch},
     pdpb::CheckPolicy,
     raft_cmdpb::{ComputeHashRequest, RaftCmdRequest},
+    raft_serverpb::RaftMessage,
 };
 use protobuf::Message;
 use raft::eraftpb;
@@ -278,6 +279,7 @@ impl_box_observer_g!(
     ConsistencyCheckObserver,
     WrappedConsistencyCheckObserver
 );
+impl_box_observer!(BoxMessageObserver, MessageObserver, WrappedMessageObserver);
 
 /// Registry contains all registered coprocessors.
 #[derive(Clone)]
@@ -296,6 +298,7 @@ where
     read_index_observers: Vec<Entry<BoxReadIndexObserver>>,
     pd_task_observers: Vec<Entry<BoxPdTaskObserver>>,
     update_safe_ts_observers: Vec<Entry<BoxUpdateSafeTsObserver>>,
+    message_observers: Vec<Entry<BoxMessageObserver>>,
     // TODO: add endpoint
 }
 
@@ -313,6 +316,7 @@ impl<E: KvEngine> Default for Registry<E> {
             read_index_observers: Default::default(),
             pd_task_observers: Default::default(),
             update_safe_ts_observers: Default::default(),
+            message_observers: Default::default(),
         }
     }
 }
@@ -381,6 +385,10 @@ impl<E: KvEngine> Registry<E> {
     pub fn register_update_safe_ts_observer(&mut self, priority: u32, qo: BoxUpdateSafeTsObserver) {
         push!(priority, qo, self.update_safe_ts_observers);
     }
+
+    pub fn register_message_observer(&mut self, priority: u32, qo: BoxMessageObserver) {
+        push!(priority, qo, self.message_observers);
+    }
 }
 
 /// A macro that loops over all observers and returns early when error is found
@@ -780,6 +788,17 @@ impl<E: KvEngine> CoprocessorHost<E> {
         true
     }
 
+    /// Returns false if the message should not be stepped later.
+    pub fn on_raft_message(&self, msg: &RaftMessage) -> bool {
+        for observer in &self.registry.message_observers {
+            let observer = observer.observer.inner();
+            if !observer.on_raft_message(msg) {
+                return false;
+            }
+        }
+        true
+    }
+
     pub fn on_flush_applied_cmd_batch(
         &self,
         max_level: ObserveLevel,
@@ -890,6 +909,7 @@ mod tests {
         OnUpdateSafeTs = 23,
         PrePersist = 24,
         PreWriteApplyState = 25,
+        OnRaftMessage = 26,
     }
 
     impl Coprocessor for TestCoprocessor {}
@@ -1132,6 +1152,14 @@ mod tests {
         }
     }
 
+    impl MessageObserver for TestCoprocessor {
+        fn on_raft_message(&self, _: &RaftMessage) -> bool {
+            self.called
+                .fetch_add(ObserverIndex::OnRaftMessage as usize, Ordering::SeqCst);
+            true
+        }
+    }
+
     macro_rules! assert_all {
         ($target:expr, $expect:expr) => {{
             for (c, e) in ($target).iter().zip($expect) {
@@ -1168,6 +1196,8 @@ mod tests {
             .register_cmd_observer(1, BoxCmdObserver::new(ob.clone()));
         host.registry
             .register_update_safe_ts_observer(1, BoxUpdateSafeTsObserver::new(ob.clone()));
+        host.registry
+            .register_message_observer(1, BoxMessageObserver::new(ob.clone()));
 
         let mut index: usize = 0;
         let region = Region::default();
@@ -1282,6 +1312,11 @@ mod tests {
         host.pre_write_apply_state(&region);
         index += ObserverIndex::PreWriteApplyState as usize;
         assert_all!([&ob.called], &[index]);
+
+        let msg = RaftMessage::default();
+        host.on_raft_message(&msg);
+        index += ObserverIndex::OnRaftMessage as usize;
+        assert_all!([&ob.called], &[index]);
     }
 
     #[test]

components/raftstore/src/coprocessor/mod.rs

@@ -26,14 +26,16 @@ mod metrics;
 pub mod region_info_accessor;
 mod split_check;
 pub mod split_observer;
+use kvproto::raft_serverpb::RaftMessage;
 
 pub use self::{
     config::{Config, ConsistencyCheckMethod},
     consistency_check::{ConsistencyCheckObserver, Raw as RawConsistencyCheckObserver},
     dispatcher::{
         BoxAdminObserver, BoxApplySnapshotObserver, BoxCmdObserver, BoxConsistencyCheckObserver,
-        BoxPdTaskObserver, BoxQueryObserver, BoxRegionChangeObserver, BoxRoleObserver,
-        BoxSplitCheckObserver, BoxUpdateSafeTsObserver, CoprocessorHost, Registry, StoreHandle,
+        BoxMessageObserver, BoxPdTaskObserver, BoxQueryObserver, BoxRegionChangeObserver,
+        BoxRoleObserver, BoxSplitCheckObserver, BoxUpdateSafeTsObserver, CoprocessorHost, Registry,
+        StoreHandle,
     },
     error::{Error, Result},
     region_info_accessor::{
@@ -269,6 +271,7 @@ pub struct RoleChange {
     /// Which peer is voted by itself.
     pub vote: u64,
     pub initialized: bool,
+    pub peer_id: u64,
 }
 
 impl RoleChange {
@@ -280,6 +283,7 @@ impl RoleChange {
             prev_lead_transferee: raft::INVALID_ID,
             vote: raft::INVALID_ID,
             initialized: true,
+            peer_id: raft::INVALID_ID,
         }
     }
 }
@@ -334,6 +338,13 @@ pub trait RegionChangeObserver: Coprocessor {
     }
 }
 
+pub trait MessageObserver: Coprocessor {
+    /// Returns false if the message should not be stepped later.
+    fn on_raft_message(&self, _: &RaftMessage) -> bool {
+        true
+    }
+}
+
 #[derive(Clone, Debug, Default)]
 pub struct Cmd {
     pub index: u64,

components/raftstore/src/store/fsm/peer.rs

@@ -610,6 +610,9 @@ where
         for m in msgs.drain(..) {
             match m {
                 PeerMsg::RaftMessage(msg) => {
+                    if !self.ctx.coprocessor_host.on_raft_message(&msg.msg) {
+                        continue;
+                    }
                     if let Err(e) = self.on_raft_message(msg) {
                         error!(%e;
                             "handle raft message err";

components/raftstore/src/store/fsm/store.rs

@@ -753,6 +753,9 @@ impl<'a, EK: KvEngine + 'static, ER: RaftEngine + 'static, T: Transport>
             match m {
                 StoreMsg::Tick(tick) => self.on_tick(tick),
                 StoreMsg::RaftMessage(msg) => {
+                    if !self.ctx.coprocessor_host.on_raft_message(&msg.msg) {
+                        continue;
+                    }
                     if let Err(e) = self.on_raft_message(msg) {
                         if matches!(&e, Error::RegionNotRegistered { .. }) {
                             // This may happen in normal cases when add-peer runs slowly

components/raftstore/src/store/peer.rs

@@ -2335,6 +2335,7 @@ where
                     prev_lead_transferee: self.lead_transferee,
                     vote: self.raft_group.raft.vote,
                     initialized: self.is_initialized(),
+                    peer_id: self.peer.get_id(),
                 },
             );
             self.cmd_epoch_checker.maybe_update_term(self.term());

components/raftstore/src/store/snap.rs

@@ -207,7 +207,9 @@ fn retry_delete_snapshot(mgr: &SnapManagerCore, key: &SnapKey, snap: &Snapshot)
     false
 }
 
-fn gen_snapshot_meta(cf_files: &[CfFile], for_balance: bool) -> RaftStoreResult<SnapshotMeta> {
+// Create a SnapshotMeta that can be later put into RaftSnapshotData or written
+// into file.
+pub fn gen_snapshot_meta(cf_files: &[CfFile], for_balance: bool) -> RaftStoreResult<SnapshotMeta> {
     let mut meta = Vec::with_capacity(cf_files.len());
     for cf_file in cf_files {
         if !SNAPSHOT_CFS.iter().any(|cf| cf_file.cf == *cf) {
@@ -663,7 +665,8 @@ impl Snapshot {
         Ok(snapshot_meta)
     }
 
-    fn set_snapshot_meta(&mut self, snapshot_meta: SnapshotMeta) -> RaftStoreResult<()> {
+    // Validate and set SnapshotMeta of this Snapshot.
+    pub fn set_snapshot_meta(&mut self, snapshot_meta: SnapshotMeta) -> RaftStoreResult<()> {
         let mut cf_file_count_from_meta: Vec<usize> = vec![];
         let mut file_count = 0;
         let mut current_cf = "";
@@ -812,8 +815,9 @@ impl Snapshot {
         }
     }
 
-    // Only called in `do_build`.
-    fn save_meta_file(&mut self) -> RaftStoreResult<()> {
+    // Save `SnapshotMeta` to file.
+    // Used in `do_build` and by external crates.
+    pub fn save_meta_file(&mut self) -> RaftStoreResult<()> {
         let v = box_try!(self.meta_file.meta.as_ref().unwrap().write_to_bytes());
         if let Some(mut f) = self.meta_file.file.take() {
             // `meta_file` could be None for this case: in `init_for_building` the snapshot
@@ -1125,6 +1129,10 @@ impl Snapshot {
         file_system::metadata(&self.meta_file.path)
     }
 
+    pub fn meta_path(&self) -> &PathBuf {
+        &self.meta_file.path
+    }
+
     pub fn total_size(&self) -> u64 {
         self.cf_files
             .iter()