adapt to API changes

prtest:full

Signed-off-by: Roman Volosatovs <rvolosatovs@riseup.net>

Author: rvolosatovs

crates/wasi/src/p3/sockets/host/ip_name_lookup.rs

@@ -1,13 +1,13 @@
 use tokio::net::lookup_host;
 use wasmtime::component::Accessor;
 
-use crate::p3::bindings::sockets::ip_name_lookup::{ErrorCode, Host, HostConcurrent};
+use crate::p3::bindings::sockets::ip_name_lookup::{ErrorCode, Host, HostWithStore};
 use crate::p3::bindings::sockets::types;
 use crate::p3::sockets::WasiSockets;
 use crate::sockets::WasiSocketsCtxView;
 use crate::sockets::util::{from_ipv4_addr, from_ipv6_addr, parse_host};
 
-impl HostConcurrent for WasiSockets {
+impl HostWithStore for WasiSockets {
     async fn resolve_addresses<U>(
         store: &Accessor<U, Self>,
         name: String,

crates/wasi/src/p3/sockets/host/types/mod.rs

@@ -2,7 +2,7 @@ use core::net::SocketAddr;
 
 use wasmtime::component::Accessor;
 
-use crate::p3::bindings::sockets::types::{Host, HostConcurrent};
+use crate::p3::bindings::sockets::types::Host;
 use crate::p3::sockets::WasiSockets;
 use crate::sockets::{SocketAddrCheck, SocketAddrUse, WasiSocketsCtxView};
 
@@ -11,8 +11,6 @@ mod udp;
 
 impl Host for WasiSocketsCtxView<'_> {}
 
-impl HostConcurrent for WasiSockets {}
-
 fn get_socket_addr_check<T>(store: &Accessor<T, WasiSockets>) -> SocketAddrCheck {
     store.with(|mut view| view.get().ctx.socket_addr_check.clone())
 }

crates/wasi/src/p3/sockets/host/types/tcp.rs

@@ -20,7 +20,7 @@ use wasmtime::component::{
 
 use crate::p3::DEFAULT_BUFFER_CAPACITY;
 use crate::p3::bindings::sockets::types::{
-    Duration, ErrorCode, HostTcpSocket, HostTcpSocketConcurrent, IpAddressFamily, IpSocketAddress,
+    Duration, ErrorCode, HostTcpSocket, HostTcpSocketWithStore, IpAddressFamily, IpSocketAddress,
     TcpSocket,
 };
 use crate::p3::sockets::WasiSockets;
@@ -192,6 +192,18 @@ impl<T> AccessorTask<T, WasiSockets, wasmtime::Result<()>> for ListenTask {
     }
 }
 
+struct ResultWriteTask {
+    result: Result<(), ErrorCode>,
+    result_tx: FutureWriter<Result<(), ErrorCode>>,
+}
+
+impl<T> AccessorTask<T, WasiSockets, wasmtime::Result<()>> for ResultWriteTask {
+    async fn run(self, store: &Accessor<T, WasiSockets>) -> wasmtime::Result<()> {
+        self.result_tx.write(store, self.result).await;
+        Ok(())
+    }
+}
+
 struct ReceiveTask {
     stream: Arc<TcpStream>,
     data_tx: StreamWriter<Cursor<BytesMut>>,
@@ -238,16 +250,22 @@ impl<T> AccessorTask<T, WasiSockets, wasmtime::Result<()>> for ReceiveTask {
                 }
             }
         };
-        self.result_tx.write(store, res).await;
         _ = self
             .stream
             .as_socketlike_view::<std::net::TcpStream>()
             .shutdown(Shutdown::Read);
+
+        // Write the result async from a separate task to ensure that all resources used by this
+        // task are freed
+        store.spawn(ResultWriteTask {
+            result: res,
+            result_tx: self.result_tx,
+        });
         Ok(())
     }
 }
 
-impl HostTcpSocketConcurrent for WasiSockets {
+impl HostTcpSocketWithStore for WasiSockets {
     async fn bind<T>(
         store: &Accessor<T, Self>,
         socket: Resource<TcpSocket>,
@@ -283,7 +301,7 @@ impl HostTcpSocketConcurrent for WasiSockets {
                 || !is_valid_remote_address(remote_address)
                 || !is_valid_address_family(ip, socket.family)
             {
-                return Ok(Err(ErrorCode::InvalidArgument));
+                return anyhow::Ok(Err(ErrorCode::InvalidArgument));
             }
             match mem::replace(&mut socket.tcp_state, TcpState::Connecting) {
                 TcpState::Default(sock) | TcpState::Bound(sock) => Ok(Ok(sock)),
@@ -433,25 +451,28 @@ impl HostTcpSocketConcurrent for WasiSockets {
             let (data_tx, data_rx) = instance
                 .stream::<_, _, BytesMut>(&mut view)
                 .context("failed to create stream")?;
-            let (result_tx, result_rx) = instance
-                .future(|| unreachable!(), &mut view)
-                .context("failed to create future")?;
             let TcpSocket { tcp_state, .. } = get_socket_mut(view.get().table, &socket)?;
             match mem::replace(tcp_state, TcpState::Closed) {
                 TcpState::Connected(stream) => {
                     *tcp_state = TcpState::Receiving(Arc::clone(&stream));
+                    let (result_tx, result_rx) = instance
+                        .future(|| unreachable!(), &mut view)
+                        .context("failed to create future")?;
                     view.spawn(ReceiveTask {
                         stream,
                         data_tx,
                         result_tx,
                     });
+                    Ok((data_rx.into(), result_rx.into()))
                 }
                 prev => {
                     *tcp_state = prev;
-                    _ = result_tx.write(store, Err(ErrorCode::InvalidState));
+                    let (_, result_rx) = instance
+                        .future(|| Err(ErrorCode::InvalidState), &mut view)
+                        .context("failed to create future")?;
+                    Ok((data_rx.into(), result_rx.into()))
                 }
-            };
-            Ok((data_rx.into(), result_rx.into()))
+            }
         })
     }
 }

crates/wasi/src/p3/sockets/host/types/udp.rs

@@ -4,7 +4,7 @@ use anyhow::Context as _;
 use wasmtime::component::{Accessor, Resource, ResourceTable};
 
 use crate::p3::bindings::sockets::types::{
-    ErrorCode, HostUdpSocket, HostUdpSocketConcurrent, IpAddressFamily, IpSocketAddress,
+    ErrorCode, HostUdpSocket, HostUdpSocketWithStore, IpAddressFamily, IpSocketAddress,
 };
 use crate::p3::sockets::WasiSockets;
 use crate::p3::sockets::udp::UdpSocket;
@@ -34,7 +34,7 @@ fn get_socket_mut<'a>(
         .context("failed to get socket resource from table")
 }
 
-impl HostUdpSocketConcurrent for WasiSockets {
+impl HostUdpSocketWithStore for WasiSockets {
     async fn bind<T>(
         store: &Accessor<T, Self>,
         socket: Resource<UdpSocket>,