rust: suppress wake notifications while destroying async tasks

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"
 ///

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);
+    }
+}

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<oneshot::Sender<()>>);
+
+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),
+        }
+    }
+}

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();
+}