try-out: atomicbool

Author: joostjager

lightning/src/util/sweep.rs

@@ -36,6 +36,7 @@ use bitcoin::{BlockHash, ScriptBuf, Transaction, Txid};
 use crate::sync::Arc;
 use core::future::Future;
 use core::ops::Deref;
+use core::sync::atomic::{AtomicBool, Ordering};
 use core::task;
 
 use super::async_poll::dummy_waker;
@@ -350,7 +351,8 @@ where
 	L::Target: Logger,
 	O::Target: OutputSpender,
 {
-	sweeper_state: Mutex<RuntimeSweeperState>,
+	sweeper_state: Mutex<SweeperState>,
+	pending_sweep: AtomicBool,
 	broadcaster: B,
 	fee_estimator: E,
 	chain_data_source: Option<F>,
@@ -380,12 +382,10 @@ where
 		output_spender: O, change_destination_source: D, kv_store: K, logger: L,
 	) -> Self {
 		let outputs = Vec::new();
-		let sweeper_state = Mutex::new(RuntimeSweeperState {
-			persistent: SweeperState { outputs, best_block },
-			sweep_pending: false,
-		});
+		let sweeper_state = Mutex::new(SweeperState { outputs, best_block });
 		Self {
 			sweeper_state,
+			pending_sweep: AtomicBool::new(false),
 			broadcaster,
 			fee_estimator,
 			chain_data_source,
@@ -427,7 +427,7 @@ where
 			return Ok(());
 		}
 
-		let state_lock = &mut self.sweeper_state.lock().unwrap().persistent;
+		let mut state_lock = self.sweeper_state.lock().unwrap();
 		for descriptor in relevant_descriptors {
 			let output_info = TrackedSpendableOutput {
 				descriptor,
@@ -444,20 +444,20 @@ where
 
 			state_lock.outputs.push(output_info);
 		}
-		self.persist_state(&state_lock).map_err(|e| {
+		self.persist_state(&*state_lock).map_err(|e| {
 			log_error!(self.logger, "Error persisting OutputSweeper: {:?}", e);
 		})
 	}
 
 	/// Returns a list of the currently tracked spendable outputs.
 	pub fn tracked_spendable_outputs(&self) -> Vec<TrackedSpendableOutput> {
-		self.sweeper_state.lock().unwrap().persistent.outputs.clone()
+		self.sweeper_state.lock().unwrap().outputs.clone()
 	}
 
 	/// Gets the latest best block which was connected either via the [`Listen`] or
 	/// [`Confirm`] interfaces.
 	pub fn current_best_block(&self) -> BestBlock {
-		self.sweeper_state.lock().unwrap().persistent.best_block
+		self.sweeper_state.lock().unwrap().best_block
 	}
 
 	/// Regenerates and broadcasts the spending transaction for any outputs that are pending
@@ -481,24 +481,29 @@ where
 			true
 		};
 
+		// Prevent concurrent sweeps.
+		if self.pending_sweep.load(Ordering::Relaxed) {
+			return Ok(());
+		}
+
 		// See if there is anything to sweep before requesting a change address.
 		{
-			let mut sweeper_state = self.sweeper_state.lock().unwrap();
+			let sweeper_state = self.sweeper_state.lock().unwrap();
 
-			// Prevent concurrent sweeping.
-			if sweeper_state.sweep_pending {
-				return Ok(());
-			}
-
-			let cur_height = sweeper_state.persistent.best_block.height;
-			let has_respends =
-				sweeper_state.persistent.outputs.iter().any(|o| filter_fn(o, cur_height));
+			let cur_height = sweeper_state.best_block.height;
+			let has_respends = sweeper_state.outputs.iter().any(|o| filter_fn(o, cur_height));
 			if !has_respends {
 				return Ok(());
 			}
+		}
 
-			// There is something to sweep. Block concurrent sweeps.
-			sweeper_state.sweep_pending = true;
+		// Mark sweep pending, if no other thread did so already.
+		if self
+			.pending_sweep
+			.compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
+			.is_err()
+		{
+			return Ok(());
 		}
 
 		// Request a new change address outside of the mutex to avoid the mutex crossing await.
@@ -509,10 +514,7 @@ where
 		{
 			let mut runtime_sweeper_state = self.sweeper_state.lock().unwrap();
 
-			// Always allow a new sweep after this spend, also in the error case.
-			runtime_sweeper_state.sweep_pending = false;
-
-			let sweeper_state = &mut runtime_sweeper_state.persistent;
+			let sweeper_state = &mut runtime_sweeper_state;
 
 			let change_destination_script = change_destination_script_result?;
 
@@ -527,6 +529,8 @@ where
 				.collect();
 
 			if respend_descriptors.is_empty() {
+				self.pending_sweep.store(false, Ordering::Release);
+
 				// It could be that a tx confirmed and there is now nothing to sweep anymore.
 				return Ok(());
 			}
@@ -545,6 +549,8 @@ where
 					spending_tx
 				},
 				Err(e) => {
+					self.pending_sweep.store(false, Ordering::Release);
+
 					log_error!(self.logger, "Error spending outputs: {:?}", e);
 					return Ok(());
 				},
@@ -570,6 +576,8 @@ where
 			self.broadcaster.broadcast_transactions(&[&spending_tx]);
 		}
 
+		self.pending_sweep.store(false, Ordering::Release);
+
 		Ok(())
 	}
 
@@ -668,22 +676,22 @@ where
 	fn filtered_block_connected(
 		&self, header: &Header, txdata: &chain::transaction::TransactionData, height: u32,
 	) {
-		let state_lock = &mut self.sweeper_state.lock().unwrap().persistent;
+		let mut state_lock = self.sweeper_state.lock().unwrap();
 		assert_eq!(state_lock.best_block.block_hash, header.prev_blockhash,
 			"Blocks must be connected in chain-order - the connected header must build on the last connected header");
 		assert_eq!(state_lock.best_block.height, height - 1,
 			"Blocks must be connected in chain-order - the connected block height must be one greater than the previous height");
 
-		self.transactions_confirmed_internal(state_lock, header, txdata, height);
-		self.best_block_updated_internal(state_lock, header, height);
+		self.transactions_confirmed_internal(&mut *state_lock, header, txdata, height);
+		self.best_block_updated_internal(&mut *state_lock, header, height);
 
-		let _ = self.persist_state(&state_lock).map_err(|e| {
+		let _ = self.persist_state(&*state_lock).map_err(|e| {
 			log_error!(self.logger, "Error persisting OutputSweeper: {:?}", e);
 		});
 	}
 
 	fn block_disconnected(&self, header: &Header, height: u32) {
-		let state_lock = &mut self.sweeper_state.lock().unwrap().persistent;
+		let mut state_lock = self.sweeper_state.lock().unwrap();
 
 		let new_height = height - 1;
 		let block_hash = header.block_hash();
@@ -721,15 +729,15 @@ where
 	fn transactions_confirmed(
 		&self, header: &Header, txdata: &chain::transaction::TransactionData, height: u32,
 	) {
-		let state_lock = &mut self.sweeper_state.lock().unwrap().persistent;
-		self.transactions_confirmed_internal(state_lock, header, txdata, height);
-		self.persist_state(state_lock).unwrap_or_else(|e| {
+		let mut state_lock = self.sweeper_state.lock().unwrap();
+		self.transactions_confirmed_internal(&mut *state_lock, header, txdata, height);
+		self.persist_state(&*state_lock).unwrap_or_else(|e| {
 			log_error!(self.logger, "Error persisting OutputSweeper: {:?}", e);
 		});
 	}
 
 	fn transaction_unconfirmed(&self, txid: &Txid) {
-		let state_lock = &mut self.sweeper_state.lock().unwrap().persistent;
+		let mut state_lock = self.sweeper_state.lock().unwrap();
 
 		// Get what height was unconfirmed.
 		let unconf_height = state_lock
@@ -746,22 +754,22 @@ where
 				.filter(|o| o.status.confirmation_height() >= Some(unconf_height))
 				.for_each(|o| o.status.unconfirmed());
 
-			self.persist_state(state_lock).unwrap_or_else(|e| {
+			self.persist_state(&*state_lock).unwrap_or_else(|e| {
 				log_error!(self.logger, "Error persisting OutputSweeper: {:?}", e);
 			});
 		}
 	}
 
 	fn best_block_updated(&self, header: &Header, height: u32) {
-		let state_lock = &mut self.sweeper_state.lock().unwrap().persistent;
-		self.best_block_updated_internal(state_lock, header, height);
-		let _ = self.persist_state(state_lock).map_err(|e| {
+		let mut state_lock = self.sweeper_state.lock().unwrap();
+		self.best_block_updated_internal(&mut *state_lock, header, height);
+		let _ = self.persist_state(&*state_lock).map_err(|e| {
 			log_error!(self.logger, "Error persisting OutputSweeper: {:?}", e);
 		});
 	}
 
 	fn get_relevant_txids(&self) -> Vec<(Txid, u32, Option<BlockHash>)> {
-		let state_lock = &self.sweeper_state.lock().unwrap().persistent;
+		let state_lock = self.sweeper_state.lock().unwrap();
 		state_lock
 			.outputs
 			.iter()
@@ -782,11 +790,6 @@ where
 	}
 }
 
-struct RuntimeSweeperState {
-	persistent: SweeperState,
-	sweep_pending: bool,
-}
-
 #[derive(Debug, Clone)]
 struct SweeperState {
 	outputs: Vec<TrackedSpendableOutput>,
@@ -849,10 +852,10 @@ where
 			}
 		}
 
-		let sweeper_state =
-			Mutex::new(RuntimeSweeperState { persistent: state, sweep_pending: false });
+		let sweeper_state = Mutex::new(state);
 		Ok(Self {
 			sweeper_state,
+			pending_sweep: AtomicBool::new(false),
 			broadcaster,
 			fee_estimator,
 			chain_data_source,
@@ -898,12 +901,12 @@ where
 			}
 		}
 
-		let sweeper_state =
-			Mutex::new(RuntimeSweeperState { persistent: state, sweep_pending: false });
+		let sweeper_state = Mutex::new(state);
 		Ok((
 			best_block,
 			OutputSweeper {
 				sweeper_state,
+				pending_sweep: AtomicBool::new(false),
 				broadcaster,
 				fee_estimator,
 				chain_data_source,