Synchronize with destroying spawned tasks more in wasip2

This commit fixes a test failure I'm running into in wasi-libc
development where spawned tasks for wasip2 are keeping objects alive in
a race condition where sometimes the task is torn down and sometimes
it's not. Specifically wasip2 is built on wasip3-style primitives for
UDP/TCP which means that futures are used, and if futures aren't
immediately ready they're resolved in a spawned Tokio task. This spawned
task can interact with UDP, for example, where exclusivity of a UDP
socket is tested via `Arc::get_mut` which will nondeterministically
return true or false depending if a previously spawned task has exited
or not. In wasi-libc this means that reconnecting a UDP socket sometimes
fails and sometimes passes because the background task may or may not
have exited.

Here this is resolved by making the `drop` methods async and then
hooking into the preexisting `cancel` method which aborts the task and
then waits on the result. This synchronizes with the task to ensure that
the state of the socket is guaranteed to be exclusive after a disconnect
and ready for another connect. This similar fix is then applied to
ip-name-lookup as well.

Author: alexcrichton

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

@@ -154,12 +154,36 @@ fn test_udp_connect_and_send(net: &Network, family: IpAddressFamily) {
     ));
 }
 
+fn test_udp_reconnect_after_pending_receive(net: &Network, family: IpAddressFamily) {
+    let unspecified_addr = IpSocketAddress::new(IpAddress::new_unspecified(family), 0);
+    let remote = IpSocketAddress::new(IpAddress::new_loopback(family), 4321);
+
+    let client = UdpSocket::new(family).unwrap();
+    client.blocking_bind(&net, unspecified_addr).unwrap();
+
+    // Connect/reconnect in a loop, and this should always succeed...
+    for _ in 0..100 {
+        {
+            let (rx, _tx) = client.stream(None).unwrap();
+            assert!(rx.receive(1).unwrap().is_empty());
+        }
+        {
+            let (rx, _tx) = client.stream(Some(remote)).unwrap();
+            assert_eq!(client.remote_address(), Ok(remote));
+            assert!(rx.receive(1).unwrap().is_empty());
+        }
+    }
+}
+
 fn main() {
     let net = Network::default();
 
     test_udp_connect_disconnect_reconnect(&net, IpAddressFamily::Ipv4);
     test_udp_connect_disconnect_reconnect(&net, IpAddressFamily::Ipv6);
 
+    test_udp_reconnect_after_pending_receive(&net, IpAddressFamily::Ipv4);
+    test_udp_reconnect_after_pending_receive(&net, IpAddressFamily::Ipv6);
+
     test_udp_disconnect_local_address(&net, IpAddressFamily::Ipv4);
     test_udp_disconnect_local_address(&net, IpAddressFamily::Ipv6);
 

crates/wasi/src/p2/bindings.rs

@@ -172,6 +172,7 @@ pub mod sync {
                 "wasi:sockets/udp.incoming-datagram-stream": super::super::sockets::udp::IncomingDatagramStream,
                 "wasi:sockets/udp.outgoing-datagram-stream": super::super::sockets::udp::OutgoingDatagramStream,
                 "wasi:sockets/udp.udp-socket": crate::p2::UdpSocket,
+                "wasi:sockets/ip-name-lookup.resolve-address-stream": crate::p2::ip_name_lookup::ResolveAddressStream,
 
                 // Error host trait from wasmtime-wasi-io is synchronous, so we can alias it
                 "wasi:io/error": wasmtime_wasi_io::bindings::wasi::io::error,
@@ -371,6 +372,8 @@ mod async_io {
             "wasi:sockets/udp.[method]udp-socket.start-bind": async | tracing | trappable,
             "wasi:sockets/udp.[method]udp-socket.stream": async | tracing | trappable,
             "wasi:sockets/udp.[drop]outgoing-datagram-stream": async | tracing | trappable,
+            "wasi:sockets/udp.[drop]incoming-datagram-stream": async | tracing | trappable,
+            "wasi:sockets/ip-name-lookup.[drop]resolve-address-stream": async | tracing | trappable,
             default: tracing | trappable,
         },
         exports: { default: async },

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

@@ -217,12 +217,12 @@ impl udp::HostIncomingDatagramStream for WasiSocketsCtxView<'_> {
         wasmtime_wasi_io::poll::subscribe(self.table, this)
     }
 
-    fn drop(&mut self, this: Resource<udp::IncomingDatagramStream>) -> Result<(), wasmtime::Error> {
-        // As in the filesystem implementation, we assume closing a socket
-        // doesn't block.
+    async fn drop(
+        &mut self,
+        this: Resource<udp::IncomingDatagramStream>,
+    ) -> Result<(), wasmtime::Error> {
         let dropped = self.table.delete(this)?;
-        drop(dropped);
-
+        dropped.finish().await;
         Ok(())
     }
 }
@@ -479,7 +479,7 @@ pub mod sync {
         }
 
         fn drop(&mut self, rep: Resource<IncomingDatagramStream>) -> wasmtime::Result<()> {
-            AsyncHostIncomingDatagramStream::drop(self, rep)
+            in_tokio(async { AsyncHostIncomingDatagramStream::drop(self, rep).await })
         }
     }
 

crates/wasi/src/p2/ip_name_lookup.rs

@@ -3,14 +3,14 @@ use crate::p2::bindings::sockets::ip_name_lookup::{Host, HostResolveAddressStrea
 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};
+use crate::sockets::{MaybeSpawned, WasiSocketsCtxView};
 use std::net::IpAddr;
 use std::vec;
 use wasmtime::Result;
 use wasmtime::component::Resource;
 use wasmtime_wasi_io::poll::{DynPollable, Pollable, subscribe};
 
-pub struct ResolveAddressStream(MaybeReady<Result<vec::IntoIter<IpAddr>, ErrorCode>>);
+pub struct ResolveAddressStream(MaybeSpawned<Result<vec::IntoIter<IpAddr>, ErrorCode>>);
 
 impl Host for WasiSocketsCtxView<'_> {
     fn resolve_addresses(
@@ -24,12 +24,12 @@ impl Host for WasiSocketsCtxView<'_> {
         _ = self.table.get(&network)?;
 
         let fut = resolve_addresses(&self.ctx, name);
-        let stream = ResolveAddressStream(MaybeReady::poll_or_spawn(async move {
+        let stream = ResolveAddressStream(MaybeSpawned::poll_or_spawn(async move {
             Ok(fut.await?.into_iter())
         }));
 
         // Attempt to surface errors immediately.
-        if let MaybeReady::Ready(Err(err)) = &stream.0 {
+        if let MaybeSpawned::Ready(Err(err)) = &stream.0 {
             return Err((*err).into());
         }
         Ok(self.table.push(stream)?)
@@ -59,8 +59,11 @@ impl HostResolveAddressStream for WasiSocketsCtxView<'_> {
         subscribe(self.table, resource)
     }
 
-    fn drop(&mut self, resource: Resource<ResolveAddressStream>) -> Result<()> {
-        self.table.delete(resource)?;
+    async fn drop(&mut self, resource: Resource<ResolveAddressStream>) -> Result<()> {
+        let stream = self.table.delete(resource)?;
+        if let MaybeSpawned::Pending(fut) = stream.0 {
+            fut.cancel().await;
+        }
         Ok(())
     }
 }
@@ -71,3 +74,47 @@ impl Pollable for ResolveAddressStream {
         std::future::poll_fn(|cx| self.0.poll_ready(cx).map(|_| ())).await
     }
 }
+
+mod sync {
+    use super::ResolveAddressStream;
+    use crate::p2::SocketError;
+    use crate::p2::bindings::sockets::network::{IpAddress, Network};
+    use crate::p2::bindings::sync::sockets::ip_name_lookup::{Host, HostResolveAddressStream};
+    use crate::runtime::in_tokio;
+    use crate::sockets::WasiSocketsCtxView;
+    use wasmtime::Result;
+    use wasmtime::component::Resource;
+    use wasmtime_wasi_io::poll::DynPollable;
+
+    impl Host for WasiSocketsCtxView<'_> {
+        fn resolve_addresses(
+            &mut self,
+            network: Resource<Network>,
+            name: String,
+        ) -> Result<Resource<ResolveAddressStream>, SocketError> {
+            <Self as super::Host>::resolve_addresses(self, network, name)
+        }
+    }
+
+    impl HostResolveAddressStream for WasiSocketsCtxView<'_> {
+        fn resolve_next_address(
+            &mut self,
+            resource: Resource<ResolveAddressStream>,
+        ) -> Result<Option<IpAddress>, SocketError> {
+            <Self as super::HostResolveAddressStream>::resolve_next_address(self, resource)
+        }
+
+        fn subscribe(
+            &mut self,
+            resource: Resource<ResolveAddressStream>,
+        ) -> Result<Resource<DynPollable>> {
+            <Self as super::HostResolveAddressStream>::subscribe(self, resource)
+        }
+
+        fn drop(&mut self, resource: Resource<ResolveAddressStream>) -> Result<()> {
+            in_tokio(<Self as super::HostResolveAddressStream>::drop(
+                self, resource,
+            ))
+        }
+    }
+}

crates/wasi/src/p2/mod.rs

@@ -329,6 +329,7 @@ pub fn add_to_linker_with_options_async<T: WasiView>(
     bindings::sockets::tcp::add_to_linker::<T, WasiSockets>(l, T::sockets)?;
     bindings::sockets::udp::add_to_linker::<T, WasiSockets>(l, T::sockets)?;
     bindings::sockets::udp_create_socket::add_to_linker::<T, WasiSockets>(l, T::sockets)?;
+    bindings::sockets::ip_name_lookup::add_to_linker::<T, WasiSockets>(l, T::sockets)?;
     Ok(())
 }
 
@@ -362,7 +363,6 @@ where
     sockets::tcp_create_socket::add_to_linker::<T, WasiSockets>(l, T::sockets)?;
     sockets::instance_network::add_to_linker::<T, WasiSockets>(l, T::sockets)?;
     sockets::network::add_to_linker::<T, WasiSockets>(l, &options.into(), T::sockets)?;
-    sockets::ip_name_lookup::add_to_linker::<T, WasiSockets>(l, T::sockets)?;
     Ok(())
 }
 
@@ -469,6 +469,7 @@ pub fn add_to_linker_with_options_sync<T: WasiView>(
     bindings::sync::sockets::tcp::add_to_linker::<T, WasiSockets>(l, T::sockets)?;
     bindings::sync::sockets::udp::add_to_linker::<T, WasiSockets>(l, T::sockets)?;
     bindings::sync::sockets::udp_create_socket::add_to_linker::<T, WasiSockets>(l, T::sockets)?;
+    bindings::sync::sockets::ip_name_lookup::add_to_linker::<T, WasiSockets>(l, T::sockets)?;
     Ok(())
 }
 

crates/wasi/src/p2/tcp.rs

@@ -5,7 +5,7 @@ use crate::p2::{
 };
 use crate::runtime::poll_now;
 use crate::sockets::{
-    MaybeReady, TcpListenStream, TcpReceiveStream, TcpSendStream, TcpSocket as P3Socket,
+    MaybeSpawned, TcpListenStream, TcpReceiveStream, TcpSendStream, TcpSocket as P3Socket,
 };
 use std::future::poll_fn;
 use std::mem;
@@ -148,8 +148,8 @@ impl From<WriteError> for StreamError {
 
 enum WriteState {
     Ready(TcpSendStream, usize),
-    Writing(MaybeReady<Result<TcpSendStream, WriteError>>),
-    Closing(MaybeReady<Result<(), WriteError>>),
+    Writing(MaybeSpawned<Result<TcpSendStream, WriteError>>),
+    Closing(MaybeSpawned<Result<(), WriteError>>),
     Closed(WriteError),
 }
 
@@ -192,7 +192,7 @@ impl WriteState {
             }
         };
 
-        *self = WriteState::Writing(MaybeReady::poll_or_spawn(async move {
+        *self = WriteState::Writing(MaybeSpawned::poll_or_spawn(async move {
             while !bytes.is_empty() {
                 match stream.write(&bytes).await {
                     Ok(n) => {
@@ -234,7 +234,7 @@ impl WriteState {
 
             // Schedule the shutdown after the current write has finished:
             WriteState::Writing(write) => {
-                WriteState::Closing(MaybeReady::poll_or_spawn(async move {
+                WriteState::Closing(MaybeSpawned::poll_or_spawn(async move {
                     _ = write.into_future().await?;
                     Ok(())
                 }))

crates/wasi/src/p2/udp.rs

@@ -3,7 +3,7 @@ use futures::TryFutureExt;
 use crate::{
     p2::bindings::sockets::network::ErrorCode,
     runtime::poll_now,
-    sockets::{MaybeReady, UdpSocket as P3Socket},
+    sockets::{MaybeSpawned, UdpSocket as P3Socket},
 };
 use std::{
     net::SocketAddr,
@@ -38,7 +38,7 @@ pub(crate) enum AsyncOperation {
 pub struct IncomingDatagramStream {
     pub(crate) inner: Arc<Mutex<P3Socket>>,
     pub(crate) connected_addr: Option<SocketAddr>,
-    pub(crate) current_recv: Option<MaybeReady<Result<(Vec<u8>, SocketAddr), ErrorCode>>>,
+    pub(crate) current_recv: Option<MaybeSpawned<Result<(Vec<u8>, SocketAddr), ErrorCode>>>,
 }
 impl IncomingDatagramStream {
     pub(crate) fn new(inner: Arc<Mutex<P3Socket>>) -> Self {
@@ -56,7 +56,7 @@ impl IncomingDatagramStream {
         if self.current_recv.is_none() {
             let connected_addr = self.connected_addr;
             let inner = self.inner.clone();
-            let recv = MaybeReady::poll_or_spawn(async move {
+            let recv = MaybeSpawned::poll_or_spawn(async move {
                 loop {
                     let fut = inner.lock().unwrap().recv();
                     let (data, addr) = fut.await?;
@@ -96,13 +96,20 @@ impl IncomingDatagramStream {
 
         self.current_recv.take().unwrap().unwrap_ready()
     }
+
+    pub(crate) async fn finish(mut self) {
+        let Some(MaybeSpawned::Pending(recv)) = self.current_recv.take() else {
+            return;
+        };
+        recv.cancel().await;
+    }
 }
 
 pub struct OutgoingDatagramStream {
     pub(crate) inner: Arc<Mutex<P3Socket>>,
     /// Number of datagrams permitted by most recent `check-send` call.
     pub(crate) check_send_permit_count: usize,
-    pub(crate) prev_send: Option<MaybeReady<Result<(), ErrorCode>>>,
+    pub(crate) prev_send: Option<MaybeSpawned<Result<(), ErrorCode>>>,
 }
 impl OutgoingDatagramStream {
     pub(crate) fn new(inner: Arc<Mutex<P3Socket>>) -> Self {
@@ -140,7 +147,7 @@ impl OutgoingDatagramStream {
 
         debug_assert!(self.prev_send.is_none());
 
-        let mut send = MaybeReady::poll_or_spawn(
+        let mut send = MaybeSpawned::poll_or_spawn(
             self.inner
                 .lock()
                 .unwrap()

crates/wasi/src/sockets/mod.rs

@@ -1,4 +1,4 @@
-use crate::runtime::poll_noop;
+use crate::runtime::{AbortOnDropJoinHandle, poll_noop};
 use core::fmt;
 use core::future::Future;
 use core::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
@@ -232,15 +232,14 @@ pub(crate) enum SocketAddressFamily {
 /// (1) polling a future for completion and
 /// (2) obtaining the output of a future
 /// into separate operations. This is a common pattern in WASI 0.2.
-pub(crate) enum MaybeReady<T> {
-    Pending(Pin<Box<dyn Future<Output = T> + Send>>),
+pub(crate) enum MaybeReady<T, F = Pin<Box<dyn Future<Output = T> + Send>>> {
+    Pending(F),
     Ready(T),
 }
-impl<T> MaybeReady<T> {
-    pub(crate) fn new(fut: impl Future<Output = T> + Send + 'static) -> Self {
-        Self::Pending(Box::pin(fut))
-    }
 
+pub(crate) type MaybeSpawned<T> = MaybeReady<T, AbortOnDropJoinHandle<T>>;
+
+impl<T> MaybeSpawned<T> {
     /// Poll the future and attempt to resolve it immediately. If the future is
     /// not ready yet, it will be moved to a background task.
     pub(crate) fn poll_or_spawn(fut: impl Future<Output = T> + Send + 'static) -> Self
@@ -253,6 +252,16 @@ impl<T> MaybeReady<T> {
             None => Self::new(crate::runtime::spawn(fut)),
         }
     }
+}
+
+impl<T, F> MaybeReady<T, F>
+where
+    F: Future<Output = T> + Unpin,
+{
+    pub(crate) fn new(fut: F) -> Self {
+        Self::Pending(fut)
+    }
+
     pub(crate) fn unwrap_ready(self) -> T {
         match self {
             Self::Ready(val) => val,
@@ -261,7 +270,7 @@ impl<T> MaybeReady<T> {
     }
     pub(crate) fn poll_ready(&mut self, cx: &mut std::task::Context<'_>) -> Poll<&mut T> {
         match self {
-            Self::Pending(fut) => match fut.as_mut().poll(cx) {
+            Self::Pending(fut) => match Pin::new(fut).as_mut().poll(cx) {
                 Poll::Ready(val) => {
                     *self = Self::Ready(val);
                     Poll::Ready(match self {

crates/wasi/src/sockets/tcp.rs

@@ -191,7 +191,7 @@ impl TcpSocket {
             unreachable!();
         };
 
-        self.tcp_state = TcpState::Connecting(MaybeReady::new(async move {
+        self.tcp_state = TcpState::Connecting(MaybeReady::new(Box::pin(async move {
             // Perform all checks before doing any syscalls.
             {
                 if !already_bound {
@@ -213,7 +213,7 @@ impl TcpSocket {
 
             let stream = sock.connect(addr).await?;
             Ok(stream)
-        }));
+        })));
 
         Ok(())
     }
@@ -569,7 +569,7 @@ impl TcpListenStream {
             let listener = self.inner.clone();
             let permissions = self.permissions.clone();
 
-            self.pending_accept = Some(MaybeReady::new(async move {
+            self.pending_accept = Some(MaybeReady::new(Box::pin(async move {
                 loop {
                     match accept(&listener).await {
                         Ok((client, addr)) => {
@@ -589,7 +589,7 @@ impl TcpListenStream {
                         }
                     }
                 }
-            }));
+            })));
         }
 
         with_ambient_tokio_runtime(|| {