net: TcpListener uses an AsyncPollable, and TcpStream uses AsyncInput/OutputStream

Author: pchickey

src/net/tcp_listener.rs

@@ -3,15 +3,17 @@ use wasi::sockets::tcp::{ErrorCode, IpAddressFamily, IpSocketAddress, TcpSocket}
 
 use crate::io;
 use crate::iter::AsyncIterator;
-use crate::runtime::Reactor;
 use std::io::ErrorKind;
 use std::net::SocketAddr;
 
 use super::TcpStream;
+use crate::runtime::AsyncPollable;
 
 /// A TCP socket server, listening for connections.
 #[derive(Debug)]
 pub struct TcpListener {
+    // Field order matters: must drop this child before parent below
+    pollable: AsyncPollable,
     socket: TcpSocket,
 }
 
@@ -40,18 +42,17 @@ impl TcpListener {
             }
             SocketAddr::V6(_) => todo!("IPv6 not yet supported in `wstd::net::TcpListener`"),
         };
-        let reactor = Reactor::current();
-
         socket
             .start_bind(&network, local_address)
             .map_err(to_io_err)?;
-        reactor.wait_for(socket.subscribe()).await;
+        let pollable = AsyncPollable::new(socket.subscribe());
+        pollable.wait_for().await;
         socket.finish_bind().map_err(to_io_err)?;
 
         socket.start_listen().map_err(to_io_err)?;
-        reactor.wait_for(socket.subscribe()).await;
+        pollable.wait_for().await;
         socket.finish_listen().map_err(to_io_err)?;
-        Ok(Self { socket })
+        Ok(Self { pollable, socket })
     }
 
     /// Returns the local socket address of this listener.
@@ -77,18 +78,12 @@ impl<'a> AsyncIterator for Incoming<'a> {
     type Item = io::Result<TcpStream>;
 
     async fn next(&mut self) -> Option<Self::Item> {
-        Reactor::current()
-            .wait_for(self.listener.socket.subscribe())
-            .await;
+        self.listener.pollable.wait_for().await;
         let (socket, input, output) = match self.listener.socket.accept().map_err(to_io_err) {
             Ok(accepted) => accepted,
             Err(err) => return Some(Err(err)),
         };
-        Some(Ok(TcpStream {
-            socket,
-            input,
-            output,
-        }))
+        Some(Ok(TcpStream::new(input, output, socket)))
     }
 }
 

src/net/tcp_stream.rs

@@ -1,23 +1,27 @@
-use std::io::Error;
+use std::cell::RefCell;
 
 use wasi::{
-    io::streams::StreamError,
-    sockets::tcp::{InputStream, OutputStream, TcpSocket},
+    io::streams::{InputStream, OutputStream},
+    sockets::tcp::TcpSocket,
 };
 
-use crate::{
-    io::{self, AsyncRead, AsyncWrite},
-    runtime::Reactor,
-};
+use crate::io::{self, AsyncInputStream, AsyncOutputStream, AsyncRead, AsyncWrite};
 
 /// A TCP stream between a local and a remote socket.
 pub struct TcpStream {
-    pub(super) input: InputStream,
-    pub(super) output: OutputStream,
-    pub(super) socket: TcpSocket,
+    input: RefCell<AsyncInputStream>,
+    output: RefCell<AsyncOutputStream>,
+    socket: TcpSocket,
 }
 
 impl TcpStream {
+    pub(crate) fn new(input: InputStream, output: OutputStream, socket: TcpSocket) -> Self {
+        TcpStream {
+            input: RefCell::new(AsyncInputStream::new(input)),
+            output: RefCell::new(AsyncOutputStream::new(output)),
+            socket,
+        }
+    }
     /// Returns the socket address of the remote peer of this TCP connection.
     pub fn peer_addr(&self) -> io::Result<String> {
         let addr = self
@@ -40,53 +44,33 @@ impl Drop for TcpStream {
 
 impl AsyncRead for TcpStream {
     async fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
-        Reactor::current().wait_for(self.input.subscribe()).await;
-        let slice = match self.input.read(buf.len() as u64) {
-            Ok(slice) => slice,
-            Err(StreamError::Closed) => return Ok(0),
-            Err(e) => return Err(to_io_err(e)),
-        };
-        let bytes_read = slice.len();
-        buf[..bytes_read].clone_from_slice(&slice);
-        Ok(bytes_read)
+        self.input.borrow_mut().read(buf).await
     }
 }
 
 impl AsyncRead for &TcpStream {
     async fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
-        Reactor::current().wait_for(self.input.subscribe()).await;
-        let slice = match self.input.read(buf.len() as u64) {
-            Ok(slice) => slice,
-            Err(StreamError::Closed) => return Ok(0),
-            Err(e) => return Err(to_io_err(e)),
-        };
-        let bytes_read = slice.len();
-        buf[..bytes_read].clone_from_slice(&slice);
-        Ok(bytes_read)
+        self.input.borrow_mut().read(buf).await
     }
 }
 
 impl AsyncWrite for TcpStream {
     async fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-        Reactor::current().wait_for(self.output.subscribe()).await;
-        self.output.write(buf).map_err(to_io_err)?;
-        Ok(buf.len())
+        self.output.borrow_mut().write(buf).await
     }
 
     async fn flush(&mut self) -> io::Result<()> {
-        self.output.flush().map_err(to_io_err)
+        self.output.borrow_mut().flush().await
     }
 }
 
 impl AsyncWrite for &TcpStream {
     async fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-        Reactor::current().wait_for(self.output.subscribe()).await;
-        self.output.write(buf).map_err(to_io_err)?;
-        Ok(buf.len())
+        self.output.borrow_mut().write(buf).await
     }
 
     async fn flush(&mut self) -> io::Result<()> {
-        self.output.flush().map_err(to_io_err)
+        self.output.borrow_mut().flush().await
     }
 }
 
@@ -125,10 +109,3 @@ impl<'a> Drop for WriteHalf<'a> {
             .shutdown(wasi::sockets::tcp::ShutdownType::Send);
     }
 }
-
-fn to_io_err(err: StreamError) -> std::io::Error {
-    match err {
-        StreamError::LastOperationFailed(err) => Error::other(err.to_debug_string()),
-        StreamError::Closed => Error::other("Stream was closed"),
-    }
-}