Merge stop_wait_group into daemon env.

src/apply_command.rs

@@ -55,7 +55,6 @@ impl<Command: ReplicableCommand> Raft<Command> {
         let rf = self.inner_state.clone();
         let condvar = self.apply_command_signal.clone();
         let snapshot_daemon = self.snapshot_daemon.clone();
-        let stop_wait_group = self.stop_wait_group.clone();
         let apply_command_daemon = move || {
             log::info!("{:?} apply command daemon running ...", me);
 
@@ -111,8 +110,6 @@ impl<Command: ReplicableCommand> Raft<Command> {
                 }
             }
             log::info!("{:?} apply command daemon done.", me);
-
-            drop(stop_wait_group);
         };
         self.daemon_env
             .watch_daemon(Daemon::ApplyCommand, apply_command_daemon);

src/daemon_env.rs

@@ -1,6 +1,7 @@
 use std::cell::RefCell;
 use std::sync::{Arc, Weak};
 
+use crossbeam_utils::sync::WaitGroup;
 use parking_lot::Mutex;
 
 #[cfg(all(not(test), feature = "integration-test"))]
@@ -33,6 +34,7 @@ macro_rules! check_or_record {
 pub(crate) struct DaemonEnv {
     data: Arc<Mutex<DaemonEnvData>>,
     thread_env: ThreadEnv,
+    stop_wait_group: Option<WaitGroup>,
 }
 
 #[derive(Debug)]
@@ -126,17 +128,30 @@ impl DaemonEnv {
         T: Send + 'static,
     {
         let thread_env = self.for_thread();
+        let stop_wait_group = self
+            .stop_wait_group
+            .clone()
+            .expect("Expecting a valid stop wait group when creating daemons");
         let thread = std::thread::Builder::new()
             .name(format!("ruaft-daemon-{:?}", daemon))
             .spawn(move || {
                 thread_env.attach();
                 func();
                 ThreadEnv::detach();
+                drop(stop_wait_group);
             })
             .expect("Creating daemon thread should never fail");
         self.data.lock().daemons.push((daemon, thread));
     }
 
+    pub fn wait_for_daemons(&mut self) {
+        if let Some(stop_wait_group) = self.stop_wait_group.take() {
+            stop_wait_group.wait();
+        } else {
+            panic!("Daemons can only be waited once")
+        }
+    }
+
     /// Makes sure that all daemons have been shutdown, no more errors can be
     /// added, checks if any error has been added, or if any daemon panicked.
     pub fn shutdown(self) {
@@ -224,7 +239,12 @@ impl DaemonEnv {
             #[cfg(all(not(test), feature = "integration-test"))]
             local_logger: thread_local_logger::get(),
         };
-        Self { data, thread_env }
+
+        Self {
+            data,
+            thread_env,
+            stop_wait_group: Some(WaitGroup::new()),
+        }
     }
 
     /// Creates a [`ThreadEnv`] that could be attached to a thread. Any code
@@ -278,6 +298,7 @@ impl ThreadEnv {
         DaemonEnv {
             data: env.data.upgrade().unwrap(),
             thread_env: env,
+            stop_wait_group: None,
         }
     }
 

src/election.rs

@@ -204,11 +204,6 @@ impl<Command: ReplicableCommand> Raft<Command> {
             }
 
             log::info!("{:?} election timer daemon done.", this.me);
-
-            let stop_wait_group = this.stop_wait_group.clone();
-            // Making sure the rest of `this` is dropped before the wait group.
-            drop(this);
-            drop(stop_wait_group);
         };
         self.daemon_env
             .watch_daemon(Daemon::ElectionTimer, election_daemon);

src/raft.rs

@@ -1,7 +1,6 @@
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 
-use crossbeam_utils::sync::WaitGroup;
 use parking_lot::{Condvar, Mutex};
 use serde_derive::{Deserialize, Serialize};
 
@@ -44,7 +43,6 @@ pub struct Raft<Command> {
     pub(crate) thread_pool: utils::ThreadPoolHolder,
 
     pub(crate) daemon_env: DaemonEnv,
-    pub(crate) stop_wait_group: WaitGroup,
 }
 
 impl<Command: ReplicableCommand> Raft<Command> {
@@ -119,7 +117,6 @@ impl<Command: ReplicableCommand> Raft<Command> {
             heartbeats_daemon: HeartbeatsDaemon::create(),
             thread_pool: utils::ThreadPoolHolder::new(thread_pool),
             daemon_env,
-            stop_wait_group: WaitGroup::new(),
         };
 
         // Running in a standalone thread.
@@ -183,7 +180,7 @@ impl<Command: ReplicableCommand> Raft<Command> {
         // DaemonEnv. The shutdown sequence is stop_wait_group -> thread_pool
         // -> DaemonEnv. The first and third cannot be combined with the second
         // in the middle.
-        self.stop_wait_group.wait();
+        self.daemon_env.wait_for_daemons();
         self.thread_pool
             .take()
             .expect(

src/snapshot.rs

@@ -142,7 +142,6 @@ impl<C: 'static + Clone + Send + serde::Serialize> Raft<C> {
         let rf = self.inner_state.clone();
         let persister = self.persister.clone();
         let snapshot_daemon = self.snapshot_daemon.clone();
-        let stop_wait_group = self.stop_wait_group.clone();
 
         log::info!("{:?} snapshot daemon running ...", me);
         let snapshot_daemon = move || loop {
@@ -155,7 +154,6 @@ impl<C: 'static + Clone + Send + serde::Serialize> Raft<C> {
                 drop(rf);
                 drop(persister);
                 drop(snapshot_daemon);
-                drop(stop_wait_group);
                 break;
             }
             if persister.state_size() >= max_state_size {

src/sync_log_entries.rs

@@ -102,11 +102,6 @@ impl<Command: ReplicableCommand> Raft<Command> {
             }
 
             log::info!("{:?} sync log entries daemon done.", this.me);
-
-            let stop_wait_group = this.stop_wait_group.clone();
-            // Making sure the rest of `this` is dropped before the wait group.
-            drop(this);
-            drop(stop_wait_group);
         };
         self.daemon_env
             .watch_daemon(Daemon::SyncLogEntries, sync_log_entry_daemon);

src/verify_authority.rs

@@ -377,7 +377,6 @@ impl<Command: 'static + Send> Raft<Command> {
         let keep_running = self.keep_running.clone();
         let this_daemon = self.verify_authority_daemon.clone();
         let rf = self.inner_state.clone();
-        let stop_wait_group = self.stop_wait_group.clone();
 
         let verify_authority_daemon = move || {
             log::info!("{:?} verify authority daemon running ...", me);
@@ -394,8 +393,6 @@ impl<Command: 'static + Send> Raft<Command> {
                 );
             }
             log::info!("{:?} verify authority daemon done.", me);
-
-            drop(stop_wait_group);
         };
         self.daemon_env
             .watch_daemon(Daemon::VerifyAuthority, verify_authority_daemon);