diff --git a/crates/guest-rust/src/lib.rs b/crates/guest-rust/src/lib.rs index 3efde07ad..e88633847 100644 --- a/crates/guest-rust/src/lib.rs +++ b/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}; diff --git a/crates/guest-rust/src/rt/async_support.rs b/crates/guest-rust/src/rt/async_support.rs index bfcc091c8..5326f419b 100644 --- a/crates/guest-rust/src/rt/async_support.rs +++ b/crates/guest-rust/src/rt/async_support.rs @@ -97,7 +97,7 @@ type BoxFuture<'a> = Pin + '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"))] diff --git a/crates/guest-rust/src/rt/async_support/spawn.rs b/crates/guest-rust/src/rt/async_support/spawn.rs index 2dff5b378..fa6542961 100644 --- a/crates/guest-rust/src/rt/async_support/spawn.rs +++ b/crates/guest-rust/src/rt/async_support/spawn.rs @@ -7,7 +7,9 @@ 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::stream::{FuturesUnordered, StreamExt}; /// Any newly-deferred work queued by calls to the `spawn` function while @@ -94,10 +96,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 +110,93 @@ 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 + 'static) { - unsafe { SPAWNED.push(Box::pin(future)) } +pub fn spawn_local(future: impl Future + 'static) -> JoinHandle { + let (sender, receiver) = oneshot::channel(); + unsafe { + SPAWNED.push(Box::pin(async move { + SpawnedFuture { + fut: future, + sender: Some(sender), + } + .await + })); + } + JoinHandle { receiver } +} + +struct SpawnedFuture { + fut: F, + sender: Option>, +} + +impl Future for SpawnedFuture +where + F: Future + 'static, +{ + type Output = (); + + fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<()> { + // SAFETY: `fut` is not moved and only used to be polled in place. All + // other fields are `Unpin`. + let inner = unsafe { self.get_unchecked_mut() }; + let sender = inner.sender.take(); + match sender { + None => Poll::Ready(()), + Some(mut sender) => { + if let Poll::Ready(()) = sender.poll_canceled(cx) { + return Poll::Ready(()); + } + // SAFETY: `fut` has not been moved. + let fut = unsafe { Pin::new_unchecked(&mut inner.fut) }; + match fut.poll(cx) { + Poll::Ready(t) => { + let _ = sender.send(t); + Poll::Ready(()) + } + Poll::Pending => { + inner.sender = Some(sender); + Poll::Pending + } + } + } + } + } +} + +/// 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 { + receiver: oneshot::Receiver, +} + +impl JoinHandle { + /// 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.receiver.close(); + } +} + +impl Future for JoinHandle { + type Output = Option; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match Pin::new(&mut self.receiver).poll(cx) { + Poll::Ready(Ok(result)) => Poll::Ready(Some(result)), + Poll::Ready(Err(_)) => Poll::Ready(None), + Poll::Pending => Poll::Pending, + } + } } diff --git a/tests/runtime/rust-spawn-and-await/runner.rs b/tests/runtime/rust-spawn-and-await/runner.rs new file mode 100644 index 000000000..d5333ee7e --- /dev/null +++ b/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>> { + let mut task = Box::pin(start()); + assert!( + task.as_mut() + .poll(&mut Context::from_waker(noop_waker_ref())) + .is_pending() + ); + task +} diff --git a/tests/runtime/rust-spawn-and-await/test.rs b/tests/runtime/rust-spawn-and-await/test.rs new file mode 100644 index 000000000..1c234cc82 --- /dev/null +++ b/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>> = const { RefCell::new(None) }; + // Send through this channel to resolve the `JoinHandle`. + static RESOLVE_CHANNEL: RefCell>> = const { RefCell::new(None) }; + // Side channel to check that the `JoinHandle` has resolved without explicitly awaiting it. + static ACK_CHANNEL: RefCell>> = 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 { + let task = HANDLE.with(|slot| slot.borrow_mut().take().unwrap()); + task.await + } + + async fn cancel_task() -> Option { + 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(); + } +} diff --git a/tests/runtime/rust-spawn-and-await/test.wit b/tests/runtime/rust-spawn-and-await/test.wit new file mode 100644 index 000000000..dce5f9b00 --- /dev/null +++ b/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; + cancel-task: async func() -> option; + resolve: async func(); + await-resolve: async func(); +} + +world test { + export i; +} + +world runner { + import i; + + export run: async func(); +}