tokio/io/util/
copy_buf.rs

1use crate::io::{AsyncBufRead, AsyncWrite};
2use std::future::Future;
3use std::io;
4use std::pin::Pin;
5use std::task::{ready, Context, Poll};
6
7cfg_io_util! {
8    /// A future that asynchronously copies the entire contents of a reader into a
9    /// writer.
10    ///
11    /// This struct is generally created by calling [`copy_buf`][copy_buf]. Please
12    /// see the documentation of `copy_buf()` for more details.
13    ///
14    /// [copy_buf]: copy_buf()
15    #[derive(Debug)]
16    #[must_use = "futures do nothing unless you `.await` or poll them"]
17    struct CopyBuf<'a, R: ?Sized, W: ?Sized> {
18        reader: &'a mut R,
19        writer: &'a mut W,
20        amt: u64,
21        need_flush: bool,
22    }
23
24    /// Asynchronously copies the entire contents of a reader into a writer.
25    ///
26    /// This function returns a future that will continuously read data from
27    /// `reader` and then write it into `writer` in a streaming fashion until
28    /// `reader` returns EOF or fails.
29    ///
30    /// On success, the total number of bytes that were copied from `reader` to
31    /// `writer` is returned.
32    ///
33    /// This is a [`tokio::io::copy`] alternative for [`AsyncBufRead`] readers
34    /// with no extra buffer allocation, since [`AsyncBufRead`] allow access
35    /// to the reader's inner buffer.
36    ///
37    /// # When to use async alternatives instead of `SyncIoBridge`
38    ///
39    /// If you are looking to use [`std::io::copy`] with a synchronous consumer
40    /// (like a `hasher` or compressor), consider using async alternatives instead of
41    /// wrapping the reader with [`SyncIoBridge`]. See the [`SyncIoBridge`]
42    /// documentation for detailed examples and guidance on hashing, compression,
43    /// and data parsing.
44    ///
45    /// [`tokio::io::copy`]: crate::io::copy
46    /// [`AsyncBufRead`]: crate::io::AsyncBufRead
47    /// [`SyncIoBridge`]: https://docs.rs/tokio-util/latest/tokio_util/io/struct.SyncIoBridge.html
48    ///
49    /// # Errors
50    ///
51    /// Errors from `poll_fill_buf` and `poll_write` are returned immediately,
52    /// except [`io::ErrorKind::Interrupted`], which is retried. Errors from
53    /// flushing the writer are returned without retrying.
54    ///
55    /// # Examples
56    ///
57    /// ```
58    /// use tokio::io;
59    ///
60    /// # async fn dox() -> std::io::Result<()> {
61    /// let mut reader: &[u8] = b"hello";
62    /// let mut writer: Vec<u8> = vec![];
63    ///
64    /// io::copy_buf(&mut reader, &mut writer).await?;
65    ///
66    /// assert_eq!(b"hello", &writer[..]);
67    /// # Ok(())
68    /// # }
69    /// ```
70    pub async fn copy_buf<'a, R, W>(reader: &'a mut R, writer: &'a mut W) -> io::Result<u64>
71    where
72        R: AsyncBufRead + Unpin + ?Sized,
73        W: AsyncWrite + Unpin + ?Sized,
74    {
75        CopyBuf {
76            reader,
77            writer,
78            amt: 0,
79            need_flush: false,
80        }.await
81    }
82}
83
84impl<R, W> Future for CopyBuf<'_, R, W>
85where
86    R: AsyncBufRead + Unpin + ?Sized,
87    W: AsyncWrite + Unpin + ?Sized,
88{
89    type Output = io::Result<u64>;
90
91    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
92        #[cfg(any(
93            feature = "fs",
94            feature = "io-std",
95            feature = "net",
96            feature = "process",
97            feature = "rt",
98            feature = "signal",
99            feature = "sync",
100            feature = "time",
101        ))]
102        // Keep track of task budget
103        let coop = ready!(crate::task::coop::poll_proceed(cx));
104        loop {
105            let me = &mut *self;
106            let buffer = match Pin::new(&mut *me.reader).poll_fill_buf(cx) {
107                Poll::Ready(Ok(buffer)) => {
108                    #[cfg(any(
109                        feature = "fs",
110                        feature = "io-std",
111                        feature = "net",
112                        feature = "process",
113                        feature = "rt",
114                        feature = "signal",
115                        feature = "sync",
116                        feature = "time",
117                    ))]
118                    coop.made_progress();
119                    buffer
120                }
121                Poll::Ready(Err(err)) if err.kind() == io::ErrorKind::Interrupted => continue,
122                Poll::Ready(Err(err)) => {
123                    #[cfg(any(
124                        feature = "fs",
125                        feature = "io-std",
126                        feature = "net",
127                        feature = "process",
128                        feature = "rt",
129                        feature = "signal",
130                        feature = "sync",
131                        feature = "time",
132                    ))]
133                    coop.made_progress();
134                    return Poll::Ready(Err(err));
135                }
136                Poll::Pending => {
137                    // Try flushing when the reader has no progress to avoid deadlock
138                    // when the reader depends on buffered writer.
139                    if me.need_flush {
140                        ready!(Pin::new(&mut *me.writer).poll_flush(cx))?;
141                        me.need_flush = false;
142                    }
143
144                    return Poll::Pending;
145                }
146            };
147
148            if buffer.is_empty() {
149                ready!(Pin::new(&mut self.writer).poll_flush(cx))?;
150                return Poll::Ready(Ok(self.amt));
151            }
152
153            let i = match ready!(Pin::new(&mut *me.writer).poll_write(cx, buffer)) {
154                Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
155                res => res?,
156            };
157            if i == 0 {
158                return Poll::Ready(Err(std::io::ErrorKind::WriteZero.into()));
159            }
160            self.amt += i as u64;
161            self.need_flush = true;
162            Pin::new(&mut *self.reader).consume(i);
163        }
164    }
165}
166
167#[cfg(test)]
168mod tests {
169    use super::*;
170
171    #[test]
172    fn assert_unpin() {
173        use std::marker::PhantomPinned;
174        crate::is_unpin::<CopyBuf<'_, PhantomPinned, PhantomPinned>>();
175    }
176}
tokio/io/util/copy_buf.rs

tokio/io/util/
copy_buf.rs
