Implement KVStore for TestStore

We implement the async `KVStore` trait for `TestStore`. This is mostly
useful in LDK Node tests, where we use `TestStore` and we now require
all stores to implement both variants.

Author: tnull

lightning/src/util/persist.rs

@@ -1257,14 +1257,13 @@ mod tests {
 
 					let monitor_name = mon.persistence_key();
 					assert_eq!(
-						persister_0
-							.kv_store
-							.list(
-								CHANNEL_MONITOR_UPDATE_PERSISTENCE_PRIMARY_NAMESPACE,
-								&monitor_name.to_string()
-							)
-							.unwrap()
-							.len() as u64,
+						KVStoreSync::list(
+							&*persister_0.kv_store,
+							CHANNEL_MONITOR_UPDATE_PERSISTENCE_PRIMARY_NAMESPACE,
+							&monitor_name.to_string()
+						)
+						.unwrap()
+						.len() as u64,
 						mon.get_latest_update_id() % persister_0_max_pending_updates,
 						"Wrong number of updates stored in persister 0",
 					);
@@ -1276,14 +1275,13 @@ mod tests {
 					assert_eq!(mon.get_latest_update_id(), $expected_update_id);
 					let monitor_name = mon.persistence_key();
 					assert_eq!(
-						persister_1
-							.kv_store
-							.list(
-								CHANNEL_MONITOR_UPDATE_PERSISTENCE_PRIMARY_NAMESPACE,
-								&monitor_name.to_string()
-							)
-							.unwrap()
-							.len() as u64,
+						KVStoreSync::list(
+							&*persister_1.kv_store,
+							CHANNEL_MONITOR_UPDATE_PERSISTENCE_PRIMARY_NAMESPACE,
+							&monitor_name.to_string()
+						)
+						.unwrap()
+						.len() as u64,
 						mon.get_latest_update_id() % persister_1_max_pending_updates,
 						"Wrong number of updates stored in persister 1",
 					);
@@ -1481,28 +1479,26 @@ mod tests {
 		let persisted_chan_data = persister_0.read_all_channel_monitors_with_updates().unwrap();
 		let (_, monitor) = &persisted_chan_data[0];
 		let monitor_name = monitor.persistence_key();
-		persister_0
-			.kv_store
-			.write(
-				CHANNEL_MONITOR_UPDATE_PERSISTENCE_PRIMARY_NAMESPACE,
-				&monitor_name.to_string(),
-				UpdateName::from(1).as_str(),
-				vec![0u8; 1],
-			)
-			.unwrap();
+		KVStoreSync::write(
+			&*persister_0.kv_store,
+			CHANNEL_MONITOR_UPDATE_PERSISTENCE_PRIMARY_NAMESPACE,
+			&monitor_name.to_string(),
+			UpdateName::from(1).as_str(),
+			vec![0u8; 1],
+		)
+		.unwrap();
 
 		// Do the stale update cleanup
 		persister_0.cleanup_stale_updates(false).unwrap();
 
 		// Confirm the stale update is unreadable/gone
-		assert!(persister_0
-			.kv_store
-			.read(
-				CHANNEL_MONITOR_UPDATE_PERSISTENCE_PRIMARY_NAMESPACE,
-				&monitor_name.to_string(),
-				UpdateName::from(1).as_str()
-			)
-			.is_err());
+		assert!(KVStoreSync::read(
+			&*persister_0.kv_store,
+			CHANNEL_MONITOR_UPDATE_PERSISTENCE_PRIMARY_NAMESPACE,
+			&monitor_name.to_string(),
+			UpdateName::from(1).as_str()
+		)
+		.is_err());
 	}
 
 	fn persist_fn<P: Deref, ChannelSigner: EcdsaChannelSigner>(_persist: P) -> bool

lightning/src/util/test_utils.rs

@@ -57,7 +57,7 @@ use crate::util::dyn_signer::{
 use crate::util::logger::{Logger, Record};
 #[cfg(feature = "std")]
 use crate::util::mut_global::MutGlobal;
-use crate::util::persist::{KVStoreSync, MonitorName};
+use crate::util::persist::{KVStore, KVStoreSync, MonitorName};
 use crate::util::ser::{Readable, ReadableArgs, Writeable, Writer};
 use crate::util::test_channel_signer::{EnforcementState, TestChannelSigner};
 
@@ -84,7 +84,10 @@ use crate::io;
 use crate::prelude::*;
 use crate::sign::{EntropySource, NodeSigner, RandomBytes, Recipient, SignerProvider};
 use crate::sync::{Arc, Mutex};
+use alloc::boxed::Box;
+use core::future::Future;
 use core::mem;
+use core::pin::Pin;
 use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
 use core::time::Duration;
 
@@ -950,6 +953,33 @@ impl TestStore {
 	}
 }
 
+impl KVStore for TestStore {
+	fn read(
+		&self, primary_namespace: &str, secondary_namespace: &str, key: &str,
+	) -> Pin<Box<dyn Future<Output = Result<Vec<u8>, io::Error>> + 'static + Send>> {
+		let res = self.read_internal(&primary_namespace, &secondary_namespace, &key);
+		Box::pin(async move { TestStoreFuture::new(res).await })
+	}
+	fn write(
+		&self, primary_namespace: &str, secondary_namespace: &str, key: &str, buf: Vec<u8>,
+	) -> Pin<Box<dyn Future<Output = Result<(), io::Error>> + 'static + Send>> {
+		let res = self.write_internal(&primary_namespace, &secondary_namespace, &key, buf);
+		Box::pin(async move { TestStoreFuture::new(res).await })
+	}
+	fn remove(
+		&self, primary_namespace: &str, secondary_namespace: &str, key: &str, lazy: bool,
+	) -> Pin<Box<dyn Future<Output = Result<(), io::Error>> + 'static + Send>> {
+		let res = self.remove_internal(&primary_namespace, &secondary_namespace, &key, lazy);
+		Box::pin(async move { TestStoreFuture::new(res).await })
+	}
+	fn list(
+		&self, primary_namespace: &str, secondary_namespace: &str,
+	) -> Pin<Box<dyn Future<Output = Result<Vec<String>, io::Error>> + 'static + Send>> {
+		let res = self.list_internal(primary_namespace, secondary_namespace);
+		Box::pin(async move { TestStoreFuture::new(res).await })
+	}
+}
+
 impl KVStoreSync for TestStore {
 	fn read(
 		&self, primary_namespace: &str, secondary_namespace: &str, key: &str,
@@ -974,6 +1004,37 @@ impl KVStoreSync for TestStore {
 	}
 }
 
+// A `Future` that returns the result only on the second poll.
+pub(crate) struct TestStoreFuture<R> {
+	res: Mutex<Option<io::Result<R>>>,
+	first_poll: AtomicBool,
+}
+
+impl<R> TestStoreFuture<R> {
+	fn new(res: io::Result<R>) -> Self {
+		let res = Mutex::new(Some(res));
+		let first_poll = AtomicBool::new(true);
+		Self { res, first_poll }
+	}
+}
+
+impl<R> Future for TestStoreFuture<R> {
+	type Output = Result<R, io::Error>;
+	fn poll(
+		self: Pin<&mut Self>, _cx: &mut core::task::Context<'_>,
+	) -> core::task::Poll<Self::Output> {
+		if self.first_poll.swap(false, Ordering::Relaxed) {
+			core::task::Poll::Pending
+		} else {
+			if let Some(res) = self.res.lock().unwrap().take() {
+				core::task::Poll::Ready(res)
+			} else {
+				unreachable!("We should never poll more than twice");
+			}
+		}
+	}
+}
+
 unsafe impl Sync for TestStore {}
 unsafe impl Send for TestStore {}