fix Pollable impl for OutgoingDatagramStream (#12629)

* fix `Pollable` impl for `OutgoingDatagramStream`

Previously, `OutgoingDatagramStream` used its `send_state` field to determine
whether the socket was ready for writing without necessarily polling the
`tokio::net::UdpSocket`.  The underlying assumption was that a freshly-bound
socket would be immediately ready for writing, but that's not true for `tokio`.
`tokio` assumes a socket is not ready for _anything_ by default and relies on
e.g. `epoll` to tell it otherwise.

In practice, this meant the `Pollable::ready` impl for `OutgoingDatagramStream`
would resolve immediately for a fresh socket, leaving the guest to race with
`tokio`'s use of `epoll`.  Usually, `tokio` won and all was well, but
occasionally it lost and the `OutgoingDatagramStream::send` call would return
`Ok(0)` (meaning "would block") leaving the guest confused since it was just
told that the socket was ready for writing.

This commit removes `SendState` entirely, relying exclusively on
`tokio::net::UdpSocket::ready` for the `Pollable` and `check_send`
implementations.  As a side effect, we no longer attempt to prevent the guest
from sending more datagrams than `check_send` indicated since the number
returned by `check_send` is only a guess anyway, and `tokio` will push back if
it needs to.  For `p3`, the whole `check_send`/`send` pattern is gone, so we
won't need to concern ourselves with it going forward.

Fixes #12612

* address review feedback

Author: dicej

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

@@ -1,7 +1,7 @@
 use test_programs::wasi::sockets::network::{
     ErrorCode, IpAddress, IpAddressFamily, IpSocketAddress, Network,
 };
-use test_programs::wasi::sockets::udp::UdpSocket;
+use test_programs::wasi::sockets::udp::{OutgoingDatagram, UdpSocket};
 
 const SOME_PORT: u16 = 47; // If the tests pass, this will never actually be connected to.
 
@@ -117,6 +117,30 @@ fn test_udp_connect_dual_stack(net: &Network) {
     ));
 }
 
+/// If `pollable.block` says a `OutgoingDatagramStream` is ready, it better be
+/// ready.
+fn test_udp_connect_and_send(net: &Network, family: IpAddressFamily) {
+    let unspecified_port = IpSocketAddress::new(IpAddress::new_loopback(family), 0);
+    let remote = IpSocketAddress::new(IpAddress::new_loopback(family), 4320);
+
+    let client = UdpSocket::new(family).unwrap();
+    client.blocking_bind(&net, unspecified_port).unwrap();
+
+    let (_, tx) = client.stream(Some(remote)).unwrap();
+    assert_eq!(client.remote_address(), Ok(remote));
+
+    tx.subscribe().block();
+
+    assert!(matches!(tx.check_send(), Ok(n) if n > 0));
+    assert!(matches!(
+        tx.send(&[OutgoingDatagram {
+            data: b"hello".into(),
+            remote_address: None
+        }]),
+        Ok(1)
+    ));
+}
+
 fn main() {
     let net = Network::default();
 
@@ -136,4 +160,11 @@ fn main() {
     test_udp_connect_without_bind(IpAddressFamily::Ipv6);
 
     test_udp_connect_dual_stack(&net);
+
+    // `wasmtime-wasi`'s use of Tokio was once flaky such that these tests would
+    // _usually_, but not always, pass, so we run them a bunch of times here:
+    for _ in 0..256 {
+        test_udp_connect_and_send(&net, IpAddressFamily::Ipv4);
+        test_udp_connect_and_send(&net, IpAddressFamily::Ipv6);
+    }
 }

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

@@ -166,26 +166,45 @@ fn test_udp_connect_dual_stack() {
     ));
 }
 
+/// A UDP socket should be immediately writable
+async fn test_udp_connect_and_send(family: IpAddressFamily) {
+    let unspecified_port = IpSocketAddress::new(IpAddress::new_loopback(family), 0);
+    let remote = IpSocketAddress::new(IpAddress::new_loopback(family), 4320);
+
+    let client = UdpSocket::create(family).unwrap();
+    client.bind(unspecified_port).unwrap();
+
+    client.connect(remote).unwrap();
+    assert_eq!(client.get_remote_address(), Ok(remote));
+
+    assert_eq!(client.send(b"hello".into(), None).await, Ok(()));
+}
+
 impl test_programs::p3::exports::wasi::cli::run::Guest for Component {
     async fn run() -> Result<(), ()> {
+        let supports_ipv6 = supports_ipv6();
+
         test_udp_connect_disconnect_reconnect(IpAddressFamily::Ipv4);
         test_udp_connect_unspec(IpAddressFamily::Ipv4);
         test_udp_connect_local_address_change(IpAddressFamily::Ipv4);
         test_udp_connect_port_0(IpAddressFamily::Ipv4);
         test_udp_connect_wrong_family(IpAddressFamily::Ipv4);
         test_udp_connect_without_bind(IpAddressFamily::Ipv4);
         test_udp_connect_with_bind(IpAddressFamily::Ipv4);
+        test_udp_connect_and_send(IpAddressFamily::Ipv4).await;
 
-        if supports_ipv6() {
+        if supports_ipv6 {
             test_udp_connect_disconnect_reconnect(IpAddressFamily::Ipv6);
             test_udp_connect_unspec(IpAddressFamily::Ipv6);
             test_udp_connect_local_address_change(IpAddressFamily::Ipv6);
             test_udp_connect_port_0(IpAddressFamily::Ipv6);
             test_udp_connect_wrong_family(IpAddressFamily::Ipv6);
             test_udp_connect_without_bind(IpAddressFamily::Ipv6);
             test_udp_connect_with_bind(IpAddressFamily::Ipv6);
+            test_udp_connect_and_send(IpAddressFamily::Ipv6).await;
             test_udp_connect_dual_stack();
         }
+
         Ok(())
     }
 }

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

@@ -1,13 +1,15 @@
 use crate::p2::bindings::sockets::network::{ErrorCode, IpAddressFamily, IpSocketAddress, Network};
 use crate::p2::bindings::sockets::udp;
-use crate::p2::udp::{IncomingDatagramStream, OutgoingDatagramStream, SendState};
+use crate::p2::udp::{IncomingDatagramStream, OutgoingDatagramStream};
 use crate::p2::{Pollable, SocketError, SocketResult};
 use crate::sockets::util::{is_valid_address_family, is_valid_remote_address};
 use crate::sockets::{
     MAX_UDP_DATAGRAM_SIZE, SocketAddrUse, SocketAddressFamily, UdpSocket, WasiSocketsCtxView,
 };
 use async_trait::async_trait;
 use std::net::SocketAddr;
+use std::pin::pin;
+use std::task::{Context, Poll, Waker};
 use tokio::io::Interest;
 use wasmtime::component::Resource;
 use wasmtime::format_err;
@@ -82,7 +84,6 @@ impl udp::HostUdpSocket for WasiSocketsCtxView<'_> {
             inner: socket.socket().clone(),
             remote_address,
             family: socket.address_family(),
-            send_state: SendState::Idle,
             socket_addr_check: socket.socket_addr_check().cloned(),
         };
 
@@ -252,9 +253,8 @@ impl udp::HostIncomingDatagramStream for WasiSocketsCtxView<'_> {
 #[async_trait]
 impl Pollable for IncomingDatagramStream {
     async fn ready(&mut self) {
-        // FIXME: Add `Interest::ERROR` when we update to tokio 1.32.
         self.inner
-            .ready(Interest::READABLE)
+            .ready(Interest::READABLE.add(Interest::ERROR))
             .await
             .expect("failed to await UDP socket readiness");
     }
@@ -264,17 +264,20 @@ impl udp::HostOutgoingDatagramStream for WasiSocketsCtxView<'_> {
     fn check_send(&mut self, this: Resource<udp::OutgoingDatagramStream>) -> SocketResult<u64> {
         let stream = self.table.get_mut(&this)?;
 
-        let permit = match stream.send_state {
-            SendState::Idle => {
-                const PERMIT: usize = 16;
-                stream.send_state = SendState::Permitted(PERMIT);
-                PERMIT
-            }
-            SendState::Permitted(n) => n,
-            SendState::Waiting => 0,
-        };
-
-        Ok(permit.try_into().unwrap())
+        Ok(
+            if let Poll::Ready(_) =
+                pin!(stream.inner.ready(Interest::WRITABLE.add(Interest::ERROR)))
+                    .poll(&mut Context::from_waker(Waker::noop()))
+            {
+                // 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
+                // looping or returning `EWOULDBLOCK`.
+                16
+            } else {
+                0
+            },
+        )
     }
 
     async fn send(
@@ -324,22 +327,6 @@ impl udp::HostOutgoingDatagramStream for WasiSocketsCtxView<'_> {
 
         let stream = self.table.get_mut(&this)?;
 
-        match stream.send_state {
-            SendState::Permitted(n) if n >= datagrams.len() => {
-                stream.send_state = SendState::Idle;
-            }
-            SendState::Permitted(_) => {
-                return Err(SocketError::trap(wasmtime::format_err!(
-                    "unpermitted: argument exceeds permitted size"
-                )));
-            }
-            SendState::Idle | SendState::Waiting => {
-                return Err(SocketError::trap(wasmtime::format_err!(
-                    "unpermitted: must call check-send first"
-                )));
-            }
-        }
-
         if datagrams.is_empty() {
             return Ok(0);
         }
@@ -354,8 +341,8 @@ impl udp::HostOutgoingDatagramStream for WasiSocketsCtxView<'_> {
                     return Ok(count);
                 }
                 Err(e) if matches!(e.downcast_ref(), Some(ErrorCode::WouldBlock)) => {
-                    stream.send_state = SendState::Waiting;
-                    return Ok(count);
+                    debug_assert!(count == 0);
+                    return Ok(0);
                 }
                 Err(e) => {
                     return Err(e);
@@ -386,17 +373,10 @@ impl udp::HostOutgoingDatagramStream for WasiSocketsCtxView<'_> {
 #[async_trait]
 impl Pollable for OutgoingDatagramStream {
     async fn ready(&mut self) {
-        match self.send_state {
-            SendState::Idle | SendState::Permitted(_) => {}
-            SendState::Waiting => {
-                // FIXME: Add `Interest::ERROR` when we update to tokio 1.32.
-                self.inner
-                    .ready(Interest::WRITABLE)
-                    .await
-                    .expect("failed to await UDP socket readiness");
-                self.send_state = SendState::Idle;
-            }
-        }
+        self.inner
+            .ready(Interest::WRITABLE.add(Interest::ERROR))
+            .await
+            .expect("failed to await UDP socket readiness");
     }
 }
 

crates/wasi/src/p2/udp.rs

@@ -18,19 +18,6 @@ pub struct OutgoingDatagramStream {
     /// Socket address family.
     pub(crate) family: SocketAddressFamily,
 
-    pub(crate) send_state: SendState,
-
     /// The check of allowed addresses
     pub(crate) socket_addr_check: Option<SocketAddrCheck>,
 }
-
-pub(crate) enum SendState {
-    /// Waiting for the API consumer to call `check-send`.
-    Idle,
-
-    /// Ready to send up to x datagrams.
-    Permitted(usize),
-
-    /// Waiting for the OS.
-    Waiting,
-}