Opt out of tokio's per-task budget in wasip2 noop polls

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,