Remove pruned LSPS2/LSPS5 peer state entries from the KVStore

Previously, we'd persist peer states to the `KVStore`, but, while we
pruned them eventually from our in-memory state, we wouldn't remove it
from the `KVStore`.

Here, we change this and regularly prune and delete peer state entries
from the `KVStore`.
Note we still prune the state-internal data on peer disconnection, but
leave removal to our (BP-driven) async `persist` calls.

Author: tnull

lightning-liquidity/src/lsps2/service.rs

@@ -514,7 +514,7 @@ impl PeerState {
 				// We abort the flow, and prune any data kept.
 				self.intercept_scid_by_channel_id.retain(|_, iscid| intercept_scid != iscid);
 				self.intercept_scid_by_user_channel_id.retain(|_, iscid| intercept_scid != iscid);
-				// TODO: Remove peer state entry from the KVStore
+				self.needs_persist |= true;
 				return false;
 			}
 			true
@@ -1662,28 +1662,41 @@ where
 	}
 
 	pub(crate) fn peer_disconnected(&self, counterparty_node_id: PublicKey) {
-		let mut outer_state_lock = self.per_peer_state.write().unwrap();
-		let is_prunable =
-			if let Some(inner_state_lock) = outer_state_lock.get(&counterparty_node_id) {
-				let mut peer_state_lock = inner_state_lock.lock().unwrap();
-				peer_state_lock.prune_expired_request_state();
-				peer_state_lock.is_prunable()
-			} else {
-				return;
-			};
-		if is_prunable {
-			outer_state_lock.remove(&counterparty_node_id);
+		let outer_state_lock = self.per_peer_state.write().unwrap();
+		if let Some(inner_state_lock) = outer_state_lock.get(&counterparty_node_id) {
+			let mut peer_state_lock = inner_state_lock.lock().unwrap();
+			// We clean up the peer state, but leave removing the peer entry to `prune_peer_state`
+			// which also removes it from the store.
+			peer_state_lock.prune_expired_request_state();
 		}
 	}
 
 	#[allow(clippy::bool_comparison)]
-	pub(crate) fn prune_peer_state(&self) {
+	pub(crate) async fn prune_peer_state(&self) {
 		let mut outer_state_lock = self.per_peer_state.write().unwrap();
-		outer_state_lock.retain(|_, inner_state_lock| {
+		let mut need_remove = Vec::new();
+		outer_state_lock.retain(|counterparty_node_id, inner_state_lock| {
 			let mut peer_state_lock = inner_state_lock.lock().unwrap();
 			peer_state_lock.prune_expired_request_state();
-			peer_state_lock.is_prunable() == false
+			let is_prunable = peer_state_lock.is_prunable();
+			if is_prunable {
+				need_remove.push(*counterparty_node_id);
+			}
+			is_prunable == false
 		});
+
+		for counterparty_node_id in need_remove {
+			let key = counterparty_node_id.to_string();
+			let _ = self
+				.kv_store
+				.remove(
+					LIQUIDITY_MANAGER_PERSISTENCE_PRIMARY_NAMESPACE,
+					LSPS2_SERVICE_PERSISTENCE_SECONDARY_NAMESPACE,
+					&key,
+					true,
+				)
+				.await;
+		}
 	}
 }
 

lightning-liquidity/src/lsps5/service.rs

@@ -270,18 +270,48 @@ where
 		});
 
 		if should_prune {
-			outer_state_lock.retain(|client_id, peer_state| {
-				if self.client_has_open_channel(client_id) {
-					// Don't prune clients with open channels
-					return true;
-				}
-				// TODO: Remove peer state entry from the KVStore
-				!peer_state.prune_stale_webhooks(now)
-			});
+			for (_, peer_state) in outer_state_lock.iter_mut() {
+				// Prune stale webhooks, but leave removal of the peers states to prune_peer_state
+				// which will also remove it from the store.
+				peer_state.prune_stale_webhooks(now)
+			}
 			*last_pruning = Some(now);
 		}
 	}
 
+	pub(crate) async fn prune_peer_state(&self) {
+		let mut outer_state_lock = self.per_peer_state.write().unwrap();
+		let mut need_remove = Vec::new();
+
+		self.check_prune_stale_webhooks(&mut outer_state_lock);
+
+		outer_state_lock.retain(|client_id, peer_state| {
+			if self.client_has_open_channel(client_id) {
+				// Don't prune clients with open channels
+				return true;
+			}
+
+			let is_prunable = peer_state.is_prunable();
+			if is_prunable {
+				need_remove.push(*client_id);
+			}
+			!is_prunable
+		});
+
+		for counterparty_node_id in need_remove {
+			let key = counterparty_node_id.to_string();
+			let _ = self
+				.kv_store
+				.remove(
+					LIQUIDITY_MANAGER_PERSISTENCE_PRIMARY_NAMESPACE,
+					LSPS5_SERVICE_PERSISTENCE_SECONDARY_NAMESPACE,
+					&key,
+					true,
+				)
+				.await;
+		}
+	}
+
 	fn handle_set_webhook(
 		&self, counterparty_node_id: PublicKey, request_id: LSPSRequestId,
 		params: SetWebhookRequest,
@@ -733,11 +763,17 @@ impl PeerState {
 	}
 
 	// Returns whether the entire state is empty and can be pruned.
-	fn prune_stale_webhooks(&mut self, now: LSPSDateTime) -> bool {
+	fn prune_stale_webhooks(&mut self, now: LSPSDateTime) {
 		self.webhooks.retain(|(_, webhook)| {
-			now.duration_since(&webhook.last_used) < MIN_WEBHOOK_RETENTION_DAYS
+			let should_prune = now.duration_since(&webhook.last_used) >= MIN_WEBHOOK_RETENTION_DAYS;
+			if should_prune {
+				self.needs_persist |= true;
+			}
+			!should_prune
 		});
+	}
 
+	fn is_prunable(&mut self) -> bool {
 		self.webhooks.is_empty()
 	}
 }

lightning-liquidity/src/manager.rs

@@ -61,6 +61,7 @@ use bitcoin::secp256k1::PublicKey;
 use core::future::Future as StdFuture;
 use core::ops::Deref;
 use core::task;
+use core::time::Duration;
 
 const LSPS_FEATURE_BIT: usize = 729;
 
@@ -311,7 +312,9 @@ pub struct LiquidityManager<
 	service_config: Option<LiquidityServiceConfig>,
 	_client_config: Option<LiquidityClientConfig>,
 	best_block: RwLock<Option<BestBlock>>,
+	last_peer_state_pruning: Mutex<Option<Duration>>,
 	_chain_source: Option<C>,
+	time_provider: TP,
 	pending_msgs_or_needs_persist_notifier: Arc<Notifier>,
 }
 
@@ -461,7 +464,7 @@ where
 					kv_store.clone(),
 					node_signer,
 					lsps5_service_config.clone(),
-					time_provider,
+					time_provider.clone(),
 				))
 			} else {
 				None
@@ -512,6 +515,8 @@ where
 			None
 		};
 
+		let last_peer_state_pruning = Mutex::new(None);
+
 		Ok(Self {
 			pending_messages,
 			pending_events,
@@ -529,7 +534,9 @@ where
 			service_config,
 			_client_config: client_config,
 			best_block: RwLock::new(chain_params.map(|chain_params| chain_params.best_block)),
+			last_peer_state_pruning,
 			_chain_source: chain_source,
+			time_provider,
 			pending_msgs_or_needs_persist_notifier,
 		})
 	}
@@ -650,14 +657,32 @@ where
 	/// This will be regularly called by LDK's background processor if necessary and only needs to
 	/// be called manually if it's not utilized.
 	pub async fn persist(&self) -> Result<(), lightning::io::Error> {
+		let should_prune_state = {
+			const PRUNE_INTERVAL: Duration = Duration::from_secs(600);
+			let mut last_peer_state_pruning_lock = self.last_peer_state_pruning.lock().unwrap();
+			let now = self.time_provider.duration_since_epoch();
+			if last_peer_state_pruning_lock.map_or(true, |l| l + PRUNE_INTERVAL < now) {
+				*last_peer_state_pruning_lock = Some(now);
+				true
+			} else {
+				false
+			}
+		};
+
 		// TODO: We should eventually persist in parallel.
 		self.pending_events.persist().await?;
 
 		if let Some(lsps2_service_handler) = self.lsps2_service_handler.as_ref() {
+			if should_prune_state {
+				lsps2_service_handler.prune_peer_state().await;
+			}
 			lsps2_service_handler.persist().await?;
 		}
 
 		if let Some(lsps5_service_handler) = self.lsps5_service_handler.as_ref() {
+			if should_prune_state {
+				lsps5_service_handler.prune_peer_state().await;
+			}
 			lsps5_service_handler.persist().await?;
 		}
 
@@ -1015,9 +1040,6 @@ where
 		*self.best_block.write().unwrap() = Some(new_best_block);
 
 		// TODO: Call best_block_updated on all sub-modules that require it, e.g., LSPS1MessageHandler.
-		if let Some(lsps2_service_handler) = self.lsps2_service_handler.as_ref() {
-			lsps2_service_handler.prune_peer_state();
-		}
 	}
 
 	fn get_relevant_txids(&self) -> Vec<(bitcoin::Txid, u32, Option<bitcoin::BlockHash>)> {