tokio/io/util/copy.rs
1use crate::io::{AsyncRead, AsyncWrite, ReadBuf};
2
3use std::future::Future;
4use std::io;
5use std::pin::Pin;
6use std::task::{ready, Context, Poll};
7
8#[derive(Debug)]
9pub(super) struct CopyBuffer {
10 read_done: bool,
11 need_flush: bool,
12 pos: usize,
13 cap: usize,
14 amt: u64,
15 buf: Box<[u8]>,
16}
17
18impl CopyBuffer {
19 pub(super) fn new(buf_size: usize) -> Self {
20 Self {
21 read_done: false,
22 need_flush: false,
23 pos: 0,
24 cap: 0,
25 amt: 0,
26 buf: vec![0; buf_size].into_boxed_slice(),
27 }
28 }
29
30 fn poll_fill_buf<R>(
31 &mut self,
32 cx: &mut Context<'_>,
33 mut reader: Pin<&mut R>,
34 ) -> Poll<io::Result<()>>
35 where
36 R: AsyncRead + ?Sized,
37 {
38 let me = &mut *self;
39 let mut buf = ReadBuf::new(&mut me.buf);
40 buf.set_filled(me.cap);
41
42 let res = loop {
43 match reader.as_mut().poll_read(cx, &mut buf) {
44 Poll::Ready(Err(e)) if e.kind() == io::ErrorKind::Interrupted => continue,
45 res => break res,
46 }
47 };
48 if let Poll::Ready(Ok(())) = res {
49 let filled_len = buf.filled().len();
50 me.read_done = me.cap == filled_len;
51 me.cap = filled_len;
52 }
53 res
54 }
55
56 fn poll_write_buf<R, W>(
57 &mut self,
58 cx: &mut Context<'_>,
59 mut reader: Pin<&mut R>,
60 mut writer: Pin<&mut W>,
61 ) -> Poll<io::Result<usize>>
62 where
63 R: AsyncRead + ?Sized,
64 W: AsyncWrite + ?Sized,
65 {
66 let me = &mut *self;
67 loop {
68 match writer.as_mut().poll_write(cx, &me.buf[me.pos..me.cap]) {
69 Poll::Ready(Err(e)) if e.kind() == io::ErrorKind::Interrupted => continue,
70 Poll::Pending => {
71 // Top up the buffer towards full if we can read a bit more
72 // data - this should improve the chances of a large write
73 if !me.read_done && me.cap < me.buf.len() {
74 ready!(me.poll_fill_buf(cx, reader.as_mut()))?;
75 }
76 return Poll::Pending;
77 }
78 res => return res,
79 }
80 }
81 }
82
83 pub(super) fn poll_copy<R, W>(
84 &mut self,
85 cx: &mut Context<'_>,
86 mut reader: Pin<&mut R>,
87 mut writer: Pin<&mut W>,
88 ) -> Poll<io::Result<u64>>
89 where
90 R: AsyncRead + ?Sized,
91 W: AsyncWrite + ?Sized,
92 {
93 ready!(crate::trace::trace_leaf(cx));
94 #[cfg(any(
95 feature = "fs",
96 feature = "io-std",
97 feature = "net",
98 feature = "process",
99 feature = "rt",
100 feature = "signal",
101 feature = "sync",
102 feature = "time",
103 ))]
104 // Keep track of task budget
105 let coop = ready!(crate::task::coop::poll_proceed(cx));
106 loop {
107 // If there is some space left in our buffer, then we try to read some
108 // data to continue, thus maximizing the chances of a large write.
109 if self.cap < self.buf.len() && !self.read_done {
110 match self.poll_fill_buf(cx, reader.as_mut()) {
111 Poll::Ready(Ok(())) => {
112 #[cfg(any(
113 feature = "fs",
114 feature = "io-std",
115 feature = "net",
116 feature = "process",
117 feature = "rt",
118 feature = "signal",
119 feature = "sync",
120 feature = "time",
121 ))]
122 coop.made_progress();
123 }
124 Poll::Ready(Err(err)) => {
125 #[cfg(any(
126 feature = "fs",
127 feature = "io-std",
128 feature = "net",
129 feature = "process",
130 feature = "rt",
131 feature = "signal",
132 feature = "sync",
133 feature = "time",
134 ))]
135 coop.made_progress();
136 return Poll::Ready(Err(err));
137 }
138 Poll::Pending => {
139 // Ignore pending reads when our buffer is not empty, because
140 // we can try to write data immediately.
141 if self.pos == self.cap {
142 // Try flushing when the reader has no progress to avoid deadlock
143 // when the reader depends on buffered writer.
144 if self.need_flush {
145 ready!(writer.as_mut().poll_flush(cx))?;
146 #[cfg(any(
147 feature = "fs",
148 feature = "io-std",
149 feature = "net",
150 feature = "process",
151 feature = "rt",
152 feature = "signal",
153 feature = "sync",
154 feature = "time",
155 ))]
156 coop.made_progress();
157 self.need_flush = false;
158 }
159
160 return Poll::Pending;
161 }
162 }
163 }
164 }
165
166 // If our buffer has some data, let's write it out!
167 while self.pos < self.cap {
168 let i = ready!(self.poll_write_buf(cx, reader.as_mut(), writer.as_mut()))?;
169 #[cfg(any(
170 feature = "fs",
171 feature = "io-std",
172 feature = "net",
173 feature = "process",
174 feature = "rt",
175 feature = "signal",
176 feature = "sync",
177 feature = "time",
178 ))]
179 coop.made_progress();
180 if i == 0 {
181 return Poll::Ready(Err(io::Error::new(
182 io::ErrorKind::WriteZero,
183 "write zero byte into writer",
184 )));
185 } else {
186 self.pos += i;
187 self.amt += i as u64;
188 self.need_flush = true;
189 }
190 }
191
192 // If pos larger than cap, this loop will never stop.
193 // In particular, user's wrong poll_write implementation returning
194 // incorrect written length may lead to thread blocking.
195 debug_assert!(
196 self.pos <= self.cap,
197 "writer returned length larger than input slice"
198 );
199
200 // All data has been written, the buffer can be considered empty again
201 self.pos = 0;
202 self.cap = 0;
203
204 // If we've written all the data and we've seen EOF, flush out the
205 // data and finish the transfer.
206 if self.read_done {
207 ready!(writer.as_mut().poll_flush(cx))?;
208 #[cfg(any(
209 feature = "fs",
210 feature = "io-std",
211 feature = "net",
212 feature = "process",
213 feature = "rt",
214 feature = "signal",
215 feature = "sync",
216 feature = "time",
217 ))]
218 coop.made_progress();
219 return Poll::Ready(Ok(self.amt));
220 }
221 }
222 }
223}
224
225/// A future that asynchronously copies the entire contents of a reader into a
226/// writer.
227#[derive(Debug)]
228#[must_use = "futures do nothing unless you `.await` or poll them"]
229struct Copy<'a, R: ?Sized, W: ?Sized> {
230 reader: &'a mut R,
231 writer: &'a mut W,
232 buf: CopyBuffer,
233}
234
235cfg_io_util! {
236 /// Asynchronously copies the entire contents of a reader into a writer.
237 ///
238 /// This function returns a future that will continuously read data from
239 /// `reader` and then write it into `writer` in a streaming fashion until
240 /// `reader` returns EOF or fails.
241 ///
242 /// On success, the total number of bytes that were copied from `reader` to
243 /// `writer` is returned.
244 ///
245 /// This is an asynchronous version of [`std::io::copy`][std].
246 ///
247 /// A heap-allocated copy buffer with 8 KB is created to take data from the
248 /// reader to the writer, check [`copy_buf`] if you want an alternative for
249 /// [`AsyncBufRead`]. You can use `copy_buf` with [`BufReader`] to change the
250 /// buffer capacity.
251 ///
252 /// # When to use async alternatives instead of `SyncIoBridge`
253 ///
254 /// If you are looking to use [`std::io::copy`] with a synchronous consumer
255 /// (like a `hasher` or compressor), consider using async alternatives instead of
256 /// wrapping the reader with [`SyncIoBridge`].
257 /// See the [`SyncIoBridge`] documentation for detailed examples and guidance.
258 ///
259 /// [std]: std::io::copy
260 /// [`copy_buf`]: crate::io::copy_buf
261 /// [`AsyncBufRead`]: crate::io::AsyncBufRead
262 /// [`BufReader`]: crate::io::BufReader
263 /// [`SyncIoBridge`]: https://docs.rs/tokio-util/latest/tokio_util/io/struct.SyncIoBridge.html
264 ///
265 /// # Errors
266 ///
267 /// The returned future will return an error immediately if any call to
268 /// `poll_read` or `poll_write` returns an error other than
269 /// [`io::ErrorKind::Interrupted`], which is retried. Errors from flushing
270 /// the writer are returned without retrying.
271 ///
272 /// # Examples
273 ///
274 /// ```
275 /// use tokio::io;
276 ///
277 /// # async fn dox() -> std::io::Result<()> {
278 /// let mut reader: &[u8] = b"hello";
279 /// let mut writer: Vec<u8> = vec![];
280 ///
281 /// io::copy(&mut reader, &mut writer).await?;
282 ///
283 /// assert_eq!(&b"hello"[..], &writer[..]);
284 /// # Ok(())
285 /// # }
286 /// ```
287 pub async fn copy<'a, R, W>(reader: &'a mut R, writer: &'a mut W) -> io::Result<u64>
288 where
289 R: AsyncRead + Unpin + ?Sized,
290 W: AsyncWrite + Unpin + ?Sized,
291 {
292 Copy {
293 reader,
294 writer,
295 buf: CopyBuffer::new(super::DEFAULT_BUF_SIZE)
296 }.await
297 }
298}
299
300impl<R, W> Future for Copy<'_, R, W>
301where
302 R: AsyncRead + Unpin + ?Sized,
303 W: AsyncWrite + Unpin + ?Sized,
304{
305 type Output = io::Result<u64>;
306
307 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> {
308 let me = &mut *self;
309
310 me.buf
311 .poll_copy(cx, Pin::new(&mut *me.reader), Pin::new(&mut *me.writer))
312 }
313}