diff --git a/crates/guest-rust/src/rt/async_support.rs b/crates/guest-rust/src/rt/async_support.rs index bee1128f8..72ad6f2a8 100644 --- a/crates/guest-rust/src/rt/async_support.rs +++ b/crates/guest-rust/src/rt/async_support.rs @@ -158,8 +158,8 @@ struct CabiWaitable { unsafe impl Send for CabiWaitable {} -impl TaskState<'_> { - fn new(future: BoxFuture<'_>) -> TaskState<'_> { +impl<'a> TaskState<'a> { + fn new(future: impl Future + 'a) -> TaskState<'a> { let shared = Arc::new(SharedTaskState::default()); TaskState { waker: shared.clone().into(), @@ -183,14 +183,14 @@ impl TaskState<'_> { EVENT_STREAM_WRITE => rtdebug!("EVENT_STREAM_WRITE({event1:#x}, {event2:#x})"), EVENT_FUTURE_READ => rtdebug!("EVENT_FUTURE_READ({event1:#x}, {event2:#x})"), EVENT_FUTURE_WRITE => rtdebug!("EVENT_FUTURE_WRITE({event1:#x}, {event2:#x})"), + + // Delegate cancellation to the set of spawned tasks that this + // component-model task is maintaining. The cancel here is a signal + // for the "root computation", but everything else spawned, if any, + // gets to continue along its merry way. EVENT_CANCEL => { rtdebug!("EVENT_CANCEL"); - - // Cancellation is mapped to destruction in Rust, so return a - // code/bool indicating we're done. The caller will then - // appropriately deallocate this `TaskState` which will - // transitively run all destructors. - return CallbackCode::Exit; + self.with_p3_task_set(|me| me.tasks.cancel_root()); } _ => unreachable!(), } @@ -209,7 +209,7 @@ impl TaskState<'_> { // `waker.sleep_state` above to ensure that if a waker is woken it // won't actually signal our inter-task stream since we're already // in the process of handling the future. - if event0 != EVENT_NONE { + if event0 != EVENT_NONE && event0 != EVENT_CANCEL { me.deliver_waitable_event(event1, event2) } @@ -617,7 +617,7 @@ mod task_state { pub fn start_task(task: impl Future + 'static) -> i32 { // Allocate a new `TaskState` which will track all state necessary for // our exported task. - let state = Box::into_raw(Box::new(TaskState::new(Box::pin(task)))); + let state = Box::into_raw(Box::new(TaskState::new(task))); // Store our `TaskState` into our context-local-storage slot and then // pretend we got EVENT_NONE to kick off everything. @@ -671,9 +671,9 @@ pub unsafe fn callback(event0: u32, event1: u32, event2: u32) -> u32 { /// to async-lowered imports as necessary. pub fn block_on(future: impl Future) -> T { let mut result = None; - let mut state = TaskState::new(Box::pin(async { + let mut state = TaskState::new(async { result = Some(future.await); - })); + }); let mut event = (EVENT_NONE, 0, 0); loop { match state.callback(event.0, event.1, event.2) { diff --git a/crates/guest-rust/src/rt/async_support/spawn.rs b/crates/guest-rust/src/rt/async_support/spawn.rs index 031281f0e..d378fef9c 100644 --- a/crates/guest-rust/src/rt/async_support/spawn.rs +++ b/crates/guest-rust/src/rt/async_support/spawn.rs @@ -22,12 +22,24 @@ static mut SPAWNED: Vec = Vec::new(); #[derive(Default)] pub struct Tasks<'a> { tasks: FuturesUnordered>, + root: Option, } impl<'a> Tasks<'a> { - pub fn new(root: BoxFuture<'a>) -> Tasks<'a> { + pub fn new(root: impl Future + 'a) -> Tasks<'a> { + let (abort, registration) = AbortHandle::new_pair(); + let root: BoxFuture<'a> = Box::pin(async move { + let _ = Abortable::new(root, registration).await; + }); Tasks { tasks: [root].into_iter().collect(), + root: Some(abort), + } + } + + pub fn cancel_root(&mut self) { + if let Some(root) = self.root.take() { + root.abort(); } } diff --git a/crates/guest-rust/src/rt/async_support/spawn_disabled.rs b/crates/guest-rust/src/rt/async_support/spawn_disabled.rs index 414b746af..ca4b5599d 100644 --- a/crates/guest-rust/src/rt/async_support/spawn_disabled.rs +++ b/crates/guest-rust/src/rt/async_support/spawn_disabled.rs @@ -1,4 +1,6 @@ use crate::rt::async_support::BoxFuture; +use alloc::boxed::Box; +use core::future::Future; use core::task::{Context, Poll}; #[derive(Default)] @@ -7,8 +9,14 @@ pub struct Tasks<'a> { } impl<'a> Tasks<'a> { - pub fn new(root: BoxFuture<'a>) -> Tasks<'a> { - Tasks { future: Some(root) } + pub fn new(root: impl Future + 'a) -> Tasks<'a> { + Tasks { + future: Some(Box::pin(root)), + } + } + + pub fn cancel_root(&mut self) { + self.future = None; } pub fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll<()> { diff --git a/tests/runtime/rust-cancel-wake/runner.rs b/tests/runtime/rust-cancel-wake/runner.rs index 5fdc6fb1c..66651d236 100644 --- a/tests/runtime/rust-cancel-wake/runner.rs +++ b/tests/runtime/rust-cancel-wake/runner.rs @@ -36,6 +36,6 @@ impl Guest for Component { yield_async().await; let status = get_status(); assert!(status.dropped); - assert_eq!(status.polls, 2); + assert_eq!(status.polls, 3); } } diff --git a/tests/runtime/rust-cancel-wake/test.rs b/tests/runtime/rust-cancel-wake/test.rs index 4f929d30b..00dd4cd26 100644 --- a/tests/runtime/rust-cancel-wake/test.rs +++ b/tests/runtime/rust-cancel-wake/test.rs @@ -37,7 +37,6 @@ impl Guest for Component { }) .await .unwrap(); - unreachable!("the cancelled waiter must not resume"); }); // FuturesUnordered destroys the most recently polled future first. diff --git a/tests/runtime/rust-spawn-and-await/runner.rs b/tests/runtime/rust-spawn-and-await/runner.rs index d5333ee7e..ed9461fad 100644 --- a/tests/runtime/rust-spawn-and-await/runner.rs +++ b/tests/runtime/rust-spawn-and-await/runner.rs @@ -35,12 +35,13 @@ impl Guest for Component { 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. + // Check that cancelling the CM-async task which spawned a task only + // cancels the export's own future, and the spawned task keeps running + // and can still be awaited. let cm_task = start_task(); drop(cm_task); - assert_eq!(await_task().await, None); resolve().await; + assert_eq!(await_task().await, Some(42)); // Check that cancelling a `JoinHandle` returns None after the CM-async task // has been terminated.