futures_util/io/
buf_writer.rs1use 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 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 pub fn new(inner: W) -> Self {
46 Self::with_capacity(DEFAULT_BUF_SIZE, inner)
47 }
48
49 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 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 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 pub fn buffer(&self) -> &[u8] {
106 &self.buf
107 }
108
109 pub(super) fn capacity(&self) -> usize {
111 self.buf.capacity()
112 }
113
114 #[inline]
116 pub(super) fn spare_capacity(&self) -> usize {
117 self.buf.capacity() - self.buf.len()
118 }
119
120 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 unsafe {
132 self.write_to_buffer_unchecked(&buf[..amt_to_buffer]);
133 }
134
135 amt_to_buffer
136 }
137
138 #[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 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}