Merge pull request #240 from Berrysoft/dev/async-stream

feat(io): compat async stream for futures

Author: George-Miao

compio-io/Cargo.toml

@@ -14,6 +14,7 @@ repository = { workspace = true }
 compio-buf = { workspace = true, features = ["arrayvec"] }
 futures-util = { workspace = true }
 paste = { workspace = true }
+pin-project-lite = { version = "0.2.14", optional = true }
 
 [dev-dependencies]
 compio-runtime = { workspace = true }
@@ -22,7 +23,7 @@ tokio = { workspace = true, features = ["macros", "rt"] }
 
 [features]
 default = []
-compat = []
+compat = ["futures-util/io", "dep:pin-project-lite"]
 
 # Nightly features
 allocator_api = ["compio-buf/allocator_api"]

compio-io/src/buffer.rs

@@ -162,11 +162,20 @@ impl Buffer {
     ///
     /// https://github.com/compio-rs/compio/issues/209
     pub async fn flush_to(&mut self, writer: &mut impl AsyncWrite) -> IoResult<usize> {
+        if self.slice().is_empty() {
+            return Ok(0);
+        }
         let mut total = 0;
         loop {
             let written = self
                 .with(|inner| async { writer.write(inner.into_slice()).await.into_inner() })
                 .await?;
+            if written == 0 {
+                return Err(std::io::Error::new(
+                    std::io::ErrorKind::UnexpectedEof,
+                    "cannot flush all buffer data",
+                ));
+            }
             total += written;
             if self.advance(written) {
                 break;

compio-io/src/compat.rs

@@ -1,8 +1,14 @@
 //! Compat wrappers for interop with other crates.
 
-use std::io::{self, BufRead, Read, Write};
+use std::{
+    future::Future,
+    io::{self, BufRead, Read, Write},
+    pin::Pin,
+    task::{Context, Poll},
+};
 
 use compio_buf::{BufResult, IntoInner, IoBuf, IoBufMut, SetBufInit};
+use pin_project_lite::pin_project;
 
 use crate::{buffer::Buffer, util::DEFAULT_BUF_SIZE};
 
@@ -158,3 +164,181 @@ impl<S: crate::AsyncWrite> SyncStream<S> {
         Ok(len)
     }
 }
+
+type PinBoxFuture<T> = Pin<Box<dyn Future<Output = T>>>;
+
+pin_project! {
+    /// A stream wrapper for [`futures_util::io`] traits.
+    pub struct AsyncStream<S> {
+        #[pin]
+        inner: SyncStream<S>,
+        read_future: Option<PinBoxFuture<io::Result<usize>>>,
+        write_future: Option<PinBoxFuture<io::Result<usize>>>,
+        shutdown_future: Option<PinBoxFuture<io::Result<()>>>,
+    }
+}
+
+impl<S> AsyncStream<S> {
+    /// Create [`AsyncStream`] with the stream and default buffer size.
+    pub fn new(stream: S) -> Self {
+        Self::new_impl(SyncStream::new(stream))
+    }
+
+    /// Create [`AsyncStream`] with the stream and buffer size.
+    pub fn with_capacity(cap: usize, stream: S) -> Self {
+        Self::new_impl(SyncStream::with_capacity(cap, stream))
+    }
+
+    fn new_impl(inner: SyncStream<S>) -> Self {
+        Self {
+            inner,
+            read_future: None,
+            write_future: None,
+            shutdown_future: None,
+        }
+    }
+
+    /// Get the reference of the inner stream.
+    pub fn get_ref(&self) -> &S {
+        self.inner.get_ref()
+    }
+}
+
+macro_rules! poll_future {
+    ($f:expr, $cx:expr, $e:expr) => {{
+        let mut future = match $f.take() {
+            Some(f) => f,
+            None => Box::pin($e),
+        };
+        let f = future.as_mut();
+        match f.poll($cx) {
+            Poll::Pending => {
+                $f.replace(future);
+                return Poll::Pending;
+            }
+            Poll::Ready(res) => res,
+        }
+    }};
+}
+
+macro_rules! poll_future_would_block {
+    ($f:expr, $cx:expr, $e:expr, $io:expr) => {{
+        if let Some(mut f) = $f.take() {
+            if f.as_mut().poll($cx).is_pending() {
+                $f.replace(f);
+                return Poll::Pending;
+            }
+        }
+
+        match $io {
+            Ok(len) => Poll::Ready(Ok(len)),
+            Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
+                $f.replace(Box::pin($e));
+                $cx.waker().wake_by_ref();
+                Poll::Pending
+            }
+            Err(e) => Poll::Ready(Err(e)),
+        }
+    }};
+}
+
+impl<S: crate::AsyncRead + 'static> futures_util::AsyncRead for AsyncStream<S> {
+    fn poll_read(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut [u8],
+    ) -> Poll<io::Result<usize>> {
+        let this = self.project();
+        // Safety:
+        // - The futures won't live longer than the stream.
+        // - `self` is pinned.
+        // - The inner stream won't be moved.
+        let inner: &'static mut SyncStream<S> =
+            unsafe { &mut *(this.inner.get_unchecked_mut() as *mut _) };
+
+        poll_future_would_block!(
+            this.read_future,
+            cx,
+            inner.fill_read_buf(),
+            io::Read::read(inner, buf)
+        )
+    }
+}
+
+impl<S: crate::AsyncRead + 'static> futures_util::AsyncBufRead for AsyncStream<S> {
+    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
+        let this = self.project();
+
+        let inner: &'static mut SyncStream<S> =
+            unsafe { &mut *(this.inner.get_unchecked_mut() as *mut _) };
+        poll_future_would_block!(
+            this.read_future,
+            cx,
+            inner.fill_read_buf(),
+            // Safety: anyway the slice won't be used after free.
+            io::BufRead::fill_buf(inner).map(|slice| unsafe { &*(slice as *const _) })
+        )
+    }
+
+    fn consume(self: Pin<&mut Self>, amt: usize) {
+        let this = self.project();
+
+        let inner: &'static mut SyncStream<S> =
+            unsafe { &mut *(this.inner.get_unchecked_mut() as *mut _) };
+        inner.consume(amt)
+    }
+}
+
+impl<S: crate::AsyncWrite + 'static> futures_util::AsyncWrite for AsyncStream<S> {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<io::Result<usize>> {
+        let this = self.project();
+
+        if this.shutdown_future.is_some() {
+            debug_assert!(this.write_future.is_none());
+            return Poll::Pending;
+        }
+
+        let inner: &'static mut SyncStream<S> =
+            unsafe { &mut *(this.inner.get_unchecked_mut() as *mut _) };
+        poll_future_would_block!(
+            this.write_future,
+            cx,
+            inner.flush_write_buf(),
+            io::Write::write(inner, buf)
+        )
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        let this = self.project();
+
+        if this.shutdown_future.is_some() {
+            debug_assert!(this.write_future.is_none());
+            return Poll::Pending;
+        }
+
+        let inner: &'static mut SyncStream<S> =
+            unsafe { &mut *(this.inner.get_unchecked_mut() as *mut _) };
+        let res = poll_future!(this.write_future, cx, inner.flush_write_buf());
+        Poll::Ready(res.map(|_| ()))
+    }
+
+    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        let this = self.project();
+
+        // Avoid shutdown on flush because the inner buffer might be passed to the
+        // driver.
+        if this.write_future.is_some() {
+            debug_assert!(this.shutdown_future.is_none());
+            return Poll::Pending;
+        }
+
+        let inner: &'static mut SyncStream<S> =
+            unsafe { &mut *(this.inner.get_unchecked_mut() as *mut _) };
+        let res = poll_future!(this.shutdown_future, cx, inner.get_mut().shutdown());
+        Poll::Ready(res)
+    }
+}

compio-io/tests/compat.rs

@@ -0,0 +1,74 @@
+use std::io::Cursor;
+
+use compio_io::compat::AsyncStream;
+use futures_util::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt};
+
+#[tokio::test]
+async fn async_compat_read() {
+    let src = &[1u8, 1, 4, 5, 1, 4, 1, 9, 1, 9, 8, 1, 0][..];
+    let mut stream = AsyncStream::new(src);
+
+    let mut buf = [0; 6];
+    let len = stream.read(&mut buf).await.unwrap();
+
+    assert_eq!(len, 6);
+    assert_eq!(buf, [1, 1, 4, 5, 1, 4]);
+
+    let mut buf = [0; 20];
+    let len = stream.read(&mut buf).await.unwrap();
+    assert_eq!(len, 7);
+    assert_eq!(&buf[..7], [1, 9, 1, 9, 8, 1, 0]);
+}
+
+#[tokio::test]
+async fn async_compat_bufread() {
+    let src = &[1u8, 1, 4, 5, 1, 4, 1, 9, 1, 9, 8, 1, 0][..];
+    let mut stream = AsyncStream::new(src);
+
+    let slice = stream.fill_buf().await.unwrap();
+    assert_eq!(slice, [1, 1, 4, 5, 1, 4, 1, 9, 1, 9, 8, 1, 0]);
+    stream.consume_unpin(6);
+
+    let mut buf = [0; 7];
+    let len = stream.read(&mut buf).await.unwrap();
+
+    assert_eq!(len, 7);
+    assert_eq!(buf, [1, 9, 1, 9, 8, 1, 0]);
+}
+
+#[tokio::test]
+async fn async_compat_write() {
+    let dst = Cursor::new([0u8; 10]);
+    let mut stream = AsyncStream::new(dst);
+
+    let len = stream.write(&[1, 1, 4, 5, 1, 4]).await.unwrap();
+    stream.flush().await.unwrap();
+
+    assert_eq!(len, 6);
+    assert_eq!(stream.get_ref().position(), 6);
+    assert_eq!(stream.get_ref().get_ref(), &[1, 1, 4, 5, 1, 4, 0, 0, 0, 0]);
+
+    let dst = Cursor::new([0u8; 10]);
+    let mut stream = AsyncStream::with_capacity(10, dst);
+    let len = stream
+        .write(&[1, 1, 4, 5, 1, 4, 1, 9, 1, 9, 8, 1, 0])
+        .await
+        .unwrap();
+    assert_eq!(len, 10);
+
+    stream.flush().await.unwrap();
+    assert_eq!(stream.get_ref().get_ref(), &[1, 1, 4, 5, 1, 4, 1, 9, 1, 9]);
+}
+
+#[tokio::test]
+async fn async_compat_flush_fail() {
+    let dst = Cursor::new([0u8; 10]);
+    let mut stream = AsyncStream::new(dst);
+    let len = stream
+        .write(&[1, 1, 4, 5, 1, 4, 1, 9, 1, 9, 8, 1, 0])
+        .await
+        .unwrap();
+    assert_eq!(len, 13);
+    let err = stream.flush().await.unwrap_err();
+    assert_eq!(err.kind(), std::io::ErrorKind::UnexpectedEof);
+}