fix: reset region detector after migration (#8234)

Signed-off-by: WenyXu <wenymedia@gmail.com>

src/common/meta/src/ddl.rs

@@ -76,6 +76,9 @@ pub trait RegionFailureDetectorController: Send + Sync {
     /// Registers failure detectors for the given identifiers.
     async fn register_failure_detectors(&self, detecting_regions: Vec<DetectingRegion>);
 
+    /// Resets failure detectors for the given identifiers.
+    async fn reset_failure_detectors(&self, detecting_regions: Vec<DetectingRegion>);
+
     /// Deregisters failure detectors for the given identifiers.
     async fn deregister_failure_detectors(&self, detecting_regions: Vec<DetectingRegion>);
 }
@@ -88,6 +91,8 @@ pub struct NoopRegionFailureDetectorControl;
 impl RegionFailureDetectorController for NoopRegionFailureDetectorControl {
     async fn register_failure_detectors(&self, _detecting_regions: Vec<DetectingRegion>) {}
 
+    async fn reset_failure_detectors(&self, _detecting_regions: Vec<DetectingRegion>) {}
+
     async fn deregister_failure_detectors(&self, _detecting_regions: Vec<DetectingRegion>) {}
 }
 

src/meta-srv/src/failure_detector.rs

@@ -160,6 +160,11 @@ impl PhiAccrualFailureDetector {
     pub(crate) fn acceptable_heartbeat_pause_millis(&self) -> u32 {
         self.acceptable_heartbeat_pause_millis
     }
+
+    #[cfg(test)]
+    pub(crate) fn last_heartbeat_millis(&self) -> Option<i64> {
+        self.last_heartbeat_millis
+    }
 }
 
 /// Calculation of phi, derived from the Cumulative distribution function for

src/meta-srv/src/procedure/region_migration.rs

@@ -628,6 +628,23 @@ impl Context {
         );
     }
 
+    /// Notifies the RegionSupervisor to reset failure detectors of candidate regions.
+    pub async fn reset_failure_detectors_for_candidate_regions(&self) {
+        let datanode_id = self.persistent_ctx.to_peer.id;
+        let region_ids = &self.persistent_ctx.region_ids;
+        let detecting_regions = region_ids
+            .iter()
+            .map(|region_id| (datanode_id, *region_id))
+            .collect::<Vec<_>>();
+        self.region_failure_detector_controller
+            .reset_failure_detectors(detecting_regions)
+            .await;
+        info!(
+            "Reset failure detectors after migration success for datanode {}, regions {:?}",
+            datanode_id, region_ids
+        );
+    }
+
     /// Notifies the RegionSupervisor to deregister failure detectors.
     ///
     /// The original failure detectors won't be removed once the procedure was triggered.
@@ -645,11 +662,11 @@ impl Context {
             .await;
     }
 
-    /// Notifies the RegionSupervisor to deregister failure detectors for the candidate region on the destination peer.
+    /// Notifies the RegionSupervisor to deregister failure detectors for the candidate regions on the destination peer.
     ///
-    /// The candidate region may be created on the destination peer,
-    /// so we need to deregister the failure detectors for the candidate region if the procedure is aborted.
-    pub async fn deregister_failure_detectors_for_candidate_region(&self) {
+    /// The candidate regions may be created on the destination peer,
+    /// so we need to deregister the failure detectors for the candidate regions if the procedure is aborted.
+    pub async fn deregister_failure_detectors_for_candidate_regions(&self) {
         let to_peer_id = self.persistent_ctx.to_peer.id;
         let region_ids = &self.persistent_ctx.region_ids;
         let detecting_regions = region_ids
@@ -867,7 +884,7 @@ impl RegionMigrationProcedure {
             }
         }
         self.context
-            .deregister_failure_detectors_for_candidate_region()
+            .deregister_failure_detectors_for_candidate_regions()
             .await;
         self.context.register_failure_detectors().await;
 
@@ -919,7 +936,7 @@ impl Procedure for RegionMigrationProcedure {
                     // Consumes the opening region guard before deregistering the failure detectors.
                     self.context.volatile_ctx.opening_region_guards.clear();
                     self.context
-                        .deregister_failure_detectors_for_candidate_region()
+                        .deregister_failure_detectors_for_candidate_regions()
                         .await;
                     error!(
                         e;

src/meta-srv/src/procedure/region_migration/update_metadata/upgrade_candidate_region.rs

@@ -215,6 +215,7 @@ impl UpdateMetadata {
         }
 
         ctx.deregister_failure_detectors().await;
+        ctx.reset_failure_detectors_for_candidate_regions().await;
         // Consumes the guard.
         ctx.volatile_ctx.opening_region_guards.clear();
 

src/meta-srv/src/region/failure_detector.rs

@@ -73,6 +73,19 @@ impl RegionFailureDetector {
         })
     }
 
+    /// Resets the [`PhiAccrualFailureDetector`] for the specified [`DetectingRegion`] with the provided timestamp.
+    /// If a detector already exists for the region, it is reset. Otherwise, a new
+    /// detector is created and initialized with the provided timestamp.
+    pub(crate) fn reset_region_failure_detector(
+        &self,
+        detecting_region: DetectingRegion,
+        ts_millis: i64,
+    ) {
+        let mut detector = PhiAccrualFailureDetector::from_options(self.options);
+        detector.heartbeat(ts_millis);
+        self.detectors.insert(detecting_region, detector);
+    }
+
     /// Returns a [FailureDetectorEntry] iterator.
     pub(crate) fn iter(&self) -> impl Iterator<Item = FailureDetectorEntry<'_>> + '_ {
         self.detectors
@@ -142,4 +155,36 @@ mod tests {
         container.clear();
         assert!(container.is_empty());
     }
+
+    #[test]
+    fn test_reset_region_failure_detector() {
+        let container = RegionFailureDetector::new(Default::default());
+        let detecting_region = (2, RegionId::new(1, 1));
+        let first_heartbeat = 1_000;
+
+        {
+            let mut detector = container.region_failure_detector(detecting_region);
+            detector.heartbeat(first_heartbeat);
+            assert_eq!(Some(first_heartbeat), detector.last_heartbeat_millis());
+        }
+
+        let reset_at = 30_000;
+        container.reset_region_failure_detector(detecting_region, reset_at);
+
+        {
+            let detector = container.region_failure_detector(detecting_region);
+            assert_eq!(Some(reset_at), detector.last_heartbeat_millis());
+        }
+
+        let new_detecting_region = (3, RegionId::new(1, 1));
+        assert!(!container.contains(&new_detecting_region));
+
+        container.reset_region_failure_detector(new_detecting_region, reset_at);
+
+        assert!(container.contains(&new_detecting_region));
+        {
+            let detector = container.region_failure_detector(new_detecting_region);
+            assert_eq!(Some(reset_at), detector.last_heartbeat_millis());
+        }
+    }
 }

src/meta-srv/src/region/supervisor.rs

@@ -83,6 +83,7 @@ impl From<&Stat> for DatanodeHeartbeat {
 /// - `Tick`: This event is used to trigger region failure detection periodically.
 /// - `InitializeAllRegions`: This event is used to initialize all region failure detectors.
 /// - `RegisterFailureDetectors`: This event is used to register failure detectors for regions.
+/// - `ResetFailureDetectors`: This event is used to reset failure detectors for regions.
 /// - `DeregisterFailureDetectors`: This event is used to deregister failure detectors for regions.
 /// - `HeartbeatArrived`: This event presents the metasrv received [`DatanodeHeartbeat`] from the datanodes.
 /// - `Clear`: This event is used to reset the state of the supervisor, typically used
@@ -95,6 +96,7 @@ pub(crate) enum Event {
     InitializeAllRegions(tokio::sync::oneshot::Sender<()>),
     RegisterFailureDetectors(Vec<DetectingRegion>),
     DeregisterFailureDetectors(Vec<DetectingRegion>),
+    ResetFailureDetectors(Vec<DetectingRegion>),
     HeartbeatArrived(DatanodeHeartbeat),
     Clear,
     #[cfg(test)]
@@ -112,6 +114,9 @@ impl Debug for Event {
                 .debug_tuple("RegisterFailureDetectors")
                 .field(arg0)
                 .finish(),
+            Self::ResetFailureDetectors(arg0) => {
+                f.debug_tuple("ResetFailureDetectors").field(arg0).finish()
+            }
             Self::DeregisterFailureDetectors(arg0) => f
                 .debug_tuple("DeregisterFailureDetectors")
                 .field(arg0)
@@ -319,6 +324,16 @@ impl RegionFailureDetectorController for RegionFailureDetectorControl {
         }
     }
 
+    async fn reset_failure_detectors(&self, detecting_regions: Vec<DetectingRegion>) {
+        if let Err(err) = self
+            .sender
+            .send(Event::ResetFailureDetectors(detecting_regions))
+            .await
+        {
+            error!(err; "RegionSupervisor has stop receiving heartbeat.");
+        }
+    }
+
     async fn deregister_failure_detectors(&self, detecting_regions: Vec<DetectingRegion>) {
         if let Err(err) = self
             .sender
@@ -429,6 +444,9 @@ impl RegionSupervisor {
                 Event::RegisterFailureDetectors(detecting_regions) => {
                     self.register_failure_detectors(detecting_regions).await
                 }
+                Event::ResetFailureDetectors(detecting_regions) => {
+                    self.reset_failure_detectors(detecting_regions).await
+                }
                 Event::DeregisterFailureDetectors(detecting_regions) => {
                     self.deregister_failure_detectors(detecting_regions).await
                 }
@@ -508,6 +526,14 @@ impl RegionSupervisor {
         }
     }
 
+    async fn reset_failure_detectors(&mut self, detecting_regions: Vec<DetectingRegion>) {
+        let ts_millis = current_time_millis();
+        for region in detecting_regions {
+            self.failure_detector
+                .reset_region_failure_detector(region, ts_millis);
+        }
+    }
+
     async fn deregister_failure_detectors(&mut self, detecting_regions: Vec<DetectingRegion>) {
         for region in detecting_regions {
             self.failure_detector.remove(&region);