Yield instead of sleep in component-async-tests (#12567)

* Yield instead of sleep in component-async-tests

Use cooperative yields instead of sleeps to make tests more
deterministic and also avoid them unnecessarily taking up test
parallelism by sleeping. Yielding should have the same effect in terms
of testing by exercising behavior returning `Pending` in futures, so
there's no expected loss in test coverage here.

* Yield fewer times

Author: alexcrichton

crates/misc/component-async-tests/src/lib.rs

@@ -9,10 +9,10 @@ pub mod resource_stream;
 pub mod round_trip;
 pub mod round_trip_direct;
 pub mod round_trip_many;
-pub mod sleep;
 pub mod transmit;
 pub mod util;
-pub mod yield_host;
+pub mod yield_;
+pub mod yield_runner;
 
 /// Host implementation, usable primarily by tests
 pub struct Ctx {

crates/misc/component-async-tests/src/round_trip.rs

@@ -1,5 +1,5 @@
 use super::Ctx;
-use std::time::Duration;
+use crate::util::yield_times;
 use wasmtime::component::Accessor;
 
 pub mod bindings {
@@ -20,7 +20,7 @@ pub mod non_concurrent_export_bindings {
 
 impl bindings::local::local::baz::HostWithStore for Ctx {
     async fn foo<T>(_: &Accessor<T, Self>, s: String) -> wasmtime::Result<String> {
-        crate::util::sleep(Duration::from_millis(10)).await;
+        yield_times(10).await;
         Ok(format!("{s} - entered host - exited host"))
     }
 }

crates/misc/component-async-tests/src/round_trip_direct.rs

@@ -1,5 +1,5 @@
 use super::Ctx;
-use std::time::Duration;
+use crate::util::yield_times;
 use wasmtime::component::Accessor;
 
 pub mod bindings {
@@ -11,7 +11,7 @@ pub mod bindings {
 
 impl bindings::RoundTripDirectImportsWithStore for Ctx {
     async fn foo<T>(_: &Accessor<T, Self>, s: String) -> String {
-        crate::util::sleep(Duration::from_millis(10)).await;
+        yield_times(5).await;
         format!("{s} - entered host - exited host")
     }
 }

crates/misc/component-async-tests/src/round_trip_many.rs

@@ -1,5 +1,5 @@
 use super::Ctx;
-use std::time::Duration;
+use crate::util::yield_times;
 use wasmtime::Result;
 use wasmtime::component::Accessor;
 
@@ -41,7 +41,7 @@ impl bindings::local::local::many::HostWithStore for Ctx {
         Option<Stuff>,
         Result<Stuff, ()>,
     ) {
-        crate::util::sleep(Duration::from_millis(10)).await;
+        yield_times(5).await;
         (
             format!("{a} - entered host - exited host"),
             b,

crates/misc/component-async-tests/src/sleep.rs

@@ -3,7 +3,6 @@ use std::{
     marker::PhantomData,
     pin::Pin,
     task::{Context, Poll},
-    thread,
 };
 use wasmtime::Result;
 use wasmtime::{
@@ -14,24 +13,9 @@ use wasmtime::{
     },
 };
 
-pub async fn sleep(duration: std::time::Duration) {
-    if cfg!(miri) {
-        // TODO: We should be able to use `tokio::time::sleep` here, but as of
-        // this writing the miri-compatible version of `wasmtime-fiber` uses
-        // threads behind the scenes, which means thread-local storage is not
-        // preserved when we switch fibers, and that confuses Tokio.  If we ever
-        // fix that we can stop using our own, special version of `sleep` and
-        // switch back to the Tokio version.
-
-        let (tx, rx) = oneshot::channel();
-        let handle = thread::spawn(move || {
-            thread::sleep(duration);
-            _ = tx.send(());
-        });
-        _ = rx.await;
-        _ = handle.join();
-    } else {
-        tokio::time::sleep(duration).await;
+pub async fn yield_times(n: usize) {
+    for _ in 0..n {
+        tokio::task::yield_now().await;
     }
 }
 

crates/misc/component-async-tests/src/util.rs

@@ -0,0 +1,15 @@
+use super::Ctx;
+use wasmtime::component::Accessor;
+
+wasmtime::component::bindgen!({
+    path: "wit",
+    world: "yield-host",
+});
+
+impl local::local::yield_::HostWithStore for Ctx {
+    async fn yield_times<T>(_: &Accessor<T, Self>, times: u64) {
+        crate::util::yield_times(times.try_into().unwrap()).await;
+    }
+}
+
+impl local::local::yield_::Host for Ctx {}

crates/misc/component-async-tests/src/yield_.rs

@@ -8,7 +8,7 @@ use wasmtime::component::{Accessor, Resource};
 pub mod bindings {
     wasmtime::component::bindgen!({
         path: "wit",
-        world: "yield-host",
+        world: "yield-runner",
         imports: { default: trappable },
         with: {
             "local:local/ready.thing": super::Thing,

…/component-async-tests/src/yield_host.rs …omponent-async-tests/src/yield_runner.rscrates/misc/component-async-tests/src/yield_host.rs renamed to crates/misc/component-async-tests/src/yield_runner.rs crates/misc/component-async-tests/src/yield_host.rs renamed to crates/misc/component-async-tests/src/yield_runner.rs

@@ -1,20 +1,19 @@
 use super::util::test_run;
 use crate::scenario::util::{config, make_component};
 use component_async_tests::util;
-use component_async_tests::{Ctx, sleep};
+use component_async_tests::{Ctx, yield_};
 use std::future;
 use std::pin::pin;
 use std::task::Poll;
-use std::time::Duration;
 use wasmtime::Result;
 use wasmtime::component::{Accessor, Linker, ResourceTable};
 use wasmtime::{AsContextMut, Engine, Store, StoreContextMut};
 use wasmtime_wasi::WasiCtxBuilder;
 
-mod sleep_post_return {
+mod yield_post_return {
     wasmtime::component::bindgen!({
         path: "wit",
-        world: "sleep-post-return-callee",
+        world: "yield-post-return-callee",
         exports: { default: task_exit },
     });
 }
@@ -36,29 +35,29 @@ pub async fn async_post_return_caller() -> Result<()> {
 }
 
 #[tokio::test]
-pub async fn async_sleep_post_return_caller() -> Result<()> {
-    test_sleep_post_return(&[
-        test_programs_artifacts::ASYNC_SLEEP_POST_RETURN_CALLER_COMPONENT,
-        test_programs_artifacts::ASYNC_SLEEP_POST_RETURN_CALLEE_COMPONENT,
+pub async fn async_yield_post_return_caller() -> Result<()> {
+    test_yield_post_return(&[
+        test_programs_artifacts::ASYNC_YIELD_POST_RETURN_CALLER_COMPONENT,
+        test_programs_artifacts::ASYNC_YIELD_POST_RETURN_CALLEE_COMPONENT,
     ])
     .await
 }
 
 #[tokio::test]
-pub async fn async_sleep_post_return_callee() -> Result<()> {
-    test_sleep_post_return(&[test_programs_artifacts::ASYNC_SLEEP_POST_RETURN_CALLEE_COMPONENT])
+pub async fn async_yield_post_return_callee() -> Result<()> {
+    test_yield_post_return(&[test_programs_artifacts::ASYNC_YIELD_POST_RETURN_CALLEE_COMPONENT])
         .await
 }
 
-async fn test_sleep_post_return(components: &[&str]) -> Result<()> {
+async fn test_yield_post_return(components: &[&str]) -> Result<()> {
     let engine = Engine::new(&config())?;
 
     let component = make_component(&engine, components).await?;
 
     let mut linker = Linker::new(&engine);
 
     wasmtime_wasi::p2::add_to_linker_async(&mut linker)?;
-    sleep::local::local::sleep::add_to_linker::<_, Ctx>(&mut linker, |ctx| ctx)?;
+    yield_::local::local::yield_::add_to_linker::<_, Ctx>(&mut linker, |ctx| ctx)?;
 
     let mut store = Store::new(
         &engine,
@@ -69,19 +68,19 @@ async fn test_sleep_post_return(components: &[&str]) -> Result<()> {
         },
     );
 
-    let guest = sleep_post_return::SleepPostReturnCallee::instantiate_async(
+    let guest = yield_post_return::YieldPostReturnCallee::instantiate_async(
         &mut store, &component, &linker,
     )
     .await?;
 
     async fn run_with(
         accessor: &Accessor<Ctx>,
-        guest: &sleep_post_return::SleepPostReturnCallee,
+        guest: &yield_post_return::YieldPostReturnCallee,
     ) -> Result<()> {
-        // This function should return immediately, then sleep the specified
-        // number of milliseconds after returning, and then finally exit.
+        // This function should return immediately, then yield the specified
+        // number of times after returning, and then finally exit.
         let exit = guest
-            .local_local_sleep_post_return()
+            .local_local_yield_post_return()
             .call_run(accessor, 100)
             .await?
             .1;
@@ -93,7 +92,7 @@ async fn test_sleep_post_return(components: &[&str]) -> Result<()> {
 
     async fn run(
         store: StoreContextMut<'_, Ctx>,
-        guest: &sleep_post_return::SleepPostReturnCallee,
+        guest: &yield_post_return::YieldPostReturnCallee,
     ) -> Result<()> {
         store
             .run_concurrent(async |accessor| {
@@ -104,7 +103,7 @@ async fn test_sleep_post_return(components: &[&str]) -> Result<()> {
                 // outstanding guest or host tasks to poll for a while, trusting
                 // that we'll resolve the future independently, with or without
                 // giving it more work to do.
-                util::sleep(Duration::from_millis(100)).await;
+                util::yield_times(10).await;
 
                 run_with(accessor, guest).await?;