Add some more async tests (#1287)

* Add some more async tests

I went through various parts of the guest runtime and sprinkled
`if true { panic!() }` throughout. Anything that wasn't hit got a test.
This led to a few issues in wasip3-prototyping:

* bytecodealliance/wasip3-prototyping#137
* bytecodealliance/wasip3-prototyping#138
* bytecodealliance/wasip3-prototyping#139
* bytecodealliance/wasip3-prototyping#140

so not all tests are fully enabled yet but they should get enabled as
bugs get fixed.

Additionally this starts to use `CALLBACK_CODE_YIELD`. That's disabled
for callbacks due to bytecodealliance/wasip3-prototyping#140 but it's
enabled for `block_on` where it translates to a `waitable-set.poll`
instead of a `waitable-set.wait`.

* Fix fallback build of `poll`

* Adjust poll signature

* Fix native signature

Author: alexcrichton

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

@@ -129,14 +129,7 @@ impl FutureState {
             self.deliver_waitable_event(event1, event2)
         }
 
-        loop {
-            match self.poll() {
-                // TODO: don't re-loop here once a host supports
-                // `CALLBACK_CODE_YIELD`.
-                CALLBACK_CODE_YIELD => {}
-                other => break other,
-            }
-        }
+        self.poll()
     }
 
     /// Deliver the `code` event to the `waitable` store within our map. This
@@ -358,12 +351,24 @@ pub unsafe fn callback(event0: u32, event1: u32, event2: u32) -> u32 {
     // our future so deallocate it. Otherwise put our future back in
     // context-local storage and forward the code.
     unsafe {
-        let rc = (*state).callback(event0, event1, event2);
+        let rc = match (*state).callback(event0, event1, event2) {
+            // FIXME(wasip3-prototyping#140) this seems to break tests in
+            // that repo. Handle this return code by re-running our callback
+            // until it stops yielding.
+            CALLBACK_CODE_YIELD => loop {
+                match (*state).callback(EVENT_NONE, 0, 0) {
+                    CALLBACK_CODE_YIELD => {}
+                    other => break other,
+                }
+            },
+            other => other,
+        };
         if rc == CALLBACK_CODE_EXIT {
             drop(Box::from_raw(state));
         } else {
             context_set(state.cast());
         }
+        rtdebug!(" => (cb) {rc:#x}");
         rc
     }
 }
@@ -389,6 +394,7 @@ pub fn block_on<T: 'static>(future: impl Future<Output = T> + 'static) -> T {
     loop {
         match state.callback(event.0, event.1, event.2) {
             CALLBACK_CODE_EXIT => break rx.try_recv().unwrap().unwrap(),
+            CALLBACK_CODE_YIELD => event = state.waitable_set.as_ref().unwrap().poll(),
             _ => event = state.waitable_set.as_ref().unwrap().wait(),
         }
     }

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

@@ -437,6 +437,7 @@ impl<T> fmt::Display for FutureWriteError<T> {
 impl<T> std::error::Error for FutureWriteError<T> {}
 
 /// Result of [`FutureWrite::cancel`].
+#[derive(Debug)]
 pub enum FutureWriteCancel<T: 'static> {
     /// The cancel request raced with the receipt of the sent value, and the
     /// value was actually sent. Neither the value nor the writer are made

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

@@ -149,8 +149,8 @@ unsafe impl<T: Subtask> WaitableOp for SubtaskOps<T> {
         trap_because_of_future_cancel()
     }
 
-    fn result_into_cancel(_result: Self::Result) -> Self::Cancel {
-        todo!()
+    fn result_into_cancel(result: Self::Result) -> Self::Cancel {
+        drop(result);
     }
 }
 

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

@@ -322,6 +322,15 @@ where
         if let Some(cx) = cx {
             let handle = S::in_progress_waitable(in_progress);
             self.register_waker(handle, cx);
+        } else {
+            // This should not be dynamically reachable and, if it were, it may
+            // mean that this needs to be re-thought and/or the caller should be
+            // adjusted. Conservatively panic for now to defer fleshing this out
+            // for later.
+            panic!(
+                "unexpected poll to completion in a non-future way (no context) \
+                 and this operation is still pending"
+            );
         }
         Poll::Pending
     }

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

@@ -1,26 +1,51 @@
 //! Low-level FFI-like bindings around `waitable-set` in the canonical ABI.
 
+use super::EVENT_NONE;
 use std::num::NonZeroU32;
 
 pub struct WaitableSet(NonZeroU32);
 
 impl WaitableSet {
     pub fn new() -> WaitableSet {
-        WaitableSet(NonZeroU32::new(unsafe { new() }).unwrap())
+        let ret = WaitableSet(NonZeroU32::new(unsafe { new() }).unwrap());
+        rtdebug!("waitable-set.new() = {}", ret.0.get());
+        ret
     }
 
     pub fn join(&self, waitable: u32) {
+        rtdebug!("waitable-set.join({waitable}, {})", self.0.get());
         unsafe { join(waitable, self.0.get()) }
     }
 
     pub fn remove_waitable_from_all_sets(waitable: u32) {
+        rtdebug!("waitable-set.join({waitable}, 0)");
         unsafe { join(waitable, 0) }
     }
 
     pub fn wait(&self) -> (u32, u32, u32) {
         unsafe {
             let mut payload = [0; 2];
             let event0 = wait(self.0.get(), &mut payload);
+            rtdebug!(
+                "waitable-set.wait({}) = ({event0}, {:#x}, {:#x})",
+                self.0.get(),
+                payload[0],
+                payload[1],
+            );
+            (event0, payload[0], payload[1])
+        }
+    }
+
+    pub fn poll(&self) -> (u32, u32, u32) {
+        unsafe {
+            let mut payload = [0; 2];
+            let event0 = poll(self.0.get(), &mut payload);
+            rtdebug!(
+                "waitable-set.poll({}) = ({event0}, {:#x}, {:#x})",
+                self.0.get(),
+                payload[0],
+                payload[1],
+            );
             (event0, payload[0], payload[1])
         }
     }
@@ -33,6 +58,7 @@ impl WaitableSet {
 impl Drop for WaitableSet {
     fn drop(&mut self) {
         unsafe {
+            rtdebug!("waitable-set.drop({})", self.0.get());
             drop(self.0.get());
         }
     }
@@ -54,6 +80,10 @@ unsafe fn join(_: u32, _: u32) {
 unsafe fn wait(_: u32, _: *mut [u32; 2]) -> u32 {
     unreachable!();
 }
+#[cfg(not(target_arch = "wasm32"))]
+unsafe fn poll(_: u32, _: *mut [u32; 2]) -> u32 {
+    unreachable!();
+}
 
 #[cfg(target_arch = "wasm32")]
 #[link(wasm_import_module = "$root")]
@@ -66,4 +96,6 @@ extern "C" {
     fn join(waitable: u32, set: u32);
     #[link_name = "[waitable-set-wait]"]
     fn wait(_: u32, _: *mut [u32; 2]) -> u32;
+    #[link_name = "[waitable-set-poll]"]
+    fn poll(_: u32, _: *mut [u32; 2]) -> u32;
 }

crates/guest-rust/src/lib.rs

@@ -895,6 +895,7 @@ pub mod rt {
 
 #[cfg(feature = "async")]
 pub use wit_bindgen_rt::async_support::{
-    AbiBuffer, FutureRead, FutureReader, FutureWrite, FutureWriter, StreamRead, StreamReader,
-    StreamResult, StreamWrite, StreamWriter,
+    block_on, spawn, AbiBuffer, FutureRead, FutureReader, FutureWrite, FutureWriteCancel,
+    FutureWriteError, FutureWriter, StreamRead, StreamReader, StreamResult, StreamWrite,
+    StreamWriter,
 };

tests/runtime-async/async/future-cancel-read/runner.rs

@@ -0,0 +1,17 @@
+include!(env!("BINDINGS"));
+
+use crate::my::test::i::*;
+
+fn main() {
+    wit_bindgen::block_on(async {
+        let (tx, rx) = wit_future::new();
+        cancel_before_read(rx).await;
+        drop(tx);
+
+        let (tx, rx) = wit_future::new();
+        cancel_after_read(rx).await;
+        drop(tx);
+
+        start_read_then_cancel().await;
+    });
+}

tests/runtime-async/async/future-cancel-read/test.rs

@@ -0,0 +1,43 @@
+use wit_bindgen::FutureReader;
+
+use futures::task::noop_waker_ref;
+use std::future::{Future, IntoFuture};
+use std::task::Context;
+
+include!(env!("BINDINGS"));
+
+struct Component;
+
+export!(Component);
+
+impl crate::exports::my::test::i::Guest for Component {
+    async fn cancel_before_read(x: FutureReader<u32>) {
+        let mut read = Box::pin(x.into_future());
+        let reader = read.as_mut().cancel().unwrap_err();
+        drop(reader);
+    }
+
+    async fn cancel_after_read(x: FutureReader<u32>) {
+        let mut read = Box::pin(x.into_future());
+        assert!(read
+            .as_mut()
+            .poll(&mut Context::from_waker(noop_waker_ref()))
+            .is_pending());
+        let reader = read.as_mut().cancel().unwrap_err();
+        drop(reader);
+    }
+
+    async fn start_read_then_cancel() {
+        // FIXME(wasip3-prototyping#137)
+        if false {
+            let (tx, rx) = wit_future::new::<u32>();
+            let mut read = Box::pin(rx.into_future());
+            assert!(read
+                .as_mut()
+                .poll(&mut Context::from_waker(noop_waker_ref()))
+                .is_pending());
+            drop(tx);
+            assert!(read.as_mut().cancel().unwrap().is_none());
+        }
+    }
+}

tests/runtime-async/async/future-cancel-read/test.wit

@@ -0,0 +1,15 @@
+package my:test;
+
+interface i {
+  cancel-before-read: async func(x: future<u32>);
+  cancel-after-read: async func(x: future<u32>);
+  start-read-then-cancel: async func();
+}
+
+world test {
+  export i;
+}
+
+world runner {
+  import i;
+}

tests/runtime-async/async/future-cancel-write/runner.rs

@@ -0,0 +1,78 @@
+include!(env!("BINDINGS"));
+
+use crate::my::test::i::{read_and_drop, take_then_close};
+use futures::task::noop_waker_ref;
+use std::future::Future;
+use std::task::Context;
+use wit_bindgen::FutureWriteCancel;
+
+fn main() {
+    wit_bindgen::block_on(async {
+        // cancel from the other end
+        let (tx, rx) = wit_future::new();
+        let f1 = async { tx.write("hello".into()).await };
+        let f2 = async { take_then_close(rx) };
+        let (result, ()) = futures::join!(f1, f2);
+        assert_eq!(result.unwrap_err().value, "hello");
+
+        // cancel before we actually hit the intrinsic
+        let (tx, _rx) = wit_future::new::<String>();
+        let mut future = Box::pin(tx.write("hello2".into()));
+        let tx = match future.as_mut().cancel() {
+            FutureWriteCancel::Cancelled(val, tx) => {
+                assert_eq!(val, "hello2");
+                tx
+            }
+            _ => unreachable!(),
+        };
+
+        // cancel after we hit the intrinsic
+        let mut future = Box::pin(tx.write("hello3".into()));
+        assert!(future
+            .as_mut()
+            .poll(&mut Context::from_waker(noop_waker_ref()))
+            .is_pending());
+        match future.as_mut().cancel() {
+            FutureWriteCancel::Cancelled(val, _) => {
+                assert_eq!(val, "hello3");
+            }
+            _ => unreachable!(),
+        };
+
+        // cancel after we hit the intrinsic and then close the other end
+        //
+        // FIXME(wasip3-prototyping#137)
+        if false {
+            let (tx, rx) = wit_future::new::<String>();
+            let mut future = Box::pin(tx.write("hello3".into()));
+            assert!(future
+                .as_mut()
+                .poll(&mut Context::from_waker(noop_waker_ref()))
+                .is_pending());
+            drop(rx);
+            match future.as_mut().cancel() {
+                FutureWriteCancel::Closed(val) => assert_eq!(val, "hello3"),
+                other => panic!("expected closed, got: {other:?}"),
+            };
+        }
+
+        // Start a write, wait for it to be pending, then go complete the write
+        // in some async work, then cancel it and witness that it was written,
+        // not cancelled.
+        //
+        // FIXME(wasip3-prototyping#138)
+        if false {
+            let (tx, rx) = wit_future::new::<String>();
+            let mut future = Box::pin(tx.write("hello3".into()));
+            assert!(future
+                .as_mut()
+                .poll(&mut Context::from_waker(noop_waker_ref()))
+                .is_pending());
+            read_and_drop(rx).await;
+            match future.as_mut().cancel() {
+                FutureWriteCancel::AlreadySent => {}
+                other => panic!("expected sent, got: {other:?}"),
+            };
+        }
+    });
+}