Merge pull request #210 from Berrysoft/fix/flush-error-safe

fix(io): make buffer flush error safe

Author: Berrysoft

compio-io/src/buffer.rs

@@ -7,9 +7,9 @@ use std::{
     future::Future,
 };
 
-use compio_buf::{BufResult, IoBuf, IoBufMut, SetBufInit};
+use compio_buf::{BufResult, IntoInner, IoBuf, IoBufMut, SetBufInit, Slice};
 
-use crate::util::MISSING_BUF;
+use crate::{util::MISSING_BUF, AsyncWrite, IoResult};
 
 pub struct Inner {
     buf: Vec<u8>,
@@ -33,6 +33,12 @@ impl Inner {
     fn slice(&self) -> &[u8] {
         &self.buf[self.pos..]
     }
+
+    #[inline]
+    pub(crate) fn into_slice(self) -> Slice<Self> {
+        let pos = self.pos;
+        IoBuf::slice(self, pos..)
+    }
 }
 
 impl IoBuf for Inner {
@@ -96,6 +102,7 @@ impl Buffer {
 
     /// If the inner buffer is empty.
     #[inline]
+    #[allow(unused)]
     pub fn is_empty(&self) -> bool {
         self.inner().as_slice().is_empty()
     }
@@ -115,6 +122,7 @@ impl Buffer {
     }
 
     /// The buffer needs to be flushed
+    #[inline]
     pub fn need_flush(&self) -> bool {
         // TODO: Better way to determine if we need to flush the buffer
         let buf = self.buf();
@@ -129,7 +137,7 @@ impl Buffer {
 
     /// Execute a funcition with ownership of the buffer, and restore the buffer
     /// afterwards
-    pub async fn with<R, Fut, F>(&mut self, func: F) -> std::io::Result<R>
+    pub async fn with<R, Fut, F>(&mut self, func: F) -> IoResult<R>
     where
         Fut: Future<Output = BufResult<R, Inner>>,
         F: FnOnce(Inner) -> Fut,
@@ -150,6 +158,24 @@ impl Buffer {
         res
     }
 
+    /// Wrapper to flush the buffer to a writer with error safety.
+    ///
+    /// https://github.com/compio-rs/compio/issues/209
+    pub async fn flush_to(&mut self, writer: &mut impl AsyncWrite) -> IoResult<usize> {
+        let mut total = 0;
+        loop {
+            let written = self
+                .with(|inner| async { writer.write(inner.into_slice()).await.into_inner() })
+                .await?;
+            total += written;
+            if self.advance(written) {
+                break;
+            }
+        }
+        self.reset();
+        Ok(total)
+    }
+
     /// Mark some bytes as read by advancing the progress tracker, return a
     /// `bool` indicating if all bytes are read.
     #[inline]

compio-io/src/compat.rs

@@ -4,7 +4,7 @@ use std::io::{self, BufRead, Read, Write};
 
 use compio_buf::{BufResult, IntoInner, IoBuf, IoBufMut, SetBufInit};
 
-use crate::{buffer::Buffer, util::DEFAULT_BUF_SIZE, AsyncWriteExt};
+use crate::{buffer::Buffer, util::DEFAULT_BUF_SIZE};
 
 /// A wrapper for [`AsyncRead`](crate::AsyncRead) +
 /// [`AsyncWrite`](crate::AsyncWrite), providing sync traits impl. The sync
@@ -153,15 +153,7 @@ 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(|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();
+        let len = self.write_buffer.flush_to(stream).await?;
         stream.flush().await?;
         Ok(len)
     }

compio-io/src/util/mod.rs

@@ -1,7 +1,6 @@
 //! IO related utilities functions for ease of use.
 
 mod take;
-use compio_buf::{BufResult, IoBuf};
 pub use take::Take;
 
 mod null;
@@ -13,7 +12,7 @@ pub use repeat::{repeat, Repeat};
 mod internal;
 pub(crate) use internal::*;
 
-use crate::{buffer::Buffer, AsyncRead, AsyncWrite, AsyncWriteExt, IoResult};
+use crate::{buffer::Buffer, AsyncRead, AsyncWrite, IoResult};
 
 /// Asynchronously copies the entire contents of a reader into a writer.
 ///
@@ -40,18 +39,8 @@ 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| async {
-                    let len = w.buf_len();
-                    let BufResult(res, w) = writer.write_all(w).await;
-                    BufResult(res.map(|()| len), w)
-                })
-                .await?;
+            let written = buf.flush_to(writer).await?;
             total += written;
-
-            if buf.advance(written) {
-                buf.reset()
-            }
         }
 
         if read == 0 {

compio-io/src/write/buf.rs

@@ -1,11 +1,11 @@
 use std::future::ready;
 
-use compio_buf::{buf_try, BufResult, IntoInner, IoBuf};
+use compio_buf::{buf_try, BufResult, IntoInner, IoBuf, IoVectoredBuf};
 
 use crate::{
     buffer::Buffer,
     util::{slice_to_buf, DEFAULT_BUF_SIZE},
-    AsyncWrite, AsyncWriteExt, IoResult,
+    AsyncWrite, IoResult,
 };
 
 /// Wraps a writer and buffers its output.
@@ -48,8 +48,21 @@ impl<W> BufWriter<W> {
     }
 }
 
+impl<W: AsyncWrite> BufWriter<W> {
+    async fn flush_if_needed(&mut self) -> IoResult<()> {
+        if self.buf.need_flush() {
+            self.flush().await?;
+        }
+        Ok(())
+    }
+}
+
 impl<W: AsyncWrite> AsyncWrite for BufWriter<W> {
-    async fn write<T: IoBuf>(&mut self, mut buf: T) -> compio_buf::BufResult<usize, T> {
+    async fn write<T: IoBuf>(&mut self, mut buf: T) -> BufResult<usize, T> {
+        // The previous flush may error because disk full. We need to make the buffer
+        // all-done before writing new data to it.
+        (_, buf) = buf_try!(self.flush_if_needed().await, buf);
+
         let written = self
             .buf
             .with_sync(|w| {
@@ -60,17 +73,14 @@ impl<W: AsyncWrite> AsyncWrite for BufWriter<W> {
             })
             .expect("Closure always return Ok");
 
-        if self.buf.need_flush() {
-            (_, buf) = buf_try!(self.flush().await, buf);
-        }
+        (_, buf) = buf_try!(self.flush_if_needed().await, buf);
 
         BufResult(Ok(written), buf)
     }
 
-    async fn write_vectored<T: compio_buf::IoVectoredBuf>(
-        &mut self,
-        mut buf: T,
-    ) -> compio_buf::BufResult<usize, T> {
+    async fn write_vectored<T: IoVectoredBuf>(&mut self, mut buf: T) -> BufResult<usize, T> {
+        (_, buf) = buf_try!(self.flush_if_needed().await, buf);
+
         let written = self
             .buf
             .with(|mut w| {
@@ -90,18 +100,15 @@ impl<W: AsyncWrite> AsyncWrite for BufWriter<W> {
             .await
             .expect("Closure always return Ok");
 
-        if self.buf.need_flush() {
-            (_, buf) = buf_try!(self.flush().await, buf);
-        }
+        (_, buf) = buf_try!(self.flush_if_needed().await, buf);
 
         BufResult(Ok(written), buf)
     }
 
     async fn flush(&mut self) -> IoResult<()> {
         let Self { writer, buf } = self;
 
-        buf.with(|w| writer.write_all(w)).await?;
-        buf.reset();
+        buf.flush_to(writer).await?;
 
         Ok(())
     }