Fix: stop tick task when shutting down Raft

Signed-off-by: Matthias Wahl <[email protected]>

Author: mfelsche

openraft/src/core/tick.rs

@@ -6,8 +6,12 @@ use std::sync::Arc;
 use std::time::Duration;
 
 use tokio::sync::mpsc;
+use tokio::task::JoinHandle;
 use tokio::time::sleep_until;
 use tokio::time::Instant;
+use tracing::Instrument;
+use tracing::Level;
+use tracing::Span;
 
 use crate::raft::RaftMsg;
 use crate::NodeId;
@@ -55,11 +59,12 @@ where
     tx: mpsc::UnboundedSender<RaftMsg<C, N, S>>,
 
     /// Emit event or not
-    running: Arc<AtomicBool>,
+    enabled: Arc<AtomicBool>,
 }
 
 pub(crate) struct TickHandle {
-    running: Arc<AtomicBool>,
+    enabled: Arc<AtomicBool>,
+    join_handle: JoinHandle<()>,
 }
 
 impl<C, N, S> Tick<C, N, S>
@@ -68,12 +73,16 @@ where
     N: RaftNetworkFactory<C>,
     S: RaftStorage<C>,
 {
-    pub(crate) fn new(interval: Duration, tx: mpsc::UnboundedSender<RaftMsg<C, N, S>>, enabled: bool) -> Self {
-        Tick {
+    pub(crate) fn spawn(interval: Duration, tx: mpsc::UnboundedSender<RaftMsg<C, N, S>>, enabled: bool) -> TickHandle {
+        let enabled = Arc::new(AtomicBool::from(enabled));
+        let this = Self {
             interval,
-            running: Arc::new(AtomicBool::from(enabled)),
+            enabled: enabled.clone(),
             tx,
-        }
+        };
+        let join_handle =
+            tokio::spawn(this.tick_loop().instrument(tracing::span!(parent: &Span::current(), Level::DEBUG, "tick")));
+        TickHandle { enabled, join_handle }
     }
 
     pub(crate) async fn tick_loop(self) {
@@ -84,7 +93,7 @@ where
             let at = Instant::now() + self.interval;
             sleep_until(at).await;
 
-            if !self.running.load(Ordering::Relaxed) {
+            if !self.enabled.load(Ordering::Relaxed) {
                 i -= 1;
                 continue;
             }
@@ -97,17 +106,14 @@ where
             }
         }
     }
-
-    /// Return a handle to control the ticker.
-    pub(crate) fn get_handle(&self) -> TickHandle {
-        TickHandle {
-            running: self.running.clone(),
-        }
-    }
 }
 
 impl TickHandle {
     pub(crate) fn enable(&self, enabled: bool) {
-        self.running.store(enabled, Ordering::Relaxed);
+        self.enabled.store(enabled, Ordering::Relaxed);
+    }
+
+    pub(crate) async fn shutdown(&self) {
+        self.join_handle.abort();
     }
 }

openraft/src/raft.rs

@@ -13,9 +13,7 @@ use tokio::sync::watch;
 use tokio::sync::Mutex;
 use tokio::task::JoinError;
 use tokio::task::JoinHandle;
-use tracing::Instrument;
 use tracing::Level;
-use tracing::Span;
 
 use crate::config::Config;
 use crate::config::RuntimeConfig;
@@ -196,16 +194,12 @@ impl<C: RaftTypeConfig, N: RaftNetworkFactory<C>, S: RaftStorage<C>> Raft<C, N,
         let (tx_metrics, rx_metrics) = watch::channel(RaftMetrics::new_initial(id));
         let (tx_shutdown, rx_shutdown) = oneshot::channel();
 
-        let tick = Tick::new(
+        let tick_handle = Tick::spawn(
             Duration::from_millis(config.heartbeat_interval * 3 / 2),
             tx_api.clone(),
             config.enable_tick,
         );
 
-        let tick_handle = tick.get_handle();
-        let _tick_join_handle =
-            tokio::spawn(tick.tick_loop().instrument(tracing::span!(parent: &Span::current(), Level::DEBUG, "tick")));
-
         let runtime_config = Arc::new(RuntimeConfig::new(&config));
 
         let core_handle = RaftCore::spawn(
@@ -771,8 +765,8 @@ impl<C: RaftTypeConfig, N: RaftNetworkFactory<C>, S: RaftStorage<C>> Raft<C, N,
             let send_res = tx.send(());
             tracing::info!("sending shutdown signal to RaftCore, sending res: {:?}", send_res);
         }
-
         self.join_core_task().await;
+        self.inner.tick_handle.shutdown().await;
 
         // TODO(xp): API change: replace `JoinError` with `Fatal`,
         //           to let the caller know the return value of RaftCore task.