try out async

Author: joostjager

lightning/src/util/sweep.rs

@@ -36,6 +36,7 @@ use bitcoin::{BlockHash, ScriptBuf, Transaction, Txid};
 
 use core::future::Future;
 use core::ops::Deref;
+use core::pin::Pin;
 use core::sync::atomic::{AtomicBool, Ordering};
 use core::task;
 
@@ -414,7 +415,7 @@ where
 	/// Returns `Err` on persistence failure, in which case the call may be safely retried.
 	///
 	/// [`Event::SpendableOutputs`]: crate::events::Event::SpendableOutputs
-	pub fn track_spendable_outputs(
+	pub async fn track_spendable_outputs(
 		&self, output_descriptors: Vec<SpendableOutputDescriptor>, channel_id: Option<ChannelId>,
 		exclude_static_outputs: bool, delay_until_height: Option<u32>,
 	) -> Result<(), ()> {
@@ -430,24 +431,32 @@ where
 			return Ok(());
 		}
 
-		let mut state_lock = self.sweeper_state.lock().unwrap();
-		for descriptor in relevant_descriptors {
-			let output_info = TrackedSpendableOutput {
-				descriptor,
-				channel_id,
-				status: OutputSpendStatus::PendingInitialBroadcast {
-					delayed_until_height: delay_until_height,
-				},
-			};
-
-			let mut outputs = state_lock.persistent.outputs.iter();
-			if outputs.find(|o| o.descriptor == output_info.descriptor).is_some() {
-				continue;
-			}
+		let persist_fut;
+		{
+			let mut state_lock = self.sweeper_state.lock().unwrap();
+			for descriptor in relevant_descriptors {
+				let output_info = TrackedSpendableOutput {
+					descriptor,
+					channel_id,
+					status: OutputSpendStatus::PendingInitialBroadcast {
+						delayed_until_height: delay_until_height,
+					},
+				};
+
+				let mut outputs = state_lock.persistent.outputs.iter();
+				if outputs.find(|o| o.descriptor == output_info.descriptor).is_some() {
+					continue;
+				}
 
-			state_lock.persistent.outputs.push(output_info);
+				state_lock.persistent.outputs.push(output_info);
+			}
+			persist_fut = self.persist_state(&state_lock.persistent);
+			state_lock.dirty = false;
 		}
-		self.flush_state(&mut state_lock).map_err(|e| {
+
+		persist_fut.await.map_err(|e| {
+			self.sweeper_state.lock().unwrap().dirty = true;
+
 			log_error!(self.logger, "Error persisting OutputSweeper: {:?}", e);
 		})
 	}
@@ -476,12 +485,20 @@ where
 		}
 
 		let result = {
-			self.regenerate_and_broadcast_spend_if_necessary_internal().await?;
+			let dirty = self.regenerate_and_broadcast_spend_if_necessary_internal().await?;
 
 			// If there is still dirty state, we need to persist it.
-			let mut sweeper_state = self.sweeper_state.lock().unwrap();
-			if sweeper_state.dirty {
-				self.flush_state(&mut sweeper_state).map_err(|e| {
+			if dirty {
+				let persist_fut;
+				{
+					let mut sweeper_state = self.sweeper_state.lock().unwrap();
+					persist_fut = self.persist_state(&sweeper_state.persistent);
+					sweeper_state.dirty = false;
+				}
+
+				persist_fut.await.map_err(|e| {
+					self.sweeper_state.lock().unwrap().dirty = true;
+
 					log_error!(self.logger, "Error persisting OutputSweeper: {:?}", e);
 				})
 			} else {
@@ -496,7 +513,7 @@ where
 	}
 
 	/// Regenerates and broadcasts the spending transaction for any outputs that are pending
-	async fn regenerate_and_broadcast_spend_if_necessary_internal(&self) -> Result<(), ()> {
+	async fn regenerate_and_broadcast_spend_if_necessary_internal(&self) -> Result<bool, ()> {
 		let filter_fn = |o: &TrackedSpendableOutput, cur_height: u32| {
 			if o.status.is_confirmed() {
 				// Don't rebroadcast confirmed txs.
@@ -524,7 +541,7 @@ where
 			let has_respends =
 				sweeper_state.persistent.outputs.iter().any(|o| filter_fn(o, cur_height));
 			if !has_respends {
-				return Ok(());
+				return Ok(sweeper_state.dirty);
 			}
 		}
 
@@ -533,6 +550,7 @@ where
 			self.change_destination_source.get_change_destination_script().await?;
 
 		// Sweep the outputs.
+		let persist_fut;
 		{
 			let mut sweeper_state = self.sweeper_state.lock().unwrap();
 
@@ -549,7 +567,7 @@ where
 
 			if respend_descriptors.is_empty() {
 				// It could be that a tx confirmed and there is now nothing to sweep anymore.
-				return Ok(());
+				return Ok(sweeper_state.dirty);
 			}
 
 			let spending_tx = self
@@ -581,14 +599,18 @@ where
 				output_info.status.broadcast(cur_hash, cur_height, spending_tx.clone());
 			}
 
-			self.flush_state(&mut sweeper_state).map_err(|e| {
-				log_error!(self.logger, "Error persisting OutputSweeper: {:?}", e);
-			})?;
-
+			persist_fut = self.persist_state(&sweeper_state.persistent);
+			sweeper_state.dirty = false;
 			self.broadcaster.broadcast_transactions(&[&spending_tx]);
 		}
 
-		Ok(())
+		persist_fut.await.map_err(|e| {
+			self.sweeper_state.lock().unwrap().dirty = true;
+
+			log_error!(self.logger, "Error persisting OutputSweeper: {:?}", e);
+		})?;
+
+		Ok(false)
 	}
 
 	fn prune_confirmed_outputs(&self, sweeper_state: &mut SweeperState) {
@@ -613,29 +635,19 @@ where
 		sweeper_state.dirty = true;
 	}
 
-	/// Flushes the current state to the persistence layer and marks the state as clean.
-	fn flush_state(&self, sweeper_state: &mut SweeperState) -> Result<(), io::Error> {
-		self.kv_store
-			.write(
-				OUTPUT_SWEEPER_PERSISTENCE_PRIMARY_NAMESPACE,
-				OUTPUT_SWEEPER_PERSISTENCE_SECONDARY_NAMESPACE,
-				OUTPUT_SWEEPER_PERSISTENCE_KEY,
-				&sweeper_state.persistent.encode(),
-			)
-			.map_err(|e| {
-				log_error!(
-					self.logger,
-					"Write for key {}/{}/{} failed due to: {}",
-					OUTPUT_SWEEPER_PERSISTENCE_PRIMARY_NAMESPACE,
-					OUTPUT_SWEEPER_PERSISTENCE_SECONDARY_NAMESPACE,
-					OUTPUT_SWEEPER_PERSISTENCE_KEY,
-					e
-				);
-				e
-			})
-			.map(|_| {
-				sweeper_state.dirty = false;
-			})
+	fn persist_state<'a>(
+		&self, sweeper_state: &PersistentSweeperState,
+	) -> Pin<Box<dyn Future<Output = Result<(), io::Error>> + 'a + Send>> {
+		let encoded = &sweeper_state.encode();
+
+		let result = self.kv_store.write(
+			OUTPUT_SWEEPER_PERSISTENCE_PRIMARY_NAMESPACE,
+			OUTPUT_SWEEPER_PERSISTENCE_SECONDARY_NAMESPACE,
+			OUTPUT_SWEEPER_PERSISTENCE_KEY,
+			encoded,
+		);
+
+		Box::pin(async move { result })
 	}
 
 	fn spend_outputs(