Add LSPS2ServiceHandler persistence

We add simple `persist` call to `LSPS2ServiceHandler` that sequentially
persist all the peer states under a key that encodes their node id.

Author: tnull

lightning-liquidity/src/lib.rs

@@ -65,6 +65,7 @@ pub mod lsps2;
 pub mod lsps5;
 mod manager;
 pub mod message_queue;
+pub mod persist;
 #[allow(dead_code)]
 #[allow(unused_imports)]
 mod sync;

lightning-liquidity/src/lsps2/service.rs

@@ -11,6 +11,7 @@
 
 use alloc::string::{String, ToString};
 use alloc::vec::Vec;
+use lightning::util::persist::KVStore;
 
 use core::cmp::Ordering as CmpOrdering;
 use core::ops::Deref;
@@ -28,6 +29,9 @@ use crate::lsps2::utils::{
 	compute_opening_fee, is_expired_opening_fee_params, is_valid_opening_fee_params,
 };
 use crate::message_queue::{MessageQueue, MessageQueueNotifierGuard};
+use crate::persist::{
+	LIQUIDITY_MANAGER_PERSISTENCE_PRIMARY_NAMESPACE, LSPS2_SERVICE_PERSISTENCE_SECONDARY_NAMESPACE,
+};
 use crate::prelude::hash_map::Entry;
 use crate::prelude::{new_hash_map, HashMap};
 use crate::sync::{Arc, Mutex, MutexGuard, RwLock};
@@ -38,6 +42,7 @@ use lightning::ln::msgs::{ErrorAction, LightningError};
 use lightning::ln::types::ChannelId;
 use lightning::util::errors::APIError;
 use lightning::util::logger::Level;
+use lightning::util::ser::Writeable;
 use lightning::{impl_writeable_tlv_based, impl_writeable_tlv_based_enum};
 
 use lightning_types::payment::PaymentHash;
@@ -564,11 +569,13 @@ macro_rules! get_or_insert_peer_state_entry {
 }
 
 /// The main object allowing to send and receive bLIP-52 / LSPS2 messages.
-pub struct LSPS2ServiceHandler<CM: Deref>
+pub struct LSPS2ServiceHandler<CM: Deref, K: Deref + Clone>
 where
 	CM::Target: AChannelManager,
+	K::Target: KVStore,
 {
 	channel_manager: CM,
+	kv_store: K,
 	pending_messages: Arc<MessageQueue>,
 	pending_events: Arc<EventQueue>,
 	per_peer_state: RwLock<HashMap<PublicKey, Mutex<PeerState>>>,
@@ -578,14 +585,15 @@ where
 	config: LSPS2ServiceConfig,
 }
 
-impl<CM: Deref> LSPS2ServiceHandler<CM>
+impl<CM: Deref, K: Deref + Clone> LSPS2ServiceHandler<CM, K>
 where
 	CM::Target: AChannelManager,
+	K::Target: KVStore,
 {
 	/// Constructs a `LSPS2ServiceHandler`.
 	pub(crate) fn new(
 		pending_messages: Arc<MessageQueue>, pending_events: Arc<EventQueue>, channel_manager: CM,
-		config: LSPS2ServiceConfig,
+		kv_store: K, config: LSPS2ServiceConfig,
 	) -> Self {
 		Self {
 			pending_messages,
@@ -595,6 +603,7 @@ where
 			peer_by_channel_id: RwLock::new(new_hash_map()),
 			total_pending_requests: AtomicUsize::new(0),
 			channel_manager,
+			kv_store,
 			config,
 		}
 	}
@@ -1442,6 +1451,50 @@ where
 		);
 	}
 
+	async fn persist_peer_state(
+		&self, counterparty_node_id: PublicKey,
+	) -> Result<(), lightning::io::Error> {
+		let fut = {
+			let outer_state_lock = self.per_peer_state.read().unwrap();
+			let encoded = match outer_state_lock.get(&counterparty_node_id) {
+				None => {
+					let err = lightning::io::Error::new(
+						lightning::io::ErrorKind::Other,
+						"Failed to get peer entry",
+					);
+					return Err(err);
+				},
+				Some(entry) => entry.lock().unwrap().encode(),
+			};
+			let key = counterparty_node_id.to_string();
+
+			self.kv_store.write(
+				LIQUIDITY_MANAGER_PERSISTENCE_PRIMARY_NAMESPACE,
+				LSPS2_SERVICE_PERSISTENCE_SECONDARY_NAMESPACE,
+				&key,
+				encoded,
+			)
+		};
+
+		fut.await
+	}
+
+	pub(crate) async fn persist(&self) -> Result<(), lightning::io::Error> {
+		// 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 need_persist: Vec<PublicKey> = {
+			let outer_state_lock = self.per_peer_state.read().unwrap();
+			outer_state_lock.iter().filter_map(|(k, v)| Some(*k)).collect()
+		};
+
+		for counterparty_node_id in need_persist.into_iter() {
+			self.persist_peer_state(counterparty_node_id).await?;
+		}
+
+		Ok(())
+	}
+
 	pub(crate) fn peer_disconnected(&self, counterparty_node_id: PublicKey) {
 		let mut outer_state_lock = self.per_peer_state.write().unwrap();
 		let is_prunable =
@@ -1468,9 +1521,10 @@ where
 	}
 }
 
-impl<CM: Deref> LSPSProtocolMessageHandler for LSPS2ServiceHandler<CM>
+impl<CM: Deref, K: Deref + Clone> LSPSProtocolMessageHandler for LSPS2ServiceHandler<CM, K>
 where
 	CM::Target: AChannelManager,
+	K::Target: KVStore,
 {
 	type ProtocolMessage = LSPS2Message;
 	const PROTOCOL_NUMBER: Option<u16> = Some(2);

lightning-liquidity/src/manager.rs

@@ -7,6 +7,7 @@
 // You may not use this file except in accordance with one or both of these
 // licenses.
 
+use alloc::boxed::Box;
 use alloc::string::ToString;
 use alloc::vec::Vec;
 
@@ -34,6 +35,7 @@ use crate::lsps2::msgs::LSPS2Message;
 use crate::lsps2::service::{LSPS2ServiceConfig, LSPS2ServiceHandler};
 use crate::prelude::{new_hash_map, new_hash_set, HashMap, HashSet};
 use crate::sync::{Arc, Mutex, RwLock};
+use crate::utils::async_poll::dummy_waker;
 #[cfg(feature = "time")]
 use crate::utils::time::DefaultTimeProvider;
 use crate::utils::time::TimeProvider;
@@ -53,7 +55,9 @@ use lightning_types::features::{InitFeatures, NodeFeatures};
 
 use bitcoin::secp256k1::PublicKey;
 
+use core::future::Future as StdFuture;
 use core::ops::Deref;
+use core::task;
 
 const LSPS_FEATURE_BIT: usize = 729;
 
@@ -297,15 +301,14 @@ pub struct LiquidityManager<
 	#[cfg(lsps1_service)]
 	lsps1_service_handler: Option<LSPS1ServiceHandler<ES, CM, C>>,
 	lsps1_client_handler: Option<LSPS1ClientHandler<ES>>,
-	lsps2_service_handler: Option<LSPS2ServiceHandler<CM>>,
+	lsps2_service_handler: Option<LSPS2ServiceHandler<CM, K>>,
 	lsps2_client_handler: Option<LSPS2ClientHandler<ES>>,
 	lsps5_service_handler: Option<LSPS5ServiceHandler<CM, NS, TP>>,
 	lsps5_client_handler: Option<LSPS5ClientHandler<ES>>,
 	service_config: Option<LiquidityServiceConfig>,
 	_client_config: Option<LiquidityClientConfig>,
 	best_block: RwLock<Option<BestBlock>>,
 	_chain_source: Option<C>,
-	kv_store: K,
 }
 
 #[cfg(feature = "time")]
@@ -392,14 +395,15 @@ where
 		let lsps2_service_handler = service_config.as_ref().and_then(|config| {
 			config.lsps2_service_config.as_ref().map(|config| {
 				if let Some(number) =
-					<LSPS2ServiceHandler<CM> as LSPSProtocolMessageHandler>::PROTOCOL_NUMBER
+					<LSPS2ServiceHandler<CM, K> as LSPSProtocolMessageHandler>::PROTOCOL_NUMBER
 				{
 					supported_protocols.push(number);
 				}
 				LSPS2ServiceHandler::new(
 					Arc::clone(&pending_messages),
 					Arc::clone(&pending_events),
 					channel_manager.clone(),
+					kv_store.clone(),
 					config.clone(),
 				)
 			})
@@ -495,7 +499,6 @@ where
 			_client_config: client_config,
 			best_block: RwLock::new(chain_params.map(|chain_params| chain_params.best_block)),
 			_chain_source: chain_source,
-			kv_store,
 		}
 	}
 
@@ -534,8 +537,8 @@ where
 
 	/// Returns a reference to the LSPS2 server-side handler.
 	///
-	/// The returned handler allows to initiate the LSPS2 service-side flow.
-	pub fn lsps2_service_handler(&self) -> Option<&LSPS2ServiceHandler<CM>> {
+	/// The returned hendler allows to initiate the LSPS2 service-side flow.
+	pub fn lsps2_service_handler(&self) -> Option<&LSPS2ServiceHandler<CM, K>> {
 		self.lsps2_service_handler.as_ref()
 	}
 
@@ -610,6 +613,19 @@ where
 		self.pending_events.get_and_clear_pending_events()
 	}
 
+	/// Persists the state of the service handlers towards the given [`KVStore`] implementation.
+	///
+	/// 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> {
+		// TODO: We should eventually persist in parallel.
+		if let Some(lsps2_service_handler) = self.lsps2_service_handler.as_ref() {
+			lsps2_service_handler.persist().await?;
+		}
+
+		Ok(())
+	}
+
 	fn handle_lsps_message(
 		&self, msg: LSPSMessage, sender_node_id: &PublicKey,
 	) -> Result<(), lightning::ln::msgs::LightningError> {
@@ -1110,7 +1126,9 @@ where
 	/// Returns a reference to the LSPS2 server-side handler.
 	///
 	/// Wraps [`LiquidityManager::lsps2_service_handler`].
-	pub fn lsps2_service_handler(&self) -> Option<&LSPS2ServiceHandler<CM>> {
+	pub fn lsps2_service_handler(
+		&self,
+	) -> Option<&LSPS2ServiceHandler<CM, Arc<KVStoreSyncWrapper<KS>>>> {
 		self.inner.lsps2_service_handler()
 	}
 
@@ -1163,6 +1181,21 @@ where
 	pub fn get_and_clear_pending_events(&self) -> Vec<LiquidityEvent> {
 		self.inner.get_and_clear_pending_events()
 	}
+
+	/// Persists the state of the service handlers towards the given [`KVStoreSync`] implementation.
+	///
+	/// Wraps [`LiquidityManager::persist`].
+	pub fn persist(&self) -> Result<(), lightning::io::Error> {
+		let mut waker = dummy_waker();
+		let mut ctx = task::Context::from_waker(&mut waker);
+		match Box::pin(self.inner.persist()).as_mut().poll(&mut ctx) {
+			task::Poll::Ready(result) => result,
+			task::Poll::Pending => {
+				// In a sync context, we can't wait for the future to complete.
+				unreachable!("LiquidityManager::persist should not be pending in a sync context");
+			},
+		}
+	}
 }
 
 impl<

lightning-liquidity/src/persist.rs

@@ -0,0 +1,20 @@
+// This file is Copyright its original authors, visible in version control
+// history.
+//
+// This file is licensed under the Apache License, Version 2.0 <LICENSE-APACHE
+// or http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your option.
+// You may not use this file except in accordance with one or both of these
+// licenses.
+
+//! Types and utils for persistence.
+
+/// The primary namespace under which the [`LiquidityManager`] will be persisted.
+///
+/// [`LiquidityManager`]: crate::LiquidityManager
+pub const LIQUIDITY_MANAGER_PERSISTENCE_PRIMARY_NAMESPACE: &str = "liquidity_manager";
+
+/// The secondary namespace under which the [`LSPS2ServiceHandler`] data will be persisted.
+///
+/// [`LSPS2ServiceHandler`]: crate::lsps2::service::LSPS2ServiceHandler
+pub const LSPS2_SERVICE_PERSISTENCE_SECONDARY_NAMESPACE: &str = "lsps2_service";