futures_util/io/
buf_writer.rs

1use super::DEFAULT_BUF_SIZE;
2use futures_core::ready;
3use futures_core::task::{Context, Poll};
4use futures_io::{AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite, IoSlice, SeekFrom};
5use pin_project_lite::pin_project;
6use std::fmt;
7use std::io::{self, Write};
8use std::pin::Pin;
9use std::ptr;
10use std::vec::Vec;
11
12pin_project! {
13    /// Wraps a writer and buffers its output.
14    ///
15    /// It can be excessively inefficient to work directly with something that
16    /// implements [`AsyncWrite`]. A `BufWriter` keeps an in-memory buffer of data and
17    /// writes it to an underlying writer in large, infrequent batches.
18    ///
19    /// `BufWriter` can improve the speed of programs that make *small* and
20    /// *repeated* write calls to the same file or network socket. It does not
21    /// help when writing very large amounts at once, or writing just one or a few
22    /// times. It also provides no advantage when writing to a destination that is
23    /// in memory, like a `Vec<u8>`.
24    ///
25    /// When the `BufWriter` is dropped, the contents of its buffer will be
26    /// discarded. Creating multiple instances of a `BufWriter` on the same
27    /// stream can cause data loss. If you need to write out the contents of its
28    /// buffer, you must manually call flush before the writer is dropped.
29    ///
30    /// [`AsyncWrite`]: futures_io::AsyncWrite
31    /// [`flush`]: super::AsyncWriteExt::flush
32    ///
33    // TODO: Examples
34    pub struct BufWriter<W> {
35        #[pin]
36        inner: W,
37        buf: Vec<u8>,
38        written: usize,
39    }
40}
41
42impl<W: AsyncWrite> BufWriter<W> {
43    /// Creates a new `BufWriter` with a default buffer capacity. The default is currently 8 KB,
44    /// but may change in the future.
45    pub fn new(inner: W) -> Self {
46        Self::with_capacity(DEFAULT_BUF_SIZE, inner)
47    }
48
49    /// Creates a new `BufWriter` with the specified buffer capacity.
50    pub fn with_capacity(cap: usize, inner: W) -> Self {
51        Self { inner, buf: Vec::with_capacity(cap), written: 0 }
52    }
53
54    pub(super) fn flush_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
55        let mut this = self.project();
56
57        let len = this.buf.len();
58        let mut ret = Ok(());
59        while *this.written < len {
60            match ready!(this.inner.as_mut().poll_write(cx, &this.buf[*this.written..])) {
61                Ok(0) => {
62                    ret = Err(io::Error::new(
63                        io::ErrorKind::WriteZero,
64                        "failed to write the buffered data",
65                    ));
66                    break;
67                }
68                Ok(n) => *this.written += n,
69                Err(e) => {
70                    ret = Err(e);
71                    break;
72                }
73            }
74        }
75        if *this.written > 0 {
76            this.buf.drain(..*this.written);
77        }
78        *this.written = 0;
79        Poll::Ready(ret)
80    }
81
82    /// Write directly using `inner`, bypassing buffering
83    pub(super) fn inner_poll_write(
84        self: Pin<&mut Self>,
85        cx: &mut Context<'_>,
86        buf: &[u8],
87    ) -> Poll<io::Result<usize>> {
88        self.project().inner.poll_write(cx, buf)
89    }
90
91    /// Write directly using `inner`, bypassing buffering
92    pub(super) fn inner_poll_write_vectored(
93        self: Pin<&mut Self>,
94        cx: &mut Context<'_>,
95        bufs: &[IoSlice<'_>],
96    ) -> Poll<io::Result<usize>> {
97        self.project().inner.poll_write_vectored(cx, bufs)
98    }
99}
100
101impl<W> BufWriter<W> {
102    delegate_access_inner!(inner, W, ());
103
104    /// Returns a reference to the internally buffered data.
105    pub fn buffer(&self) -> &[u8] {
106        &self.buf
107    }
108
109    /// Capacity of `buf`. how many chars can be held in buffer
110    pub(super) fn capacity(&self) -> usize {
111        self.buf.capacity()
112    }
113
114    /// Remaining number of bytes to reach `buf` 's capacity
115    #[inline]
116    pub(super) fn spare_capacity(&self) -> usize {
117        self.buf.capacity() - self.buf.len()
118    }
119
120    /// Write a byte slice directly into buffer
121    ///
122    /// Will truncate the number of bytes written to `spare_capacity()` so you want to
123    /// calculate the size of your slice to avoid losing bytes
124    ///
125    /// Based on `std::io::BufWriter`
126    pub(super) fn write_to_buf(self: Pin<&mut Self>, buf: &[u8]) -> usize {
127        let available = self.spare_capacity();
128        let amt_to_buffer = available.min(buf.len());
129
130        // SAFETY: `amt_to_buffer` is <= buffer's spare capacity by construction.
131        unsafe {
132            self.write_to_buffer_unchecked(&buf[..amt_to_buffer]);
133        }
134
135        amt_to_buffer
136    }
137
138    /// Write byte slice directly into `self.buf`
139    ///
140    /// Based on `std::io::BufWriter`
141    #[inline]
142    unsafe fn write_to_buffer_unchecked(self: Pin<&mut Self>, buf: &[u8]) {
143        debug_assert!(buf.len() <= self.spare_capacity());
144        let this = self.project();
145        let old_len = this.buf.len();
146        let buf_len = buf.len();
147        let src = buf.as_ptr();
148        unsafe {
149            let dst = this.buf.as_mut_ptr().add(old_len);
150            ptr::copy_nonoverlapping(src, dst, buf_len);
151            this.buf.set_len(old_len + buf_len);
152        }
153    }
154}
155
156impl<W: AsyncWrite> AsyncWrite for BufWriter<W> {
157    fn poll_write(
158        mut self: Pin<&mut Self>,
159        cx: &mut Context<'_>,
160        buf: &[u8],
161    ) -> Poll<io::Result<usize>> {
162        if self.buf.len() + buf.len() > self.buf.capacity() {
163            ready!(self.as_mut().flush_buf(cx))?;
164        }
165        if buf.len() >= self.buf.capacity() {
166            self.project().inner.poll_write(cx, buf)
167        } else {
168            Poll::Ready(self.project().buf.write(buf))
169        }
170    }
171
172    fn poll_write_vectored(
173        mut self: Pin<&mut Self>,
174        cx: &mut Context<'_>,
175        bufs: &[IoSlice<'_>],
176    ) -> Poll<io::Result<usize>> {
177        let total_len = bufs.iter().map(|b| b.len()).sum::<usize>();
178        if self.buf.len() + total_len > self.buf.capacity() {
179            ready!(self.as_mut().flush_buf(cx))?;
180        }
181        if total_len >= self.buf.capacity() {
182            self.project().inner.poll_write_vectored(cx, bufs)
183        } else {
184            Poll::Ready(self.project().buf.write_vectored(bufs))
185        }
186    }
187
188    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
189        ready!(self.as_mut().flush_buf(cx))?;
190        self.project().inner.poll_flush(cx)
191    }
192
193    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
194        ready!(self.as_mut().flush_buf(cx))?;
195        self.project().inner.poll_close(cx)
196    }
197}
198
199impl<W: AsyncRead> AsyncRead for BufWriter<W> {
200    delegate_async_read!(inner);
201}
202
203impl<W: AsyncBufRead> AsyncBufRead for BufWriter<W> {
204    delegate_async_buf_read!(inner);
205}
206
207impl<W: fmt::Debug> fmt::Debug for BufWriter<W> {
208    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
209        f.debug_struct("BufWriter")
210            .field("writer", &self.inner)
211            .field("buffer", &format_args!("{}/{}", self.buf.len(), self.buf.capacity()))
212            .field("written", &self.written)
213            .finish()
214    }
215}
216
217impl<W: AsyncWrite + AsyncSeek> AsyncSeek for BufWriter<W> {
218    /// Seek to the offset, in bytes, in the underlying writer.
219    ///
220    /// Seeking always writes out the internal buffer before seeking.
221    fn poll_seek(
222        mut self: Pin<&mut Self>,
223        cx: &mut Context<'_>,
224        pos: SeekFrom,
225    ) -> Poll<io::Result<u64>> {
226        ready!(self.as_mut().flush_buf(cx))?;
227        self.project().inner.poll_seek(cx, pos)
228    }
229}