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&mut DebugStruct<'_, '_>.debug_struct("BufWriter")
210 .field("writer", &self.inner)
211 .field("buffer", &format_args!("{}/{}", self.buf.len(), self.buf.capacity()))
212 .field(name:"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}
230