Merge pull request #94 from lightningdevkit/vss-orange-optimizations

Parallelize Cashu cache loading

Author: benthecarman

orange-sdk/src/dyn_store.rs

@@ -14,6 +14,11 @@ use std::sync::Arc;
 
 use ldk_node::lightning::io;
 use ldk_node::lightning::util::persist::KVStore;
+use tokio::task::JoinSet;
+
+/// Matches the connection capacity used by the VSS HTTP client. Keeping the
+/// limit here also prevents large wallets from spawning one task per record.
+const MAX_CONCURRENT_READS: usize = 10;
 
 /// Object-safe view of a `KVStore` backend. Async methods return boxed futures so the trait
 /// can be used through `dyn`.
@@ -119,3 +124,172 @@ impl KVStore for LdkNodeStore {
 		self.0.list_async(p, s)
 	}
 }
+
+/// Reads a set of keys concurrently while preserving the input order.
+///
+/// Storage formats expose record collections as a list followed by individual
+/// reads. For a remote store, performing those reads serially adds one network
+/// round trip per record. This helper bounds that fan-out to the VSS connection
+/// pool size and retains the same fail-fast I/O behavior as a serial loop.
+pub(crate) async fn read_keys_bounded(
+	store: Arc<dyn DynStore>, primary_namespace: &str, secondary_namespace: &str, keys: Vec<String>,
+) -> Result<Vec<(String, Vec<u8>)>, io::Error> {
+	if keys.is_empty() {
+		return Ok(Vec::new());
+	}
+
+	type ReadResult = (usize, String, Result<Vec<u8>, io::Error>);
+
+	let primary_namespace = primary_namespace.to_owned();
+	let secondary_namespace = secondary_namespace.to_owned();
+	let mut pending = keys.into_iter().enumerate();
+	let mut reads: JoinSet<ReadResult> = JoinSet::new();
+	let mut results = Vec::with_capacity(pending.len());
+	results.resize_with(pending.len(), || None);
+
+	let spawn_read = |reads: &mut JoinSet<ReadResult>, index, key: String| {
+		let store = Arc::clone(&store);
+		let primary_namespace = primary_namespace.clone();
+		let secondary_namespace = secondary_namespace.clone();
+		reads.spawn(async move {
+			let data = store.read_async(&primary_namespace, &secondary_namespace, &key).await;
+			(index, key, data)
+		});
+	};
+
+	for _ in 0..MAX_CONCURRENT_READS {
+		if let Some((index, key)) = pending.next() {
+			spawn_read(&mut reads, index, key);
+		} else {
+			break;
+		}
+	}
+
+	while let Some(result) = reads.join_next().await {
+		let (index, key, data) = result.map_err(|e| {
+			io::Error::new(io::ErrorKind::Other, format!("store read task failed: {e}"))
+		})?;
+		results[index] = Some((key, data?));
+
+		if let Some((index, key)) = pending.next() {
+			spawn_read(&mut reads, index, key);
+		}
+	}
+
+	Ok(results
+		.into_iter()
+		.map(|result| result.expect("every bounded store read must complete"))
+		.collect())
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use std::future::ready;
+	use std::sync::atomic::{AtomicUsize, Ordering};
+	use tokio::sync::{Semaphore, mpsc};
+
+	struct ControlledStore {
+		entered: mpsc::UnboundedSender<String>,
+		release: Arc<Semaphore>,
+		active: Arc<AtomicUsize>,
+		max_active: Arc<AtomicUsize>,
+		fail_key: Option<String>,
+	}
+
+	impl KVStore for ControlledStore {
+		fn read(
+			&self, _primary_namespace: &str, _secondary_namespace: &str, key: &str,
+		) -> impl Future<Output = Result<Vec<u8>, io::Error>> + Send + 'static {
+			let key = key.to_owned();
+			let entered = self.entered.clone();
+			let release = Arc::clone(&self.release);
+			let active = Arc::clone(&self.active);
+			let max_active = Arc::clone(&self.max_active);
+			let should_fail = self.fail_key.as_ref() == Some(&key);
+			async move {
+				let active_count = active.fetch_add(1, Ordering::SeqCst) + 1;
+				max_active.fetch_max(active_count, Ordering::SeqCst);
+				let _ = entered.send(key.clone());
+				let _permit =
+					release.acquire_owned().await.expect("test semaphore must remain open");
+				active.fetch_sub(1, Ordering::SeqCst);
+
+				if should_fail {
+					Err(io::Error::new(io::ErrorKind::InvalidData, "controlled read failure"))
+				} else {
+					Ok(key.into_bytes())
+				}
+			}
+		}
+
+		fn write(
+			&self, _primary_namespace: &str, _secondary_namespace: &str, _key: &str, _buf: Vec<u8>,
+		) -> impl Future<Output = Result<(), io::Error>> + Send + 'static {
+			ready(Ok(()))
+		}
+
+		fn remove(
+			&self, _primary_namespace: &str, _secondary_namespace: &str, _key: &str, _lazy: bool,
+		) -> impl Future<Output = Result<(), io::Error>> + Send + 'static {
+			ready(Ok(()))
+		}
+
+		fn list(
+			&self, _primary_namespace: &str, _secondary_namespace: &str,
+		) -> impl Future<Output = Result<Vec<String>, io::Error>> + Send + 'static {
+			ready(Ok(Vec::new()))
+		}
+	}
+
+	#[tokio::test]
+	async fn bounded_reads_limit_concurrency_and_preserve_order() {
+		let (entered, mut entered_rx) = mpsc::unbounded_channel();
+		let release = Arc::new(Semaphore::new(0));
+		let max_active = Arc::new(AtomicUsize::new(0));
+		let store: Arc<dyn DynStore> = Arc::new(ControlledStore {
+			entered,
+			release: Arc::clone(&release),
+			active: Arc::new(AtomicUsize::new(0)),
+			max_active: Arc::clone(&max_active),
+			fail_key: None,
+		});
+		let keys: Vec<_> = (0..12).rev().map(|index| format!("key-{index:02}")).collect();
+
+		let reads = tokio::spawn(read_keys_bounded(store, "primary", "secondary", keys.clone()));
+		for _ in 0..MAX_CONCURRENT_READS {
+			entered_rx.recv().await.expect("bounded read must start");
+		}
+		assert_eq!(max_active.load(Ordering::SeqCst), MAX_CONCURRENT_READS);
+		assert!(entered_rx.try_recv().is_err());
+
+		release.add_permits(keys.len());
+		let records = reads.await.unwrap().unwrap();
+		let result_keys: Vec<_> = records.iter().map(|(key, _)| key.clone()).collect();
+		assert_eq!(result_keys, keys);
+		assert!(records.iter().all(|(key, data)| key.as_bytes() == data));
+		assert_eq!(max_active.load(Ordering::SeqCst), MAX_CONCURRENT_READS);
+	}
+
+	#[tokio::test]
+	async fn bounded_reads_propagate_store_errors() {
+		let (entered, _entered_rx) = mpsc::unbounded_channel();
+		let store: Arc<dyn DynStore> = Arc::new(ControlledStore {
+			entered,
+			release: Arc::new(Semaphore::new(2)),
+			active: Arc::new(AtomicUsize::new(0)),
+			max_active: Arc::new(AtomicUsize::new(0)),
+			fail_key: Some("bad".to_owned()),
+		});
+
+		let error = read_keys_bounded(
+			store,
+			"primary",
+			"secondary",
+			vec!["good".to_owned(), "bad".to_owned()],
+		)
+		.await
+		.unwrap_err();
+		assert_eq!(error.kind(), io::ErrorKind::InvalidData);
+	}
+}

orange-sdk/src/store.rs

@@ -13,7 +13,7 @@
 
 use bitcoin_payment_instructions::amount::Amount;
 
-use crate::dyn_store::DynStore;
+use crate::dyn_store::{DynStore, read_keys_bounded};
 use ldk_node::bitcoin::Txid;
 use ldk_node::bitcoin::hex::{DisplayHex, FromHex};
 use ldk_node::lightning::io;
@@ -458,11 +458,11 @@ impl TxMetadataStore {
 			.await
 			.expect("We do not allow reads to fail");
 		let mut tx_metadata = HashMap::with_capacity(keys.len());
-		for key in keys {
-			let data_bytes =
-				KVStore::read(store.as_ref(), STORE_PRIMARY_KEY, STORE_SECONDARY_KEY, &key)
-					.await
-					.expect("We do not allow reads to fail");
+		let records =
+			read_keys_bounded(Arc::clone(&store), STORE_PRIMARY_KEY, STORE_SECONDARY_KEY, keys)
+				.await
+				.expect("We do not allow reads to fail");
+		for (key, data_bytes) in records {
 			let key =
 				PaymentId::from_str(&key).expect("Invalid key in transaction metadata storage");
 			let data = Readable::read(&mut &data_bytes[..])

orange-sdk/src/trusted_wallet/cashu/cashu_store.rs

@@ -3,7 +3,7 @@ use std::fmt::Debug;
 use std::str::FromStr;
 use std::sync::{Arc, RwLock};
 
-use crate::dyn_store::DynStore;
+use crate::dyn_store::{DynStore, read_keys_bounded};
 use async_trait::async_trait;
 use cdk::cdk_database::WalletDatabase;
 use cdk::wallet::types::WalletSaga;
@@ -180,15 +180,20 @@ impl CashuKvDatabase {
 	}
 
 	async fn load_caches(&self) -> Result<(), DatabaseError> {
-		// Load mints cache
-		if let Ok(mints) = self.load_mints_from_store().await {
+		// These use independent namespaces and can be restored concurrently. This
+		// lets remote stores pipeline the requests while synchronous stores retain
+		// their existing behavior.
+		let (mints, proofs) =
+			tokio::join!(self.load_mints_from_store(), self.load_proofs_from_store());
+
+		if let Ok(mints) = mints {
 			let mut cache = self.mints_cache.write().unwrap();
 			*cache = mints;
 		}
 
 		// Proof errors must not be hidden: presenting an empty cache for an unreadable
 		// proof set could make the wallet report an incorrect balance.
-		let proofs = self.load_proofs_from_store().await?;
+		let proofs = proofs?;
 		let mut cache = self.proofs_cache.write().unwrap();
 		*cache = proofs;
 
@@ -641,12 +646,12 @@ impl WalletDatabase<cdk::cdk_database::Error> for CashuKvDatabase {
 			.await
 			.map_err(DatabaseError::Io)?;
 
-		let mut quotes = Vec::with_capacity(keys.len());
-		for key in keys {
-			let data = KVStore::read(self.store.as_ref(), CASHU_PRIMARY_KEY, MINT_QUOTES_KEY, &key)
+		let records =
+			read_keys_bounded(Arc::clone(&self.store), CASHU_PRIMARY_KEY, MINT_QUOTES_KEY, keys)
 				.await
 				.map_err(DatabaseError::Io)?;
-
+		let mut quotes = Vec::with_capacity(records.len());
+		for (_, data) in records {
 			if !data.is_empty() {
 				let quote: MintQuote = serde_json::from_slice(&data)
 					.map_err(|e| DatabaseError::Serialization(e.to_string()))?;