split mutex

Author: joostjager

lightning/src/util/sweep.rs

@@ -28,7 +28,7 @@ use crate::{impl_writeable_tlv_based, log_debug, log_error};
 use bitcoin::block::Header;
 use bitcoin::locktime::absolute::LockTime;
 use bitcoin::secp256k1::Secp256k1;
-use bitcoin::{BlockHash, Transaction, Txid};
+use bitcoin::{BlockHash, ScriptBuf, Transaction, Txid};
 
 use core::ops::Deref;
 
@@ -372,7 +372,7 @@ where
 		output_spender: O, change_destination_source: D, kv_store: K, logger: L,
 	) -> Self {
 		let outputs = Vec::new();
-		let sweeper_state = Mutex::new(SweeperState { outputs, best_block });
+		let sweeper_state = Mutex::new(SweeperState { outputs, best_block, sweep_pending: false });
 		Self {
 			sweeper_state,
 			broadcaster,
@@ -451,14 +451,71 @@ where
 
 	/// Regenerates and broadcasts the spending transaction for any outputs that are pending
 	pub async fn regenerate_and_broadcast_spend_if_necessary_locked(&self) -> Result<(), ()> {
-		let mut sweeper_state = self.sweeper_state.lock().unwrap();
-		self.regenerate_and_broadcast_spend_if_necessary(&mut *sweeper_state).await
+		// Collect spendable output descriptors.
+		let respend_descriptors_clones: Vec<SpendableOutputDescriptor>;
+		let respend_descriptors: Vec<&SpendableOutputDescriptor>;
+		{
+			let mut sweeper_state = self.sweeper_state.lock().unwrap();
+
+			// Prevent concurrent sweeping.
+			if sweeper_state.sweep_pending {
+				return Ok(());
+			}
+
+			let cur_height = sweeper_state.best_block.height;
+			let filter_fn = |o: &TrackedSpendableOutput| {
+				if o.status.is_confirmed() {
+					// Don't rebroadcast confirmed txs.
+					return false;
+				}
+
+				if o.status.is_delayed(cur_height) {
+					// Don't generate and broadcast if still delayed
+					return false;
+				}
+
+				if o.status.latest_broadcast_height() >= Some(cur_height) {
+					// Only broadcast once per block height.
+					return false;
+				}
+
+				true
+			};
+
+			// Clone first, otherwise we can't take references that outlive the lock.
+			respend_descriptors_clones =
+				sweeper_state.outputs.iter().filter(|o| filter_fn(*o)).map(|o| o.descriptor.clone()).collect();
+
+			respend_descriptors = respend_descriptors_clones.iter().collect();
+
+			if respend_descriptors.is_empty() {
+				// Nothing to do.
+				return Ok(());
+			}
+
+			// There is something to sweep. Block concurrent sweeps.
+			sweeper_state.sweep_pending = true;
+		}
+
+		// Request a new change address outside of the mutex to avoid the mutex crossing await.
+		let change_destination_script_result = self.change_destination_source.get_change_destination_script().await;
+
+		// Sweep the outputs.
+		{
+			let mut sweeper_state = self.sweeper_state.lock().unwrap();
+
+			// Always allow a new sweep after this spend.
+			sweeper_state.sweep_pending = false;
+
+			let change_destination_script = change_destination_script_result?;
+			self.regenerate_and_broadcast_spend_if_necessary(&mut *sweeper_state, respend_descriptors, change_destination_script).await
+		}
 	}
 
 	async fn regenerate_and_broadcast_spend_if_necessary(
-		&self, sweeper_state: &mut SweeperState,
+		&self, sweeper_state: &mut SweeperState, respend_descriptors: Vec<&SpendableOutputDescriptor>, change_destination_script: ScriptBuf,
 	) -> Result<(), ()> {
-		let spending_tx_opt = self.regenerate_spend_if_necessary(sweeper_state).await;
+		let spending_tx_opt = self.regenerate_spend_if_necessary(sweeper_state, respend_descriptors, change_destination_script).await;
 		if let Some(spending_tx) = spending_tx_opt {
 			self.persist_state(&*sweeper_state).map_err(|e| {
 				log_error!(self.logger, "Error persisting OutputSweeper: {:?}", e);
@@ -471,38 +528,12 @@ where
 	}
 
 	async fn regenerate_spend_if_necessary(
-		&self, sweeper_state: &mut SweeperState,
+		&self, sweeper_state: &mut SweeperState, respend_descriptors: Vec<&SpendableOutputDescriptor>, change_destination_script: ScriptBuf,
 	) -> Option<Transaction> {
 		let cur_height = sweeper_state.best_block.height;
 		let cur_hash = sweeper_state.best_block.block_hash;
-		let filter_fn = |o: &TrackedSpendableOutput| {
-			if o.status.is_confirmed() {
-				// Don't rebroadcast confirmed txs.
-				return false;
-			}
 
-			if o.status.is_delayed(cur_height) {
-				// Don't generate and broadcast if still delayed
-				return false;
-			}
-
-			if o.status.latest_broadcast_height() >= Some(cur_height) {
-				// Only broadcast once per block height.
-				return false;
-			}
-
-			true
-		};
-
-		let respend_descriptors: Vec<&SpendableOutputDescriptor> =
-			sweeper_state.outputs.iter().filter(|o| filter_fn(*o)).map(|o| &o.descriptor).collect();
-
-		if respend_descriptors.is_empty() {
-			// Nothing to do.
-			return None;
-		}
-
-		let spending_tx = match self.spend_outputs(&*sweeper_state, respend_descriptors).await {
+		let spending_tx = match self.spend_outputs(&*sweeper_state, &respend_descriptors, change_destination_script) {
 			Ok(spending_tx) => {
 				log_debug!(
 					self.logger,
@@ -517,10 +548,14 @@ where
 			},
 		};
 
-		// As we didn't modify the state so far, the same filter_fn yields the same elements as
-		// above.
-		let respend_outputs = sweeper_state.outputs.iter_mut().filter(|o| filter_fn(&**o));
-		for output_info in respend_outputs {
+		// Watch outputs and update status.
+		for output_info in respend_descriptors {
+			let output_info = sweeper_state
+				.outputs
+				.iter_mut()
+				.find(|o| o.descriptor == *output_info)
+				.expect("We should have found the output info");
+
 			if let Some(filter) = self.chain_data_source.as_ref() {
 				let watched_output = output_info.to_watched_output(cur_hash);
 				filter.register_output(watched_output);
@@ -573,17 +608,15 @@ where
 			})
 	}
 
-	async fn spend_outputs(
-		&self, sweeper_state: &SweeperState, descriptors: Vec<&SpendableOutputDescriptor>,
+	fn spend_outputs(
+		&self, sweeper_state: &SweeperState, descriptors: &Vec<&SpendableOutputDescriptor>, change_destination_script: ScriptBuf,
 	) -> Result<Transaction, ()> {
 		let tx_feerate =
 			self.fee_estimator.get_est_sat_per_1000_weight(ConfirmationTarget::OutputSpendingFee);
-		let change_destination_script =
-			self.change_destination_source.get_change_destination_script().await?;
 		let cur_height = sweeper_state.best_block.height;
 		let locktime = Some(LockTime::from_height(cur_height).unwrap_or(LockTime::ZERO));
 		self.output_spender.spend_spendable_outputs(
-			&descriptors,
+			descriptors,
 			Vec::new(),
 			change_destination_script,
 			tx_feerate,
@@ -746,11 +779,15 @@ where
 struct SweeperState {
 	outputs: Vec<TrackedSpendableOutput>,
 	best_block: BestBlock,
+	sweep_pending: bool,
 }
 
 impl_writeable_tlv_based!(SweeperState, {
 	(0, outputs, required_vec),
 	(2, best_block, required),
+
+	// TODO: Do not persist this field.
+	(4, sweep_pending, required),
 });
 
 /// A `enum` signalling to the [`OutputSweeper`] that it should delay spending an output until a