Joel Galenson | 7a98302 | 2021-05-19 15:35:43 -0700 | [diff] [blame^] | 1 | use super::DEFAULT_BUF_SIZE; |
Haibo Huang | 5bf8796 | 2021-02-09 17:18:40 -0800 | [diff] [blame] | 2 | use futures_core::ready; |
Jason Macnak | c417d3b | 2020-04-06 10:30:28 -0700 | [diff] [blame] | 3 | use futures_core::task::{Context, Poll}; |
Haibo Huang | 52627c8 | 2020-05-08 19:26:17 -0700 | [diff] [blame] | 4 | use futures_io::{AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite, IoSlice, SeekFrom}; |
Haibo Huang | 5bf8796 | 2021-02-09 17:18:40 -0800 | [diff] [blame] | 5 | use pin_project_lite::pin_project; |
Jason Macnak | c417d3b | 2020-04-06 10:30:28 -0700 | [diff] [blame] | 6 | use std::fmt; |
| 7 | use std::io::{self, Write}; |
| 8 | use std::pin::Pin; |
Jason Macnak | c417d3b | 2020-04-06 10:30:28 -0700 | [diff] [blame] | 9 | |
Haibo Huang | 5bf8796 | 2021-02-09 17:18:40 -0800 | [diff] [blame] | 10 | pin_project! { |
| 11 | /// Wraps a writer and buffers its output. |
| 12 | /// |
| 13 | /// It can be excessively inefficient to work directly with something that |
| 14 | /// implements [`AsyncWrite`]. A `BufWriter` keeps an in-memory buffer of data and |
| 15 | /// writes it to an underlying writer in large, infrequent batches. |
| 16 | /// |
| 17 | /// `BufWriter` can improve the speed of programs that make *small* and |
| 18 | /// *repeated* write calls to the same file or network socket. It does not |
| 19 | /// help when writing very large amounts at once, or writing just one or a few |
| 20 | /// times. It also provides no advantage when writing to a destination that is |
| 21 | /// in memory, like a `Vec<u8>`. |
| 22 | /// |
| 23 | /// When the `BufWriter` is dropped, the contents of its buffer will be |
| 24 | /// discarded. Creating multiple instances of a `BufWriter` on the same |
| 25 | /// stream can cause data loss. If you need to write out the contents of its |
| 26 | /// buffer, you must manually call flush before the writer is dropped. |
| 27 | /// |
| 28 | /// [`AsyncWrite`]: futures_io::AsyncWrite |
| 29 | /// [`flush`]: super::AsyncWriteExt::flush |
| 30 | /// |
| 31 | // TODO: Examples |
| 32 | pub struct BufWriter<W> { |
| 33 | #[pin] |
| 34 | inner: W, |
| 35 | buf: Vec<u8>, |
| 36 | written: usize, |
| 37 | } |
Jason Macnak | c417d3b | 2020-04-06 10:30:28 -0700 | [diff] [blame] | 38 | } |
| 39 | |
Jason Macnak | c417d3b | 2020-04-06 10:30:28 -0700 | [diff] [blame] | 40 | impl<W: AsyncWrite> BufWriter<W> { |
| 41 | /// Creates a new `BufWriter` with a default buffer capacity. The default is currently 8 KB, |
| 42 | /// but may change in the future. |
| 43 | pub fn new(inner: W) -> Self { |
| 44 | Self::with_capacity(DEFAULT_BUF_SIZE, inner) |
| 45 | } |
| 46 | |
| 47 | /// Creates a new `BufWriter` with the specified buffer capacity. |
| 48 | pub fn with_capacity(cap: usize, inner: W) -> Self { |
Joel Galenson | 7a98302 | 2021-05-19 15:35:43 -0700 | [diff] [blame^] | 49 | Self { inner, buf: Vec::with_capacity(cap), written: 0 } |
Jason Macnak | c417d3b | 2020-04-06 10:30:28 -0700 | [diff] [blame] | 50 | } |
| 51 | |
| 52 | fn flush_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { |
Chih-Hung Hsieh | d45e96e | 2020-10-25 23:16:22 -0700 | [diff] [blame] | 53 | let mut this = self.project(); |
Jason Macnak | c417d3b | 2020-04-06 10:30:28 -0700 | [diff] [blame] | 54 | |
Chih-Hung Hsieh | d45e96e | 2020-10-25 23:16:22 -0700 | [diff] [blame] | 55 | let len = this.buf.len(); |
Jason Macnak | c417d3b | 2020-04-06 10:30:28 -0700 | [diff] [blame] | 56 | let mut ret = Ok(()); |
Chih-Hung Hsieh | d45e96e | 2020-10-25 23:16:22 -0700 | [diff] [blame] | 57 | while *this.written < len { |
| 58 | match ready!(this.inner.as_mut().poll_write(cx, &this.buf[*this.written..])) { |
Jason Macnak | c417d3b | 2020-04-06 10:30:28 -0700 | [diff] [blame] | 59 | Ok(0) => { |
| 60 | ret = Err(io::Error::new( |
| 61 | io::ErrorKind::WriteZero, |
| 62 | "failed to write the buffered data", |
| 63 | )); |
| 64 | break; |
| 65 | } |
Chih-Hung Hsieh | d45e96e | 2020-10-25 23:16:22 -0700 | [diff] [blame] | 66 | Ok(n) => *this.written += n, |
Jason Macnak | c417d3b | 2020-04-06 10:30:28 -0700 | [diff] [blame] | 67 | Err(e) => { |
| 68 | ret = Err(e); |
| 69 | break; |
| 70 | } |
| 71 | } |
| 72 | } |
Chih-Hung Hsieh | d45e96e | 2020-10-25 23:16:22 -0700 | [diff] [blame] | 73 | if *this.written > 0 { |
| 74 | this.buf.drain(..*this.written); |
Jason Macnak | c417d3b | 2020-04-06 10:30:28 -0700 | [diff] [blame] | 75 | } |
Chih-Hung Hsieh | d45e96e | 2020-10-25 23:16:22 -0700 | [diff] [blame] | 76 | *this.written = 0; |
Jason Macnak | c417d3b | 2020-04-06 10:30:28 -0700 | [diff] [blame] | 77 | Poll::Ready(ret) |
| 78 | } |
| 79 | |
Haibo Huang | 52627c8 | 2020-05-08 19:26:17 -0700 | [diff] [blame] | 80 | delegate_access_inner!(inner, W, ()); |
Jason Macnak | c417d3b | 2020-04-06 10:30:28 -0700 | [diff] [blame] | 81 | |
| 82 | /// Returns a reference to the internally buffered data. |
| 83 | pub fn buffer(&self) -> &[u8] { |
| 84 | &self.buf |
| 85 | } |
| 86 | } |
| 87 | |
| 88 | impl<W: AsyncWrite> AsyncWrite for BufWriter<W> { |
| 89 | fn poll_write( |
| 90 | mut self: Pin<&mut Self>, |
| 91 | cx: &mut Context<'_>, |
| 92 | buf: &[u8], |
| 93 | ) -> Poll<io::Result<usize>> { |
| 94 | if self.buf.len() + buf.len() > self.buf.capacity() { |
| 95 | ready!(self.as_mut().flush_buf(cx))?; |
| 96 | } |
| 97 | if buf.len() >= self.buf.capacity() { |
Haibo Huang | 52627c8 | 2020-05-08 19:26:17 -0700 | [diff] [blame] | 98 | self.project().inner.poll_write(cx, buf) |
Jason Macnak | c417d3b | 2020-04-06 10:30:28 -0700 | [diff] [blame] | 99 | } else { |
Haibo Huang | 52627c8 | 2020-05-08 19:26:17 -0700 | [diff] [blame] | 100 | Poll::Ready(self.project().buf.write(buf)) |
Jason Macnak | c417d3b | 2020-04-06 10:30:28 -0700 | [diff] [blame] | 101 | } |
| 102 | } |
| 103 | |
| 104 | fn poll_write_vectored( |
| 105 | mut self: Pin<&mut Self>, |
| 106 | cx: &mut Context<'_>, |
| 107 | bufs: &[IoSlice<'_>], |
| 108 | ) -> Poll<io::Result<usize>> { |
| 109 | let total_len = bufs.iter().map(|b| b.len()).sum::<usize>(); |
| 110 | if self.buf.len() + total_len > self.buf.capacity() { |
| 111 | ready!(self.as_mut().flush_buf(cx))?; |
| 112 | } |
| 113 | if total_len >= self.buf.capacity() { |
Haibo Huang | 52627c8 | 2020-05-08 19:26:17 -0700 | [diff] [blame] | 114 | self.project().inner.poll_write_vectored(cx, bufs) |
Jason Macnak | c417d3b | 2020-04-06 10:30:28 -0700 | [diff] [blame] | 115 | } else { |
Haibo Huang | 52627c8 | 2020-05-08 19:26:17 -0700 | [diff] [blame] | 116 | Poll::Ready(self.project().buf.write_vectored(bufs)) |
Jason Macnak | c417d3b | 2020-04-06 10:30:28 -0700 | [diff] [blame] | 117 | } |
| 118 | } |
| 119 | |
| 120 | fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { |
| 121 | ready!(self.as_mut().flush_buf(cx))?; |
Haibo Huang | 52627c8 | 2020-05-08 19:26:17 -0700 | [diff] [blame] | 122 | self.project().inner.poll_flush(cx) |
Jason Macnak | c417d3b | 2020-04-06 10:30:28 -0700 | [diff] [blame] | 123 | } |
| 124 | |
| 125 | fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { |
| 126 | ready!(self.as_mut().flush_buf(cx))?; |
Haibo Huang | 52627c8 | 2020-05-08 19:26:17 -0700 | [diff] [blame] | 127 | self.project().inner.poll_close(cx) |
Jason Macnak | c417d3b | 2020-04-06 10:30:28 -0700 | [diff] [blame] | 128 | } |
| 129 | } |
| 130 | |
| 131 | impl<W: AsyncRead> AsyncRead for BufWriter<W> { |
Haibo Huang | 52627c8 | 2020-05-08 19:26:17 -0700 | [diff] [blame] | 132 | delegate_async_read!(inner); |
Jason Macnak | c417d3b | 2020-04-06 10:30:28 -0700 | [diff] [blame] | 133 | } |
| 134 | |
| 135 | impl<W: AsyncBufRead> AsyncBufRead for BufWriter<W> { |
Haibo Huang | 52627c8 | 2020-05-08 19:26:17 -0700 | [diff] [blame] | 136 | delegate_async_buf_read!(inner); |
Jason Macnak | c417d3b | 2020-04-06 10:30:28 -0700 | [diff] [blame] | 137 | } |
| 138 | |
| 139 | impl<W: fmt::Debug> fmt::Debug for BufWriter<W> { |
| 140 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
| 141 | f.debug_struct("BufWriter") |
| 142 | .field("writer", &self.inner) |
| 143 | .field("buffer", &format_args!("{}/{}", self.buf.len(), self.buf.capacity())) |
| 144 | .field("written", &self.written) |
| 145 | .finish() |
| 146 | } |
| 147 | } |
| 148 | |
| 149 | impl<W: AsyncWrite + AsyncSeek> AsyncSeek for BufWriter<W> { |
| 150 | /// Seek to the offset, in bytes, in the underlying writer. |
| 151 | /// |
| 152 | /// Seeking always writes out the internal buffer before seeking. |
| 153 | fn poll_seek( |
| 154 | mut self: Pin<&mut Self>, |
| 155 | cx: &mut Context<'_>, |
| 156 | pos: SeekFrom, |
| 157 | ) -> Poll<io::Result<u64>> { |
| 158 | ready!(self.as_mut().flush_buf(cx))?; |
Haibo Huang | 52627c8 | 2020-05-08 19:26:17 -0700 | [diff] [blame] | 159 | self.project().inner.poll_seek(cx, pos) |
Jason Macnak | c417d3b | 2020-04-06 10:30:28 -0700 | [diff] [blame] | 160 | } |
| 161 | } |