yield to executor when polling with zero-duration wait (#13085)

* yield to executor when polling with zero-duration wait

This addresses #13040 for `wasmtime-wasi`'s WASIp2 implementation by calling
`tokio::task::yield_now` in `<Deadline as Pollable>::ready` when `self` is a
`Deadline::Past`.  Previously, it did nothing, which had the effect of starving
other pollables that rely on yielding control to `mio` in order to progress,
e.g. when the guest polls with a zero timeout in a busy loop.

This also addresses that issue for `wasmtime`'s `thread.yield` implementation by
modifying the `poll_until` event loop in `concurrent.rs` to yield to the
executor prior to executing any low-priority tasks.  This works because
`thread.yield` operates by queueing a low-priority task to resume the calling
fiber just prior to suspending it.

Note that, because the `wasmtime` crate does not depend on Tokio, we can't
directly call `tokio::task::yield_now`.  Instead, we return a `Future` which,
when polled for the first time, wakes the `Context`'s `Waker` and returns
`Poll::Pending`; then it returns `Poll::Ready(())` for subsequent polls.  That's
enough to ensure `mio` has a chance to run when using `tokio`.  Once we have a
mechanism to allow the embedder to configure a custom yield function, we'll be
able to use that instead.

Fixes #13040

* remove obsolete assertions from `p2_sleep` test

Now that zero-duration `monotonic_clock` waits yield to the executor, they
aren't ready immediately, so the assertions in this test are no longer
appropriate.

* address review feedback

- Expand comments regarding why we yield to the runtime in `wasmtime` and `wasmtime-wasi`
- Move `yield_to_executor` to `StoreOpaque::yield_now` and remove `vm/async_yield` in favor of the new function
- Revert `take_low_priority` performance tweak (will open a separate PR for that)
- Inline `collect_work_itmes_to_run` for clarity

Author: dicej

crates/test-programs/src/bin/p2_sleep.rs

@@ -18,16 +18,10 @@ fn sleep_0ms() {
     let p = monotonic_clock::subscribe_instant(monotonic_clock::now());
     p.block();
     let p = monotonic_clock::subscribe_duration(0);
-    assert!(
-        p.ready(),
-        "timer subscription with duration 0 is ready immediately"
-    );
+    p.block();
 }
 
 fn sleep_backwards_in_time() {
     let p = monotonic_clock::subscribe_instant(monotonic_clock::now() - 1);
-    assert!(
-        p.ready(),
-        "timer subscription for instant which has passed is ready immediately"
-    );
+    p.block();
 }

crates/test-programs/src/bin/p2_tcp_busy_poll.rs

@@ -0,0 +1,72 @@
+use test_programs::sockets::supports_ipv6;
+use test_programs::wasi::clocks::monotonic_clock;
+use test_programs::wasi::io::poll;
+use test_programs::wasi::sockets::network::{
+    IpAddressFamily, IpSocketAddress, Ipv4SocketAddress, Ipv6SocketAddress, Network,
+};
+use test_programs::wasi::sockets::tcp::TcpSocket;
+
+// Historically, `wasmtime-wasi` had a bug such that polling with a
+// zero-duration (i.e. always-ready) wait would starve the host executor and
+// prevent e.g. socket readiness from being delivered.  Here we verify that such
+// starvation does not happen.
+fn test_tcp_busy_poll(family: IpAddressFamily, address: IpSocketAddress) {
+    let zero_wait = monotonic_clock::subscribe_duration(0);
+
+    let net = Network::default();
+
+    let listener = TcpSocket::new(family).unwrap();
+    listener.blocking_bind(&net, address).unwrap();
+    listener.set_listen_backlog_size(32).unwrap();
+    listener.blocking_listen().unwrap();
+
+    let address = listener.local_address().unwrap();
+
+    let message = b"Hello, world!";
+
+    for _ in 0..100 {
+        let client = TcpSocket::new(family).unwrap();
+        let (_rx, tx) = client.blocking_connect(&net, address).unwrap();
+        tx.blocking_write_util(message).unwrap();
+
+        let (_accepted, rx, _tx) = listener.blocking_accept().unwrap();
+        let rx_ready = rx.subscribe();
+        let mut counter = 0;
+        loop {
+            if counter > 1_000_000 {
+                panic!("socket still not ready!");
+            }
+
+            let ready = poll::poll(&[&zero_wait, &rx_ready]);
+            if ready.contains(&1) {
+                break;
+            }
+            counter += 1;
+        }
+
+        let data = rx.read(message.len().try_into().unwrap()).unwrap();
+        assert_eq!(data, message); // Not guaranteed to work but should work in practice.
+    }
+}
+
+fn main() {
+    test_tcp_busy_poll(
+        IpAddressFamily::Ipv4,
+        IpSocketAddress::Ipv4(Ipv4SocketAddress {
+            port: 0,
+            address: (127, 0, 0, 1),
+        }),
+    );
+
+    if supports_ipv6() {
+        test_tcp_busy_poll(
+            IpAddressFamily::Ipv6,
+            IpSocketAddress::Ipv6(Ipv6SocketAddress {
+                port: 0,
+                address: (0, 0, 0, 0, 0, 0, 0, 1),
+                flow_info: 0,
+                scope_id: 0,
+            }),
+        );
+    }
+}

crates/test-programs/src/bin/p3_sockets_tcp_busy_poll.rs

@@ -0,0 +1,123 @@
+use futures::join;
+use std::ptr;
+use test_programs::async_::{
+    BLOCKED, EVENT_NONE, EVENT_STREAM_READ, thread_yield, waitable_join, waitable_set_drop,
+    waitable_set_new, waitable_set_poll,
+};
+use test_programs::p3::wasi::sockets::types::{
+    IpAddressFamily, IpSocketAddress, Ipv4SocketAddress, Ipv6SocketAddress, TcpSocket,
+};
+use test_programs::p3::wit_stream;
+use test_programs::sockets::supports_ipv6;
+use wit_bindgen::StreamResult;
+
+#[link(wasm_import_module = "wasi:http/types@0.3.0-rc-2026-03-15")]
+unsafe extern "C" {
+    #[link_name = "[async-lower][stream-read-0][static]request.new"]
+    fn stream_read(_: u32, _: *mut u8, _: usize) -> u32;
+}
+
+// Historically, `wasmtime-wasi` had a bug such that polling would starve the
+// host executor and prevent e.g. socket readiness from being delivered.  Here
+// we verify that such starvation does not happen.
+async fn test_tcp_busy_poll(family: IpAddressFamily, address: IpSocketAddress) {
+    let listener = TcpSocket::create(family).unwrap();
+    listener.bind(address).unwrap();
+    listener.set_listen_backlog_size(32).unwrap();
+    let mut accept = listener.listen().unwrap();
+
+    let address = listener.get_local_address().unwrap();
+
+    let message = b"Hello, world!";
+
+    for _ in 0..100 {
+        let client = TcpSocket::create(family).unwrap();
+        client.connect(address).await.unwrap();
+        let (mut data_tx, data_rx) = wit_stream::new();
+        join!(
+            async {
+                client.send(data_rx).await.unwrap();
+            },
+            async {
+                let remaining = data_tx.write_all(message.into()).await;
+                assert!(remaining.is_empty());
+                drop(data_tx);
+            }
+        );
+
+        let sock = accept.next().await.unwrap();
+        let (mut data_rx, fut) = sock.receive();
+
+        unsafe {
+            if (stream_read(data_rx.handle(), ptr::null_mut(), 0)) == BLOCKED {
+                let set = waitable_set_new();
+                waitable_join(data_rx.handle(), set);
+                let mut counter = 0;
+                loop {
+                    if counter > 1_000_000 {
+                        panic!("socket still not ready!");
+                    }
+
+                    let (mut event, mut waitable, _) = waitable_set_poll(set);
+                    if event == EVENT_NONE {
+                        // `waitable-set.poll` does not yield, so we
+                        // call `thread.yield` to give the executor a
+                        // chance to run, then poll one more time.
+                        assert!(thread_yield() == 0);
+                        (event, waitable, _) = waitable_set_poll(set);
+                    }
+                    if event == EVENT_STREAM_READ && waitable == data_rx.handle() {
+                        waitable_join(data_rx.handle(), 0);
+                        waitable_set_drop(set);
+                        break;
+                    }
+                    counter += 1;
+                }
+            }
+        }
+
+        let (result, data) = data_rx.read(Vec::with_capacity(100)).await;
+        assert_eq!(result, StreamResult::Complete(message.len()));
+        assert_eq!(data, message); // Not guaranteed to work but should work in practice.
+
+        let (result, data) = data_rx.read(Vec::with_capacity(1)).await;
+        assert_eq!(result, StreamResult::Dropped);
+        assert_eq!(data, []);
+
+        fut.await.unwrap();
+    }
+}
+
+struct Component;
+
+test_programs::p3::export!(Component);
+
+impl test_programs::p3::exports::wasi::cli::run::Guest for Component {
+    async fn run() -> Result<(), ()> {
+        test_tcp_busy_poll(
+            IpAddressFamily::Ipv4,
+            IpSocketAddress::Ipv4(Ipv4SocketAddress {
+                port: 0,
+                address: (127, 0, 0, 1),
+            }),
+        )
+        .await;
+
+        if supports_ipv6() {
+            test_tcp_busy_poll(
+                IpAddressFamily::Ipv6,
+                IpSocketAddress::Ipv6(Ipv6SocketAddress {
+                    port: 0,
+                    address: (0, 0, 0, 0, 0, 0, 0, 1),
+                    flow_info: 0,
+                    scope_id: 0,
+                }),
+            )
+            .await;
+        }
+
+        Ok(())
+    }
+}
+
+fn main() {}

crates/wasi/src/p2/host/clocks.rs

@@ -124,7 +124,26 @@ enum Deadline {
 impl Pollable for Deadline {
     async fn ready(&mut self) {
         match self {
-            Deadline::Past => {}
+            Deadline::Past => {
+                // It is important we yield to Tokio here; otherwise we risk
+                // starving `mio` such that it is unable to signal readiness for
+                // other pollables (e.g. TCP sockets) when the guest is polling
+                // in a busy loop.
+                //
+                // This is somewhat of a hack to ensure that
+                // `wasmtime-wasi-io`'s implementation of `wasi:io/poll` does
+                // not starve `mio` when the guest calls `wasi:io/poll#poll` in
+                // a busy loop with a zero timeout.  It relies on the guest
+                // using the most natural approach to making a non-blocking call
+                // to `wasi:io/poll#poll`, which is to include a zero-duration
+                // `monotonic_clock::subscribe_{instant,duration}` in the list
+                // of pollables.  That's what `wasi-libc`'s `poll(2)`
+                // implementation does as of this writing, for example.  There
+                // are hypothetically other ways to generate a pollable that's
+                // always immediately ready, which this hack doesn't cover, but
+                // we consider this sufficient for now.
+                tokio::task::yield_now().await
+            }
             Deadline::Instant(instant) => tokio::time::sleep_until(*instant).await,
             Deadline::Never => std::future::pending().await,
         }

crates/wasi/tests/all/p2/async_.rs

@@ -307,6 +307,10 @@ async fn p2_tcp_listen() {
     run(P2_TCP_LISTEN_COMPONENT, false).await.unwrap()
 }
 #[test_log::test(tokio::test(flavor = "multi_thread"))]
+async fn p2_tcp_busy_poll() {
+    run(P2_TCP_BUSY_POLL_COMPONENT, false).await.unwrap()
+}
+#[test_log::test(tokio::test(flavor = "multi_thread"))]
 async fn p2_udp_sockopts() {
     run(P2_UDP_SOCKOPTS_COMPONENT, false).await.unwrap()
 }

crates/wasi/tests/all/p2/sync.rs

@@ -293,6 +293,10 @@ fn p2_tcp_listen() {
     run(P2_TCP_LISTEN_COMPONENT, false).unwrap()
 }
 #[test_log::test]
+fn p2_tcp_busy_poll() {
+    run(P2_TCP_BUSY_POLL_COMPONENT, false).unwrap()
+}
+#[test_log::test]
 fn p2_udp_sockopts() {
     run(P2_UDP_SOCKOPTS_COMPONENT, false).unwrap()
 }

crates/wasi/tests/all/p3/mod.rs

@@ -115,6 +115,11 @@ async fn p3_sockets_tcp_streams() -> wasmtime::Result<()> {
     run(P3_SOCKETS_TCP_STREAMS_COMPONENT).await
 }
 
+#[test_log::test(tokio::test(flavor = "multi_thread"))]
+async fn p3_sockets_tcp_busy_poll() -> wasmtime::Result<()> {
+    run(P3_SOCKETS_TCP_BUSY_POLL_COMPONENT).await
+}
+
 #[test_log::test(tokio::test(flavor = "multi_thread"))]
 async fn p3_sockets_udp_bind() -> wasmtime::Result<()> {
     run(P3_SOCKETS_UDP_BIND_COMPONENT).await

crates/wasmtime/src/runtime/component/concurrent.rs

@@ -1196,8 +1196,12 @@ impl<T> StoreContextMut<'_, T> {
 
             enum PollResult<R> {
                 Complete(R),
-                ProcessWork(Vec<WorkItem>),
+                ProcessWork {
+                    ready: Vec<WorkItem>,
+                    low_priority: bool,
+                },
             }
+
             let result = future::poll_fn(|cx| {
                 // First, poll the future we were passed as an argument and
                 // return immediately if it's ready.
@@ -1220,13 +1224,23 @@ impl<T> StoreContextMut<'_, T> {
                     Poll::Pending => Poll::Pending,
                 };
 
-                // Next, collect the next batch of work items to process, if any.
-                // This will be either all of the high-priority work items, or if
-                // there are none, a single low-priority work item.
+                // Next, collect the next batch of work items to process, if
+                // any.  This will be either all of the high-priority work
+                // items, or if there are none, a single low-priority work item.
                 let state = reset.store.0.concurrent_state_mut();
-                let ready = state.collect_work_items_to_run();
+                let mut ready = mem::take(&mut state.high_priority);
+                let mut low_priority = false;
+                if ready.is_empty() {
+                    if let Some(item) = state.low_priority.pop_back() {
+                        ready.push(item);
+                        low_priority = true;
+                    }
+                }
                 if !ready.is_empty() {
-                    return Poll::Ready(Ok(PollResult::ProcessWork(ready)));
+                    return Poll::Ready(Ok(PollResult::ProcessWork {
+                        ready,
+                        low_priority,
+                    }));
                 }
 
                 // Finally, if we have nothing else to do right now, determine what to do
@@ -1239,7 +1253,10 @@ impl<T> StoreContextMut<'_, T> {
                         // successfully, so we return now and continue
                         // the outer loop in case there is another one
                         // ready to complete.
-                        Poll::Ready(Ok(PollResult::ProcessWork(Vec::new())))
+                        Poll::Ready(Ok(PollResult::ProcessWork {
+                            ready: Vec::new(),
+                            low_priority: false,
+                        }))
                     }
                     Poll::Ready(false) => {
                         // Poll the future we were passed one last time
@@ -1287,7 +1304,10 @@ impl<T> StoreContextMut<'_, T> {
                 PollResult::Complete(value) => break Ok(value),
                 // The future we were passed has not yet completed, so handle
                 // any work items and then loop again.
-                PollResult::ProcessWork(ready) => {
+                PollResult::ProcessWork {
+                    ready,
+                    low_priority,
+                } => {
                     struct Dispose<'a, T: 'static, I: Iterator<Item = WorkItem>> {
                         store: StoreContextMut<'a, T>,
                         ready: I,
@@ -1312,6 +1332,31 @@ impl<T> StoreContextMut<'_, T> {
                         ready: ready.into_iter(),
                     };
 
+                    // If we're about to run a low-priority task, first yield to
+                    // the executor.  This ensures that it won't be starved of
+                    // the ability to e.g. update the readiness of sockets,
+                    // etc. which the guest may be using `thread.yield` along
+                    // with `waitable-set.poll` to monitor in a CPU-heavy loop.
+                    //
+                    // This works for e.g. `thread.yield` and callbacks which
+                    // return `CALLBACK_CODE_YIELD` because we queue a low
+                    // priority item to resume the task (i.e. resume the thread
+                    // or call the callback, respectively) just prior to
+                    // suspending it.  Indeed, as of this writing those are the
+                    // _only_ situations we queue low-priority tasks.
+                    // Therefore, we interpret the guest's request to yield as
+                    // meaning "yield to other guest tasks _and_/_or_ host
+                    // operations such as updating socket readiness", the latter
+                    // being the async runtime's responsibility.
+                    //
+                    // In the future, if this ends up causing measurable
+                    // performance issues, this could be optimized such that we
+                    // only yield periodically (e.g. for batches of low priority
+                    // items) and not for each and every idividual item.
+                    if low_priority {
+                        dispose.store.0.yield_now().await
+                    }
+
                     while let Some(item) = dispose.ready.next() {
                         dispose
                             .store
@@ -4972,19 +5017,6 @@ impl ConcurrentState {
         }
     }
 
-    /// Collect the next set of work items to run. This will be either all
-    /// high-priority items, or a single low-priority item if there are no
-    /// high-priority items.
-    fn collect_work_items_to_run(&mut self) -> Vec<WorkItem> {
-        let mut ready = mem::take(&mut self.high_priority);
-        if ready.is_empty() {
-            if let Some(item) = self.low_priority.pop_back() {
-                ready.push(item);
-            }
-        }
-        ready
-    }
-
     fn push<V: Send + Sync + 'static>(
         &mut self,
         value: V,