1 | use bytes::BufMut; |
2 | use std::future::Future; |
3 | use std::io; |
4 | use std::pin::Pin; |
5 | use std::task::{Context, Poll}; |
6 | use tokio::io::AsyncRead; |
7 | |
8 | /// Read data from an `AsyncRead` into an implementer of the [`BufMut`] trait. |
9 | /// |
10 | /// [`BufMut`]: bytes::BufMut |
11 | /// |
12 | /// # Example |
13 | /// |
14 | /// ``` |
15 | /// use bytes::{Bytes, BytesMut}; |
16 | /// use tokio_stream as stream; |
17 | /// use tokio::io::Result; |
18 | /// use tokio_util::io::{StreamReader, read_buf}; |
19 | /// # #[tokio::main] |
20 | /// # async fn main() -> std::io::Result<()> { |
21 | /// |
22 | /// // Create a reader from an iterator. This particular reader will always be |
23 | /// // ready. |
24 | /// let mut read = StreamReader::new(stream::iter(vec![Result::Ok(Bytes::from_static(&[0, 1, 2, 3]))])); |
25 | /// |
26 | /// let mut buf = BytesMut::new(); |
27 | /// let mut reads = 0; |
28 | /// |
29 | /// loop { |
30 | /// reads += 1; |
31 | /// let n = read_buf(&mut read, &mut buf).await?; |
32 | /// |
33 | /// if n == 0 { |
34 | /// break; |
35 | /// } |
36 | /// } |
37 | /// |
38 | /// // one or more reads might be necessary. |
39 | /// assert!(reads >= 1); |
40 | /// assert_eq!(&buf[..], &[0, 1, 2, 3]); |
41 | /// # Ok(()) |
42 | /// # } |
43 | /// ``` |
44 | pub async fn read_buf<R, B>(read: &mut R, buf: &mut B) -> io::Result<usize> |
45 | where |
46 | R: AsyncRead + Unpin, |
47 | B: BufMut, |
48 | { |
49 | return ReadBufFn(read, buf).await; |
50 | |
51 | struct ReadBufFn<'a, R, B>(&'a mut R, &'a mut B); |
52 | |
53 | impl<'a, R, B> Future for ReadBufFn<'a, R, B> |
54 | where |
55 | R: AsyncRead + Unpin, |
56 | B: BufMut, |
57 | { |
58 | type Output = io::Result<usize>; |
59 | |
60 | fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { |
61 | let this = &mut *self; |
62 | crate::util::poll_read_buf(Pin::new(this.0), cx, this.1) |
63 | } |
64 | } |
65 | } |
66 | |