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/wasi/src/p2/bindings.rs

@@ -371,6 +371,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,11 +217,15 @@ impl udp::HostIncomingDatagramStream for WasiSocketsCtxView<'_> {
         wasmtime_wasi_io::poll::subscribe(self.table, this)
     }
 
-    fn drop(&mut self, this: Resource<udp::IncomingDatagramStream>) -> Result<(), wasmtime::Error> {
+    async fn drop(
+        &mut self,
+        this: Resource<udp::IncomingDatagramStream>,
+    ) -> Result<(), wasmtime::Error> {
         // As in the filesystem implementation, we assume closing a socket
         // doesn't block.
         let dropped = self.table.delete(this)?;
-        drop(dropped);
+
+        dropped.finish().await;
 
         Ok(())
     }
@@ -481,7 +485,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

@@ -2,15 +2,15 @@ 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::sockets::ip_name_lookup::resolve_addresses;
-use crate::sockets::{MaybeReady, WasiSocketsCtxView, noop_cx};
+use crate::sockets::{MaybeSpawned, WasiSocketsCtxView, noop_cx};
 use std::net::IpAddr;
 use std::task::Poll;
 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(())
     }
 }

crates/wasi/src/p2/tcp.rs

@@ -4,7 +4,7 @@ use crate::p2::{
     DynInputStream, DynOutputStream, InputStream, OutputStream, Pollable, SocketResult, StreamError,
 };
 use crate::sockets::{
-    MaybeReady, TcpListenStream, TcpReceiveStream, TcpSendStream, TcpSocket as P3Socket, noop_cx,
+    MaybeSpawned, TcpListenStream, TcpReceiveStream, TcpSendStream, TcpSocket as P3Socket, noop_cx,
 };
 use std::future::poll_fn;
 use std::mem;
@@ -147,8 +147,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),
 }
 
@@ -191,7 +191,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) => {
@@ -233,7 +233,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

@@ -2,7 +2,7 @@ use futures::TryFutureExt;
 
 use crate::{
     p2::bindings::sockets::network::ErrorCode,
-    sockets::{MaybeReady, UdpSocket as P3Socket, noop_cx},
+    sockets::{MaybeSpawned, UdpSocket as P3Socket, noop_cx},
 };
 use std::{
     net::SocketAddr,
@@ -37,7 +37,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 {
@@ -55,7 +55,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,3 +1,4 @@
+use crate::runtime::AbortOnDropJoinHandle;
 use core::fmt;
 use core::future::Future;
 use core::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
@@ -231,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
@@ -252,6 +252,16 @@ impl<T> MaybeReady<T> {
             Poll::Pending => 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,
@@ -260,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(|| {