| use crate::io::util::DEFAULT_BUF_SIZE; |
| use crate::io::{AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite, ReadBuf}; |
| |
| use pin_project_lite::pin_project; |
| use std::fmt; |
| use std::io::{self, IoSlice, SeekFrom, Write}; |
| use std::pin::Pin; |
| use std::task::{Context, Poll}; |
| |
| pin_project! { |
| /// Wraps a writer and buffers its output. |
| /// |
| /// It can be excessively inefficient to work directly with something that |
| /// implements [`AsyncWrite`]. A `BufWriter` keeps an in-memory buffer of data and |
| /// writes it to an underlying writer in large, infrequent batches. |
| /// |
| /// `BufWriter` can improve the speed of programs that make *small* and |
| /// *repeated* write calls to the same file or network socket. It does not |
| /// help when writing very large amounts at once, or writing just one or a few |
| /// times. It also provides no advantage when writing to a destination that is |
| /// in memory, like a `Vec<u8>`. |
| /// |
| /// When the `BufWriter` is dropped, the contents of its buffer will be |
| /// discarded. Creating multiple instances of a `BufWriter` on the same |
| /// stream can cause data loss. If you need to write out the contents of its |
| /// buffer, you must manually call flush before the writer is dropped. |
| /// |
| /// [`AsyncWrite`]: AsyncWrite |
| /// [`flush`]: super::AsyncWriteExt::flush |
| /// |
| #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] |
| pub struct BufWriter<W> { |
| #[pin] |
| pub(super) inner: W, |
| pub(super) buf: Vec<u8>, |
| pub(super) written: usize, |
| pub(super) seek_state: SeekState, |
| } |
| } |
| |
| impl<W: AsyncWrite> 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::with_capacity(cap), |
| written: 0, |
| seek_state: SeekState::Init, |
| } |
| } |
| |
| fn flush_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { |
| let mut me = self.project(); |
| |
| let len = me.buf.len(); |
| let mut ret = Ok(()); |
| while *me.written < len { |
| match ready!(me.inner.as_mut().poll_write(cx, &me.buf[*me.written..])) { |
| Ok(0) => { |
| ret = Err(io::Error::new( |
| io::ErrorKind::WriteZero, |
| "failed to write the buffered data", |
| )); |
| break; |
| } |
| Ok(n) => *me.written += n, |
| Err(e) => { |
| ret = Err(e); |
| break; |
| } |
| } |
| } |
| if *me.written > 0 { |
| me.buf.drain(..*me.written); |
| } |
| *me.written = 0; |
| Poll::Ready(ret) |
| } |
| |
| /// 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 |
| } |
| |
| /// Returns a reference to the internally buffered data. |
| pub fn buffer(&self) -> &[u8] { |
| &self.buf |
| } |
| } |
| |
| impl<W: AsyncWrite> AsyncWrite for BufWriter<W> { |
| fn poll_write( |
| mut self: Pin<&mut Self>, |
| cx: &mut Context<'_>, |
| buf: &[u8], |
| ) -> Poll<io::Result<usize>> { |
| if self.buf.len() + buf.len() > self.buf.capacity() { |
| ready!(self.as_mut().flush_buf(cx))?; |
| } |
| |
| let me = self.project(); |
| if buf.len() >= me.buf.capacity() { |
| me.inner.poll_write(cx, buf) |
| } else { |
| Poll::Ready(me.buf.write(buf)) |
| } |
| } |
| |
| fn poll_write_vectored( |
| mut self: Pin<&mut Self>, |
| cx: &mut Context<'_>, |
| mut bufs: &[IoSlice<'_>], |
| ) -> Poll<io::Result<usize>> { |
| if self.inner.is_write_vectored() { |
| let total_len = bufs |
| .iter() |
| .fold(0usize, |acc, b| acc.saturating_add(b.len())); |
| if total_len > self.buf.capacity() - self.buf.len() { |
| ready!(self.as_mut().flush_buf(cx))?; |
| } |
| let me = self.as_mut().project(); |
| if total_len >= me.buf.capacity() { |
| // It's more efficient to pass the slices directly to the |
| // underlying writer than to buffer them. |
| // The case when the total_len calculation saturates at |
| // usize::MAX is also handled here. |
| me.inner.poll_write_vectored(cx, bufs) |
| } else { |
| bufs.iter().for_each(|b| me.buf.extend_from_slice(b)); |
| Poll::Ready(Ok(total_len)) |
| } |
| } else { |
| // Remove empty buffers at the beginning of bufs. |
| while bufs.first().map(|buf| buf.len()) == Some(0) { |
| bufs = &bufs[1..]; |
| } |
| if bufs.is_empty() { |
| return Poll::Ready(Ok(0)); |
| } |
| // Flush if the first buffer doesn't fit. |
| let first_len = bufs[0].len(); |
| if first_len > self.buf.capacity() - self.buf.len() { |
| ready!(self.as_mut().flush_buf(cx))?; |
| debug_assert!(self.buf.is_empty()); |
| } |
| let me = self.as_mut().project(); |
| if first_len >= me.buf.capacity() { |
| // The slice is at least as large as the buffering capacity, |
| // so it's better to write it directly, bypassing the buffer. |
| debug_assert!(me.buf.is_empty()); |
| return me.inner.poll_write(cx, &bufs[0]); |
| } else { |
| me.buf.extend_from_slice(&bufs[0]); |
| bufs = &bufs[1..]; |
| } |
| let mut total_written = first_len; |
| debug_assert!(total_written != 0); |
| // Append the buffers that fit in the internal buffer. |
| for buf in bufs { |
| if buf.len() > me.buf.capacity() - me.buf.len() { |
| break; |
| } else { |
| me.buf.extend_from_slice(buf); |
| total_written += buf.len(); |
| } |
| } |
| Poll::Ready(Ok(total_written)) |
| } |
| } |
| |
| fn is_write_vectored(&self) -> bool { |
| true |
| } |
| |
| fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { |
| ready!(self.as_mut().flush_buf(cx))?; |
| self.get_pin_mut().poll_flush(cx) |
| } |
| |
| fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { |
| ready!(self.as_mut().flush_buf(cx))?; |
| self.get_pin_mut().poll_shutdown(cx) |
| } |
| } |
| |
| #[derive(Debug, Clone, Copy)] |
| pub(super) enum SeekState { |
| /// start_seek has not been called. |
| Init, |
| /// start_seek has been called, but poll_complete has not yet been called. |
| Start(SeekFrom), |
| /// Waiting for completion of poll_complete. |
| Pending, |
| } |
| |
| /// Seek to the offset, in bytes, in the underlying writer. |
| /// |
| /// Seeking always writes out the internal buffer before seeking. |
| impl<W: AsyncWrite + AsyncSeek> AsyncSeek for BufWriter<W> { |
| fn start_seek(self: Pin<&mut Self>, pos: SeekFrom) -> io::Result<()> { |
| // We need to flush the internal buffer before seeking. |
| // It receives a `Context` and returns a `Poll`, so it cannot be called |
| // inside `start_seek`. |
| *self.project().seek_state = SeekState::Start(pos); |
| Ok(()) |
| } |
| |
| fn poll_complete(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> { |
| let pos = match self.seek_state { |
| SeekState::Init => { |
| return self.project().inner.poll_complete(cx); |
| } |
| SeekState::Start(pos) => Some(pos), |
| SeekState::Pending => None, |
| }; |
| |
| // Flush the internal buffer before seeking. |
| ready!(self.as_mut().flush_buf(cx))?; |
| |
| let mut me = self.project(); |
| if let Some(pos) = pos { |
| // Ensure previous seeks have finished before starting a new one |
| ready!(me.inner.as_mut().poll_complete(cx))?; |
| if let Err(e) = me.inner.as_mut().start_seek(pos) { |
| *me.seek_state = SeekState::Init; |
| return Poll::Ready(Err(e)); |
| } |
| } |
| match me.inner.poll_complete(cx) { |
| Poll::Ready(res) => { |
| *me.seek_state = SeekState::Init; |
| Poll::Ready(res) |
| } |
| Poll::Pending => { |
| *me.seek_state = SeekState::Pending; |
| Poll::Pending |
| } |
| } |
| } |
| } |
| |
| impl<W: AsyncWrite + AsyncRead> AsyncRead for BufWriter<W> { |
| fn poll_read( |
| self: Pin<&mut Self>, |
| cx: &mut Context<'_>, |
| buf: &mut ReadBuf<'_>, |
| ) -> Poll<io::Result<()>> { |
| self.get_pin_mut().poll_read(cx, buf) |
| } |
| } |
| |
| impl<W: AsyncWrite + AsyncBufRead> AsyncBufRead for BufWriter<W> { |
| fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> { |
| self.get_pin_mut().poll_fill_buf(cx) |
| } |
| |
| fn consume(self: Pin<&mut Self>, amt: usize) { |
| self.get_pin_mut().consume(amt) |
| } |
| } |
| |
| 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.buf.len(), self.buf.capacity()), |
| ) |
| .field("written", &self.written) |
| .finish() |
| } |
| } |
| |
| #[cfg(test)] |
| mod tests { |
| use super::*; |
| |
| #[test] |
| fn assert_unpin() { |
| crate::is_unpin::<BufWriter<()>>(); |
| } |
| } |