refactor: Improve SPV shutdown handling with CancellationToken (#187)

Author: xdustinface

dash-spv-ffi/Cargo.toml

@@ -17,6 +17,7 @@ dashcore = { path = "../dash", package = "dashcore" }
 libc = "0.2"
 once_cell = "1.19"
 tokio = { version = "1", features = ["full"] }
+tokio-util = "0.7"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 log = "0.4"

dash-spv-ffi/src/client.rs

@@ -17,11 +17,12 @@ use std::collections::HashMap;
 use std::ffi::{CStr, CString};
 use std::os::raw::{c_char, c_void};
 use std::str::FromStr;
-use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::{Arc, Mutex};
 use std::time::Duration;
 use tokio::runtime::Runtime;
 use tokio::sync::mpsc::{error::TryRecvError, UnboundedReceiver};
+use tokio_util::sync::CancellationToken;
 
 /// Global callback registry for thread-safe callback management
 static CALLBACK_REGISTRY: Lazy<Arc<Mutex<CallbackRegistry>>> =
@@ -104,12 +105,6 @@ struct SyncCallbackData {
     _marker: std::marker::PhantomData<()>,
 }
 
-async fn wait_for_shutdown_signal(signal: Arc<AtomicBool>) {
-    while !signal.load(Ordering::Relaxed) {
-        tokio::time::sleep(Duration::from_millis(50)).await;
-    }
-}
-
 /// FFIDashSpvClient structure
 type InnerClient = DashSpvClient<
     key_wallet_manager::wallet_manager::WalletManager<
@@ -126,7 +121,7 @@ pub struct FFIDashSpvClient {
     event_callbacks: Arc<Mutex<FFIEventCallbacks>>,
     active_threads: Arc<Mutex<Vec<std::thread::JoinHandle<()>>>>,
     sync_callbacks: Arc<Mutex<Option<SyncCallbackData>>>,
-    shutdown_signal: Arc<AtomicBool>,
+    shutdown_token: CancellationToken,
     // Stored event receiver for pull-based draining (no background thread by default)
     event_rx: Arc<Mutex<Option<UnboundedReceiver<dash_spv::types::SpvEvent>>>>,
 }
@@ -197,7 +192,7 @@ pub unsafe extern "C" fn dash_spv_ffi_client_new(
                 event_callbacks: Arc::new(Mutex::new(FFIEventCallbacks::default())),
                 active_threads: Arc::new(Mutex::new(Vec::new())),
                 sync_callbacks: Arc::new(Mutex::new(None)),
-                shutdown_signal: Arc::new(AtomicBool::new(false)),
+                shutdown_token: CancellationToken::new(),
                 event_rx: Arc::new(Mutex::new(None)),
             };
             Box::into_raw(Box::new(ffi_client))
@@ -378,8 +373,8 @@ pub unsafe extern "C" fn dash_spv_ffi_client_drain_events(client: *mut FFIDashSp
     FFIErrorCode::Success as i32
 }
 
-fn stop_client_internal(client: &FFIDashSpvClient) -> Result<(), dash_spv::SpvError> {
-    client.shutdown_signal.store(true, Ordering::Relaxed);
+fn stop_client_internal(client: &mut FFIDashSpvClient) -> Result<(), dash_spv::SpvError> {
+    client.shutdown_token.cancel();
 
     // Ensure callbacks are cleared so no further progress/completion notifications fire.
     {
@@ -411,7 +406,7 @@ fn stop_client_internal(client: &FFIDashSpvClient) -> Result<(), dash_spv::SpvEr
         res
     });
 
-    client.shutdown_signal.store(false, Ordering::Relaxed);
+    client.shutdown_token = CancellationToken::new();
 
     result
 }
@@ -525,7 +520,7 @@ pub unsafe extern "C" fn dash_spv_ffi_client_start(client: *mut FFIDashSpvClient
 pub unsafe extern "C" fn dash_spv_ffi_client_stop(client: *mut FFIDashSpvClient) -> i32 {
     null_check!(client);
 
-    let client = &(*client);
+    let client = &mut (*client);
     match stop_client_internal(client) {
         Ok(()) => FFIErrorCode::Success as i32,
         Err(e) => {
@@ -785,7 +780,6 @@ pub unsafe extern "C" fn dash_spv_ffi_client_sync_to_tip_with_progress(
     let inner = client.inner.clone();
     let runtime = client.runtime.clone();
     let sync_callbacks = client.sync_callbacks.clone();
-    let shutdown_signal = client.shutdown_signal.clone();
 
     // Take progress receiver from client
     let progress_receiver = {
@@ -797,7 +791,7 @@ pub unsafe extern "C" fn dash_spv_ffi_client_sync_to_tip_with_progress(
     if let Some(mut receiver) = progress_receiver {
         let runtime_handle = runtime.handle().clone();
         let sync_callbacks_clone = sync_callbacks.clone();
-        let shutdown_signal_clone = shutdown_signal.clone();
+        let shutdown_token_monitor = client.shutdown_token.clone();
 
         let handle = std::thread::spawn(move || {
             runtime_handle.block_on(async move {
@@ -859,7 +853,7 @@ pub unsafe extern "C" fn dash_spv_ffi_client_sync_to_tip_with_progress(
                                 None => break,
                             }
                         }
-                        _ = wait_for_shutdown_signal(shutdown_signal_clone.clone()) => {
+                        _ = shutdown_token_monitor.cancelled() => {
                             break;
                         }
                     }
@@ -874,15 +868,12 @@ pub unsafe extern "C" fn dash_spv_ffi_client_sync_to_tip_with_progress(
     // Spawn sync task in a separate thread with safe callback access
     let runtime_handle = runtime.handle().clone();
     let sync_callbacks_clone = sync_callbacks.clone();
-    let shutdown_signal_for_thread = shutdown_signal.clone();
-    let stop_triggered_for_thread = Arc::new(AtomicBool::new(false));
+    let shutdown_token_sync = client.shutdown_token.clone();
     let sync_handle = std::thread::spawn(move || {
-        let stop_triggered_for_callback = stop_triggered_for_thread.clone();
+        let shutdown_token_callback = shutdown_token_sync.clone();
         // Run monitoring loop
         let monitor_result = runtime_handle.block_on({
             let inner = inner.clone();
-            let shutdown_signal_for_thread = shutdown_signal_for_thread.clone();
-            let stop_triggered_for_thread = stop_triggered_for_callback.clone();
             async move {
                 let mut spv_client = {
                     let mut guard = inner.lock().unwrap();
@@ -903,8 +894,7 @@ pub unsafe extern "C" fn dash_spv_ffi_client_sync_to_tip_with_progress(
                         Ok(inner) => inner,
                         Err(_) => Ok(()),
                     },
-                    _ = wait_for_shutdown_signal(shutdown_signal_for_thread.clone()) => {
-                        stop_triggered_for_thread.store(true, Ordering::Relaxed);
+                    _ = shutdown_token_sync.cancelled() => {
                         abort_handle.abort();
                         match monitor_future.as_mut().await {
                             Ok(inner) => inner,
@@ -930,7 +920,7 @@ pub unsafe extern "C" fn dash_spv_ffi_client_sync_to_tip_with_progress(
                     ..
                 }) = registry.unregister(callback_data.callback_id)
                 {
-                    if stop_triggered_for_callback.load(Ordering::Relaxed) {
+                    if shutdown_token_callback.is_cancelled() {
                         let msg = CString::new("Sync stopped by request").unwrap_or_else(|_| {
                             CString::new("Sync stopped").expect("hardcoded string is safe")
                         });
@@ -984,7 +974,7 @@ pub unsafe extern "C" fn dash_spv_ffi_client_sync_to_tip_with_progress(
 pub unsafe extern "C" fn dash_spv_ffi_client_cancel_sync(client: *mut FFIDashSpvClient) -> i32 {
     null_check!(client);
 
-    let client = &(*client);
+    let client = &mut (*client);
 
     match stop_client_internal(client) {
         Ok(()) => FFIErrorCode::Success as i32,
@@ -1318,8 +1308,8 @@ pub unsafe extern "C" fn dash_spv_ffi_client_destroy(client: *mut FFIDashSpvClie
     if !client.is_null() {
         let client = Box::from_raw(client);
 
-        // Set shutdown signal to stop all threads
-        client.shutdown_signal.store(true, Ordering::Relaxed);
+        // Cancel shutdown token to stop all threads
+        client.shutdown_token.cancel();
 
         // Clean up any registered callbacks
         if let Some(ref callback_data) = *client.sync_callbacks.lock().unwrap() {

dash-spv/Cargo.toml

@@ -23,6 +23,7 @@ clap = { version = "4.0", features = ["derive"] }
 
 # Async runtime
 tokio = { version = "1.0", features = ["full"] }
+tokio-util = "0.7"
 async-trait = "0.1"
 
 # Error handling

dash-spv/src/network/manager.rs

@@ -3,7 +3,7 @@
 use std::collections::{HashMap, HashSet};
 use std::net::SocketAddr;
 use std::path::PathBuf;
-use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 use std::time::{Duration, SystemTime};
 use tokio::sync::{mpsc, Mutex};
@@ -14,6 +14,7 @@ use async_trait::async_trait;
 use dashcore::network::constants::ServiceFlags;
 use dashcore::network::message::NetworkMessage;
 use dashcore::Network;
+use tokio_util::sync::CancellationToken;
 
 use crate::client::config::MempoolStrategy;
 use crate::client::ClientConfig;
@@ -43,8 +44,8 @@ pub struct PeerNetworkManager {
     reputation_manager: Arc<PeerReputationManager>,
     /// Network type
     network: Network,
-    /// Shutdown signal
-    shutdown: Arc<AtomicBool>,
+    /// Shutdown token
+    shutdown_token: CancellationToken,
     /// Channel for incoming messages
     message_tx: mpsc::Sender<(SocketAddr, NetworkMessage)>,
     message_rx: Arc<Mutex<mpsc::Receiver<(SocketAddr, NetworkMessage)>>>,
@@ -109,7 +110,7 @@ impl PeerNetworkManager {
             peer_store: Arc::new(peer_store),
             reputation_manager,
             network: config.network,
-            shutdown: Arc::new(AtomicBool::new(false)),
+            shutdown_token: CancellationToken::new(),
             message_tx,
             message_rx: Arc::new(Mutex::new(message_rx)),
             tasks: Arc::new(Mutex::new(JoinSet::new())),
@@ -204,7 +205,7 @@ impl PeerNetworkManager {
         let network = self.network;
         let message_tx = self.message_tx.clone();
         let addrv2_handler = self.addrv2_handler.clone();
-        let shutdown = self.shutdown.clone();
+        let shutdown_token = self.shutdown_token.clone();
         let reputation_manager = self.reputation_manager.clone();
         let mempool_strategy = self.mempool_strategy;
         let user_agent = self.user_agent.clone();
@@ -245,7 +246,7 @@ impl PeerNetworkManager {
                                 pool.clone(),
                                 message_tx,
                                 addrv2_handler,
-                                shutdown,
+                                shutdown_token,
                                 reputation_manager.clone(),
                                 connected_peer_count.clone(),
                             )
@@ -287,19 +288,19 @@ impl PeerNetworkManager {
         pool: Arc<ConnectionPool>,
         message_tx: mpsc::Sender<(SocketAddr, NetworkMessage)>,
         addrv2_handler: Arc<AddrV2Handler>,
-        shutdown: Arc<AtomicBool>,
+        shutdown_token: CancellationToken,
         reputation_manager: Arc<PeerReputationManager>,
         connected_peer_count: Arc<AtomicUsize>,
     ) {
         tokio::spawn(async move {
             log::debug!("Starting peer reader loop for {}", addr);
             let mut loop_iteration = 0;
 
-            while !shutdown.load(Ordering::Relaxed) {
+            loop {
                 loop_iteration += 1;
 
                 // Check shutdown signal first with detailed logging
-                if shutdown.load(Ordering::Relaxed) {
+                if shutdown_token.is_cancelled() {
                     log::info!("Breaking peer reader loop for {} - shutdown signal received (iteration {})", addr, loop_iteration);
                     break;
                 }
@@ -326,7 +327,15 @@ impl PeerNetworkManager {
 
                     // Now get write lock only for the duration of the read
                     let mut conn_guard = conn.write().await;
-                    conn_guard.receive_message().await
+                    tokio::select! {
+                        message = conn_guard.receive_message() => {
+                            message
+                        },
+                        _ = shutdown_token.cancelled() => {
+                            log::info!("Breaking peer reader loop for {} - shutdown signal received while reading (iteration {})", addr, loop_iteration);
+                            break;
+                        }
+                    }
                 };
 
                 match msg_result {
@@ -576,7 +585,7 @@ impl PeerNetworkManager {
         let pool = self.pool.clone();
         let discovery = self.discovery.clone();
         let network = self.network;
-        let shutdown = self.shutdown.clone();
+        let shutdown_token = self.shutdown_token.clone();
         let addrv2_handler = self.addrv2_handler.clone();
         let peer_store = self.peer_store.clone();
         let reputation_manager = self.reputation_manager.clone();
@@ -599,7 +608,7 @@ impl PeerNetworkManager {
 
         let mut tasks = self.tasks.lock().await;
         tasks.spawn(async move {
-            while !shutdown.load(Ordering::Relaxed) {
+            while !shutdown_token.is_cancelled() {
                 // Clean up disconnected peers
                 pool.cleanup_disconnected().await;
 
@@ -612,7 +621,13 @@ impl PeerNetworkManager {
                     for addr in initial_peers.iter() {
                         if !pool.is_connected(addr).await && !pool.is_connecting(addr).await {
                             log::info!("Reconnecting to exclusive peer: {}", addr);
-                            connect_fn(*addr).await;
+                            tokio::select! {
+                                _= connect_fn(*addr) => {},
+                                _ = shutdown_token.cancelled() => {
+                                    log::info!("Maintenance loop shutting down during connection attempt (exclusive)");
+                                    break;
+                                }
+                            }
                         }
                     }
                 } else {
@@ -642,7 +657,13 @@ impl PeerNetworkManager {
 
                         for addr in best_peers {
                             if !pool.is_connected(&addr).await && !pool.is_connecting(&addr).await {
-                                connect_fn(addr).await;
+                                tokio::select! {
+                                    _= connect_fn(addr) => {},
+                                    _ = shutdown_token.cancelled() => {
+                                        log::info!("Maintenance loop shutting down during connection attempt (min peers)");
+                                        break;
+                                    }
+                                }
                                 attempted += 1;
                                 if attempted >= needed {
                                     break;
@@ -661,11 +682,23 @@ impl PeerNetworkManager {
                                 });
                             if elapsed >= DNS_DISCOVERY_DELAY {
                                 log::info!("Using DNS discovery after {}s delay", elapsed.as_secs());
-                                let dns_peers = discovery.discover_peers(network).await;
+                                let dns_peers = tokio::select! {
+                                    peers = discovery.discover_peers(network) => peers,
+                                    _ = shutdown_token.cancelled() => {
+                                        log::info!("Maintenance loop shutting down during DNS discovery");
+                                        break;
+                                    }
+                                };
                                 let mut dns_attempted = 0;
                                 for addr in dns_peers.into_iter() {
                                     if !pool.is_connected(&addr).await && !pool.is_connecting(&addr).await {
-                                        connect_fn(addr).await;
+                                        tokio::select! {
+                                            _= connect_fn(addr) => {},
+                                            _ = shutdown_token.cancelled() => {
+                                                log::info!("Maintenance loop shutting down during connection attempt (dns)");
+                                                break;
+                                            }
+                                        }
                                         dns_attempted += 1;
                                         if dns_attempted >= needed {
                                             break;
@@ -719,7 +752,15 @@ impl PeerNetworkManager {
                     }
                 }
 
-                time::sleep(MAINTENANCE_INTERVAL).await;
+                tokio::select! {
+                    _ = time::sleep(MAINTENANCE_INTERVAL) => {
+                        log::debug!("Maintenance interval elapsed");
+                    }
+                    _ = shutdown_token.cancelled() => {
+                        log::info!("Maintenance loop shutting down");
+                        break;
+                    }
+                }
             }
         });
     }
@@ -942,7 +983,7 @@ impl PeerNetworkManager {
     /// Shutdown the network manager
     pub async fn shutdown(&self) {
         log::info!("Shutting down peer network manager");
-        self.shutdown.store(true, Ordering::Relaxed);
+        self.shutdown_token.cancel();
 
         // Save known peers before shutdown
         let addresses = self.addrv2_handler.get_addresses_for_peer(MAX_ADDR_TO_STORE).await;
@@ -983,7 +1024,7 @@ impl Clone for PeerNetworkManager {
             peer_store: self.peer_store.clone(),
             reputation_manager: self.reputation_manager.clone(),
             network: self.network,
-            shutdown: self.shutdown.clone(),
+            shutdown_token: self.shutdown_token.clone(),
             message_tx: self.message_tx.clone(),
             message_rx: self.message_rx.clone(),
             tasks: self.tasks.clone(),