Opt out of tokio's per-task budget in wasip2 noop polls (#14265)

This commit updates various calls of `noop_waker` and `poll_noop`
throughout wasip2 to ensure they're all using
`tokio::task::unconstrained` to specifically opt-out of Tokio's default
task-budgeting mechanism. This is necessary to get some various tests
I'm working on wasi-libc working on wasip2 because otherwise the test
exhausts the task budget pretty quickly and then results in pretty
surprising I/O behavior. One example is that a TCP socket ready for
writing with `check-write` is subsequently not actually ready for
writing because the task budget ran out.

To the best of my ability to understand what I'm running into is the
same as what's run into in #14174. I'm not entirely sure though, but I'm
going to conservatively say...

Closes #14174

Author: alexcrichton

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

@@ -110,19 +110,62 @@ fn test_tcp_shutdown_should_not_lose_data(net: &Network, family: IpAddressFamily
     });
 }
 
+// Once a stream is writable it should in theory always be writable...
+fn test_tcp_check_write_should_not_be_rate_limited(net: &Network, family: IpAddressFamily) {
+    setup(net, family, |_server, client| {
+        client.output.subscribe().block();
+        assert!(client.output.check_write().unwrap() > 0);
+
+        for i in 0..1000 {
+            let permit = client.output.check_write().unwrap();
+            assert!(
+                permit > 0,
+                "check-write reported 0 bytes on an idle stream after {i} calls"
+            );
+        }
+    });
+}
+
+// Assert that repeated small writes are sent to the kernel's buffer, and the
+// 300 writes here should surely be much larger than any kernel buffer...
+fn test_tcp_nonblocking_write_loop(net: &Network, family: IpAddressFamily) {
+    setup(net, family, |server, client| {
+        client.output.subscribe().block();
+
+        const COUNT: usize = 300;
+        for i in 0..COUNT {
+            let permit = client.output.check_write().unwrap();
+            assert!(
+                permit > 0,
+                "check-write reported 0 bytes after {i} nonblocking writes"
+            );
+            client.output.write(b"x").unwrap();
+        }
+
+        // Everything written above must arrive at the peer.
+        client.socket.shutdown(ShutdownType::Send).unwrap();
+        let received = server.input.blocking_read_to_end().unwrap();
+        assert_eq!(received.len(), COUNT);
+    });
+}
+
 fn main() {
     let net = Network::default();
 
     test_tcp_input_stream_should_be_closed_by_remote_shutdown(&net, IpAddressFamily::Ipv4);
     test_tcp_input_stream_should_be_closed_by_local_shutdown(&net, IpAddressFamily::Ipv4);
     test_tcp_output_stream_should_be_closed_by_local_shutdown(&net, IpAddressFamily::Ipv4);
     test_tcp_shutdown_should_not_lose_data(&net, IpAddressFamily::Ipv4);
+    test_tcp_check_write_should_not_be_rate_limited(&net, IpAddressFamily::Ipv4);
+    test_tcp_nonblocking_write_loop(&net, IpAddressFamily::Ipv4);
 
     if supports_ipv6() {
         test_tcp_input_stream_should_be_closed_by_remote_shutdown(&net, IpAddressFamily::Ipv6);
         test_tcp_input_stream_should_be_closed_by_local_shutdown(&net, IpAddressFamily::Ipv6);
         test_tcp_output_stream_should_be_closed_by_local_shutdown(&net, IpAddressFamily::Ipv6);
         test_tcp_shutdown_should_not_lose_data(&net, IpAddressFamily::Ipv6);
+        test_tcp_check_write_should_not_be_rate_limited(&net, IpAddressFamily::Ipv6);
+        test_tcp_nonblocking_write_loop(&net, IpAddressFamily::Ipv6);
     }
 }
 

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

@@ -6,9 +6,9 @@ use crate::p2::{
     },
     tcp::AsyncOperation,
 };
-use crate::sockets::{WasiSocketsCtxView, noop_cx};
+use crate::runtime::poll_now;
+use crate::sockets::WasiSocketsCtxView;
 use std::net::SocketAddr;
-use std::task::Poll;
 use wasmtime::component::Resource;
 use wasmtime_wasi_io::{
     poll::DynPollable,
@@ -80,7 +80,7 @@ impl crate::p2::host::tcp::tcp::HostTcpSocket for WasiSocketsCtxView<'_> {
             return Err(ErrorCode::NotInProgress.into());
         };
 
-        let Poll::Ready(result) = socket.inner.poll_finish_connect(&mut noop_cx()) else {
+        let Some(result) = poll_now(|cx| socket.inner.poll_finish_connect(cx)) else {
             return Err(ErrorCode::WouldBlock.into());
         };
         socket.in_progress_operation = None;
@@ -134,9 +134,9 @@ impl crate::p2::host::tcp::tcp::HostTcpSocket for WasiSocketsCtxView<'_> {
             return Err(ErrorCode::InvalidState.into());
         };
 
-        let accepted = match listener.poll_accept(&mut noop_cx()) {
-            Poll::Pending => return Err(ErrorCode::WouldBlock.into()),
-            Poll::Ready(accepted) => accepted,
+        let accepted = match poll_now(|cx| listener.poll_accept(cx)) {
+            Some(accepted) => accepted,
+            None => return Err(ErrorCode::WouldBlock.into()),
         };
         let mut tcp_socket = TcpSocket::new(accepted);
         let (input, output) = tcp_socket.take_streams()?;

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

@@ -2,11 +2,11 @@ use crate::p2::bindings::sockets::network::{ErrorCode, IpAddressFamily, IpSocket
 use crate::p2::bindings::sockets::udp;
 use crate::p2::udp::{AsyncOperation, IncomingDatagramStream, OutgoingDatagramStream};
 use crate::p2::{Pollable, SocketError, SocketResult, UdpSocket};
+use crate::runtime::poll_now;
 use crate::sockets::{SocketAddressFamily, WasiSocketsCtxView};
 use async_trait::async_trait;
 use std::future::poll_fn;
 use std::net::SocketAddr;
-use std::task::{Context, Poll, Waker};
 use wasmtime::component::Resource;
 use wasmtime::format_err;
 use wasmtime_wasi_io::poll::DynPollable;
@@ -238,9 +238,7 @@ impl udp::HostOutgoingDatagramStream for WasiSocketsCtxView<'_> {
     fn check_send(&mut self, this: Resource<udp::OutgoingDatagramStream>) -> SocketResult<u64> {
         let stream = self.table.get_mut(&this)?;
 
-        let count = if let Poll::Ready(()) =
-            stream.poll_send_ready(&mut Context::from_waker(Waker::noop()))
-        {
+        let count = if let Some(()) = poll_now(|cx| stream.poll_send_ready(cx)) {
             // We don't know how many Tokio will accept, so we make up a
             // reasonable number here.  If we're wrong and `send` returns
             // `Ok(0)`, the guest will just have to deal with that, e.g. by

crates/wasi/src/p2/ip_name_lookup.rs

@@ -1,10 +1,10 @@
 use crate::p2::SocketError;
 use crate::p2::bindings::sockets::ip_name_lookup::{Host, HostResolveAddressStream};
 use crate::p2::bindings::sockets::network::{ErrorCode, IpAddress, Network};
+use crate::runtime::poll_now;
 use crate::sockets::ip_name_lookup::resolve_addresses;
-use crate::sockets::{MaybeReady, WasiSocketsCtxView, noop_cx};
+use crate::sockets::{MaybeReady, WasiSocketsCtxView};
 use std::net::IpAddr;
-use std::task::Poll;
 use std::vec;
 use wasmtime::Result;
 use wasmtime::component::Resource;
@@ -42,7 +42,7 @@ impl HostResolveAddressStream for WasiSocketsCtxView<'_> {
         resource: Resource<ResolveAddressStream>,
     ) -> Result<Option<IpAddress>, SocketError> {
         let stream: &mut ResolveAddressStream = self.table.get_mut(&resource)?;
-        let Poll::Ready(result) = stream.0.poll_ready(&mut noop_cx()) else {
+        let Some(result) = poll_now(|cx| stream.0.poll_ready(cx)) else {
             return Err(ErrorCode::WouldBlock.into());
         };
 

crates/wasi/src/p2/tcp.rs

@@ -3,8 +3,9 @@ use crate::p2::bindings::sockets::network::ErrorCode;
 use crate::p2::{
     DynInputStream, DynOutputStream, InputStream, OutputStream, Pollable, SocketResult, StreamError,
 };
+use crate::runtime::poll_now;
 use crate::sockets::{
-    MaybeReady, TcpListenStream, TcpReceiveStream, TcpSendStream, TcpSocket as P3Socket, noop_cx,
+    MaybeReady, TcpListenStream, TcpReceiveStream, TcpSendStream, TcpSocket as P3Socket,
 };
 use std::future::poll_fn;
 use std::mem;
@@ -79,14 +80,14 @@ impl ReadState {
             return Ok(bytes::Bytes::new());
         }
         let mut buf = bytes::BytesMut::zeroed(size.min(crate::MAX_READ_SIZE_ALLOC));
-        let n = match stream.poll_read(&mut noop_cx(), &mut buf) {
-            Poll::Pending => 0,
-            Poll::Ready(Ok(0)) => {
+        let n = match poll_now(|cx| stream.poll_read(cx, &mut buf)) {
+            None => 0,
+            Some(Ok(0)) => {
                 *self = ReadState::Closed;
                 return Err(StreamError::Closed);
             }
-            Poll::Ready(Ok(n)) => n,
-            Poll::Ready(Err(e)) => {
+            Some(Ok(n)) => n,
+            Some(Err(e)) => {
                 *self = ReadState::Closed;
                 return Err(StreamError::LastOperationFailed(e.into()));
             }
@@ -158,13 +159,13 @@ impl WriteState {
     }
 
     fn check_write(&mut self) -> StreamResult<usize> {
-        match self.poll_ready(&mut noop_cx()) {
-            Poll::Pending => Ok(0),
-            Poll::Ready(Ok((_, permit))) => {
+        match poll_now(|cx| self.poll_ready(cx)) {
+            None => Ok(0),
+            Some(Ok((_, permit))) => {
                 *permit = MAX_READ_SIZE_ALLOC;
                 Ok(*permit)
             }
-            Poll::Ready(Err(e)) => Err(e),
+            Some(Err(e)) => Err(e),
         }
     }
 
@@ -210,9 +211,9 @@ impl WriteState {
         }));
 
         // Attempt to finish the write, surfacing potential errors immediately:
-        match self.poll_ready(&mut noop_cx()) {
-            Poll::Pending | Poll::Ready(Ok(_)) => Ok(()),
-            Poll::Ready(Err(e)) => Err(e),
+        match poll_now(|cx| self.poll_ready(cx)) {
+            None | Some(Ok(_)) => Ok(()),
+            Some(Err(e)) => Err(e),
         }
     }
 

crates/wasi/src/p2/udp.rs

@@ -2,7 +2,8 @@ use futures::TryFutureExt;
 
 use crate::{
     p2::bindings::sockets::network::ErrorCode,
-    sockets::{MaybeReady, UdpSocket as P3Socket, noop_cx},
+    runtime::poll_now,
+    sockets::{MaybeReady, UdpSocket as P3Socket},
 };
 use std::{
     net::SocketAddr,
@@ -89,8 +90,7 @@ impl IncomingDatagramStream {
     }
 
     pub(crate) fn try_recv(&mut self) -> Result<(Vec<u8>, SocketAddr), ErrorCode> {
-        let noop_cx = &mut noop_cx();
-        if self.poll_recv_ready(noop_cx).is_pending() {
+        if poll_now(|cx| self.poll_recv_ready(cx)).is_none() {
             return Err(ErrorCode::WouldBlock);
         }
 
@@ -128,7 +128,7 @@ impl OutgoingDatagramStream {
         addr: Option<std::net::SocketAddr>,
     ) -> Result<(), ErrorCode> {
         if let Some(send) = &mut self.prev_send {
-            if !send.poll_ready(&mut noop_cx()).is_ready() {
+            if poll_now(|cx| send.poll_ready(cx)).is_none() {
                 return Err(ErrorCode::WouldBlock);
             }
 
@@ -147,7 +147,7 @@ impl OutgoingDatagramStream {
                 .send(data, addr)
                 .map_err(|e| e.into()),
         );
-        if send.poll_ready(&mut noop_cx()).is_ready() {
+        if poll_now(|cx| send.poll_ready(cx)).is_some() {
             send.unwrap_ready()
         } else {
             self.prev_send = Some(send);

crates/wasi/src/runtime.rs

@@ -19,8 +19,8 @@
 //! Each of these facilities should be used by dependencies of wasmtime-wasi
 //! which when implementing component bindings.
 
-use std::future::Future;
-use std::pin::Pin;
+use std::future::{Future, poll_fn};
+use std::pin::{Pin, pin};
 use std::sync::LazyLock;
 use std::task::{Context, Poll, Waker};
 
@@ -170,6 +170,37 @@ pub fn with_ambient_tokio_runtime<R>(f: impl FnOnce() -> R) -> R {
     }
 }
 
+/// Invokes `f` with a noop waker, intended to be used in a `poll*`-style
+/// function, and then returns whether the operation was ready or not.
+///
+/// This uses a "noop waker" to perform the poll and additionally handles
+/// details such as opting-out of Tokio's cooperative task budget.
+pub(crate) fn poll_now<T>(f: impl FnOnce(&mut Context<'_>) -> Poll<T>) -> Option<T> {
+    let mut f = Some(f);
+    // Note that this specifically opts-out of tokio's task budget. This is the
+    // implementation primitive for a number of wasip2 operations where the
+    // "true blocking operation", `poll`, happens at some future time w.r.t. an
+    // in-guest event loop. If tokio's task budget is enabled then the guest
+    // will quickly run out of budget and this function (which bottoms out in
+    // tokio primitives within `f`) will return `Pending`. This can pessimize
+    // guest loops that expect to be able to perform more work
+    // per-turn-of-the-event-loop and otherwise provide surprising behavior
+    // because the "everything is pending now" cliff can happen at seemingly
+    // arbitrary points.
+    let fut = pin!(tokio::task::unconstrained(poll_fn(move |cx| {
+        match f.take() {
+            Some(f) => f(cx),
+            None => Poll::Pending,
+        }
+    })));
+
+    let mut context = Context::from_waker(Waker::noop());
+    match fut.poll(&mut context) {
+        Poll::Ready(result) => Some(result),
+        Poll::Pending => None,
+    }
+}
+
 /// Attempts to get the result of a `future`.
 ///
 /// This function does not block and will poll the provided future once. If the
@@ -181,9 +212,5 @@ pub fn poll_noop<F>(future: Pin<&mut F>) -> Option<F::Output>
 where
     F: Future,
 {
-    let mut task = Context::from_waker(Waker::noop());
-    match future.poll(&mut task) {
-        Poll::Ready(result) => Some(result),
-        Poll::Pending => None,
-    }
+    poll_now(|cx| future.poll(cx))
 }

crates/wasi/src/sockets/mod.rs

@@ -1,3 +1,4 @@
+use crate::runtime::poll_noop;
 use core::fmt;
 use core::future::Future;
 use core::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
@@ -247,9 +248,9 @@ impl<T> MaybeReady<T> {
         T: Send + 'static,
     {
         let mut fut = Box::pin(fut);
-        match crate::runtime::with_ambient_tokio_runtime(|| fut.as_mut().poll(&mut noop_cx())) {
-            Poll::Ready(val) => Self::Ready(val),
-            Poll::Pending => Self::new(crate::runtime::spawn(fut)),
+        match crate::runtime::with_ambient_tokio_runtime(|| poll_noop(fut.as_mut())) {
+            Some(val) => Self::Ready(val),
+            None => Self::new(crate::runtime::spawn(fut)),
         }
     }
     pub(crate) fn unwrap_ready(self) -> T {
@@ -281,10 +282,6 @@ impl<T> MaybeReady<T> {
     }
 }
 
-pub(crate) fn noop_cx() -> std::task::Context<'static> {
-    std::task::Context::from_waker(futures::task::noop_waker_ref())
-}
-
 #[derive(Clone, Copy, Debug)]
 pub enum ErrorCode {
     AccessDenied,