f Introduce TestStoreFuture delaying readiness

Author: tnull

lightning/src/util/test_utils.rs

@@ -875,7 +875,11 @@ impl KVStore for TestStore {
 		let secondary_namespace = secondary_namespace.to_string();
 		let key = key.to_string();
 		let inner = Arc::clone(&self.inner);
-		Box::pin(async move { inner.read_internal(&primary_namespace, &secondary_namespace, &key) })
+		let fut =
+			Box::pin(
+				async move { inner.read_internal(&primary_namespace, &secondary_namespace, &key) },
+			);
+		Box::pin(async move { TestStoreFuture::new(Box::pin(fut)).await })
 	}
 	fn write(
 		&self, primary_namespace: &str, secondary_namespace: &str, key: &str, buf: Vec<u8>,
@@ -884,9 +888,10 @@ impl KVStore for TestStore {
 		let secondary_namespace = secondary_namespace.to_string();
 		let key = key.to_string();
 		let inner = Arc::clone(&self.inner);
-		Box::pin(async move {
+		let fut = Box::pin(async move {
 			inner.write_internal(&primary_namespace, &secondary_namespace, &key, buf)
-		})
+		});
+		Box::pin(async move { TestStoreFuture::new(Box::pin(fut)).await })
 	}
 	fn remove(
 		&self, primary_namespace: &str, secondary_namespace: &str, key: &str, lazy: bool,
@@ -895,17 +900,20 @@ impl KVStore for TestStore {
 		let secondary_namespace = secondary_namespace.to_string();
 		let key = key.to_string();
 		let inner = Arc::clone(&self.inner);
-		Box::pin(async move {
+		let fut = Box::pin(async move {
 			inner.remove_internal(&primary_namespace, &secondary_namespace, &key, lazy)
-		})
+		});
+		Box::pin(async move { TestStoreFuture::new(Box::pin(fut)).await })
 	}
 	fn list(
 		&self, primary_namespace: &str, secondary_namespace: &str,
 	) -> Pin<Box<dyn Future<Output = Result<Vec<String>, io::Error>> + 'static + Send>> {
 		let primary_namespace = primary_namespace.to_string();
 		let secondary_namespace = secondary_namespace.to_string();
 		let inner = Arc::clone(&self.inner);
-		Box::pin(async move { inner.list_internal(&primary_namespace, &secondary_namespace) })
+		let fut =
+			Box::pin(async move { inner.list_internal(&primary_namespace, &secondary_namespace) });
+		Box::pin(async move { TestStoreFuture::new(Box::pin(fut)).await })
 	}
 }
 
@@ -933,6 +941,37 @@ impl KVStoreSync for TestStore {
 	}
 }
 
+// A `Future` wrapper that returns `Pending` once before actually polling the (likely sync) future.
+pub(crate) struct TestStoreFuture<F: Future<Output = Result<R, io::Error>> + Unpin, R> {
+	future: Mutex<Option<F>>,
+	first_poll: AtomicBool,
+}
+
+impl<F: Future<Output = Result<R, io::Error>> + Unpin, R> TestStoreFuture<F, R> {
+	fn new(fut: F) -> Self {
+		let future = Mutex::new(Some(fut));
+		let first_poll = AtomicBool::new(true);
+		Self { future, first_poll }
+	}
+}
+
+impl<F: Future<Output = Result<R, io::Error>> + Unpin, R> Future for TestStoreFuture<F, 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(mut fut) = self.future.lock().unwrap().take() {
+				Pin::new(&mut fut).poll(cx)
+			} else {
+				unreachable!("We should never poll more than twice");
+			}
+		}
+	}
+}
+
 struct TestStoreInner {
 	persisted_bytes: Mutex<HashMap<String, HashMap<String, Vec<u8>>>>,
 	read_only: bool,