Update blocking yield signature, bind async yield (#1294)

* Update blocking yield signature, bind async yield

This commit updates the signature for `task_yield` and additionally
renames it to `yield_blocking`. The dual of this, `yield_async`, is also
added as a helper.

* Invert yield definition

Author: alexcrichton

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

@@ -431,21 +431,76 @@ pub fn block_on<T: 'static>(future: impl Future<Output = T> + 'static) -> T {
 /// Call the `yield` canonical built-in function.
 ///
 /// This yields control to the host temporarily, allowing other tasks to make
-/// progress.  It's a good idea to call this inside a busy loop which does not
-/// otherwise ever yield control the the host.
-pub fn task_yield() {
+/// progress. It's a good idea to call this inside a busy loop which does not
+/// otherwise ever yield control the host.
+///
+/// Note that this function is a blocking function, not an `async` function.
+/// That means that this is not an async yield which allows other tasks in this
+/// component to progress, but instead this will block the current function
+/// until the host gets back around to returning from this yield. Asynchronous
+/// functions should probably use [`yield_async`] instead.
+///
+/// # Return Value
+///
+/// This function returns a `bool` which indicates whether execution should
+/// continue after this yield point. A return value of `true` means that the
+/// task was not cancelled and execution should continue. A return value of
+/// `false`, however, means that the task was cancelled while it was suspended
+/// at this yield point. The caller should return back and exit from the task
+/// ASAP in this situation.
+pub fn yield_blocking() -> bool {
     #[cfg(not(target_arch = "wasm32"))]
-    unsafe fn yield_() {
+    unsafe fn yield_() -> bool {
         unreachable!();
     }
 
     #[cfg(target_arch = "wasm32")]
     #[link(wasm_import_module = "$root")]
     extern "C" {
         #[link_name = "[yield]"]
-        fn yield_();
+        fn yield_() -> bool;
     }
-    unsafe { yield_() }
+    // Note that the return value from the raw intrinsic is inverted, the
+    // canonical ABI returns "did this task get cancelled" while this function
+    // works as "should work continue going".
+    unsafe { !yield_() }
+}
+
+/// The asynchronous counterpart to [`yield_blocking`].
+///
+/// This function does not block the current task but instead gives the
+/// Rust-level executor a chance to yield control back to the host temporarily.
+/// This means that other Rust-level tasks may also be able to progress during
+/// this yield operation.
+///
+/// # Return Value
+///
+/// Unlike [`yield_blocking`] this function does not return anything. If this
+/// component task is cancelled while paused at this yield point then the future
+/// will be dropped and a Rust-level destructor will take over and clean up the
+/// task. It's not necessary to do anything with the return value of this
+/// function other than ensuring that you `.await` the function call.
+pub async fn yield_async() {
+    #[derive(Default)]
+    struct Yield {
+        yielded: bool,
+    }
+
+    impl Future for Yield {
+        type Output = ();
+
+        fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<()> {
+            if self.yielded {
+                Poll::Ready(())
+            } else {
+                self.yielded = true;
+                context.waker().wake_by_ref();
+                Poll::Pending
+            }
+        }
+    }
+
+    Yield::default().await;
 }
 
 /// Call the `backpressure.set` canonical built-in function.

crates/guest-rust/src/lib.rs

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

tests/runtime-async/async/cancel-import/runner.rs

@@ -3,8 +3,8 @@ include!(env!("BINDINGS"));
 use crate::my::test::i::*;
 use futures::task::noop_waker_ref;
 use std::future::Future;
-use std::pin::Pin;
-use std::task::{Context, Poll};
+use std::task::Context;
+use wit_bindgen::yield_async;
 
 fn main() {
     println!("test cancelling an import in progress");
@@ -79,7 +79,7 @@ fn main() {
         // Let the subtask's completion notification make its way to our task
         // here.
         for _ in 0..5 {
-            yield_().await;
+            yield_async().await;
         }
 
         // Now cancel the import, despite it actually being done. This should
@@ -114,7 +114,7 @@ fn main() {
         // notification that it's entered the "STARTED" state.
         backpressure_set(false);
         for _ in 0..5 {
-            yield_().await;
+            yield_async().await;
         }
 
         // Now cancel the `starting_import`. This should correctly pick up the
@@ -131,26 +131,3 @@ fn main() {
         started_import.await;
     });
 }
-
-async fn yield_() {
-    #[derive(Default)]
-    struct Yield {
-        yielded: bool,
-    }
-
-    impl Future for Yield {
-        type Output = ();
-
-        fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<()> {
-            if self.yielded {
-                Poll::Ready(())
-            } else {
-                self.yielded = true;
-                context.waker().wake_by_ref();
-                Poll::Pending
-            }
-        }
-    }
-
-    Yield::default().await;
-}

tests/runtime-async/async/pending-import/runner.rs

@@ -3,8 +3,8 @@ include!(env!("BINDINGS"));
 use crate::my::test::i::*;
 use futures::task::noop_waker_ref;
 use std::future::Future;
-use std::pin::Pin;
-use std::task::{Context, Poll};
+use std::task::Context;
+use wit_bindgen::yield_async;
 
 fn main() {
     // Test that Rust-level polling twice works.
@@ -37,31 +37,8 @@ fn main() {
         tx.write(()).await.unwrap();
 
         for _ in 0..5 {
-            yield_().await;
+            yield_async().await;
         }
         drop(import);
     });
 }
-
-async fn yield_() {
-    #[derive(Default)]
-    struct Yield {
-        yielded: bool,
-    }
-
-    impl Future for Yield {
-        type Output = ();
-
-        fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<()> {
-            if self.yielded {
-                Poll::Ready(())
-            } else {
-                self.yielded = true;
-                context.waker().wake_by_ref();
-                Poll::Pending
-            }
-        }
-    }
-
-    Yield::default().await;
-}

tests/runtime-async/async/simple-pending-import/test.rs

@@ -1,40 +1,15 @@
 include!(env!("BINDINGS"));
 
+use wit_bindgen::yield_async;
+
 struct Component;
 
 export!(Component);
 
 impl crate::exports::a::b::i::Guest for Component {
     async fn f() {
         for _ in 0..10 {
-            yield_().await;
-        }
-    }
-}
-
-async fn yield_() {
-    use std::future::Future;
-    use std::pin::Pin;
-    use std::task::{Context, Poll};
-
-    #[derive(Default)]
-    struct Yield {
-        yielded: bool,
-    }
-
-    impl Future for Yield {
-        type Output = ();
-
-        fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<()> {
-            if self.yielded {
-                Poll::Ready(())
-            } else {
-                self.yielded = true;
-                context.waker().wake_by_ref();
-                Poll::Pending
-            }
+            yield_async().await;
         }
     }
-
-    Yield::default().await;
 }