f: update_state with callback

Author: joostjager

lightning/src/util/sweep.rs

@@ -470,8 +470,7 @@ where
 			return Ok(());
 		}
 
-		let persist_fut = {
-			let mut state_lock = self.sweeper_state.lock().unwrap();
+		self.update_state(|state_lock| -> Result<((), bool), ()> {
 			for descriptor in relevant_descriptors {
 				let output_info = TrackedSpendableOutput {
 					descriptor,
@@ -491,17 +490,12 @@ where
 				}
 
 				state_lock.outputs.push(output_info);
+				state_lock.dirty = true;
 			}
 
-			state_lock.dirty = false;
-			self.persist_state(&state_lock)
-		};
-
-		persist_fut.await.map_err(|e| {
-			self.sweeper_state.lock().unwrap().dirty = true;
-
-			log_error!(self.logger, "Error persisting OutputSweeper: {:?}", e);
+			Ok(((), true))
 		})
+		.await
 	}
 
 	/// Returns a list of the currently tracked spendable outputs.
@@ -557,29 +551,18 @@ where
 		};
 
 		// See if there is anything to sweep before requesting a change address.
-		let (persist_fut, has_respends) = {
-			let mut sweeper_state = self.sweeper_state.lock().unwrap();
-
-			let cur_height = sweeper_state.best_block.height;
-			let has_respends = sweeper_state.outputs.iter().any(|o| filter_fn(o, cur_height));
-			// If there is nothing to sweep, we still persist the state if it is dirty.
-			let fut_opt = if !has_respends && sweeper_state.dirty {
-				sweeper_state.dirty = false;
-				Some(self.persist_state(&sweeper_state))
-			} else {
-				None
-			};
-
-			(fut_opt, has_respends)
-		};
+		let has_respends = self
+			.update_state(|sweeper_state| -> Result<(bool, bool), ()> {
+				let cur_height = sweeper_state.best_block.height;
+				let has_respends = sweeper_state.outputs.iter().any(|o| filter_fn(o, cur_height));
 
-		if let Some(persist_fut) = persist_fut {
-			persist_fut.await.map_err(|e| {
-				self.sweeper_state.lock().unwrap().dirty = true;
+				// If there are respends, we don't need to persist a dirty state already now and can postpone it until
+				// after the sweep.
+				let persist_if_dirty = !has_respends;
 
-				log_error!(self.logger, "Error persisting OutputSweeper: {:?}", e);
-			})?;
-		};
+				Ok((has_respends, persist_if_dirty))
+			})
+			.await?;
 
 		if !has_respends {
 			return Ok(());
@@ -590,67 +573,56 @@ where
 			self.change_destination_source.get_change_destination_script().await?;
 
 		// Sweep the outputs.
-		let spending_tx;
-		let persist_fut = {
-			let mut sweeper_state = self.sweeper_state.lock().unwrap();
-
-			let cur_height = sweeper_state.best_block.height;
-			let cur_hash = sweeper_state.best_block.block_hash;
+		let spending_tx = self
+			.update_state(|sweeper_state| -> Result<(Option<Transaction>, bool), ()> {
+				let cur_height = sweeper_state.best_block.height;
+				let cur_hash = sweeper_state.best_block.block_hash;
 
-			let respend_descriptors: Vec<&SpendableOutputDescriptor> = sweeper_state
-				.outputs
-				.iter()
-				.filter(|o| filter_fn(*o, cur_height))
-				.map(|o| &o.descriptor)
-				.collect();
-
-			// Exit if there is nothing to spend anymore and there also is no need to persist the state.
-			if respend_descriptors.is_empty() && !sweeper_state.dirty {
-				return Ok(());
-			}
-
-			// Generate the spending transaction and broadcast it.
-			spending_tx = if !respend_descriptors.is_empty() {
-				let spending_tx = self
-					.spend_outputs(&sweeper_state, &respend_descriptors, change_destination_script)
-					.map_err(|e| {
-						log_error!(self.logger, "Error spending outputs: {:?}", e);
-					})?;
-
-				log_debug!(
-					self.logger,
-					"Generating and broadcasting sweeping transaction {}",
-					spending_tx.compute_txid()
-				);
+				let respend_descriptors: Vec<&SpendableOutputDescriptor> = sweeper_state
+					.outputs
+					.iter()
+					.filter(|o| filter_fn(*o, cur_height))
+					.map(|o| &o.descriptor)
+					.collect();
+
+				// Generate the spending transaction and broadcast it.
+				if !respend_descriptors.is_empty() {
+					let spending_tx = self
+						.spend_outputs(
+							&sweeper_state,
+							&respend_descriptors,
+							change_destination_script,
+						)
+						.map_err(|e| {
+							log_error!(self.logger, "Error spending outputs: {:?}", e);
+						})?;
+
+					log_debug!(
+						self.logger,
+						"Generating and broadcasting sweeping transaction {}",
+						spending_tx.compute_txid()
+					);
 
-				// 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, cur_height));
-				for output_info in respend_outputs {
-					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);
+					// 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, cur_height));
+					for output_info in respend_outputs {
+						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);
+						}
+
+						output_info.status.broadcast(cur_hash, cur_height, spending_tx.clone());
+						sweeper_state.dirty = true;
 					}
 
-					output_info.status.broadcast(cur_hash, cur_height, spending_tx.clone());
+					Ok((Some(spending_tx), true))
+				} else {
+					Ok((None, true))
 				}
-
-				Some(spending_tx)
-			} else {
-				None
-			};
-
-			// Either the state was already dirty or we modified it above, so we persist it.
-			sweeper_state.dirty = false;
-			self.persist_state(&sweeper_state)
-		};
-
-		persist_fut.await.map_err(|e| {
-			self.sweeper_state.lock().unwrap().dirty = true;
-
-			log_error!(self.logger, "Error persisting OutputSweeper: {:?}", e);
-		})?;
+			})
+			.await?;
 
 		// Persistence completely successfully. If we have a spending transaction, we broadcast it.
 		if let Some(spending_tx) = spending_tx {
@@ -695,6 +667,37 @@ where
 		)
 	}
 
+	/// Updates the sweeper state by executing the given callback. Persists the state afterwards if it is marked dirty,
+	/// but only if persist_if_dirty is also true. Returning false for persist_if_dirty allows the callback to postpone
+	/// persisting a potentially dirty state.
+	async fn update_state<X>(
+		&self, callback: impl FnOnce(&mut SweeperState) -> Result<(X, bool), ()>,
+	) -> Result<X, ()> {
+		let (fut, res);
+		{
+			let mut state_lock = self.sweeper_state.lock().unwrap();
+
+			let persist_if_dirty;
+			(res, persist_if_dirty) = callback(&mut state_lock)?;
+
+			if !state_lock.dirty || !persist_if_dirty {
+				return Ok(res);
+			}
+
+			state_lock.dirty = false;
+
+			fut = self.persist_state(&state_lock);
+		};
+
+		fut.await.map_err(|e| {
+			self.sweeper_state.lock().unwrap().dirty = true;
+
+			log_error!(self.logger, "Error persisting OutputSweeper: {:?}", e);
+		})?;
+
+		Ok(res)
+	}
+
 	fn spend_outputs(
 		&self, sweeper_state: &SweeperState, descriptors: &[&SpendableOutputDescriptor],
 		change_destination_script: ScriptBuf,