guest-rust: Implement a Task handle for spawned futures

crates/guest-rust/src/lib.rs

@@ -908,8 +908,6 @@ pub mod resource;
 
 #[cfg(feature = "inter-task-wakeup")]
 pub use rt::async_support::UnitStreamOps;
-#[cfg(feature = "async-spawn")]
-pub use rt::async_support::spawn_local;
 #[cfg(feature = "async")]
 pub use rt::async_support::{
     AbiBuffer, FutureOps, FutureRead, FutureReader, FutureWrite, FutureWriteCancel,
@@ -918,3 +916,5 @@ pub use rt::async_support::{
     StreamRead, StreamReader, StreamResult, StreamWrite, StreamWriter, backpressure_dec,
     backpressure_inc, block_on, yield_async, yield_blocking,
 };
+#[cfg(feature = "async-spawn")]
+pub use rt::async_support::{JoinHandle, spawn_local};

crates/guest-rust/src/rt/async_support.rs

@@ -97,7 +97,7 @@ type BoxFuture<'a> = Pin<Box<dyn Future<Output = ()> + 'a>>;
 #[cfg(feature = "async-spawn")]
 mod spawn;
 #[cfg(feature = "async-spawn")]
-pub use spawn::spawn_local;
+pub use spawn::{JoinHandle, spawn_local};
 #[cfg(not(feature = "async-spawn"))]
 mod spawn_disabled;
 #[cfg(not(feature = "async-spawn"))]

crates/guest-rust/src/rt/async_support/spawn.rs

@@ -7,7 +7,10 @@ use crate::rt::async_support::BoxFuture;
 use alloc::boxed::Box;
 use alloc::vec::Vec;
 use core::future::Future;
+use core::pin::Pin;
 use core::task::{Context, Poll};
+use futures::channel::oneshot;
+use futures::future::{AbortHandle, Abortable, Aborted};
 use futures::stream::{FuturesUnordered, StreamExt};
 
 /// Any newly-deferred work queued by calls to the `spawn` function while
@@ -94,10 +97,10 @@ impl<'a> Tasks<'a> {
 ///   computations executing within a [`block_on`] call, however, the spawned
 ///   tasks will be executed within that scope. This notably means that for
 ///   [`block_on`] spawned tasks will prevent the [`block_on`] function from
-///   returning, even if a value is available to return.
-///
-/// * There is no handle returned to the spawned task meaning that it cannot be
-///   cancelled or monitored.
+///   returning, even if a value is available to return. If `spawn_local` is
+///   called within a component-model async task which is then terminated (e.g.
+///   by the host) before the future resolves, awating the `JoinHandle` will return
+///   `None`.
 ///
 /// * The task spawned here is executed *concurrently*, not in *parallel*. This
 ///   means that while one future is being polled no other future can be polled
@@ -108,8 +111,52 @@ impl<'a> Tasks<'a> {
 /// exported async function has produced a value this can be used to continue to
 /// execute some more code before the component model async task exits.
 ///
+/// # Cancellation
+///
+/// Dropping the resulting [`JoinHandle`] will detach the task and allow it to
+/// continue running in the background. It can also be explicitly cancelled with
+/// [`JoinHandle::cancel`].
+///
 /// [`block_on`]: crate::block_on
 /// [#1305]: https://github.com/bytecodealliance/wit-bindgen/issues/1305
-pub fn spawn_local(future: impl Future<Output = ()> + 'static) {
-    unsafe { SPAWNED.push(Box::pin(future)) }
+pub fn spawn_local<T: 'static>(future: impl Future<Output = T> + 'static) -> JoinHandle<T> {
+    let (sender, receiver) = oneshot::channel();
+    let (abort, registration) = AbortHandle::new_pair();
+    unsafe {
+        SPAWNED.push(Box::pin(async move {
+            let _ = sender.send(Abortable::new(future, registration).await);
+        }));
+    }
+    JoinHandle { receiver, abort }
+}
+
+/// A handle to a spawned task which can be awaited for its result.
+///
+/// Dropping the handle allows the task to continue running in the background.
+/// Awaiting the handle returns `None` if the task was cancelled or otherwise
+/// terminated without producing a result.
+pub struct JoinHandle<T> {
+    receiver: oneshot::Receiver<Result<T, Aborted>>,
+    abort: AbortHandle,
+}
+
+impl<T> JoinHandle<T> {
+    /// Cancels the spawned task. It is possible the task resolved and produced
+    /// a result before cancelation completed and the [`JoinHandle`] can still
+    /// be awaited to check for that case.
+    pub fn cancel(&mut self) {
+        self.abort.abort();
+    }
+}
+
+impl<T> Future for JoinHandle<T> {
+    type Output = Option<T>;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T>> {
+        match Pin::new(&mut self.receiver).poll(cx) {
+            Poll::Ready(Ok(Ok(result))) => Poll::Ready(Some(result)),
+            Poll::Ready(Ok(Err(_)) | Err(_)) => Poll::Ready(None),
+            Poll::Pending => Poll::Pending,
+        }
+    }
 }

tests/runtime/rust-spawn-and-await/runner.rs

@@ -0,0 +1,62 @@
+//@ wasmtime-flags = '-Wcomponent-model-async'
+
+include!(env!("BINDINGS"));
+
+use crate::test::rust_spawn_and_await::i::{
+    await_resolve, await_task, cancel_task, resolve, start,
+};
+use futures::task::noop_waker_ref;
+use std::future::Future;
+use std::pin::Pin;
+use std::task::Context;
+
+struct Component;
+
+export!(Component);
+
+impl Guest for Component {
+    async fn run() {
+        // Awaiting a `JoinHandle` works.
+        let _cm_task = start_task();
+        resolve().await;
+        let result = await_task().await;
+        assert_eq!(result, Some(42));
+
+        // Cancelling a `JoinHandle` before it completes returns `None`.
+        let _cm_task = start_task();
+        let result = cancel_task().await;
+        resolve().await;
+        assert_eq!(result, None);
+
+        // Cancelling a `JoinHandle` after it completes returns the result anyway.
+        let _cm_task = start_task();
+        resolve().await;
+        await_resolve().await;
+        let result = cancel_task().await;
+        assert_eq!(result, Some(42));
+
+        // Check that awaiting a `JoinHandle` returns None after the CM-async task has
+        // been terminated.
+        let cm_task = start_task();
+        drop(cm_task);
+        assert_eq!(await_task().await, None);
+        resolve().await;
+
+        // Check that cancelling a `JoinHandle` returns None after the CM-async task
+        // has been terminated.
+        let cm_task = start_task();
+        drop(cm_task);
+        assert_eq!(cancel_task().await, None);
+        resolve().await;
+    }
+}
+
+fn start_task() -> Pin<Box<dyn Future<Output = ()>>> {
+    let mut task = Box::pin(start());
+    assert!(
+        task.as_mut()
+            .poll(&mut Context::from_waker(noop_waker_ref()))
+            .is_pending()
+    );
+    task
+}

tests/runtime/rust-spawn-and-await/test.rs

@@ -0,0 +1,58 @@
+include!(env!("BINDINGS"));
+
+use futures::channel::oneshot;
+use std::cell::RefCell;
+use wit_bindgen::{JoinHandle, spawn_local};
+
+struct Component;
+
+export!(Component);
+
+std::thread_local! {
+    static HANDLE: RefCell<Option<JoinHandle<u32>>> = const { RefCell::new(None) };
+    // Send through this channel to resolve the `JoinHandle`.
+    static RESOLVE_CHANNEL: RefCell<Option<oneshot::Sender<()>>> = const { RefCell::new(None) };
+    // Side channel to check that the `JoinHandle` has resolved without explicitly awaiting it.
+    static ACK_CHANNEL: RefCell<Option<oneshot::Receiver<()>>> = const { RefCell::new(None) };
+}
+
+impl crate::exports::test::rust_spawn_and_await::i::Guest for Component {
+    async fn start() {
+        let (tx, rx) = oneshot::channel();
+        let (ack_tx, ack_rx) = oneshot::channel();
+        let task = spawn_local(async {
+            rx.await.unwrap();
+            let _ = ack_tx.send(());
+            42
+        });
+        HANDLE.with(|slot| assert!(slot.replace(Some(task)).is_none()));
+        RESOLVE_CHANNEL.with(|slot| assert!(slot.replace(Some(tx)).is_none()));
+        ACK_CHANNEL.with(|slot| slot.replace(Some(ack_rx)));
+        std::future::pending::<()>().await;
+    }
+
+    async fn await_task() -> Option<u32> {
+        let task = HANDLE.with(|slot| slot.borrow_mut().take().unwrap());
+        task.await
+    }
+
+    async fn cancel_task() -> Option<u32> {
+        let mut task = HANDLE.with(|slot| slot.borrow_mut().take().unwrap());
+        task.cancel();
+        task.await
+    }
+
+    async fn resolve() {
+        let channel = RESOLVE_CHANNEL.with(|slot| slot.borrow_mut().take().unwrap());
+        // Ignore error when trying to resolve the `JoinHandle` because some tests
+        // cancel it before it completes.
+        let _ = channel.send(());
+    }
+
+    async fn await_resolve() {
+        let channel = ACK_CHANNEL.with(|slot| slot.borrow_mut().take().unwrap());
+        // Ignore error when trying to resolve the `JoinHandle` because some tests
+        // cancel it before it completes.
+        channel.await.unwrap();
+    }
+}

tests/runtime/rust-spawn-and-await/test.wit

@@ -0,0 +1,20 @@
+//@ async = true
+package test:rust-spawn-and-await;
+
+interface i {
+  start: async func();
+  await-task: async func() -> option<u32>;
+  cancel-task: async func() -> option<u32>;
+  resolve: async func();
+  await-resolve: async func();
+}
+
+world test {
+  export i;
+}
+
+world runner {
+  import i;
+
+  export run: async func();
+}