Wait on all background tasks to finish (or abort)

Previously, we'd only wait for the background processor tasks to
successfully finish. It turned out that this could lead to races when
the other background tasks took too long to shutdown. Here, we attempt
to wait on all background tasks shutting down for a bit, before moving
on.

Author: tnull

src/builder.rs

@@ -1632,11 +1632,15 @@ fn build_with_store_internal(
 
 	let (stop_sender, _) = tokio::sync::watch::channel(());
 	let background_processor_task = Mutex::new(None);
+	let background_tasks = Mutex::new(None);
+	let cancellable_background_tasks = Mutex::new(None);
 
 	Ok(Node {
 		runtime,
 		stop_sender,
 		background_processor_task,
+		background_tasks,
+		cancellable_background_tasks,
 		config,
 		wallet,
 		chain_source,

src/config.rs

@@ -79,6 +79,9 @@ pub(crate) const LDK_WALLET_SYNC_TIMEOUT_SECS: u64 = 10;
 // The timeout after which we give up waiting on LDK's event handler to exit on shutdown.
 pub(crate) const LDK_EVENT_HANDLER_SHUTDOWN_TIMEOUT_SECS: u64 = 30;
 
+// The timeout after which we give up waiting on a background task to exit on shutdown.
+pub(crate) const BACKGROUND_TASK_SHUTDOWN_TIMEOUT_SECS: u64 = 5;
+
 // The timeout after which we abort a fee rate cache update operation.
 pub(crate) const FEE_RATE_CACHE_UPDATE_TIMEOUT_SECS: u64 = 5;
 

src/lib.rs

@@ -127,8 +127,8 @@ pub use builder::NodeBuilder as Builder;
 use chain::ChainSource;
 use config::{
 	default_user_config, may_announce_channel, ChannelConfig, Config,
-	LDK_EVENT_HANDLER_SHUTDOWN_TIMEOUT_SECS, NODE_ANN_BCAST_INTERVAL, PEER_RECONNECTION_INTERVAL,
-	RGS_SYNC_INTERVAL,
+	BACKGROUND_TASK_SHUTDOWN_TIMEOUT_SECS, LDK_EVENT_HANDLER_SHUTDOWN_TIMEOUT_SECS,
+	NODE_ANN_BCAST_INTERVAL, PEER_RECONNECTION_INTERVAL, RGS_SYNC_INTERVAL,
 };
 use connection::ConnectionManager;
 use event::{EventHandler, EventQueue};
@@ -179,6 +179,8 @@ pub struct Node {
 	runtime: Arc<RwLock<Option<Arc<tokio::runtime::Runtime>>>>,
 	stop_sender: tokio::sync::watch::Sender<()>,
 	background_processor_task: Mutex<Option<tokio::task::JoinHandle<()>>>,
+	background_tasks: Mutex<Option<tokio::task::JoinSet<()>>>,
+	cancellable_background_tasks: Mutex<Option<tokio::task::JoinSet<()>>>,
 	config: Arc<Config>,
 	wallet: Arc<Wallet>,
 	chain_source: Arc<ChainSource>,
@@ -232,6 +234,10 @@ impl Node {
 			return Err(Error::AlreadyRunning);
 		}
 
+		let mut background_tasks = tokio::task::JoinSet::new();
+		let mut cancellable_background_tasks = tokio::task::JoinSet::new();
+		let runtime_handle = runtime.handle();
+
 		log_info!(
 			self.logger,
 			"Starting up LDK Node with node ID {} on network: {}",
@@ -258,19 +264,27 @@ impl Node {
 		let sync_cman = Arc::clone(&self.channel_manager);
 		let sync_cmon = Arc::clone(&self.chain_monitor);
 		let sync_sweeper = Arc::clone(&self.output_sweeper);
-		runtime.spawn(async move {
-			chain_source
-				.continuously_sync_wallets(stop_sync_receiver, sync_cman, sync_cmon, sync_sweeper)
-				.await;
-		});
+		background_tasks.spawn_on(
+			async move {
+				chain_source
+					.continuously_sync_wallets(
+						stop_sync_receiver,
+						sync_cman,
+						sync_cmon,
+						sync_sweeper,
+					)
+					.await;
+			},
+			runtime_handle,
+		);
 
 		if self.gossip_source.is_rgs() {
 			let gossip_source = Arc::clone(&self.gossip_source);
 			let gossip_sync_store = Arc::clone(&self.kv_store);
 			let gossip_sync_logger = Arc::clone(&self.logger);
 			let gossip_node_metrics = Arc::clone(&self.node_metrics);
 			let mut stop_gossip_sync = self.stop_sender.subscribe();
-			runtime.spawn(async move {
+			cancellable_background_tasks.spawn_on(async move {
 				let mut interval = tokio::time::interval(RGS_SYNC_INTERVAL);
 				loop {
 					tokio::select! {
@@ -311,7 +325,7 @@ impl Node {
 						}
 					}
 				}
-			});
+			}, runtime_handle);
 		}
 
 		if let Some(listening_addresses) = &self.config.listening_addresses {
@@ -337,7 +351,7 @@ impl Node {
 				bind_addrs.extend(resolved_address);
 			}
 
-			runtime.spawn(async move {
+			cancellable_background_tasks.spawn_on(async move {
 				{
 				let listener =
 					tokio::net::TcpListener::bind(&*bind_addrs).await
@@ -356,7 +370,7 @@ impl Node {
 						_ = stop_listen.changed() => {
 							log_debug!(
 								listening_logger,
-								"Stopping listening to inbound connections.",
+								"Stopping listening to inbound connections."
 							);
 							break;
 						}
@@ -375,7 +389,7 @@ impl Node {
 				}
 
 				listening_indicator.store(false, Ordering::Release);
-			});
+			}, runtime_handle);
 		}
 
 		// Regularly reconnect to persisted peers.
@@ -384,15 +398,15 @@ impl Node {
 		let connect_logger = Arc::clone(&self.logger);
 		let connect_peer_store = Arc::clone(&self.peer_store);
 		let mut stop_connect = self.stop_sender.subscribe();
-		runtime.spawn(async move {
+		cancellable_background_tasks.spawn_on(async move {
 			let mut interval = tokio::time::interval(PEER_RECONNECTION_INTERVAL);
 			interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
 			loop {
 				tokio::select! {
 						_ = stop_connect.changed() => {
 							log_debug!(
 								connect_logger,
-								"Stopping reconnecting known peers.",
+								"Stopping reconnecting known peers."
 							);
 							return;
 						}
@@ -412,7 +426,7 @@ impl Node {
 						}
 				}
 			}
-		});
+		}, runtime_handle);
 
 		// Regularly broadcast node announcements.
 		let bcast_cm = Arc::clone(&self.channel_manager);
@@ -424,7 +438,7 @@ impl Node {
 		let mut stop_bcast = self.stop_sender.subscribe();
 		let node_alias = self.config.node_alias.clone();
 		if may_announce_channel(&self.config).is_ok() {
-			runtime.spawn(async move {
+			cancellable_background_tasks.spawn_on(async move {
 				// We check every 30 secs whether our last broadcast is NODE_ANN_BCAST_INTERVAL away.
 				#[cfg(not(test))]
 				let mut interval = tokio::time::interval(Duration::from_secs(30));
@@ -495,14 +509,15 @@ impl Node {
 						}
 					}
 				}
-			});
+			}, runtime_handle);
 		}
 
 		let stop_tx_bcast = self.stop_sender.subscribe();
 		let chain_source = Arc::clone(&self.chain_source);
-		runtime.spawn(async move {
-			chain_source.continuously_process_broadcast_queue(stop_tx_bcast).await
-		});
+		cancellable_background_tasks.spawn_on(
+			async move { chain_source.continuously_process_broadcast_queue(stop_tx_bcast).await },
+			runtime_handle,
+		);
 
 		let bump_tx_event_handler = Arc::new(BumpTransactionEventHandler::new(
 			Arc::clone(&self.tx_broadcaster),
@@ -587,24 +602,33 @@ impl Node {
 			let mut stop_liquidity_handler = self.stop_sender.subscribe();
 			let liquidity_handler = Arc::clone(&liquidity_source);
 			let liquidity_logger = Arc::clone(&self.logger);
-			runtime.spawn(async move {
-				loop {
-					tokio::select! {
-						_ = stop_liquidity_handler.changed() => {
-							log_debug!(
-								liquidity_logger,
-								"Stopping processing liquidity events.",
-							);
-							return;
+			background_tasks.spawn_on(
+				async move {
+					loop {
+						tokio::select! {
+							_ = stop_liquidity_handler.changed() => {
+								log_debug!(
+									liquidity_logger,
+									"Stopping processing liquidity events.",
+								);
+								return;
+							}
+							_ = liquidity_handler.handle_next_event() => {}
 						}
-						_ = liquidity_handler.handle_next_event() => {}
 					}
-				}
-			});
+				},
+				runtime_handle,
+			);
 		}
 
 		*runtime_lock = Some(runtime);
 
+		debug_assert!(self.background_tasks.lock().unwrap().is_none());
+		*self.background_tasks.lock().unwrap() = Some(background_tasks);
+
+		debug_assert!(self.cancellable_background_tasks.lock().unwrap().is_none());
+		*self.cancellable_background_tasks.lock().unwrap() = Some(cancellable_background_tasks);
+
 		log_info!(self.logger, "Startup complete.");
 		Ok(())
 	}
@@ -643,6 +667,53 @@ impl Node {
 		self.chain_source.stop();
 		log_debug!(self.logger, "Stopped chain sources.");
 
+		// Cancel cancellable background tasks
+		if let Some(mut tasks) = self.cancellable_background_tasks.lock().unwrap().take() {
+			tasks.abort_all();
+		} else {
+			debug_assert!(false, "Expected some cancellable background tasks");
+		};
+
+		// Wait until non-cancellable background tasks (mod LDK's background processor) are done.
+		let runtime_2 = Arc::clone(&runtime);
+		if let Some(mut tasks) = self.background_tasks.lock().unwrap().take() {
+			tokio::task::block_in_place(move || {
+				runtime_2.block_on(async {
+					loop {
+						let timeout_fut = tokio::time::timeout(
+							Duration::from_secs(BACKGROUND_TASK_SHUTDOWN_TIMEOUT_SECS),
+							tasks.join_next_with_id(),
+						);
+						match timeout_fut.await {
+							Ok(Some(Ok((id, _)))) => {
+								log_trace!(self.logger, "Stopped background task with id {}", id);
+							},
+							Ok(Some(Err(e))) => {
+								tasks.abort_all();
+								log_trace!(self.logger, "Stopping background task failed: {}", e);
+								break;
+							},
+							Ok(None) => {
+								log_debug!(self.logger, "Stopped all background tasks");
+								break;
+							},
+							Err(e) => {
+								tasks.abort_all();
+								log_error!(
+									self.logger,
+									"Stopping background task timed out: {}",
+									e
+								);
+								break;
+							},
+						}
+					}
+				})
+			});
+		} else {
+			debug_assert!(false, "Expected some background tasks");
+		};
+
 		// Wait until background processing stopped, at least until a timeout is reached.
 		if let Some(background_processor_task) =
 			self.background_processor_task.lock().unwrap().take()
@@ -676,7 +747,9 @@ impl Node {
 					log_error!(self.logger, "Stopping event handling timed out: {}", e);
 				},
 			}
-		}
+		} else {
+			debug_assert!(false, "Expected a background processing task");
+		};
 
 		#[cfg(tokio_unstable)]
 		{