Parallelize Orange store record reads

Bound list-backed reads to the VSS connection capacity so populated
wallets do not add one remote round trip per transaction metadata or
Cashu quote record. Preserve record order and existing storage formats.

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;
@@ -646,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()))?;