Add sweeper state dirty flag

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),
@@ -450,7 +452,7 @@ where
 
 			state_lock.persistent.outputs.push(output_info);
 		}
-		self.persist_state(&state_lock).map_err(|e| {
+		self.persist_state(&mut state_lock).map_err(|e| {
 			log_error!(self.logger, "Error persisting OutputSweeper: {:?}", e);
 		})
 	}
@@ -478,7 +480,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 {
+				return Ok(());
+			}
+
+			self.persist_state(&mut sweeper_state).map_err(|e| {
+				log_error!(self.logger, "Error persisting OutputSweeper: {:?}", e);
+			})
+		};
 
 		// Release the pending sweep flag again, regardless of result.
 		self.pending_sweep.store(false, Ordering::Release);
@@ -571,7 +585,7 @@ where
 				output_info.status.broadcast(cur_hash, cur_height, spending_tx.clone());
 			}
 
-			self.persist_state(&sweeper_state).map_err(|e| {
+			self.persist_state(&mut sweeper_state).map_err(|e| {
 				log_error!(self.logger, "Error persisting OutputSweeper: {:?}", e);
 			})?;
 
@@ -599,9 +613,11 @@ where
 			}
 			true
 		});
+
+		sweeper_state.dirty = true;
 	}
 
-	fn persist_state(&self, sweeper_state: &SweeperState) -> Result<(), io::Error> {
+	fn persist_state(&self, sweeper_state: &mut SweeperState) -> Result<(), io::Error> {
 		self.kv_store
 			.write(
 				OUTPUT_SWEEPER_PERSISTENCE_PRIMARY_NAMESPACE,
@@ -620,6 +636,9 @@ where
 				);
 				e
 			})
+			.map(|_| {
+				sweeper_state.dirty = false;
+			})
 	}
 
 	fn spend_outputs(
@@ -652,13 +671,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;
 	}
 }
 
@@ -682,12 +705,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) {
@@ -709,9 +728,7 @@ where
 			}
 		}
 
-		self.persist_state(&*state_lock).unwrap_or_else(|e| {
-			log_error!(self.logger, "Error persisting OutputSweeper: {:?}", e);
-		});
+		state_lock.dirty = true;
 	}
 }
 
@@ -731,9 +748,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) {
@@ -756,24 +770,18 @@ 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>)> {
-		let state_lock = self.sweeper_state.lock().unwrap();
+		let state_lock = &self.sweeper_state.lock().unwrap().persistent;
 		state_lock
-			.persistent
 			.outputs
 			.iter()
 			.filter_map(|o| match o.status {
@@ -796,6 +804,7 @@ where
 #[derive(Debug)]
 struct SweeperState {
 	persistent: PersistentSweeperState,
+	dirty: bool,
 }
 
 #[derive(Debug, Clone)]
@@ -860,7 +869,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),
@@ -909,7 +918,7 @@ where
 			}
 		}
 
-		let sweeper_state = Mutex::new(SweeperState { persistent: state });
+		let sweeper_state = Mutex::new(SweeperState { persistent: state, dirty: false });
 		Ok((
 			best_block,
 			OutputSweeper {