Persist sweeper state as part of background process

To prepare for an async kv store trait that must be awaited, this commit
moves the kv store calls from the chain notification handlers to the
background process. It uses a dirty flag to communicate that there is
something to persist. The block height is part of the persisted data. If
that data does not make it to disk, the chain notifications are replayed
after restart.

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.flush_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 {
+				self.flush_state(&mut sweeper_state).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);
@@ -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.flush_state(&mut sweeper_state).map_err(|e| {
 				log_error!(self.logger, "Error persisting OutputSweeper: {:?}", e);
 			})?;
 
@@ -599,9 +613,12 @@ where
 			}
 			true
 		});
+
+		sweeper_state.dirty = true;
 	}
 
-	fn persist_state(&self, sweeper_state: &SweeperState) -> Result<(), io::Error> {
+	/// 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,
@@ -620,6 +637,9 @@ where
 				);
 				e
 			})
+			.map(|_| {
+				sweeper_state.dirty = false;
+			})
 	}
 
 	fn spend_outputs(
@@ -652,13 +672,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 +706,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 +729,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 +749,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,18 +771,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>)> {
@@ -796,6 +806,7 @@ where
 #[derive(Debug)]
 struct SweeperState {
 	persistent: PersistentSweeperState,
+	dirty: bool,
 }
 
 #[derive(Debug, Clone)]
@@ -860,7 +871,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 +920,7 @@ where
 			}
 		}
 
-		let sweeper_state = Mutex::new(SweeperState { persistent: state });
+		let sweeper_state = Mutex::new(SweeperState { persistent: state, dirty: false });
 		Ok((
 			best_block,
 			OutputSweeper {