Ensure mutual exclusion in LSPS2/5 persistence

There are various race conditions between `persist` calls on the
LSPS2 and LSPS5 services. Thus, we simply ensure that no two
`persist` calls can be in-flight at the same time, ensuring that
any new state updates between them will ultimately be persisted by
re-persisting from the top in the originally-running `persist` call
if another is started before it finishes.

Author: TheBlueMatt

lightning-liquidity/src/lsps2/service.rs

@@ -593,6 +593,7 @@ where
 	peer_by_channel_id: RwLock<HashMap<ChannelId, PublicKey>>,
 	total_pending_requests: AtomicUsize,
 	config: LSPS2ServiceConfig,
+	persistence_in_flight: AtomicUsize,
 }
 
 impl<CM: Deref, K: Deref + Clone> LSPS2ServiceHandler<CM, K>
@@ -640,6 +641,7 @@ where
 			peer_by_intercept_scid: RwLock::new(peer_by_intercept_scid),
 			peer_by_channel_id: RwLock::new(peer_by_channel_id),
 			total_pending_requests: AtomicUsize::new(0),
+			persistence_in_flight: AtomicUsize::new(0),
 			channel_manager,
 			kv_store,
 			config,
@@ -1645,64 +1647,80 @@ where
 		// introduce some batching to upper-bound the number of requests inflight at any given
 		// time.
 
-		let mut need_remove = Vec::new();
-		let mut need_persist = Vec::new();
+		if self.persistence_in_flight.fetch_add(1, Ordering::AcqRel) > 0 {
+			// If we're not the first event processor to get here, just return early, the increment
+			// we just did will be treated as "go around again" at the end.
+			return Ok(());
+		}
 
-		{
-			// First build a list of peers to persist and prune with the read lock. This allows
-			// us to avoid the write lock unless we actually need to remove a node.
-			let outer_state_lock = self.per_peer_state.read().unwrap();
-			for (counterparty_node_id, inner_state_lock) in outer_state_lock.iter() {
-				let mut peer_state_lock = inner_state_lock.lock().unwrap();
-				peer_state_lock.prune_expired_request_state();
-				let is_prunable = peer_state_lock.is_prunable();
-				if is_prunable {
-					need_remove.push(*counterparty_node_id);
-				} else if peer_state_lock.needs_persist {
-					need_persist.push(*counterparty_node_id);
+		loop {
+			let mut need_remove = Vec::new();
+			let mut need_persist = Vec::new();
+
+			{
+				// First build a list of peers to persist and prune with the read lock. This allows
+				// us to avoid the write lock unless we actually need to remove a node.
+				let outer_state_lock = self.per_peer_state.read().unwrap();
+				for (counterparty_node_id, inner_state_lock) in outer_state_lock.iter() {
+					let mut peer_state_lock = inner_state_lock.lock().unwrap();
+					peer_state_lock.prune_expired_request_state();
+					let is_prunable = peer_state_lock.is_prunable();
+					if is_prunable {
+						need_remove.push(*counterparty_node_id);
+					} else if peer_state_lock.needs_persist {
+						need_persist.push(*counterparty_node_id);
+					}
 				}
 			}
-		}
 
-		for counterparty_node_id in need_persist.into_iter() {
-			debug_assert!(!need_remove.contains(&counterparty_node_id));
-			self.persist_peer_state(counterparty_node_id).await?;
-		}
+			for counterparty_node_id in need_persist.into_iter() {
+				debug_assert!(!need_remove.contains(&counterparty_node_id));
+				self.persist_peer_state(counterparty_node_id).await?;
+			}
 
-		for counterparty_node_id in need_remove {
-			let mut future_opt = None;
-			{
-				// We need to take the `per_peer_state` write lock to remove an entry, but also
-				// have to hold it until after the `remove` call returns (but not through
-				// future completion) to ensure that writes for the peer's state are
-				// well-ordered with other `persist_peer_state` calls even across the removal
-				// itself.
-				let mut per_peer_state = self.per_peer_state.write().unwrap();
-				if let Entry::Occupied(mut entry) = per_peer_state.entry(counterparty_node_id) {
-					let state = entry.get_mut().get_mut().unwrap();
-					if state.is_prunable() {
-						entry.remove();
-						let key = counterparty_node_id.to_string();
-						future_opt = Some(self.kv_store.remove(
-							LIQUIDITY_MANAGER_PERSISTENCE_PRIMARY_NAMESPACE,
-							LSPS2_SERVICE_PERSISTENCE_SECONDARY_NAMESPACE,
-							&key,
-						));
+			for counterparty_node_id in need_remove {
+				let mut future_opt = None;
+				{
+					// We need to take the `per_peer_state` write lock to remove an entry, but also
+					// have to hold it until after the `remove` call returns (but not through
+					// future completion) to ensure that writes for the peer's state are
+					// well-ordered with other `persist_peer_state` calls even across the removal
+					// itself.
+					let mut per_peer_state = self.per_peer_state.write().unwrap();
+					if let Entry::Occupied(mut entry) = per_peer_state.entry(counterparty_node_id) {
+						let state = entry.get_mut().get_mut().unwrap();
+						if state.is_prunable() {
+							entry.remove();
+							let key = counterparty_node_id.to_string();
+							future_opt = Some(self.kv_store.remove(
+								LIQUIDITY_MANAGER_PERSISTENCE_PRIMARY_NAMESPACE,
+								LSPS2_SERVICE_PERSISTENCE_SECONDARY_NAMESPACE,
+								&key,
+							));
+						} else {
+							// If the peer got new state, force a re-persist of the current state.
+							state.needs_persist = true;
+						}
 					} else {
-						// If the peer got new state, force a re-persist of the current state.
-						state.needs_persist = true;
+						// This should never happen, we can only have one `persist` call
+						// in-progress at once and map entries are only removed by it.
+						debug_assert!(false);
 					}
+				}
+				if let Some(future) = future_opt {
+					future.await?;
 				} else {
-					// This should never happen, we can only have one `persist` call
-					// in-progress at once and map entries are only removed by it.
-					debug_assert!(false);
+					self.persist_peer_state(counterparty_node_id).await?;
 				}
 			}
-			if let Some(future) = future_opt {
-				future.await?;
-			} else {
-				self.persist_peer_state(counterparty_node_id).await?;
+
+			if self.persistence_in_flight.fetch_sub(1, Ordering::AcqRel) != 1 {
+				// If another thread incremented the state while we were running we should go
+				// around again, but only once.
+				self.persistence_in_flight.store(1, Ordering::Release);
+				continue;
 			}
+			break;
 		}
 
 		Ok(())

lightning-liquidity/src/lsps5/service.rs

@@ -36,6 +36,7 @@ use lightning::util::persist::KVStore;
 use lightning::util::ser::Writeable;
 
 use core::ops::Deref;
+use core::sync::atomic::{AtomicUsize, Ordering};
 use core::time::Duration;
 
 use alloc::string::String;
@@ -140,6 +141,7 @@ where
 	node_signer: NS,
 	kv_store: K,
 	last_pruning: Mutex<Option<LSPSDateTime>>,
+	persistence_in_flight: AtomicUsize,
 }
 
 impl<CM: Deref, NS: Deref, K: Deref + Clone, TP: Deref> LSPS5ServiceHandler<CM, NS, K, TP>
@@ -167,6 +169,7 @@ where
 			node_signer,
 			kv_store,
 			last_pruning: Mutex::new(None),
+			persistence_in_flight: AtomicUsize::new(0),
 		}
 	}
 
@@ -245,63 +248,80 @@ where
 		// TODO: We should eventually persist in parallel, however, when we do, we probably want to
 		// introduce some batching to upper-bound the number of requests inflight at any given
 		// time.
-		let mut need_remove = Vec::new();
-		let mut need_persist = Vec::new();
-
-		self.check_prune_stale_webhooks(&mut self.per_peer_state.write().unwrap());
-		{
-			let outer_state_lock = self.per_peer_state.read().unwrap();
-
-			for (client_id, peer_state) in outer_state_lock.iter() {
-				let is_prunable = peer_state.is_prunable();
-				let has_open_channel = self.client_has_open_channel(client_id);
-				if is_prunable && !has_open_channel {
-					need_remove.push(*client_id);
-				} else if peer_state.needs_persist {
-					need_persist.push(*client_id);
-				}
-			}
-		}
 
-		for client_id in need_persist.into_iter() {
-			debug_assert!(!need_remove.contains(&client_id));
-			self.persist_peer_state(client_id).await?;
+		if self.persistence_in_flight.fetch_add(1, Ordering::AcqRel) > 0 {
+			// If we're not the first event processor to get here, just return early, the increment
+			// we just did will be treated as "go around again" at the end.
+			return Ok(());
 		}
 
-		for client_id in need_remove {
-			let mut future_opt = None;
+		loop {
+			let mut need_remove = Vec::new();
+			let mut need_persist = Vec::new();
+
+			self.check_prune_stale_webhooks(&mut self.per_peer_state.write().unwrap());
 			{
-				// We need to take the `per_peer_state` write lock to remove an entry, but also
-				// have to hold it until after the `remove` call returns (but not through
-				// future completion) to ensure that writes for the peer's state are
-				// well-ordered with other `persist_peer_state` calls even across the removal
-				// itself.
-				let mut per_peer_state = self.per_peer_state.write().unwrap();
-				if let Entry::Occupied(mut entry) = per_peer_state.entry(client_id) {
-					let state = entry.get_mut();
-					if state.is_prunable() && !self.client_has_open_channel(&client_id) {
-						entry.remove();
-						let key = client_id.to_string();
-						future_opt = Some(self.kv_store.remove(
-							LIQUIDITY_MANAGER_PERSISTENCE_PRIMARY_NAMESPACE,
-							LSPS5_SERVICE_PERSISTENCE_SECONDARY_NAMESPACE,
-							&key,
-						));
+				let outer_state_lock = self.per_peer_state.read().unwrap();
+
+				for (client_id, peer_state) in outer_state_lock.iter() {
+					let is_prunable = peer_state.is_prunable();
+					let has_open_channel = self.client_has_open_channel(client_id);
+					if is_prunable && !has_open_channel {
+						need_remove.push(*client_id);
+					} else if peer_state.needs_persist {
+						need_persist.push(*client_id);
+					}
+				}
+			}
+
+			for client_id in need_persist.into_iter() {
+				debug_assert!(!need_remove.contains(&client_id));
+				self.persist_peer_state(client_id).await?;
+			}
+
+			for client_id in need_remove {
+				let mut future_opt = None;
+				{
+					// We need to take the `per_peer_state` write lock to remove an entry, but also
+					// have to hold it until after the `remove` call returns (but not through
+					// future completion) to ensure that writes for the peer's state are
+					// well-ordered with other `persist_peer_state` calls even across the removal
+					// itself.
+					let mut per_peer_state = self.per_peer_state.write().unwrap();
+					if let Entry::Occupied(mut entry) = per_peer_state.entry(client_id) {
+						let state = entry.get_mut();
+						if state.is_prunable() && !self.client_has_open_channel(&client_id) {
+							entry.remove();
+							let key = client_id.to_string();
+							future_opt = Some(self.kv_store.remove(
+								LIQUIDITY_MANAGER_PERSISTENCE_PRIMARY_NAMESPACE,
+								LSPS5_SERVICE_PERSISTENCE_SECONDARY_NAMESPACE,
+								&key,
+							));
+						} else {
+							// If the peer was re-added, force a re-persist of the current state.
+							state.needs_persist = true;
+						}
 					} else {
-						// If the peer was re-added, force a re-persist of the current state.
-						state.needs_persist = true;
+						// This should never happen, we can only have one `persist` call
+						// in-progress at once and map entries are only removed by it.
+						debug_assert!(false);
 					}
+				}
+				if let Some(future) = future_opt {
+					future.await?;
 				} else {
-					// This should never happen, we can only have one `persist` call
-					// in-progress at once and map entries are only removed by it.
-					debug_assert!(false);
+					self.persist_peer_state(client_id).await?;
 				}
 			}
-			if let Some(future) = future_opt {
-				future.await?;
-			} else {
-				self.persist_peer_state(client_id).await?;
+
+			if self.persistence_in_flight.fetch_sub(1, Ordering::AcqRel) != 1 {
+				// If another thread incremented the state while we were running we should go
+				// around again, but only once.
+				self.persistence_in_flight.store(1, Ordering::Release);
+				continue;
 			}
+			break;
 		}
 
 		Ok(())