|
| 1 | +use std::io; |
| 2 | + |
| 3 | +use bytes::Bytes; |
| 4 | + |
| 5 | +use crate::VortexReadAt; |
| 6 | + |
| 7 | +/// A stateful asynchronous reader that wraps an internal [stateless reader][VortexReadAt]. |
| 8 | +/// |
| 9 | +/// Read operations will advance the cursor. |
| 10 | +#[derive(Clone)] |
| 11 | +pub struct VortexBufReader<R> { |
| 12 | + inner: R, |
| 13 | + pos: u64, |
| 14 | +} |
| 15 | + |
| 16 | +impl<R> VortexBufReader<R> { |
| 17 | + /// Create a new buffered reader wrapping a stateless reader, with reads |
| 18 | + /// beginning at offset 0. |
| 19 | + pub fn new(inner: R) -> Self { |
| 20 | + Self { inner, pos: 0 } |
| 21 | + } |
| 22 | + |
| 23 | + /// Set the position of the next `read_bytes` call directly. |
| 24 | + /// |
| 25 | + /// Note: this method will not fail if the position is past the end of the valid range, |
| 26 | + /// the failure will occur at read time and result in an [`UnexpectedEof`][std::io::ErrorKind::UnexpectedEof] error. |
| 27 | + pub fn set_position(&mut self, pos: u64) { |
| 28 | + self.pos = pos; |
| 29 | + } |
| 30 | +} |
| 31 | + |
| 32 | +impl<R: VortexReadAt> VortexBufReader<R> { |
| 33 | + /// Perform an exactly-sized read at the current cursor position, advancing |
| 34 | + /// the cursor and returning the bytes. |
| 35 | + /// |
| 36 | + /// If there are not enough bytes available to fulfill the request, an |
| 37 | + /// [`UnexpectedEof`][std::io::ErrorKind::UnexpectedEof] error is returned. |
| 38 | + /// |
| 39 | + /// See also [`VortexReadAt::read_byte_range`]. |
| 40 | + pub async fn read_bytes(&mut self, len: u64) -> io::Result<Bytes> { |
| 41 | + let result = self.inner.read_byte_range(self.pos, len).await?; |
| 42 | + self.pos += len; |
| 43 | + Ok(result) |
| 44 | + } |
| 45 | +} |
| 46 | + |
| 47 | +#[cfg(test)] |
| 48 | +mod tests { |
| 49 | + use std::io; |
| 50 | + |
| 51 | + use bytes::Bytes; |
| 52 | + |
| 53 | + use crate::VortexBufReader; |
| 54 | + |
| 55 | + #[tokio::test] |
| 56 | + async fn test_buf_reader() { |
| 57 | + let reader = Bytes::from("0123456789".as_bytes()); |
| 58 | + let mut buf_reader = VortexBufReader::new(reader); |
| 59 | + |
| 60 | + let first2 = buf_reader.read_bytes(2).await.unwrap(); |
| 61 | + assert_eq!(first2.as_ref(), "01".as_bytes()); |
| 62 | + |
| 63 | + buf_reader.set_position(8); |
| 64 | + let last2 = buf_reader.read_bytes(2).await.unwrap(); |
| 65 | + assert_eq!(last2.as_ref(), "89".as_bytes()); |
| 66 | + } |
| 67 | + |
| 68 | + #[tokio::test] |
| 69 | + async fn test_eof() { |
| 70 | + let reader = Bytes::from("0123456789".as_bytes()); |
| 71 | + let mut buf_reader = VortexBufReader::new(reader); |
| 72 | + |
| 73 | + // Read past end of internal reader |
| 74 | + buf_reader.set_position(10); |
| 75 | + |
| 76 | + assert_eq!( |
| 77 | + buf_reader.read_bytes(1).await.unwrap_err().kind(), |
| 78 | + io::ErrorKind::UnexpectedEof, |
| 79 | + ); |
| 80 | + } |
| 81 | +} |
0 commit comments