Switch stop signal over to tokio::sync::watch::channel

Author: tnull

src/lib.rs

@@ -570,11 +570,12 @@ impl Builder {
 			}
 		};
 
-		let stop_running = Arc::new(AtomicBool::new(false));
+		let (stop_sender, stop_receiver) = tokio::sync::watch::channel(());
 
 		Arc::new(Node {
 			runtime,
-			stop_running,
+			stop_sender,
+			stop_receiver,
 			config,
 			wallet,
 			tx_sync,
@@ -599,7 +600,8 @@ impl Builder {
 /// Needs to be initialized and instantiated through [`Builder::build`].
 pub struct Node {
 	runtime: Arc<RwLock<Option<tokio::runtime::Runtime>>>,
-	stop_running: Arc<AtomicBool>,
+	stop_sender: tokio::sync::watch::Sender<()>,
+	stop_receiver: tokio::sync::watch::Receiver<()>,
 	config: Arc<Config>,
 	wallet: Arc<Wallet<bdk::database::SqliteDatabase>>,
 	tx_sync: Arc<EsploraSyncClient<Arc<FilesystemLogger>>>,
@@ -633,8 +635,6 @@ impl Node {
 
 		let runtime = tokio::runtime::Builder::new_multi_thread().enable_all().build().unwrap();
 
-		let stop_running = Arc::new(AtomicBool::new(false));
-
 		let event_handler = Arc::new(EventHandler::new(
 			Arc::clone(&self.wallet),
 			Arc::clone(&self.event_queue),
@@ -653,66 +653,76 @@ impl Node {
 		let sync_cman = Arc::clone(&self.channel_manager);
 		let sync_cmon = Arc::clone(&self.chain_monitor);
 		let sync_logger = Arc::clone(&self.logger);
-		let stop_sync = Arc::clone(&stop_running);
+		let mut stop_sync = self.stop_receiver.clone();
 
 		std::thread::spawn(move || {
 			tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap().block_on(
 				async move {
+					let mut interval = tokio::time::interval(Duration::from_secs(30));
+					interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
 					loop {
-						if stop_sync.load(Ordering::Acquire) {
-							return;
-						}
 						let now = Instant::now();
-						match wallet.sync().await {
-							Ok(()) => log_info!(
-								sync_logger,
-								"Background sync of on-chain wallet finished in {}ms.",
-								now.elapsed().as_millis()
-							),
-							Err(err) => {
-								log_error!(
-									sync_logger,
-									"Background sync of on-chain wallet failed: {}",
-									err
-								)
+						tokio::select! {
+							_ = stop_sync.changed() => {
+								return;
+							}
+							_ = interval.tick() => {
+								match wallet.sync().await {
+									Ok(()) => log_info!(
+										sync_logger,
+										"Background sync of on-chain wallet finished in {}ms.",
+										now.elapsed().as_millis()
+										),
+									Err(err) => {
+										log_error!(
+											sync_logger,
+											"Background sync of on-chain wallet failed: {}",
+											err
+											)
+									}
+								}
 							}
 						}
-						tokio::time::sleep(Duration::from_secs(20)).await;
 					}
 				},
 			);
 		});
 
 		let sync_logger = Arc::clone(&self.logger);
-		let stop_sync = Arc::clone(&stop_running);
+		let mut stop_sync = self.stop_receiver.clone();
 		runtime.spawn(async move {
+			let mut interval = tokio::time::interval(Duration::from_secs(10));
+			interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
 			loop {
-				if stop_sync.load(Ordering::Acquire) {
-					return;
-				}
 				let now = Instant::now();
-				let confirmables = vec![
-					&*sync_cman as &(dyn Confirm + Sync + Send),
-					&*sync_cmon as &(dyn Confirm + Sync + Send),
-				];
-				match tx_sync.sync(confirmables).await {
-					Ok(()) => log_info!(
-						sync_logger,
-						"Background sync of Lightning wallet finished in {}ms.",
-						now.elapsed().as_millis()
-					),
-					Err(e) => {
-						log_error!(sync_logger, "Background sync of Lightning wallet failed: {}", e)
+				tokio::select! {
+					_ = stop_sync.changed() => {
+						return;
+					}
+					_ = interval.tick() => {
+						let confirmables = vec![
+							&*sync_cman as &(dyn Confirm + Sync + Send),
+							&*sync_cmon as &(dyn Confirm + Sync + Send),
+						];
+						match tx_sync.sync(confirmables).await {
+							Ok(()) => log_info!(
+								sync_logger,
+								"Background sync of Lightning wallet finished in {}ms.",
+								now.elapsed().as_millis()
+								),
+							Err(e) => {
+								log_error!(sync_logger, "Background sync of Lightning wallet failed: {}", e)
+							}
+						}
 					}
 				}
-				tokio::time::sleep(Duration::from_secs(5)).await;
 			}
 		});
 
 		if let Some(listening_address) = &self.config.listening_address {
 			// Setup networking
 			let peer_manager_connection_handler = Arc::clone(&self.peer_manager);
-			let stop_listen = Arc::clone(&stop_running);
+			let mut stop_listen = self.stop_receiver.clone();
 			let listening_address = listening_address.clone();
 
 			let bind_addr = listening_address
@@ -727,18 +737,22 @@ impl Node {
 						"Failed to bind to listen address/port - is something else already listening on it?",
 						);
 				loop {
-					if stop_listen.load(Ordering::Acquire) {
-						return;
-					}
 					let peer_mgr = Arc::clone(&peer_manager_connection_handler);
-					let tcp_stream = listener.accept().await.unwrap().0;
-					tokio::spawn(async move {
-						lightning_net_tokio::setup_inbound(
-							Arc::clone(&peer_mgr),
-							tcp_stream.into_std().unwrap(),
-						)
-						.await;
-					});
+					tokio::select! {
+						_ = stop_listen.changed() => {
+							return;
+						}
+						res = listener.accept() => {
+							let tcp_stream = res.unwrap().0;
+							tokio::spawn(async move {
+								lightning_net_tokio::setup_inbound(
+									Arc::clone(&peer_mgr),
+									tcp_stream.into_std().unwrap(),
+									)
+									.await;
+							});
+						}
+					}
 				}
 			});
 		}
@@ -748,36 +762,38 @@ impl Node {
 		let connect_pm = Arc::clone(&self.peer_manager);
 		let connect_logger = Arc::clone(&self.logger);
 		let connect_peer_store = Arc::clone(&self.peer_store);
-		let stop_connect = Arc::clone(&stop_running);
+		let mut stop_connect = self.stop_receiver.clone();
 		runtime.spawn(async move {
 			let mut interval = tokio::time::interval(PEER_RECONNECTION_INTERVAL);
 			interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
 			loop {
-				if stop_connect.load(Ordering::Acquire) {
-					return;
-				}
-
-				interval.tick().await;
-				let pm_peers = connect_pm
-					.get_peer_node_ids()
-					.iter()
-					.map(|(peer, _addr)| *peer)
-					.collect::<Vec<_>>();
-				for node_id in connect_cm
-					.list_channels()
-					.iter()
-					.map(|chan| chan.counterparty.node_id)
-					.filter(|id| !pm_peers.contains(id))
-				{
-					if let Some(peer_info) = connect_peer_store.get_peer(&node_id) {
-						let _ = do_connect_peer(
-							peer_info.node_id,
-							peer_info.address,
-							Arc::clone(&connect_pm),
-							Arc::clone(&connect_logger),
-						)
-						.await;
-					}
+				tokio::select! {
+						_ = stop_connect.changed() => {
+							return;
+						}
+						_ = interval.tick() => {
+							let pm_peers = connect_pm
+								.get_peer_node_ids()
+								.iter()
+								.map(|(peer, _addr)| *peer)
+								.collect::<Vec<_>>();
+							for node_id in connect_cm
+								.list_channels()
+									.iter()
+									.map(|chan| chan.counterparty.node_id)
+									.filter(|id| !pm_peers.contains(id))
+									{
+										if let Some(peer_info) = connect_peer_store.get_peer(&node_id) {
+											let _ = do_connect_peer(
+												peer_info.node_id,
+												peer_info.address,
+												Arc::clone(&connect_pm),
+												Arc::clone(&connect_logger),
+												)
+												.await;
+										}
+									}
+						}
 				}
 			}
 		});
@@ -786,28 +802,32 @@ impl Node {
 		let bcast_cm = Arc::clone(&self.channel_manager);
 		let bcast_pm = Arc::clone(&self.peer_manager);
 		let bcast_config = Arc::clone(&self.config);
-		let stop_bcast = Arc::clone(&stop_running);
+		let mut stop_bcast = self.stop_receiver.clone();
 		runtime.spawn(async move {
 			let mut interval = tokio::time::interval(NODE_ANN_BCAST_INTERVAL);
 			loop {
-				if stop_bcast.load(Ordering::Acquire) {
-					return;
-				}
-
-				if !bcast_cm.list_channels().iter().any(|chan| chan.is_public) { continue; }
-
-				interval.tick().await;
-
-				if !bcast_cm.list_channels().iter().any(|chan| chan.is_public) { continue; }
+				tokio::select! {
+						_ = stop_bcast.changed() => {
+							return;
+						}
+						_ = interval.tick(), if bcast_cm.list_channels().iter().any(|chan| chan.is_public) => {
+							while bcast_pm.get_peer_node_ids().is_empty() {
+								// Sleep a bit and retry if we don't have any peers yet.
+								tokio::time::sleep(Duration::from_secs(5)).await;
+
+								// Check back if we need to stop.
+								match stop_bcast.has_changed() {
+									Ok(false) => {},
+									Ok(true) => return,
+									Err(_) => return,
+								}
+							}
 
-				while bcast_pm.get_peer_node_ids().is_empty() {
-					// Sleep a bit and retry if we don't have any peers yet.
-					tokio::time::sleep(Duration::from_secs(5)).await;
+							let addresses =
+								bcast_config.listening_address.iter().cloned().map(|a| a.0).collect();
+							bcast_pm.broadcast_node_announcement([0; 3], [0; 32], addresses);
+						}
 				}
-
-				let addresses =
-					bcast_config.listening_address.iter().cloned().map(|a| a.0).collect();
-				bcast_pm.broadcast_node_announcement([0; 3], [0; 32], addresses);
 			}
 		});
 
@@ -820,15 +840,17 @@ impl Node {
 		let background_peer_man = Arc::clone(&self.peer_manager);
 		let background_logger = Arc::clone(&self.logger);
 		let background_scorer = Arc::clone(&self.scorer);
-		let stop_background_processing = Arc::clone(&stop_running);
+		let stop_bp = self.stop_receiver.clone();
 		let sleeper = move |d| {
-			let stop = Arc::clone(&stop_background_processing);
+			let mut stop = stop_bp.clone();
 			Box::pin(async move {
-				if stop.load(Ordering::Acquire) {
-					true
-				} else {
-					tokio::time::sleep(d).await;
-					false
+				tokio::select! {
+					_ = stop.changed() => {
+						true
+					}
+					_ = tokio::time::sleep(d) => {
+						false
+					}
 				}
 			})
 		};
@@ -860,7 +882,7 @@ impl Node {
 	pub fn stop(&self) -> Result<(), Error> {
 		let runtime = self.runtime.write().unwrap().take().ok_or(Error::NotRunning)?;
 		// Stop the runtime.
-		self.stop_running.store(true, Ordering::Release);
+		self.stop_sender.send(()).expect("Failed to send stop signal");
 
 		// Stop disconnect peers.
 		self.peer_manager.disconnect_all_peers();