chore: apply suggestions from CR

GreptimeTeam · Dec 4, 2023 · 19b8000 · 19b8000
1 parent 4a16466
commit 19b8000
Showing 1 changed file with 16 additions and 19 deletions.
diff --git a/src/datanode/src/alive_keeper.rs b/src/datanode/src/alive_keeper.rs
@@ -355,7 +355,6 @@ struct CountdownTask {
 }
 
 impl CountdownTask {
-    ///
     async fn run(&mut self) {
         // 30 years. See `Instant::far_future`.
         let far_future = Instant::now() + Duration::from_secs(86400 * 365 * 30);
@@ -377,23 +376,12 @@ impl CountdownTask {
                             countdown.set(tokio::time::sleep_until(first_deadline));
                         },
                         Some(CountdownCommand::Reset((role, deadline))) => {
-                            // The first-time granted regions might be ignored because the `first_deadline` is larger than the `region_lease_timeout`.
-                            // Therefore, we set writable at the outside.
-                            // TODO(weny): Considers setting `first_deadline` to `region_lease_timeout`.
                             let _ = self.region_server.set_writable(self.region_id, role.writable());
-
-                            if countdown.deadline() < deadline {
-                                trace!(
-                                    "Reset deadline of region {region_id} to approximately {} seconds later.",
-                                    (deadline - Instant::now()).as_secs_f32(),
-                                );
-                                countdown.set(tokio::time::sleep_until(deadline));
-                            }
-                            // Else the countdown could be either:
-                            // - not started yet;
-                            // - during startup protection;
-                            // - received a lagging heartbeat message.
-                            // All can be safely ignored.
+                            trace!(
+                                "Reset deadline of region {region_id} to approximately {} seconds later.",
+                                (deadline - Instant::now()).as_secs_f32(),
+                            );
+                            countdown.set(tokio::time::sleep_until(deadline));
                         },
                         None => {
                             info!(
@@ -413,7 +401,7 @@ impl CountdownTask {
                     let _ = self.region_server.set_writable(self.region_id, false);
                     metrics::LEASE_EXPIRED_REGION
                         .with_label_values(&[&format!("{}", self.region_id)])
-                        .add(1);
+                        .inc();
                     // resets the countdown.
                     let far_future = Instant::now() + Duration::from_secs(86400 * 365 * 30);
                     countdown.as_mut().reset(far_future);
@@ -477,14 +465,23 @@ mod test {
                     region_id: region_id.as_u64(),
                     role: api::v1::meta::RegionRole::Leader.into(),
                 }],
-                Instant::now() + Duration::from_millis(500),
+                Instant::now() + Duration::from_millis(200),
             )
             .await;
         tokio::time::sleep(Duration::from_millis(100)).await;
         assert!(alive_keeper.find_handle(region_id).await.is_some());
         let deadline = alive_keeper.deadline(region_id).await.unwrap();
         assert!(deadline >= Instant::now());
         assert_eq!(engine.role(region_id).unwrap(), RegionRole::Leader);
+
+        info!("Wait for lease expired");
+        // Sleep to wait lease expired.
+        tokio::time::sleep(Duration::from_millis(200)).await;
+        assert!(alive_keeper.find_handle(region_id).await.is_some());
+        assert_eq!(engine.role(region_id).unwrap(), RegionRole::Follower);
+
+        let deadline = alive_keeper.deadline(region_id).await.unwrap();
+        assert!(deadline > Instant::now() + Duration::from_secs(86400 * 365 * 29));
     }
 
     #[tokio::test(flavor = "multi_thread")]