Skip to main content

compio_io/read/
buf.rs

1use compio_buf::{
2    BufResult, IntoInner, IoBuf, IoBufExt, IoBufMut, IoVectoredBuf, IoVectoredBufMut, buf_try,
3};
4use futures_util::FutureExt;
5
6use crate::{AsyncRead, AsyncWrite, IoResult, buffer::Buffer, util::DEFAULT_BUF_SIZE};
7/// # AsyncBufRead
8///
9/// Async read with buffered content.
10pub trait AsyncBufRead: AsyncRead {
11    /// Try fill the internal buffer with data
12    async fn fill_buf(&mut self) -> IoResult<&'_ [u8]>;
13
14    /// Mark how much data is read
15    fn consume(&mut self, amount: usize);
16}
17
18impl<A: AsyncBufRead + ?Sized> AsyncBufRead for &mut A {
19    async fn fill_buf(&mut self) -> IoResult<&'_ [u8]> {
20        (**self).fill_buf().await
21    }
22
23    fn consume(&mut self, amount: usize) {
24        (**self).consume(amount)
25    }
26}
27
28/// Wraps a reader and buffers input from [`AsyncRead`]
29///
30/// It can be excessively inefficient to work directly with a [`AsyncRead`]
31/// instance. A `BufReader<R>` performs large, infrequent reads on the
32/// underlying [`AsyncRead`] and maintains an in-memory buffer of the results.
33///
34/// `BufReader<R>` can improve the speed of programs that make *small* and
35/// *repeated* read calls to the same file or network socket. It does not
36/// help when reading very large amounts at once, or reading just one or a few
37/// times. It also provides no advantage when reading from a source that is
38/// already in memory, like a `Vec<u8>`.
39///
40/// If the underlying reader also implements [`AsyncWrite`], `BufReader<R>`
41/// forwards write operations directly to the inner writer without touching the
42/// read buffer.
43///
44/// When the `BufReader<R>` is dropped, the contents of its buffer will be
45/// discarded. Reading from the underlying reader after unwrapping the
46/// `BufReader<R>` with [`BufReader::into_inner`] can cause data loss.
47///
48/// # Caution
49///
50/// Due to the pass-by-ownership nature of completion-based IO, the buffer is
51/// passed to the inner reader when [`fill_buf`] is called. If the future
52/// returned by [`fill_buf`] is dropped before inner `read` is completed,
53/// `BufReader` will not be able to retrieve the buffer, causing panic on next
54/// [`fill_buf`] call.
55///
56/// [`fill_buf`]: #method.fill_buf
57#[derive(Debug)]
58pub struct BufReader<R> {
59    reader: R,
60    buf: Buffer,
61}
62
63impl<R> BufReader<R> {
64    /// Creates a new `BufReader` with a default buffer capacity. The default is
65    /// currently 8 KiB, but may change in the future.
66    pub fn new(reader: R) -> Self {
67        Self::with_capacity(DEFAULT_BUF_SIZE, reader)
68    }
69
70    /// Creates a new `BufReader` with the specified buffer capacity.
71    pub fn with_capacity(cap: usize, reader: R) -> Self {
72        Self {
73            reader,
74            buf: Buffer::with_capacity(cap),
75        }
76    }
77}
78
79impl<R: AsyncRead> AsyncRead for BufReader<R> {
80    async fn read<B: IoBufMut>(&mut self, buf: B) -> BufResult<usize, B> {
81        let (mut slice, buf) = buf_try!(self.fill_buf().await, buf);
82        slice.read(buf).await.map_res(|res| {
83            self.consume(res);
84            res
85        })
86    }
87
88    async fn read_vectored<V: IoVectoredBufMut>(&mut self, buf: V) -> BufResult<usize, V> {
89        let (mut slice, buf) = buf_try!(self.fill_buf().await, buf);
90        slice.read_vectored(buf).await.map_res(|res| {
91            self.consume(res);
92            res
93        })
94    }
95}
96
97impl<R: AsyncRead> AsyncBufRead for BufReader<R> {
98    async fn fill_buf(&mut self) -> IoResult<&'_ [u8]> {
99        let Self { reader, buf } = self;
100
101        if buf.all_done() {
102            buf.reset()
103        }
104
105        if buf.need_fill() {
106            buf.with(|b| {
107                let len = b.buf_len();
108                let b = b.slice(len..);
109                reader.read(b).map(IntoInner::into_inner)
110            })
111            .await?;
112        }
113
114        Ok(buf.buffer())
115    }
116
117    fn consume(&mut self, amount: usize) {
118        self.buf.advance(amount);
119    }
120}
121
122impl<R: AsyncRead + AsyncWrite> AsyncWrite for BufReader<R> {
123    async fn write<B: IoBuf>(&mut self, buf: B) -> BufResult<usize, B> {
124        self.reader.write(buf).await
125    }
126
127    async fn write_vectored<B: IoVectoredBuf>(&mut self, buf: B) -> BufResult<usize, B> {
128        self.reader.write_vectored(buf).await
129    }
130
131    async fn flush(&mut self) -> IoResult<()> {
132        self.reader.flush().await
133    }
134
135    async fn shutdown(&mut self) -> IoResult<()> {
136        self.reader.shutdown().await
137    }
138}
139
140impl<R> IntoInner for BufReader<R> {
141    type Inner = R;
142
143    fn into_inner(self) -> Self::Inner {
144        self.reader
145    }
146}