Optimize copy.

Optimize `wstd::io::copy` using the wasi-io `splice` function. To do
this, add a way to query `AsyncRead` and `AsyncWrite` implementations
for an underlying `AsyncInputStream` and `AsyncOutputStream`.

This is similar to optimizations done in `std::io::copy`, however
those optimizations make use of specialization, which isn't available
on stable Rust. That said, using explicit interfaces instead of
specialization does have the advantage that user-defined types can
opt into optimization if they're able to.

This also eliminates the `RefCell`s in `TcpStream`.

Author: sunfishcode

src/http/body.rs

@@ -139,6 +139,10 @@ impl AsyncRead for IncomingBody {
     async fn read(&mut self, out_buf: &mut [u8]) -> crate::io::Result<usize> {
         self.body_stream.read(out_buf).await
     }
+
+    fn as_async_input_stream(&self) -> Option<&AsyncInputStream> {
+        Some(&self.body_stream)
+    }
 }
 
 impl Body for IncomingBody {

src/io/copy.rs

@@ -1,11 +1,29 @@
-use crate::io::{AsyncRead, AsyncWrite};
+use crate::io::{AsyncRead, AsyncWrite, Error};
+use wasi::io::streams::StreamError;
 
 /// Copy bytes from a reader to a writer.
 pub async fn copy<R, W>(mut reader: R, mut writer: W) -> crate::io::Result<()>
 where
     R: AsyncRead,
     W: AsyncWrite,
 {
+    // Optimized path when we have an `AsyncInputStream` and an
+    // `AsyncOutputStream`.
+    if let Some(reader) = reader.as_async_input_stream() {
+        if let Some(writer) = writer.as_async_output_stream() {
+            loop {
+                match super::splice(reader, writer, u64::MAX).await {
+                    Ok(_n) => (),
+                    Err(StreamError::Closed) => return Ok(()),
+                    Err(StreamError::LastOperationFailed(err)) => {
+                        return Err(Error::other(err.to_debug_string()));
+                    }
+                }
+            }
+        }
+    }
+
+    // Unoptimized case: read the input and then write it.
     let mut buf = [0; 1024];
     'read: loop {
         let bytes_read = reader.read(&mut buf).await?;

src/io/read.rs

@@ -24,16 +24,28 @@ pub trait AsyncRead {
             n += len;
         }
     }
+
+    // If the `AsyncRead` implementation is an unbuffered wrapper around an
+    // `AsyncInputStream`, some I/O operations can be more efficient.
+    #[inline]
+    fn as_async_input_stream(&self) -> Option<&io::AsyncInputStream> {
+        None
+    }
 }
 
 impl<R: AsyncRead + ?Sized> AsyncRead for &mut R {
     #[inline]
     async fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
-        (*self).read(buf).await
+        (**self).read(buf).await
     }
 
     #[inline]
     async fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
-        (*self).read_to_end(buf).await
+        (**self).read_to_end(buf).await
+    }
+
+    #[inline]
+    fn as_async_input_stream(&self) -> Option<&io::AsyncInputStream> {
+        (**self).as_async_input_stream()
     }
 }

src/io/streams.rs

@@ -32,10 +32,8 @@ impl AsyncInputStream {
             .wait_for()
             .await;
     }
-}
-
-impl AsyncRead for AsyncInputStream {
-    async fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
+    /// Like [`AsyncRead::read`], but doesn't require a `&mut self`.
+    pub async fn read(&self, buf: &mut [u8]) -> Result<usize> {
         self.ready().await;
         // Ideally, the ABI would be able to read directly into buf. However, with the default
         // generated bindings, it returns a newly allocated vec, which we need to copy into buf.
@@ -57,6 +55,17 @@ impl AsyncRead for AsyncInputStream {
     }
 }
 
+impl AsyncRead for AsyncInputStream {
+    async fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
+        Self::read(self, buf).await
+    }
+
+    #[inline]
+    fn as_async_input_stream(&self) -> Option<&AsyncInputStream> {
+        Some(self)
+    }
+}
+
 #[derive(Debug)]
 pub struct AsyncOutputStream {
     // Lazily initialized pollable, used for lifetime of stream to check readiness.
@@ -86,10 +95,8 @@ impl AsyncOutputStream {
             .wait_for()
             .await;
     }
-}
-impl AsyncWrite for AsyncOutputStream {
-    // Required methods
-    async fn write(&mut self, buf: &[u8]) -> Result<usize> {
+    /// Like [`AsyncWrite::write`], but doesn't require a `&mut self`.
+    pub async fn write(&self, buf: &[u8]) -> Result<usize> {
         // Loops at most twice.
         loop {
             match self.stream.check_write() {
@@ -119,7 +126,8 @@ impl AsyncWrite for AsyncOutputStream {
             }
         }
     }
-    async fn flush(&mut self) -> Result<()> {
+    /// Like [`AsyncWrite::flush`], but doesn't require a `&mut self`.
+    pub async fn flush(&self) -> Result<()> {
         match self.stream.flush() {
             Ok(()) => {
                 self.ready().await;
@@ -134,3 +142,31 @@ impl AsyncWrite for AsyncOutputStream {
         }
     }
 }
+impl AsyncWrite for AsyncOutputStream {
+    // Required methods
+    async fn write(&mut self, buf: &[u8]) -> Result<usize> {
+        Self::write(self, buf).await
+    }
+    async fn flush(&mut self) -> Result<()> {
+        Self::flush(self).await
+    }
+
+    #[inline]
+    fn as_async_output_stream(&self) -> Option<&AsyncOutputStream> {
+        Some(self)
+    }
+}
+
+/// Wait for both streams to be ready and then do a WASI splice.
+pub(crate) async fn splice(
+    reader: &AsyncInputStream,
+    writer: &AsyncOutputStream,
+    len: u64,
+) -> core::result::Result<u64, StreamError> {
+    // Wait for both streams to be ready.
+    let r = reader.ready();
+    writer.ready().await;
+    r.await;
+
+    writer.stream.splice(&reader.stream, len)
+}

src/io/write.rs

@@ -16,21 +16,33 @@ pub trait AsyncWrite {
             }
         }
     }
+
+    // If the `AsyncWrite` implementation is an unbuffered wrapper around an
+    // `AsyncOutputStream`, some I/O operations can be more efficient.
+    #[inline]
+    fn as_async_output_stream(&self) -> Option<&io::AsyncOutputStream> {
+        None
+    }
 }
 
 impl<W: AsyncWrite + ?Sized> AsyncWrite for &mut W {
     #[inline]
     async fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-        (*self).write(buf).await
+        (**self).write(buf).await
     }
 
     #[inline]
     async fn flush(&mut self) -> io::Result<()> {
-        (*self).flush().await
+        (**self).flush().await
     }
 
     #[inline]
     async fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
-        (*self).write_all(buf).await
+        (**self).write_all(buf).await
+    }
+
+    #[inline]
+    fn as_async_output_stream(&self) -> Option<&io::AsyncOutputStream> {
+        (**self).as_async_output_stream()
     }
 }

src/net/tcp_stream.rs

@@ -1,24 +1,22 @@
-use std::cell::RefCell;
-
 use wasi::{
     io::streams::{InputStream, OutputStream},
     sockets::tcp::TcpSocket,
 };
 
-use crate::io::{self, AsyncInputStream, AsyncOutputStream, AsyncRead, AsyncWrite};
+use crate::io::{self, AsyncInputStream, AsyncOutputStream};
 
 /// A TCP stream between a local and a remote socket.
 pub struct TcpStream {
-    input: RefCell<AsyncInputStream>,
-    output: RefCell<AsyncOutputStream>,
+    input: AsyncInputStream,
+    output: 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)),
+            input: AsyncInputStream::new(input),
+            output: AsyncOutputStream::new(output),
             socket,
         }
     }
@@ -42,43 +40,63 @@ impl Drop for TcpStream {
     }
 }
 
-impl AsyncRead for TcpStream {
+impl io::AsyncRead for TcpStream {
     async fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
-        self.input.borrow_mut().read(buf).await
+        self.input.read(buf).await
+    }
+
+    fn as_async_input_stream(&self) -> Option<&AsyncInputStream> {
+        Some(&self.input)
     }
 }
 
-impl AsyncRead for &TcpStream {
+impl io::AsyncRead for &TcpStream {
     async fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
-        self.input.borrow_mut().read(buf).await
+        self.input.read(buf).await
+    }
+
+    fn as_async_input_stream(&self) -> Option<&AsyncInputStream> {
+        (**self).as_async_input_stream()
     }
 }
 
-impl AsyncWrite for TcpStream {
+impl io::AsyncWrite for TcpStream {
     async fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-        self.output.borrow_mut().write(buf).await
+        self.output.write(buf).await
     }
 
     async fn flush(&mut self) -> io::Result<()> {
-        self.output.borrow_mut().flush().await
+        self.output.flush().await
+    }
+
+    fn as_async_output_stream(&self) -> Option<&AsyncOutputStream> {
+        Some(&self.output)
     }
 }
 
-impl AsyncWrite for &TcpStream {
+impl io::AsyncWrite for &TcpStream {
     async fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-        self.output.borrow_mut().write(buf).await
+        self.output.write(buf).await
     }
 
     async fn flush(&mut self) -> io::Result<()> {
-        self.output.borrow_mut().flush().await
+        self.output.flush().await
+    }
+
+    fn as_async_output_stream(&self) -> Option<&AsyncOutputStream> {
+        (**self).as_async_output_stream()
     }
 }
 
 pub struct ReadHalf<'a>(&'a TcpStream);
-impl<'a> AsyncRead for ReadHalf<'a> {
+impl<'a> io::AsyncRead for ReadHalf<'a> {
     async fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
         self.0.read(buf).await
     }
+
+    fn as_async_input_stream(&self) -> Option<&AsyncInputStream> {
+        self.0.as_async_input_stream()
+    }
 }
 
 impl<'a> Drop for ReadHalf<'a> {
@@ -91,14 +109,18 @@ impl<'a> Drop for ReadHalf<'a> {
 }
 
 pub struct WriteHalf<'a>(&'a TcpStream);
-impl<'a> AsyncWrite for WriteHalf<'a> {
+impl<'a> io::AsyncWrite for WriteHalf<'a> {
     async fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
         self.0.write(buf).await
     }
 
     async fn flush(&mut self) -> io::Result<()> {
         self.0.flush().await
     }
+
+    fn as_async_output_stream(&self) -> Option<&AsyncOutputStream> {
+        self.0.as_async_output_stream()
+    }
 }
 
 impl<'a> Drop for WriteHalf<'a> {