Refactor: Extract duplicate BufWrite into separate module

src/buf_write.rs

@@ -0,0 +1,39 @@
+use std::{
+    io,
+    pin::Pin,
+    task::{Context, Poll},
+};
+
+macro_rules! impl_async_buf_write {
+    ($AsyncBufWrite:tt) => {
+        pub(crate) trait $AsyncBufWrite {
+            /// Attempt to return an internal buffer to write to, flushing data out to the inner reader if
+            /// it is full.
+            ///
+            /// On success, returns `Poll::Ready(Ok(buf))`.
+            ///
+            /// If the buffer is full and cannot be flushed, the method returns `Poll::Pending` and
+            /// arranges for the current task context (`cx`) to receive a notification when the object
+            /// becomes readable or is closed.
+            fn poll_partial_flush_buf(
+                self: Pin<&mut Self>,
+                cx: &mut Context<'_>,
+            ) -> Poll<io::Result<&mut [u8]>>;
+
+            /// Tells this buffer that `amt` bytes have been written to its buffer, so they should be
+            /// written out to the underlying IO when possible.
+            ///
+            /// This function is a lower-level call. It needs to be paired with the `poll_flush_buf` method to
+            /// function properly. This function does not perform any I/O, it simply informs this object
+            /// that some amount of its buffer, returned from `poll_flush_buf`, has been written to and should
+            /// be sent. As such, this function may do odd things if `poll_flush_buf` isn't
+            /// called before calling it.
+            ///
+            /// The `amt` must be `<=` the number of bytes in the buffer returned by `poll_flush_buf`.
+            fn produce(self: Pin<&mut Self>, amt: usize);
+        }
+    };
+}
+
+impl_async_buf_write!(AsyncBufWriteTokio);
+impl_async_buf_write!(AsyncBufWriteFuturesIo);

src/buf_writer.rs

@@ -0,0 +1,259 @@
+// Originally sourced from `futures_util::io::buf_writer`, needs to be redefined locally so that
+// the `AsyncBufWrite` impl can access its internals, and changed a bit to make it more efficient
+// with those methods.
+
+use futures_core::ready;
+use pin_project_lite::pin_project;
+use std::{
+    fmt, io,
+    pin::Pin,
+    task::{Context, Poll},
+};
+
+const DEFAULT_BUF_SIZE: usize = 8192;
+
+pin_project! {
+    pub struct BufWriter<W> {
+        #[pin]
+        inner: W,
+        buf: Box<[u8]>,
+        written: usize,
+        buffered: usize,
+    }
+}
+
+impl<W> BufWriter<W> {
+    /// Creates a new `BufWriter` with a default buffer capacity. The default is currently 8 KB,
+    /// but may change in the future.
+    pub fn new(inner: W) -> Self {
+        Self::with_capacity(DEFAULT_BUF_SIZE, inner)
+    }
+
+    /// Creates a new `BufWriter` with the specified buffer capacity.
+    pub fn with_capacity(cap: usize, inner: W) -> Self {
+        Self {
+            inner,
+            buf: vec![0; cap].into(),
+            written: 0,
+            buffered: 0,
+        }
+    }
+
+    /// Gets a reference to the underlying writer.
+    pub fn get_ref(&self) -> &W {
+        &self.inner
+    }
+
+    /// Gets a mutable reference to the underlying writer.
+    ///
+    /// It is inadvisable to directly write to the underlying writer.
+    pub fn get_mut(&mut self) -> &mut W {
+        &mut self.inner
+    }
+
+    /// Gets a pinned mutable reference to the underlying writer.
+    ///
+    /// It is inadvisable to directly write to the underlying writer.
+    pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut W> {
+        self.project().inner
+    }
+
+    /// Consumes this `BufWriter`, returning the underlying writer.
+    ///
+    /// Note that any leftover data in the internal buffer is lost.
+    pub fn into_inner(self) -> W {
+        self.inner
+    }
+}
+
+impl<W: fmt::Debug> fmt::Debug for BufWriter<W> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("BufWriter")
+            .field("writer", &self.inner)
+            .field(
+                "buffer",
+                &format_args!("{}/{}", self.buffered, self.buf.len()),
+            )
+            .field("written", &self.written)
+            .finish()
+    }
+}
+
+macro_rules! impl_traits {
+    ($partial_flush_buf:tt, $flush_buf:tt, $shutdown_fn:tt) => {
+        impl<W: AsyncWrite> BufWriter<W> {
+            fn $partial_flush_buf(
+                self: Pin<&mut Self>,
+                cx: &mut Context<'_>,
+            ) -> Poll<io::Result<()>> {
+                let mut this = self.project();
+
+                let mut ret = Ok(());
+                while *this.written < *this.buffered {
+                    match this
+                        .inner
+                        .as_mut()
+                        .poll_write(cx, &this.buf[*this.written..*this.buffered])
+                    {
+                        Poll::Pending => {
+                            break;
+                        }
+                        Poll::Ready(Ok(0)) => {
+                            ret = Err(io::Error::new(
+                                io::ErrorKind::WriteZero,
+                                "failed to write the buffered data",
+                            ));
+                            break;
+                        }
+                        Poll::Ready(Ok(n)) => *this.written += n,
+                        Poll::Ready(Err(e)) => {
+                            ret = Err(e);
+                            break;
+                        }
+                    }
+                }
+
+                if *this.written > 0 {
+                    this.buf.copy_within(*this.written..*this.buffered, 0);
+                    *this.buffered -= *this.written;
+                    *this.written = 0;
+
+                    Poll::Ready(ret)
+                } else if *this.buffered == 0 {
+                    Poll::Ready(ret)
+                } else {
+                    ret?;
+                    Poll::Pending
+                }
+            }
+
+            fn $flush_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+                let mut this = self.project();
+
+                let mut ret = Ok(());
+                while *this.written < *this.buffered {
+                    match ready!(this
+                        .inner
+                        .as_mut()
+                        .poll_write(cx, &this.buf[*this.written..*this.buffered]))
+                    {
+                        Ok(0) => {
+                            ret = Err(io::Error::new(
+                                io::ErrorKind::WriteZero,
+                                "failed to write the buffered data",
+                            ));
+                            break;
+                        }
+                        Ok(n) => *this.written += n,
+                        Err(e) => {
+                            ret = Err(e);
+                            break;
+                        }
+                    }
+                }
+                this.buf.copy_within(*this.written..*this.buffered, 0);
+                *this.buffered -= *this.written;
+                *this.written = 0;
+                Poll::Ready(ret)
+            }
+        }
+
+        impl<W: AsyncWrite> AsyncWrite for BufWriter<W> {
+            fn poll_write(
+                mut self: Pin<&mut Self>,
+                cx: &mut Context<'_>,
+                buf: &[u8],
+            ) -> Poll<io::Result<usize>> {
+                let this = self.as_mut().project();
+                if *this.buffered + buf.len() > this.buf.len() {
+                    ready!(self.as_mut().$partial_flush_buf(cx))?;
+                }
+
+                let this = self.as_mut().project();
+                if buf.len() >= this.buf.len() {
+                    if *this.buffered == 0 {
+                        this.inner.poll_write(cx, buf)
+                    } else {
+                        // The only way that `partial_flush_buf` would have returned with
+                        // `this.buffered != 0` is if it were Pending, so our waker was already queued
+                        Poll::Pending
+                    }
+                } else {
+                    let len = buf.len().min(this.buf.len() - *this.buffered);
+                    this.buf[*this.buffered..*this.buffered + len].copy_from_slice(&buf[..len]);
+                    *this.buffered += len;
+                    Poll::Ready(Ok(len))
+                }
+            }
+
+            fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+                ready!(self.as_mut().$flush_buf(cx))?;
+                self.project().inner.poll_flush(cx)
+            }
+
+            fn $shutdown_fn(
+                mut self: Pin<&mut Self>,
+                cx: &mut Context<'_>,
+            ) -> Poll<io::Result<()>> {
+                ready!(self.as_mut().$flush_buf(cx))?;
+                self.project().inner.$shutdown_fn(cx)
+            }
+        }
+
+        impl<W: AsyncWrite> AsyncBufWrite for BufWriter<W> {
+            fn poll_partial_flush_buf(
+                mut self: Pin<&mut Self>,
+                cx: &mut Context<'_>,
+            ) -> Poll<io::Result<&mut [u8]>> {
+                ready!(self.as_mut().$partial_flush_buf(cx))?;
+                let this = self.project();
+                Poll::Ready(Ok(&mut this.buf[*this.buffered..]))
+            }
+
+            fn produce(self: Pin<&mut Self>, amt: usize) {
+                let this = self.project();
+                debug_assert!(
+                    *this.buffered + amt <= this.buf.len(),
+                    "produce called with amt exceeding buffer capacity"
+                );
+                *this.buffered += amt;
+            }
+        }
+    };
+}
+
+#[cfg(feature = "tokio")]
+mod tokio_impl {
+    use super::*;
+    use crate::buf_write::AsyncBufWriteTokio as AsyncBufWrite;
+    use tokio::io::AsyncWrite;
+
+    impl_traits!(partial_flush_buf_tokio, flush_buf_tokio, poll_shutdown);
+}
+
+#[cfg(feature = "futures-io")]
+mod futures_io_impl {
+    use super::*;
+    use crate::buf_write::AsyncBufWriteFuturesIo as AsyncBufWrite;
+    use futures_io::{AsyncSeek, AsyncWrite, SeekFrom};
+
+    impl_traits!(
+        partial_flush_buf_futures_io,
+        flush_buf_futures_io,
+        poll_close
+    );
+
+    impl<W: AsyncWrite + AsyncSeek> AsyncSeek for BufWriter<W> {
+        /// Seek to the offset, in bytes, in the underlying writer.
+        ///
+        /// Seeking always writes out the internal buffer before seeking.
+        fn poll_seek(
+            mut self: Pin<&mut Self>,
+            cx: &mut Context<'_>,
+            pos: SeekFrom,
+        ) -> Poll<io::Result<u64>> {
+            ready!(self.as_mut().flush_buf_futures_io(cx))?;
+            self.project().inner.poll_seek(cx, pos)
+        }
+    }
+}