Move persist into async part of the sweeper

Prepares for making the kv store async. Otherwise it might be necessary
to use block_on in the sweeper. For block_on, a runtime would be needed.

Author: joostjager

lightning/src/util/sweep.rs

@@ -382,8 +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(SweeperState { persistent: PersistentSweeperState { outputs, best_block } });
+		let sweeper_state = Mutex::new(SweeperState {
+			persistent: PersistentSweeperState { outputs, best_block },
+			dirty: false,
+		});
 		Self {
 			sweeper_state,
 			pending_sweep: AtomicBool::new(false),
@@ -445,7 +447,7 @@ where
 
 			state_lock.persistent.outputs.push(output_info);
 		}
-		self.persist_state(&*state_lock).await.map_err(|e| {
+		self.flush_state(&mut state_lock).await.map_err(|e| {
 			log_error!(self.logger, "Error persisting OutputSweeper: {:?}", e);
 		})
 	}
@@ -473,7 +475,19 @@ where
 			return Ok(());
 		}
 
-		let result = self.regenerate_and_broadcast_spend_if_necessary_internal().await;
+		let result = {
+			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).await.map_err(|e| {
+					log_error!(self.logger, "Error persisting OutputSweeper: {:?}", e);
+				})
+			} else {
+				Ok(())
+			}
+		};
 
 		// Release the pending sweep flag again, regardless of result.
 		self.pending_sweep.store(false, Ordering::Release);
@@ -567,7 +581,7 @@ where
 				output_info.status.broadcast(cur_hash, cur_height, spending_tx.clone());
 			}
 
-			self.persist_state(&sweeper_state).await.map_err(|e| {
+			self.flush_state(&mut sweeper_state).await.map_err(|e| {
 				log_error!(self.logger, "Error persisting OutputSweeper: {:?}", e);
 			})?;
 
@@ -595,9 +609,12 @@ where
 			}
 			true
 		});
+
+		sweeper_state.dirty = true;
 	}
 
-	async fn persist_state(&self, sweeper_state: &SweeperState) -> Result<(), io::Error> {
+	/// Flushes the current state to the persistence layer and marks the state as clean.
+	async fn flush_state(&self, sweeper_state: &mut SweeperState) -> Result<(), io::Error> {
 		self.kv_store
 			.write(
 				OUTPUT_SWEEPER_PERSISTENCE_PRIMARY_NAMESPACE,
@@ -617,6 +634,9 @@ where
 				);
 				e
 			})
+			.map(|_| {
+				sweeper_state.dirty = false;
+			})
 	}
 
 	fn spend_outputs(
@@ -649,13 +669,17 @@ where
 				}
 			}
 		}
+
+		sweeper_state.dirty = true;
 	}
 
 	fn best_block_updated_internal(
 		&self, sweeper_state: &mut SweeperState, header: &Header, height: u32,
 	) {
 		sweeper_state.persistent.best_block = BestBlock::new(header.block_hash(), height);
 		self.prune_confirmed_outputs(sweeper_state);
+
+		sweeper_state.dirty = true;
 	}
 }
 
@@ -679,12 +703,8 @@ where
 		assert_eq!(state_lock.persistent.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(&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| {
-		// 	log_error!(self.logger, "Error persisting OutputSweeper: {:?}", e);
-		// });
+		self.transactions_confirmed_internal(&mut state_lock, header, txdata, height);
+		self.best_block_updated_internal(&mut state_lock, header, height);
 	}
 
 	fn block_disconnected(&self, header: &Header, height: u32) {
@@ -706,9 +726,7 @@ where
 			}
 		}
 
-		// self.persist_state(&*state_lock).unwrap_or_else(|e| {
-		// 	log_error!(self.logger, "Error persisting OutputSweeper: {:?}", e);
-		// });
+		state_lock.dirty = true;
 	}
 }
 
@@ -728,9 +746,6 @@ where
 	) {
 		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) {
@@ -753,18 +768,13 @@ 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| {
-			// 	log_error!(self.logger, "Error persisting OutputSweeper: {:?}", e);
-			// });
+			state_lock.dirty = true;
 		}
 	}
 
 	fn best_block_updated(&self, header: &Header, height: u32) {
 		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);
-		// });
+		self.best_block_updated_internal(&mut state_lock, header, height);
 	}
 
 	fn get_relevant_txids(&self) -> Vec<(Txid, u32, Option<BlockHash>)> {
@@ -793,6 +803,7 @@ where
 #[derive(Debug)]
 struct SweeperState {
 	persistent: PersistentSweeperState,
+	dirty: bool,
 }
 
 #[derive(Debug, Clone)]
@@ -857,7 +868,7 @@ where
 			}
 		}
 
-		let sweeper_state = Mutex::new(SweeperState { persistent: state });
+		let sweeper_state = Mutex::new(SweeperState { persistent: state, dirty: false });
 		Ok(Self {
 			sweeper_state,
 			pending_sweep: AtomicBool::new(false),
@@ -906,7 +917,7 @@ where
 			}
 		}
 
-		let sweeper_state = Mutex::new(SweeperState { persistent: state });
+		let sweeper_state = Mutex::new(SweeperState { persistent: state, dirty: false });
 		Ok((
 			best_block,
 			OutputSweeper {