refactor: add timeout mechanism for wait_for_XX_shutdown method

Signed-off-by: Phoeniix Zhao <[email protected]>
xline-kv · Apr 28, 2024 · 7026b00 · 7026b00
1 parent 1afbbbc
commit 7026b00
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 8 deletions.
diff --git a/crates/curp/tests/it/common/curp_group.rs b/crates/curp/tests/it/common/curp_group.rs
@@ -29,6 +29,7 @@ use tokio::{
     runtime::{Handle, Runtime},
     sync::{mpsc, watch},
     task::{block_in_place, JoinHandle},
+    time::timeout,
 };
 use tokio_stream::wrappers::TcpListenerStream;
 use tonic::transport::{Certificate, Channel, ClientTlsConfig, Endpoint, ServerTlsConfig};
@@ -332,21 +333,28 @@ impl CurpGroup {
             .all(|node| node.task_manager.is_finished())
     }
 
-    pub async fn wait_for_node_shutdown(&self, node_id: u64) {
+    pub async fn wait_for_node_shutdown(&self, node_id: u64, duration: Duration) {
         let node = self
             .nodes
             .get(&node_id)
             .expect("{node_id} should exist in nodes");
         let res = std::iter::once(node);
-        Self::wait_for_targets_shutdown(res).await;
+        timeout(duration, Self::wait_for_targets_shutdown(res))
+            .await
+            .expect("wait for group to shutdown timeout");
         assert!(
             node.task_manager.is_finished(),
             "The target node({node_id}) is not finished yet"
         );
     }
 
-    pub async fn wait_for_group_shutdown(&self) {
-        Self::wait_for_targets_shutdown(self.nodes.values()).await;
+    pub async fn wait_for_group_shutdown(&self, duration: Duration) {
+        timeout(
+            duration,
+            Self::wait_for_targets_shutdown(self.nodes.values()),
+        )
+        .await
+        .expect("wait for group to shutdown timeout");
         assert!(self.is_finished(), "The group is not finished yet");
     }
 

diff --git a/crates/curp/tests/it/server.rs b/crates/curp/tests/it/server.rs
@@ -297,7 +297,7 @@ async fn shutdown_rpc_should_shutdown_the_cluster() {
     ));
 
     let collection = collection_task.await.unwrap();
-    group.wait_for_group_shutdown().await;
+    group.wait_for_group_shutdown(Duration::from_secs(7)).await;
 
     let group = CurpGroup::new_rocks(3, tmp_path).await;
     let client = group.new_client().await;
@@ -418,7 +418,7 @@ async fn shutdown_rpc_should_shutdown_the_cluster_when_client_has_wrong_leader()
         .unwrap();
     client.propose_shutdown().await.unwrap();
 
-    group.wait_for_group_shutdown().await;
+    group.wait_for_group_shutdown(Duration::from_secs(7)).await;
 }
 
 #[tokio::test(flavor = "multi_thread")]
@@ -548,7 +548,7 @@ async fn shutdown_rpc_should_shutdown_the_cluster_when_client_has_wrong_cluster(
         .await;
     client.propose_shutdown().await.unwrap();
 
-    group.wait_for_group_shutdown().await;
+    group.wait_for_group_shutdown(Duration::from_secs(7)).await;
 }
 
 #[tokio::test(flavor = "multi_thread")]
@@ -574,7 +574,9 @@ async fn propose_conf_change_rpc_should_work_when_client_has_wrong_cluster() {
     let members = client.propose_conf_change(changes).await.unwrap();
     assert_eq!(members.len(), 3);
     assert!(members.iter().all(|m| m.id != node_id));
-    group.wait_for_node_shutdown(node_id).await;
+    group
+        .wait_for_node_shutdown(node_id, Duration::from_secs(7))
+        .await;
 }
 
 #[tokio::test(flavor = "multi_thread")]