diff --git a/Cargo.toml b/Cargo.toml index 7fc945833..e981c961f 100755 --- a/Cargo.toml +++ b/Cargo.toml @@ -70,6 +70,7 @@ base64 = { version = "0.22.1", default-features = false, features = ["std"] } getrandom = { version = "0.3", default-features = false } chrono = { version = "0.4", default-features = false, features = ["clock"] } tokio = { version = "1.39", default-features = false, features = [ "rt-multi-thread", "time", "sync", "macros", "net" ] } +tokio-util = { version = "0.7", default-features = false, features = ["rt"] } esplora-client = { version = "0.12", default-features = false, features = ["tokio", "async-https-rustls"] } electrum-client = { version = "0.25", default-features = false, features = ["proxy", "use-rustls-ring"] } libc = "0.2" diff --git a/src/chain/bitcoind.rs b/src/chain/bitcoind.rs index f857ef533..b7763a2c0 100644 --- a/src/chain/bitcoind.rs +++ b/src/chain/bitcoind.rs @@ -31,7 +31,7 @@ use lightning_block_sync::{ }; use serde::Serialize; -use super::WalletSyncStatus; +use super::{WalletSyncGuard, WalletSyncStatus}; use crate::config::{ BitcoindRestClientConfig, Config, DEFAULT_FEE_RATE_CACHE_UPDATE_TIMEOUT_SECS, DEFAULT_TX_BROADCAST_TIMEOUT_SECS, @@ -152,12 +152,14 @@ impl BitcoindChainSource { ) { // First register for the wallet polling status to make sure `Node::sync_wallets` calls // wait on the result before proceeding. - { + let initial_sync_guard = { let mut status_lock = self.wallet_polling_status.lock().expect("lock"); if status_lock.register_or_subscribe_pending_sync().is_some() { debug_assert!(false, "Sync already in progress. This should never happen."); + return; } - } + WalletSyncGuard::new(&self.wallet_polling_status, Error::WalletOperationFailed) + }; log_info!( self.logger, @@ -285,7 +287,7 @@ impl BitcoindChainSource { } // Now propagate the initial result to unblock waiting subscribers. - self.wallet_polling_status.lock().expect("lock").propagate_result_to_subscribers(Ok(())); + initial_sync_guard.complete(Ok(())); let mut chain_polling_interval = tokio::time::interval(Duration::from_secs(CHAIN_POLLING_INTERVAL_SECS)); @@ -396,6 +398,8 @@ impl BitcoindChainSource { Error::WalletOperationFailed })?; } + let sync_guard = + WalletSyncGuard::new(&self.wallet_polling_status, Error::WalletOperationFailed); let res = self .poll_and_update_listeners_inner( @@ -406,7 +410,7 @@ impl BitcoindChainSource { ) .await; - self.wallet_polling_status.lock().expect("lock").propagate_result_to_subscribers(res); + sync_guard.complete(res); res } diff --git a/src/chain/electrum.rs b/src/chain/electrum.rs index 59fa23a6c..fbbebd760 100644 --- a/src/chain/electrum.rs +++ b/src/chain/electrum.rs @@ -25,7 +25,7 @@ use lightning::chain::{Confirm, Filter, WatchedOutput}; use lightning::util::ser::Writeable; use lightning_transaction_sync::ElectrumSyncClient; -use super::WalletSyncStatus; +use super::{WalletSyncGuard, WalletSyncStatus}; use crate::config::{ clamp_full_scan_stop_gap, Config, ElectrumSyncConfig, MAX_FULL_SCAN_STOP_GAP, MIN_FULL_SCAN_STOP_GAP, @@ -113,10 +113,12 @@ impl ElectrumChainSource { Error::WalletOperationFailed })?; } + let sync_guard = + WalletSyncGuard::new(&self.onchain_wallet_sync_status, Error::WalletOperationFailed); let res = self.sync_onchain_wallet_inner(onchain_wallet).await; - self.onchain_wallet_sync_status.lock().expect("lock").propagate_result_to_subscribers(res); + sync_guard.complete(res); res } @@ -223,14 +225,13 @@ impl ElectrumChainSource { Error::TxSyncFailed })?; } + let sync_guard = + WalletSyncGuard::new(&self.lightning_wallet_sync_status, Error::TxSyncFailed); let res = self.sync_lightning_wallet_inner(channel_manager, chain_monitor, output_sweeper).await; - self.lightning_wallet_sync_status - .lock() - .expect("lock") - .propagate_result_to_subscribers(res); + sync_guard.complete(res); res } diff --git a/src/chain/esplora.rs b/src/chain/esplora.rs index 21205bd25..901636f56 100644 --- a/src/chain/esplora.rs +++ b/src/chain/esplora.rs @@ -18,7 +18,7 @@ use lightning::chain::{Confirm, Filter, WatchedOutput}; use lightning::util::ser::Writeable; use lightning_transaction_sync::EsploraSyncClient; -use super::WalletSyncStatus; +use super::{WalletSyncGuard, WalletSyncStatus}; use crate::config::{ clamp_full_scan_stop_gap, Config, EsploraSyncConfig, BDK_CLIENT_CONCURRENCY, MAX_FULL_SCAN_STOP_GAP, MIN_FULL_SCAN_STOP_GAP, @@ -125,10 +125,12 @@ impl EsploraChainSource { Error::WalletOperationFailed })?; } + let sync_guard = + WalletSyncGuard::new(&self.onchain_wallet_sync_status, Error::WalletOperationFailed); let res = self.sync_onchain_wallet_inner(onchain_wallet).await; - self.onchain_wallet_sync_status.lock().expect("lock").propagate_result_to_subscribers(res); + sync_guard.complete(res); res } @@ -275,14 +277,13 @@ impl EsploraChainSource { Error::WalletOperationFailed })?; } + let sync_guard = + WalletSyncGuard::new(&self.lightning_wallet_sync_status, Error::WalletOperationFailed); let res = self.sync_lightning_wallet_inner(channel_manager, chain_monitor, output_sweeper).await; - self.lightning_wallet_sync_status - .lock() - .expect("lock") - .propagate_result_to_subscribers(res); + sync_guard.complete(res); res } diff --git a/src/chain/mod.rs b/src/chain/mod.rs index 0f96c409f..4aeb64704 100644 --- a/src/chain/mod.rs +++ b/src/chain/mod.rs @@ -65,6 +65,34 @@ pub(crate) enum WalletSyncStatus { InProgress { subscribers: tokio::sync::broadcast::Sender> }, } +pub(crate) struct WalletSyncGuard<'a> { + status: &'a Mutex, + cancellation_error: Error, + active: bool, +} + +impl<'a> WalletSyncGuard<'a> { + pub(crate) fn new(status: &'a Mutex, cancellation_error: Error) -> Self { + Self { status, cancellation_error, active: true } + } + + pub(crate) fn complete(mut self, res: Result<(), Error>) { + self.status.lock().expect("lock").propagate_result_to_subscribers(res); + self.active = false; + } +} + +impl Drop for WalletSyncGuard<'_> { + fn drop(&mut self) { + if self.active { + self.status + .lock() + .expect("lock") + .propagate_result_to_subscribers(Err(self.cancellation_error)); + } + } +} + impl WalletSyncStatus { fn register_or_subscribe_pending_sync( &mut self, @@ -95,16 +123,7 @@ impl WalletSyncStatus { WalletSyncStatus::InProgress { subscribers } => { // A sync is in-progress, we notify subscribers. if subscribers.receiver_count() > 0 { - match subscribers.send(res) { - Ok(_) => (), - Err(e) => { - debug_assert!( - false, - "Failed to send wallet sync result to subscribers: {:?}", - e - ); - }, - } + let _ = subscribers.send(res); } *self = WalletSyncStatus::Completed; }, @@ -561,3 +580,28 @@ impl Filter for ChainSource { } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn wallet_sync_guard_resets_abandoned_sync() { + let status = Mutex::new(WalletSyncStatus::Completed); + assert!(status.lock().expect("lock").register_or_subscribe_pending_sync().is_none()); + let sync_guard = WalletSyncGuard::new(&status, Error::WalletOperationFailed); + let mut subscriber = status + .lock() + .expect("lock") + .register_or_subscribe_pending_sync() + .expect("sync subscriber"); + + drop(sync_guard); + + assert!( + matches!(*status.lock().expect("lock"), WalletSyncStatus::Completed), + "abandoned wallet sync should reset its status" + ); + assert_eq!(subscriber.try_recv(), Ok(Err(Error::WalletOperationFailed))); + } +} diff --git a/src/connection.rs b/src/connection.rs index 88135e841..dda90d0ef 100644 --- a/src/connection.rs +++ b/src/connection.rs @@ -18,12 +18,44 @@ use crate::logger::{log_debug, log_error, log_info, LdkLogger}; use crate::types::{KeysManager, PeerManager}; use crate::Error; +type PendingConnections = + Mutex>>>>; + +struct PendingConnectionGuard<'a> { + pending_connections: &'a PendingConnections, + node_id: PublicKey, + active: bool, +} + +impl<'a> PendingConnectionGuard<'a> { + fn new(pending_connections: &'a PendingConnections, node_id: PublicKey) -> Self { + Self { pending_connections, node_id, active: true } + } + + fn disarm(&mut self) { + self.active = false; + } +} + +impl Drop for PendingConnectionGuard<'_> { + fn drop(&mut self) { + if !self.active { + return; + } + let mut pending_connections = self.pending_connections.lock().expect("lock"); + if let Some(subscribers) = pending_connections.remove(&self.node_id) { + for subscriber in subscribers { + let _ = subscriber.send(Err(Error::ConnectionFailed)); + } + } + } +} + pub(crate) struct ConnectionManager where L::Target: LdkLogger, { - pending_connections: - Mutex>>>>, + pending_connections: PendingConnections, peer_manager: Arc, tor_proxy_config: Option, keys_manager: Arc, @@ -60,26 +92,29 @@ where pub(crate) async fn do_connect_peer( &self, node_id: PublicKey, addr: SocketAddress, ) -> Result<(), Error> { + // If another task is already connecting, subscribe to its result instead of starting a + // duplicate attempt. + if let Some(pending_connection_ready_receiver) = + self.register_or_subscribe_pending_connection(&node_id) + { + return pending_connection_ready_receiver.await.map_err(|e| { + debug_assert!(false, "Failed to receive connection result: {:?}", e); + log_error!(self.logger, "Failed to receive connection result: {:?}", e); + Error::ConnectionFailed + })?; + } + + let mut pending_connection = + PendingConnectionGuard::new(&self.pending_connections, node_id); let res = self.do_connect_peer_internal(node_id, addr).await; self.propagate_result_to_subscribers(&node_id, res); + pending_connection.disarm(); res } async fn do_connect_peer_internal( &self, node_id: PublicKey, addr: SocketAddress, ) -> Result<(), Error> { - // First, we check if there is already an outbound connection in flight, if so, we just - // await on the corresponding watch channel. The task driving the connection future will - // send us the result.. - let pending_ready_receiver_opt = self.register_or_subscribe_pending_connection(&node_id); - if let Some(pending_connection_ready_receiver) = pending_ready_receiver_opt { - return pending_connection_ready_receiver.await.map_err(|e| { - debug_assert!(false, "Failed to receive connection result: {:?}", e); - log_error!(self.logger, "Failed to receive connection result: {:?}", e); - Error::ConnectionFailed - })?; - } - log_info!(self.logger, "Connecting to peer: {}@{}", node_id, addr); match addr { @@ -246,6 +281,7 @@ where match pending_connections_lock.entry(*node_id) { hash_map::Entry::Occupied(mut entry) => { let (tx, rx) = tokio::sync::oneshot::channel(); + entry.get_mut().retain(|subscriber| !subscriber.is_closed()); entry.get_mut().push(tx); Some(rx) }, @@ -277,3 +313,26 @@ where } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn pending_connection_guard_notifies_subscribers_when_abandoned() { + let node_id: PublicKey = + "0279be667ef9dcbbac55a06295ce870b07029bfcdb2dce28d959f2815b16f81798".parse().unwrap(); + let pending_connections = Mutex::new(HashMap::new()); + let (sender, mut receiver) = tokio::sync::oneshot::channel(); + pending_connections.lock().expect("lock").insert(node_id, vec![sender]); + + let connection_guard = PendingConnectionGuard::new(&pending_connections, node_id); + drop(connection_guard); + + assert!( + pending_connections.lock().expect("lock").is_empty(), + "abandoned connection attempt should remove pending state" + ); + assert_eq!(receiver.try_recv(), Ok(Err(Error::ConnectionFailed))); + } +} diff --git a/src/liquidity/client/lsps1.rs b/src/liquidity/client/lsps1.rs index 6082414c1..ef6ceea63 100644 --- a/src/liquidity/client/lsps1.rs +++ b/src/liquidity/client/lsps1.rs @@ -21,8 +21,8 @@ use tokio::sync::oneshot; use crate::connection::ConnectionManager; use crate::liquidity::{ - select_lsps_for_protocol, LspConfig, LspNode, LIQUIDITY_REQUEST_TIMEOUT_SECS, - LSPS_DISCOVERY_WAIT_TIMEOUT_SECS, + select_lsps_for_protocol, LspConfig, LspNode, PendingRequest, PendingRequestGuard, + LIQUIDITY_REQUEST_TIMEOUT_SECS, LSPS_DISCOVERY_WAIT_TIMEOUT_SECS, }; use crate::logger::{log_error, log_info, LdkLogger, Logger}; use crate::runtime::Runtime; @@ -35,11 +35,11 @@ where { pub(crate) lsp_nodes: Arc>>, pub(crate) pending_opening_params_requests: - Mutex>>, + Mutex>>, pub(crate) pending_create_order_requests: - Mutex>>, + Mutex>>, pub(crate) pending_check_order_status_requests: - Mutex>>, + Mutex>>, pub(crate) discovery_done_rx: tokio::sync::watch::Receiver, pub(crate) liquidity_manager: Arc, pub(crate) logger: L, @@ -61,12 +61,17 @@ where })?; let (request_sender, request_receiver) = oneshot::channel(); - { + let _pending_request = { let mut pending_opening_params_requests_lock = self.pending_opening_params_requests.lock().expect("lock"); let request_id = client_handler.request_supported_options(lsps1_node.node_id); - pending_opening_params_requests_lock.insert(request_id, request_sender); - } + PendingRequestGuard::insert( + &self.pending_opening_params_requests, + &mut pending_opening_params_requests_lock, + request_id, + request_sender, + ) + }; tokio::time::timeout(Duration::from_secs(LIQUIDITY_REQUEST_TIMEOUT_SECS), request_receiver) .await @@ -146,7 +151,7 @@ where let (request_sender, request_receiver) = oneshot::channel(); let request_id; - { + let _pending_request = { let mut pending_create_order_requests_lock = self.pending_create_order_requests.lock().expect("lock"); request_id = client_handler.create_order( @@ -154,8 +159,13 @@ where order_params.clone(), Some(refund_address), ); - pending_create_order_requests_lock.insert(request_id.clone(), request_sender); - } + PendingRequestGuard::insert( + &self.pending_create_order_requests, + &mut pending_create_order_requests_lock, + request_id.clone(), + request_sender, + ) + }; let response = tokio::time::timeout( Duration::from_secs(LIQUIDITY_REQUEST_TIMEOUT_SECS), @@ -191,12 +201,17 @@ where })?; let (request_sender, request_receiver) = oneshot::channel(); - { + let _pending_request = { let mut pending_check_order_status_requests_lock = self.pending_check_order_status_requests.lock().expect("lock"); let request_id = client_handler.check_order_status(&lsp_node_id, order_id); - pending_check_order_status_requests_lock.insert(request_id, request_sender); - } + PendingRequestGuard::insert( + &self.pending_check_order_status_requests, + &mut pending_check_order_status_requests_lock, + request_id, + request_sender, + ) + }; let response = tokio::time::timeout( Duration::from_secs(LIQUIDITY_REQUEST_TIMEOUT_SECS), @@ -229,7 +244,7 @@ where .iter() .any(|n| n.node_id == counterparty_node_id) { - if let Some(sender) = self + if let Some(request) = self .pending_opening_params_requests .lock() .expect("lock") @@ -237,7 +252,7 @@ where { let response = LSPS1OpeningParamsResponse { supported_options }; - match sender.send(response) { + match request.sender.send(response) { Ok(()) => (), Err(_) => { log_error!( @@ -279,7 +294,7 @@ where .iter() .any(|n| n.node_id == counterparty_node_id) { - if let Some(sender) = + if let Some(request) = self.pending_create_order_requests.lock().expect("lock").remove(&request_id) { let response = LSPS1OrderStatus { @@ -290,7 +305,7 @@ where counterparty_node_id, }; - match sender.send(response) { + match request.sender.send(response) { Ok(()) => (), Err(_) => { log_error!( @@ -329,7 +344,7 @@ where .iter() .any(|n| n.node_id == counterparty_node_id) { - if let Some(sender) = self + if let Some(request) = self .pending_check_order_status_requests .lock() .expect("lock") @@ -343,7 +358,7 @@ where counterparty_node_id, }; - match sender.send(response) { + match request.sender.send(response) { Ok(()) => (), Err(_) => { log_error!( diff --git a/src/liquidity/client/lsps2.rs b/src/liquidity/client/lsps2.rs index 3033f8d82..4e6163e21 100644 --- a/src/liquidity/client/lsps2.rs +++ b/src/liquidity/client/lsps2.rs @@ -26,8 +26,8 @@ use tokio::task::JoinSet; use crate::connection::ConnectionManager; use crate::liquidity::{ - select_all_lsps_for_protocol, select_lsps_for_protocol, LspConfig, LspNode, - LIQUIDITY_REQUEST_TIMEOUT_SECS, LSPS_DISCOVERY_WAIT_TIMEOUT_SECS, + select_all_lsps_for_protocol, select_lsps_for_protocol, LspConfig, LspNode, PendingRequest, + PendingRequestGuard, LIQUIDITY_REQUEST_TIMEOUT_SECS, LSPS_DISCOVERY_WAIT_TIMEOUT_SECS, }; use crate::logger::{log_debug, log_error, log_info, LdkLogger}; use crate::payment::store::LSPS2Parameters; @@ -41,9 +41,9 @@ where { pub(crate) lsp_nodes: Arc>>, pub(crate) pending_lsps2_fee_requests: - Mutex>>, + Mutex>>, pub(crate) pending_buy_requests: - Mutex>>, + Mutex>>, pub(crate) channel_manager: Arc, pub(crate) keys_manager: Arc, pub(crate) discovery_done_rx: tokio::sync::watch::Receiver, @@ -262,13 +262,18 @@ where })?; let (fee_request_sender, fee_request_receiver) = oneshot::channel(); - { + let _pending_request = { let mut pending_fee_requests_lock = self.pending_lsps2_fee_requests.lock().expect("lock"); let request_id = client_handler.request_opening_params(lsps2_node.node_id, lsps2_node.token.clone()); - pending_fee_requests_lock.insert(request_id, fee_request_sender); - } + PendingRequestGuard::insert( + &self.pending_lsps2_fee_requests, + &mut pending_fee_requests_lock, + request_id, + fee_request_sender, + ) + }; tokio::time::timeout( Duration::from_secs(LIQUIDITY_REQUEST_TIMEOUT_SECS), @@ -298,7 +303,7 @@ where })?; let (buy_request_sender, buy_request_receiver) = oneshot::channel(); - { + let _pending_request = { let mut pending_buy_requests_lock = self.pending_buy_requests.lock().expect("lock"); let request_id = client_handler .select_opening_params(lsps2_node.node_id, amount_msat, opening_fee_params) @@ -310,8 +315,13 @@ where ); Error::LiquidityRequestFailed })?; - pending_buy_requests_lock.insert(request_id, buy_request_sender); - } + PendingRequestGuard::insert( + &self.pending_buy_requests, + &mut pending_buy_requests_lock, + request_id, + buy_request_sender, + ) + }; let buy_response = tokio::time::timeout( Duration::from_secs(LIQUIDITY_REQUEST_TIMEOUT_SECS), @@ -428,12 +438,12 @@ where .iter() .any(|n| n.node_id == counterparty_node_id) { - if let Some(sender) = + if let Some(request) = self.pending_lsps2_fee_requests.lock().expect("lock").remove(&request_id) { let response = LSPS2FeeResponse { opening_fee_params_menu }; - match sender.send(response) { + match request.sender.send(response) { Ok(()) => (), Err(_) => { log_error!( @@ -474,12 +484,12 @@ where .iter() .any(|n| n.node_id == counterparty_node_id) { - if let Some(sender) = + if let Some(request) = self.pending_buy_requests.lock().expect("lock").remove(&request_id) { let response = LSPS2BuyResponse { intercept_scid, cltv_expiry_delta }; - match sender.send(response) { + match request.sender.send(response) { Ok(()) => (), Err(_) => { log_error!( diff --git a/src/liquidity/mod.rs b/src/liquidity/mod.rs index 0eddde1ae..ffc1f878b 100644 --- a/src/liquidity/mod.rs +++ b/src/liquidity/mod.rs @@ -10,8 +10,8 @@ pub(crate) mod client; pub(crate) mod service; -use std::collections::hash_map::Entry; use std::collections::HashMap; +use std::hash::Hash; use std::ops::Deref; use std::sync::{Arc, Mutex, RwLock}; use std::time::Duration; @@ -42,6 +42,45 @@ use crate::{Config, Error}; const LIQUIDITY_REQUEST_TIMEOUT_SECS: u64 = 5; const LSPS_DISCOVERY_WAIT_TIMEOUT_SECS: u64 = 10; +pub(crate) struct PendingRequest { + token: Arc<()>, + pub(crate) sender: oneshot::Sender, + followers: Vec>, +} + +pub(crate) struct PendingRequestGuard<'a, K: Clone + Eq + Hash, T> { + pending_requests: &'a Mutex>>, + request_key: K, + token: Arc<()>, +} + +impl<'a, K: Clone + Eq + Hash, T> PendingRequestGuard<'a, K, T> { + pub(crate) fn insert( + pending_requests: &'a Mutex>>, + pending_requests_lock: &mut HashMap>, request_key: K, + sender: oneshot::Sender, + ) -> Self { + let token = Arc::new(()); + pending_requests_lock.insert( + request_key.clone(), + PendingRequest { token: Arc::clone(&token), sender, followers: Vec::new() }, + ); + Self { pending_requests, request_key, token } + } +} + +impl Drop for PendingRequestGuard<'_, K, T> { + fn drop(&mut self) { + let mut pending_requests = self.pending_requests.lock().expect("lock"); + if pending_requests + .get(&self.request_key) + .is_some_and(|request| Arc::ptr_eq(&request.token, &self.token)) + { + pending_requests.remove(&self.request_key); + } + } +} + fn select_lsps_for_protocol( lsp_nodes: &Arc>>, protocol: u16, override_node_id: Option<&PublicKey>, ) -> Option { @@ -345,7 +384,7 @@ where lsps1_client: Arc>, lsps2_client: Arc>, lsps2_service: Arc>, - pending_lsps0_discovery: Mutex>>>>, + pending_lsps0_discovery: Mutex>>>, discovery_done_tx: tokio::sync::watch::Sender, discovery_done_rx: tokio::sync::watch::Receiver, liquidity_manager: Arc, @@ -383,13 +422,14 @@ where protocols, }) => { if self.is_lsps_node(&counterparty_node_id) { - if let Some(senders) = self + if let Some(request) = self .pending_lsps0_discovery .lock() .expect("lock") .remove(&counterparty_node_id) { - for sender in senders { + let _ = request.sender.send(protocols.clone()); + for sender in request.followers { let _ = sender.send(protocols.clone()); } } else { @@ -431,23 +471,25 @@ where let lsps0_handler = self.liquidity_manager.lsps0_client_handler(); let (sender, receiver) = oneshot::channel(); - let issued_request = { + let _pending_request = { let mut pending_discovery = self.pending_lsps0_discovery.lock().expect("lock"); - match pending_discovery.entry(*node_id) { - Entry::Occupied(mut e) => { - e.get_mut().push(sender); - false - }, - Entry::Vacant(v) => { - v.insert(vec![sender]); - lsps0_handler.list_protocols(node_id); - true - }, + if let Some(request) = pending_discovery.get_mut(node_id) { + request.followers.push(sender); + None + } else { + let request_guard = PendingRequestGuard::insert( + &self.pending_lsps0_discovery, + &mut pending_discovery, + *node_id, + sender, + ); + lsps0_handler.list_protocols(node_id); + Some(request_guard) } }; - // Only the request that issued the discovery may remove the entry; a follower removing it - // would drop the in-flight request and all other waiters. + // Only the request that issued the discovery holds the guard. If it is abandoned, dropping + // the pending request also wakes all followers without risking removal of a newer request. let protocols = tokio::time::timeout(Duration::from_secs(LIQUIDITY_REQUEST_TIMEOUT_SECS), receiver) .await @@ -458,9 +500,6 @@ where node_id, e ); - if issued_request { - self.pending_lsps0_discovery.lock().expect("lock").remove(node_id); - } Error::LiquidityRequestFailed })? .map_err(|e| { @@ -470,9 +509,6 @@ where node_id, e ); - if issued_request { - self.pending_lsps0_discovery.lock().expect("lock").remove(node_id); - } Error::LiquidityRequestFailed })?; @@ -540,3 +576,83 @@ where let _ = self.discovery_done_tx.send(true); } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn pending_request_guard_removes_abandoned_request() { + let pending_requests = Mutex::new(HashMap::new()); + let (sender, receiver) = oneshot::channel::<()>(); + let request_guard = { + let mut pending_requests_lock = pending_requests.lock().expect("lock"); + PendingRequestGuard::insert( + &pending_requests, + &mut pending_requests_lock, + "request".to_owned(), + sender, + ) + }; + let (follower_sender, mut follower_receiver) = oneshot::channel(); + pending_requests + .lock() + .expect("lock") + .get_mut("request") + .expect("request should be pending") + .followers + .push(follower_sender); + + drop(receiver); + drop(request_guard); + + assert!( + pending_requests.lock().expect("lock").is_empty(), + "abandoned pending request should be removed" + ); + assert!( + matches!(follower_receiver.try_recv(), Err(oneshot::error::TryRecvError::Closed)), + "abandoned pending request should wake its followers" + ); + } + + #[test] + fn pending_request_guard_preserves_replacement() { + let pending_requests = Mutex::new(HashMap::new()); + let (old_sender, old_receiver) = oneshot::channel::<()>(); + let old_request_guard = { + let mut pending_requests_lock = pending_requests.lock().expect("lock"); + PendingRequestGuard::insert( + &pending_requests, + &mut pending_requests_lock, + "request".to_owned(), + old_sender, + ) + }; + + pending_requests.lock().expect("lock").remove("request"); + drop(old_receiver); + + let (replacement_sender, replacement_receiver) = oneshot::channel::<()>(); + let replacement_request_guard = { + let mut pending_requests_lock = pending_requests.lock().expect("lock"); + PendingRequestGuard::insert( + &pending_requests, + &mut pending_requests_lock, + "request".to_owned(), + replacement_sender, + ) + }; + + drop(old_request_guard); + assert_eq!( + pending_requests.lock().expect("lock").len(), + 1, + "an older guard should not remove a replacement request" + ); + + drop(replacement_receiver); + drop(replacement_request_guard); + assert!(pending_requests.lock().expect("lock").is_empty()); + } +} diff --git a/src/runtime.rs b/src/runtime.rs index 7e29996e6..d4ea7cea2 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -13,6 +13,8 @@ use std::time::Duration; use lightning::util::native_async::FutureSpawner; use tokio::task::{JoinHandle, JoinSet}; +use tokio_util::sync::CancellationToken; +use tokio_util::task::TaskTracker; use crate::config::{ BACKGROUND_TASK_SHUTDOWN_TIMEOUT_SECS, LDK_EVENT_HANDLER_SHUTDOWN_TIMEOUT_SECS, @@ -28,13 +30,18 @@ pub(crate) struct Runtime { } struct CancellableBackgroundTasks { - tasks: JoinSet<()>, + tasks: TaskTracker, + cancellation_token: CancellationToken, accepting_tasks: bool, } impl CancellableBackgroundTasks { fn new() -> Self { - Self { tasks: JoinSet::new(), accepting_tasks: true } + Self { + tasks: TaskTracker::new(), + cancellation_token: CancellationToken::new(), + accepting_tasks: true, + } } } @@ -109,8 +116,7 @@ impl Runtime { where F: Future + Send + 'static, { - let mut cancellable_background_tasks = - self.cancellable_background_tasks.lock().expect("lock"); + let cancellable_background_tasks = self.cancellable_background_tasks.lock().expect("lock"); if !cancellable_background_tasks.accepting_tasks { log_trace!( self.logger, @@ -122,11 +128,32 @@ impl Runtime { // Since it seems to make a difference to `tokio` (see // https://docs.rs/tokio/latest/tokio/time/fn.timeout.html#panics) we make sure the futures // are always put in an `async` / `.await` closure. - cancellable_background_tasks.tasks.spawn_on(async { future.await }, runtime_handle); + let cancellation_token = cancellable_background_tasks.cancellation_token.clone(); + // Detach the handle while the tracker continues tracking the task. + let _ = cancellable_background_tasks.tasks.spawn_on( + async move { + tokio::select! { + biased; + _ = cancellation_token.cancelled() => {}, + _ = future => {}, + } + }, + runtime_handle, + ); } pub fn allow_cancellable_background_task_spawns(&self) { - self.cancellable_background_tasks.lock().expect("lock").accepting_tasks = true; + let mut cancellable_background_tasks = + self.cancellable_background_tasks.lock().expect("lock"); + if cancellable_background_tasks.cancellation_token.is_cancelled() { + debug_assert!( + cancellable_background_tasks.tasks.is_empty(), + "Expected all cancellable background tasks to be stopped" + ); + cancellable_background_tasks.cancellation_token = CancellationToken::new(); + } + cancellable_background_tasks.tasks.reopen(); + cancellable_background_tasks.accepting_tasks = true; } pub fn spawn_background_processor_task(&self, future: F) @@ -164,15 +191,15 @@ impl Runtime { } pub fn abort_cancellable_background_tasks(&self) { - let mut tasks = { + let tasks = { let mut cancellable_background_tasks = self.cancellable_background_tasks.lock().expect("lock"); cancellable_background_tasks.accepting_tasks = false; - core::mem::take(&mut cancellable_background_tasks.tasks) + cancellable_background_tasks.tasks.close(); + cancellable_background_tasks.cancellation_token.cancel(); + cancellable_background_tasks.tasks.clone() }; - debug_assert!(tasks.len() > 0, "Expected some cancellable background_tasks"); - tasks.abort_all(); - self.block_on(async { while let Some(_) = tasks.join_next().await {} }) + self.block_on(tasks.wait()) } pub fn wait_on_background_tasks(&self) { @@ -381,19 +408,68 @@ impl FutureSpawner for RuntimeSpawner { #[cfg(test)] mod tests { - use tokio::sync::oneshot; + use tokio::sync::{mpsc, oneshot}; use super::*; + struct DropNotifier(Option>); + + impl Drop for DropNotifier { + fn drop(&mut self) { + if let Some(sender) = self.0.take() { + let _ = sender.send(()); + } + } + } + fn test_runtime() -> Runtime { Runtime::new(Arc::new(Logger::new_log_facade())).unwrap() } + #[test] + fn completed_cancellable_tasks_are_released_before_shutdown() { + const TASK_COUNT: usize = 64; + + let runtime = test_runtime(); + let (completion_sender, mut completion_receiver) = mpsc::channel(TASK_COUNT); + for _ in 0..TASK_COUNT { + let completion_sender = completion_sender.clone(); + runtime.spawn_cancellable_background_task(async move { + completion_sender.send(()).await.expect("completion receiver should be open"); + }); + } + drop(completion_sender); + + let completed_tasks_are_released = runtime.block_on(async { + for _ in 0..TASK_COUNT { + completion_receiver.recv().await.expect("cancellable task should complete"); + } + + tokio::time::timeout(Duration::from_secs(1), async { + loop { + if runtime.cancellable_background_tasks.lock().expect("lock").tasks.is_empty() { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .is_ok() + }); + + assert!( + completed_tasks_are_released, + "completed cancellable tasks should be released before shutdown" + ); + } + #[test] fn late_cancellable_spawns_are_not_polled_after_abort() { let runtime = test_runtime(); let (started_sender, started_receiver) = oneshot::channel(); + let (dropped_sender, dropped_receiver) = oneshot::channel(); runtime.spawn_cancellable_background_task(async move { + let _drop_notifier = DropNotifier(Some(dropped_sender)); let _ = started_sender.send(()); std::future::pending::<()>().await; }); @@ -402,6 +478,9 @@ mod tests { }); runtime.abort_cancellable_background_tasks(); + runtime.block_on(async { + dropped_receiver.await.expect("aborted task should be dropped before abort returns"); + }); let (late_spawn_sender, late_spawn_receiver) = oneshot::channel(); runtime.spawn_cancellable_background_task(async move {