diff --git a/crates/guest-rust/src/rt/async_support.rs b/crates/guest-rust/src/rt/async_support.rs index 1de25d29e..bee1128f8 100644 --- a/crates/guest-rust/src/rt/async_support.rs +++ b/crates/guest-rust/src/rt/async_support.rs @@ -345,6 +345,13 @@ impl TaskState<'_> { impl Drop for TaskState<'_> { fn drop(&mut self) { + // No further wakeup notifications are needed once destruction starts. + // Destructors may wake other futures in this task, so suppress those + // notifications before cancelling the inter-task stream read. + self.shared + .sleep_state + .store(SLEEP_STATE_WOKEN, Ordering::Relaxed); + // If there's an active read of the inter-task stream, go ahead and // cancel it, since we're about to drop the stream anyway. self.cancel_inter_task_stream_read(); @@ -450,7 +457,8 @@ impl SharedTaskState { /// Status for "this task is actively being polled" const SLEEP_STATE_POLLING: u32 = 0; -/// Status for "this task has a wakeup scheduled, no more action need be taken". +/// Status for "no further wakeup notification is needed", either because the +/// task has already been woken or because it is being destroyed. const SLEEP_STATE_WOKEN: u32 = 1; /// Status for "this task is not being polled and has not been woken" /// diff --git a/tests/runtime/rust-cancel-wake/runner.rs b/tests/runtime/rust-cancel-wake/runner.rs new file mode 100644 index 000000000..5fdc6fb1c --- /dev/null +++ b/tests/runtime/rust-cancel-wake/runner.rs @@ -0,0 +1,41 @@ +//@ wasmtime-flags = '-Wcomponent-model-async' + +include!(env!("BINDINGS")); + +use crate::test::rust_cancel_wake::i::*; +use futures::task::noop_waker_ref; +use std::future::Future; +use std::task::Context; +use wit_bindgen::yield_async; + +struct Component; + +export!(Component); + +impl Guest for Component { + async fn run() { + let (checkpoint, checkpoint_rx) = wit_future::new(|| unreachable!()); + let mut task = Box::pin(cancellable(checkpoint_rx)); + assert!( + task.as_mut() + .poll(&mut Context::from_waker(noop_waker_ref())) + .is_pending() + ); + + // This callback clears the wake from spawning and lets the task sleep + // on its Rust oneshot receiver before we cancel it. + checkpoint.write(()).await.unwrap(); + // Two polls: the waiter registers its waker, then the root blocks. + while get_status().polls < 2 { + yield_async().await; + } + assert_eq!(get_status().polls, 2); + + // Dropping the import delivers EVENT_CANCEL to the sleeping task. + drop(task); + yield_async().await; + let status = get_status(); + assert!(status.dropped); + assert_eq!(status.polls, 2); + } +} diff --git a/tests/runtime/rust-cancel-wake/test.rs b/tests/runtime/rust-cancel-wake/test.rs new file mode 100644 index 000000000..4f929d30b --- /dev/null +++ b/tests/runtime/rust-cancel-wake/test.rs @@ -0,0 +1,60 @@ +use futures::channel::oneshot; +use std::future::{Future, poll_fn}; +use std::pin::Pin; +use std::sync::atomic::{AtomicBool, AtomicU32, Ordering::Relaxed}; +use std::task::Poll; +use wit_bindgen::{FutureReader, spawn_local}; + +include!(env!("BINDINGS")); + +use crate::exports::test::rust_cancel_wake::i::{Guest, Status}; + +struct Component; + +export!(Component); + +static POLLS: AtomicU32 = AtomicU32::new(0); +static DROPPED: AtomicBool = AtomicBool::new(false); + +struct SendOnDrop(Option>); + +impl Drop for SendOnDrop { + fn drop(&mut self) { + // Sending wakes the spawned future's registered oneshot receiver. + self.0.take().unwrap().send(()).unwrap(); + DROPPED.store(true, Relaxed); + } +} + +impl Guest for Component { + async fn cancellable(checkpoint: FutureReader<()>) { + let (sender, mut receiver) = oneshot::channel(); + let _sender = SendOnDrop(Some(sender)); + spawn_local(async move { + poll_fn(|cx| { + POLLS.fetch_add(1, Relaxed); + Pin::new(&mut receiver).poll(cx) + }) + .await + .unwrap(); + unreachable!("the cancelled waiter must not resume"); + }); + + // FuturesUnordered destroys the most recently polled future first. + // Poll the root last so its sender wakes the still-live waiter. + checkpoint.await; + poll_fn(|_| { + POLLS.fetch_add(1, Relaxed); + Poll::<()>::Pending + }) + .await; + unreachable!("the cancelled task must not resume"); + } + + fn get_status() -> Status { + Status { + polls: POLLS.load(Relaxed), + dropped: DROPPED.load(Relaxed), + } + } +} diff --git a/tests/runtime/rust-cancel-wake/test.wit b/tests/runtime/rust-cancel-wake/test.wit new file mode 100644 index 000000000..4057fd3d8 --- /dev/null +++ b/tests/runtime/rust-cancel-wake/test.wit @@ -0,0 +1,22 @@ +//@ async = true +package test:rust-cancel-wake; + +interface i { + record status { + polls: u32, + dropped: bool, + } + + cancellable: async func(checkpoint: future); + get-status: func() -> status; +} + +world test { + export i; +} + +world runner { + import i; + + export run: async func(); +}