Allow FutureSpawner to return the result of the spawned future

`tokio::spawn` can be use both to spawn a forever-running
background task or to spawn a task which gets `poll`ed
independently and eventually returns a result which the callsite
wants.

In LDK, we have only ever needed the first, and thus didn't bother
defining a return type for `FutureSpawner::spawn`. However, in the
next commit we'll start using `FutureSpawner` in a context where we
actually do want the spawned future's result. Thus, here, we add a
result output to `FutureSpawner::spawn`, mirroring the
`tokio::spawn` API.

Author: TheBlueMatt

lightning-block-sync/src/gossip.rs

@@ -49,8 +49,12 @@ pub trait UtxoSource: BlockSource + 'static {
 pub struct TokioSpawner;
 #[cfg(feature = "tokio")]
 impl FutureSpawner for TokioSpawner {
-	fn spawn<T: Future<Output = ()> + Send + 'static>(&self, future: T) {
-		tokio::spawn(future);
+	type E = tokio::task::JoinError;
+	type SpawnedFutureResult<O> = tokio::task::JoinHandle<O>;
+	fn spawn<O: Send + 'static, F: Future<Output = O> + Send + 'static>(
+		&self, future: F,
+	) -> Self::SpawnedFutureResult<O> {
+		tokio::spawn(future)
 	}
 }
 
@@ -280,7 +284,7 @@ where
 		let gossiper = Arc::clone(&self.gossiper);
 		let block_cache = Arc::clone(&self.block_cache);
 		let pmw = Arc::clone(&self.peer_manager_wake);
-		self.spawn.spawn(async move {
+		let _not_polled = self.spawn.spawn(async move {
 			let res = Self::retrieve_utxo(source, block_cache, short_channel_id).await;
 			fut.resolve(gossiper.network_graph(), &*gossiper, res);
 			(pmw)();

lightning/src/util/native_async.rs

@@ -8,23 +8,44 @@
 //! environment.
 
 #[cfg(all(test, feature = "std"))]
-use crate::sync::Mutex;
+use crate::sync::{Arc, Mutex};
 use crate::util::async_poll::{MaybeSend, MaybeSync};
 
+#[cfg(all(test, not(feature = "std")))]
+use alloc::rc::Rc;
+
 #[cfg(all(test, not(feature = "std")))]
 use core::cell::RefCell;
+#[cfg(test)]
+use core::convert::Infallible;
 use core::future::Future;
 #[cfg(test)]
 use core::pin::Pin;
+#[cfg(test)]
+use core::task::{Context, Poll};
 
-/// A generic trait which is able to spawn futures in the background.
+/// A generic trait which is able to spawn futures to be polled in the background.
+///
+/// When the spawned future completes, the returned [`Self::SpawnedFutureResult`] should resolve
+/// with the output of the spawned future.
+///
+/// Spawned futures must be polled independently in the background even if the returned
+/// [`Self::SpawnedFutureResult`] is dropped without being polled. This matches the semantics of
+/// `tokio::spawn`.
 ///
 /// This is not exported to bindings users as async is only supported in Rust.
 pub trait FutureSpawner: MaybeSend + MaybeSync + 'static {
+	/// The error type of [`Self::SpawnedFutureResult`]. This can be used to indicate that the
+	/// spawned future was cancelled or panicked.
+	type E;
+	/// The result of [`Self::spawn`], a future which completes when the spawned future completes.
+	type SpawnedFutureResult<O>: Future<Output = Result<O, Self::E>> + Unpin;
 	/// Spawns the given future as a background task.
 	///
 	/// This method MUST NOT block on the given future immediately.
-	fn spawn<T: Future<Output = ()> + MaybeSend + 'static>(&self, future: T);
+	fn spawn<O: MaybeSend + 'static, T: Future<Output = O> + MaybeSend + 'static>(
+		&self, future: T,
+	) -> Self::SpawnedFutureResult<O>;
 }
 
 #[cfg(test)]
@@ -39,6 +60,70 @@ pub(crate) struct FutureQueue(Mutex<Vec<Pin<Box<dyn MaybeSendableFuture>>>>);
 #[cfg(all(test, not(feature = "std")))]
 pub(crate) struct FutureQueue(RefCell<Vec<Pin<Box<dyn MaybeSendableFuture>>>>);
 
+/// A simple future which can be completed later. Used to implement [`FutureQueue`].
+#[cfg(all(test, feature = "std"))]
+pub struct FutureQueueCompletion<O>(Arc<Mutex<Option<O>>>);
+#[cfg(all(test, not(feature = "std")))]
+pub struct FutureQueueCompletion<O>(Rc<RefCell<Option<O>>>);
+
+#[cfg(all(test, feature = "std"))]
+impl<O> FutureQueueCompletion<O> {
+	fn new() -> Self {
+		Self(Arc::new(Mutex::new(None)))
+	}
+
+	fn complete(&self, o: O) {
+		*self.0.lock().unwrap() = Some(o);
+	}
+}
+
+#[cfg(all(test, feature = "std"))]
+impl<O> Clone for FutureQueueCompletion<O> {
+	fn clone(&self) -> Self {
+		Self(self.0.clone())
+	}
+}
+
+#[cfg(all(test, not(feature = "std")))]
+impl<O> FutureQueueCompletion<O> {
+	fn new() -> Self {
+		Self(Rc::new(RefCell::new(None)))
+	}
+
+	fn complete(&self, o: O) {
+		*self.0.borrow_mut() = Some(o);
+	}
+}
+
+#[cfg(all(test, not(feature = "std")))]
+impl<O> Clone for FutureQueueCompletion<O> {
+	fn clone(&self) -> Self {
+		Self(self.0.clone())
+	}
+}
+
+#[cfg(all(test, feature = "std"))]
+impl<O> Future for FutureQueueCompletion<O> {
+	type Output = Result<O, Infallible>;
+	fn poll(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<O, Infallible>> {
+		match Pin::into_inner(self).0.lock().unwrap().take() {
+			None => Poll::Pending,
+			Some(o) => Poll::Ready(Ok(o)),
+		}
+	}
+}
+
+#[cfg(all(test, not(feature = "std")))]
+impl<O> Future for FutureQueueCompletion<O> {
+	type Output = Result<O, Infallible>;
+	fn poll(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<O, Infallible>> {
+		match Pin::into_inner(self).0.borrow_mut().take() {
+			None => Poll::Pending,
+			Some(o) => Poll::Ready(Ok(o)),
+		}
+	}
+}
+
 #[cfg(test)]
 impl FutureQueue {
 	pub(crate) fn new() -> Self {
@@ -74,7 +159,6 @@ impl FutureQueue {
 			futures = self.0.borrow_mut();
 		}
 		futures.retain_mut(|fut| {
-			use core::task::{Context, Poll};
 			let waker = crate::util::async_poll::dummy_waker();
 			match fut.as_mut().poll(&mut Context::from_waker(&waker)) {
 				Poll::Ready(()) => false,
@@ -86,7 +170,16 @@ impl FutureQueue {
 
 #[cfg(test)]
 impl FutureSpawner for FutureQueue {
-	fn spawn<T: Future<Output = ()> + MaybeSend + 'static>(&self, future: T) {
+	type E = Infallible;
+	type SpawnedFutureResult<O> = FutureQueueCompletion<O>;
+	fn spawn<O: MaybeSend + 'static, F: Future<Output = O> + MaybeSend + 'static>(
+		&self, f: F,
+	) -> FutureQueueCompletion<O> {
+		let completion = FutureQueueCompletion::new();
+		let compl_ref = completion.clone();
+		let future = async move {
+			compl_ref.complete(f.await);
+		};
 		#[cfg(feature = "std")]
 		{
 			self.0.lock().unwrap().push(Box::pin(future));
@@ -95,14 +188,24 @@ impl FutureSpawner for FutureQueue {
 		{
 			self.0.borrow_mut().push(Box::pin(future));
 		}
+		completion
 	}
 }
 
 #[cfg(test)]
 impl<D: core::ops::Deref<Target = FutureQueue> + MaybeSend + MaybeSync + 'static> FutureSpawner
 	for D
 {
-	fn spawn<T: Future<Output = ()> + MaybeSend + 'static>(&self, future: T) {
+	type E = Infallible;
+	type SpawnedFutureResult<O> = FutureQueueCompletion<O>;
+	fn spawn<O: MaybeSend + 'static, F: Future<Output = O> + MaybeSend + 'static>(
+		&self, f: F,
+	) -> FutureQueueCompletion<O> {
+		let completion = FutureQueueCompletion::new();
+		let compl_ref = completion.clone();
+		let future = async move {
+			compl_ref.complete(f.await);
+		};
 		#[cfg(feature = "std")]
 		{
 			self.0.lock().unwrap().push(Box::pin(future));
@@ -111,5 +214,6 @@ impl<D: core::ops::Deref<Target = FutureQueue> + MaybeSend + MaybeSync + 'static
 		{
 			self.0.borrow_mut().push(Box::pin(future));
 		}
+		completion
 	}
 }

lightning/src/util/persist.rs

@@ -16,6 +16,7 @@ use alloc::sync::Arc;
 use bitcoin::hashes::hex::FromHex;
 use bitcoin::{BlockHash, Txid};
 
+use core::convert::Infallible;
 use core::future::Future;
 use core::mem;
 use core::ops::Deref;
@@ -491,7 +492,11 @@ where
 
 struct PanicingSpawner;
 impl FutureSpawner for PanicingSpawner {
-	fn spawn<T: Future<Output = ()> + MaybeSend + 'static>(&self, _: T) {
+	type E = Infallible;
+	type SpawnedFutureResult<O> = Box<dyn Future<Output = Result<O, Infallible>> + Unpin>;
+	fn spawn<O, T: Future<Output = O> + MaybeSend + 'static>(
+		&self, _: T,
+	) -> Self::SpawnedFutureResult<O> {
 		unreachable!();
 	}
 }
@@ -959,7 +964,7 @@ where
 		let future = inner.persist_new_channel(monitor_name, monitor);
 		let channel_id = monitor.channel_id();
 		let completion = (monitor.channel_id(), monitor.get_latest_update_id());
-		self.0.future_spawner.spawn(async move {
+		let _runs_free = self.0.future_spawner.spawn(async move {
 			match future.await {
 				Ok(()) => {
 					inner.async_completed_updates.lock().unwrap().push(completion);
@@ -991,7 +996,7 @@ where
 			None
 		};
 		let inner = Arc::clone(&self.0);
-		self.0.future_spawner.spawn(async move {
+		let _runs_free = self.0.future_spawner.spawn(async move {
 			match future.await {
 				Ok(()) => if let Some(completion) = completion {
 					inner.async_completed_updates.lock().unwrap().push(completion);
@@ -1009,7 +1014,7 @@ where
 
 	pub(crate) fn spawn_async_archive_persisted_channel(&self, monitor_name: MonitorName) {
 		let inner = Arc::clone(&self.0);
-		self.0.future_spawner.spawn(async move {
+		let _runs_free = self.0.future_spawner.spawn(async move {
 			inner.archive_persisted_channel(monitor_name).await;
 		});
 	}