refactor: remove sleep_secs method when waiting the cluster to shutdown

Signed-off-by: Phoeniix Zhao <[email protected]>

crates/curp/tests/it/common/curp_group.rs

@@ -25,7 +25,7 @@ use engine::{
     Engine, EngineType, MemorySnapshotAllocator, RocksSnapshotAllocator, Snapshot,
     SnapshotAllocator,
 };
-use futures::{future::join_all, Future};
+use futures::{future::join_all, stream::FuturesUnordered, Future};
 use itertools::Itertools;
 use tokio::{
     net::TcpListener,
@@ -339,6 +339,42 @@ impl CurpGroup {
             .all(|node| node.task_manager.is_finished())
     }
 
+    pub async fn wait_for_node_shutdown(&self, node_id: u64) {
+        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;
+        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;
+        assert!(self.is_finished(), "The group is not finished yet");
+    }
+
+    async fn wait_for_targets_shutdown(targets: impl Iterator<Item = &CurpNode>) {
+        let final_tasks: [TaskName; 3] = [
+            TaskName::WatchTask,
+            TaskName::ConfChange,
+            TaskName::LogPersist,
+        ];
+        let listeners = targets
+            .flat_map(|node| {
+                final_tasks
+                    .iter()
+                    .map(|task| node.task_manager.get_shutdown_listener(task.to_owned()))
+                    .collect::<Vec<_>>()
+            })
+            .collect::<Vec<_>>();
+        let waiters: Vec<_> = listeners.iter().map(|l| l.wait()).collect();
+        futures::future::join_all(waiters.into_iter()).await;
+    }
+
     async fn stop(&mut self) {
         debug!("curp group stopping");
 

crates/curp/tests/it/server.rs

@@ -297,8 +297,7 @@ async fn shutdown_rpc_should_shutdown_the_cluster() {
     ));
 
     let collection = collection_task.await.unwrap();
-    sleep_secs(1).await; // wait for the cluster to shutdown
-    assert!(group.is_finished());
+    group.wait_for_group_shutdown().await;
 
     let group = CurpGroup::new_rocks(3, tmp_path).await;
     let client = group.new_client().await;
@@ -419,8 +418,7 @@ async fn shutdown_rpc_should_shutdown_the_cluster_when_client_has_wrong_leader()
         .unwrap();
     client.propose_shutdown().await.unwrap();
 
-    sleep_secs(7).await; // wait for the cluster to shutdown
-    assert!(group.is_finished());
+    group.wait_for_group_shutdown().await;
 }
 
 #[tokio::test(flavor = "multi_thread")]
@@ -550,8 +548,7 @@ async fn shutdown_rpc_should_shutdown_the_cluster_when_client_has_wrong_cluster(
         .await;
     client.propose_shutdown().await.unwrap();
 
-    sleep_secs(7).await; // wait for the cluster to shutdown
-    assert!(group.is_finished());
+    group.wait_for_group_shutdown().await;
 }
 
 #[tokio::test(flavor = "multi_thread")]
@@ -577,13 +574,7 @@ 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));
-    sleep_secs(7).await;
-    assert!(group
-        .nodes
-        .get(&node_id)
-        .unwrap()
-        .task_manager
-        .is_finished());
+    group.wait_for_node_shutdown(node_id).await;
 }
 
 #[tokio::test(flavor = "multi_thread")]