feat(io): return () in write_all

Berrysoft · Berrysoft · commit 8f7ea0ca50f2 · 2024-02-04T13:57:29.000+08:00
diff --git a/compio-io/src/compat.rs b/compio-io/src/compat.rs
@@ -153,7 +153,14 @@ impl<S: crate::AsyncWrite> SyncStream<S> {
     /// Flush all data in the write buffer.
     pub async fn flush_write_buf(&mut self) -> io::Result<usize> {
         let stream = &mut self.stream;
-        let len = self.write_buffer.with(|b| stream.write_all(b)).await?;
+        let len = self
+            .write_buffer
+            .with(|w| async {
+                let len = w.buf_len();
+                let BufResult(res, w) = stream.write_all(w).await;
+                BufResult(res.map(|()| len), w)
+            })
+            .await?;
         self.write_buffer.reset();
         stream.flush().await?;
         Ok(len)
diff --git a/compio-io/src/util/mod.rs b/compio-io/src/util/mod.rs
@@ -1,6 +1,7 @@
 //! IO related utilities functions for ease of use.
 
 mod take;
+use compio_buf::{BufResult, IoBuf};
 pub use take::Take;
 
 mod null;
@@ -39,13 +40,13 @@ pub async fn copy<'a, R: AsyncRead, W: AsyncWrite>(
 
         // When EOF is reached, we are terminating, so flush before that
         if read == 0 || buf.need_flush() {
-            let written = buf.with(|w| writer.write_all(w)).await?;
-            if written == 0 {
-                return Err(std::io::Error::new(
-                    std::io::ErrorKind::WriteZero,
-                    "0 byte was written into the writer",
-                ));
-            }
+            let written = buf
+                .with(|w| async {
+                    let len = w.buf_len();
+                    let BufResult(res, w) = writer.write_all(w).await;
+                    BufResult(res.map(|()| len), w)
+                })
+                .await?;
             total += written;
 
             if buf.advance(written) {
diff --git a/compio-io/src/write/ext.rs b/compio-io/src/write/ext.rs
@@ -11,23 +11,21 @@ macro_rules! write_scalar {
                 use ::compio_buf::{arrayvec::ArrayVec, BufResult};
 
                 const LEN: usize = ::std::mem::size_of::<$t>();
-                let BufResult(len, _) = self
+                let BufResult(res, _) = self
                     .write_all(ArrayVec::<u8, LEN>::from(num.$be()))
                     .await;
-                assert_eq!(len?, LEN, "`write_all` returned unexpected length");
-                Ok(())
+                res
             }
 
             #[doc = concat!("Write a little endian `", stringify!($t), "` into the underlying writer.")]
             async fn [< write_ $t _le >](&mut self, num: $t) -> IoResult<()> {
                 use ::compio_buf::{arrayvec::ArrayVec, BufResult};
 
                 const LEN: usize = ::std::mem::size_of::<$t>();
-                let BufResult(len, _) = self
+                let BufResult(res, _) = self
                     .write_all(ArrayVec::<u8, LEN>::from(num.$le()))
                     .await;
-                assert_eq!(len?, LEN, "`write_all` returned unexpected length");
-                Ok(())
+                res
             }
         }
     };
@@ -57,24 +55,42 @@ macro_rules! loop_write_all {
                 BufResult(Err(ref e), buf) if e.kind() == ::std::io::ErrorKind::Interrupted => {
                     $buf = buf;
                 }
-                res => return res,
+                BufResult(Err(e), buf) => return BufResult(Err(e), buf),
             }
         }
 
-        return BufResult(Ok($needle), $buf);
+        return BufResult(Ok(()), $buf);
     };
 }
 
 macro_rules! loop_write_vectored {
-    (
-        $buf:ident,
-        $tracker:ident :
-        $tracker_ty:ty,
-        $iter:ident,loop
-        $read_expr:expr
-    ) => {
-        loop_write_vectored!($buf, $tracker: $tracker_ty, res, $iter, loop $read_expr, break None)
-    };
+    ($buf:ident, $tracker:ident : $tracker_ty:ty, $iter:ident,loop $read_expr:expr) => {{
+        let mut $iter = match $buf.owned_iter() {
+            Ok(buf) => buf,
+            Err(buf) => return BufResult(Ok(()), buf),
+        };
+        let mut $tracker: $tracker_ty = 0;
+
+        loop {
+            let len = $iter.buf_len();
+            if len == 0 {
+                continue;
+            }
+
+            match $read_expr.await {
+                BufResult(Ok(()), ret) => {
+                    $iter = ret;
+                    $tracker += len as $tracker_ty;
+                }
+                BufResult(Err(e), $iter) => return BufResult(Err(e), $iter.into_inner()),
+            };
+
+            match $iter.next() {
+                Ok(next) => $iter = next,
+                Err(buf) => return BufResult(Ok(()), buf),
+            }
+        }
+    }};
     (
         $buf:ident,
         $tracker:ident :
@@ -130,7 +146,7 @@ pub trait AsyncWriteExt: AsyncWrite {
     }
 
     /// Write the entire contents of a buffer into this writer.
-    async fn write_all<T: IoBuf>(&mut self, mut buf: T) -> BufResult<usize, T> {
+    async fn write_all<T: IoBuf>(&mut self, mut buf: T) -> BufResult<(), T> {
         loop_write_all!(
             buf,
             buf.buf_len(),
@@ -142,8 +158,8 @@ pub trait AsyncWriteExt: AsyncWrite {
     /// Write the entire contents of a buffer into this writer. Like
     /// [`AsyncWrite::write_vectored`], except that it tries to write the entire
     /// contents of the buffer into this writer.
-    async fn write_vectored_all<T: IoVectoredBuf>(&mut self, buf: T) -> BufResult<usize, T> {
-        loop_write_vectored!(buf, total: usize, iter, loop self.write_all(iter))
+    async fn write_vectored_all<T: IoVectoredBuf>(&mut self, buf: T) -> BufResult<(), T> {
+        loop_write_vectored!(buf, _total: usize, iter, loop self.write_all(iter))
     }
 
     write_scalar!(u8, to_be_bytes, to_le_bytes);
@@ -168,7 +184,7 @@ impl<A: AsyncWrite + ?Sized> AsyncWriteExt for A {}
 pub trait AsyncWriteAtExt: AsyncWriteAt {
     /// Like [`AsyncWriteAt::write_at`], except that it tries to write the
     /// entire contents of the buffer into this writer.
-    async fn write_all_at<T: IoBuf>(&mut self, mut buf: T, pos: u64) -> BufResult<usize, T> {
+    async fn write_all_at<T: IoBuf>(&mut self, mut buf: T, pos: u64) -> BufResult<(), T> {
         loop_write_all!(
             buf,
             buf.buf_len(),
@@ -183,7 +199,7 @@ pub trait AsyncWriteAtExt: AsyncWriteAt {
         &mut self,
         buf: T,
         pos: u64,
-    ) -> BufResult<usize, T> {
+    ) -> BufResult<(), T> {
         loop_write_vectored!(buf, total: u64, iter, loop self.write_all_at(iter, pos + total))
     }
 }
diff --git a/compio-io/tests/io.rs b/compio-io/tests/io.rs
@@ -280,12 +280,10 @@ impl AsyncWriteAt for WriteOne {
 async fn write_all() {
     let mut dst = WriteOne(vec![]);
 
-    let (len, _) = dst.write_all([1, 1, 4, 5, 1, 4]).await.unwrap();
-    assert_eq!(len, 6);
+    let ((), _) = dst.write_all([1, 1, 4, 5, 1, 4]).await.unwrap();
     assert_eq!(dst.0, [1, 1, 4, 5, 1, 4]);
 
-    let (len, _) = dst.write_all_at([114, 114, 114], 2).await.unwrap();
-    assert_eq!(len, 3);
+    let ((), _) = dst.write_all_at([114, 114, 114], 2).await.unwrap();
     assert_eq!(dst.0, [1, 1, 114, 114, 114, 4]);
 }
 
diff --git a/compio-net/tests/unix_stream.rs b/compio-net/tests/unix_stream.rs
@@ -14,8 +14,7 @@ async fn accept_read_write() -> std::io::Result<()> {
     let mut client = UnixStream::connect(&sock_path)?;
     let (mut server, _) = listener.accept().await?;
 
-    let write_len = client.write_all("hello").await.0?;
-    assert_eq!(write_len, 5);
+    client.write_all("hello").await.0?;
     drop(client);
 
     let buf = Vec::with_capacity(5);
