Skip to main content

lace_rt/
record_stream.rs

1//! Reading a Lace-010 record stream (`.lrs`) from a byte reader.
2
3use std::io::{ErrorKind, Read};
4
5use lace_sans_io::runtime_support::io_budget::IoReservation;
6use lace_sans_io::runtime_support::record_stream::RecordStreamDecoder;
7
8use crate::io::IO;
9use crate::{Error, LaceHandle, Record, Result};
10
11/// The validated records of one record stream, in stream order; see
12/// [`read_record_stream`].
13pub struct RecordStream<R> {
14    reader: R,
15    /// `None` once the stream has ended or failed.
16    decoder: Option<RecordStreamDecoder>,
17    buffer: Box<[u8]>,
18    start: usize,
19    end: usize,
20    _reservation: IoReservation,
21}
22
23/// Read the records of a record stream from `reader`, one at a time.
24///
25/// Each record is validated as it arrives and received like an ILTP
26/// response: into `lace`'s record staging, or into memory when its store has
27/// none. A record whose canonical length exceeds what remains of `max_bytes`
28/// fails at its start, before any of its bytes are held. An invalid,
29/// oversized or truncated record, or a failed read, ends the stream with one
30/// error and discards its partial bytes. The reader's chunk buffer is
31/// reserved from the IO budget until the stream drops.
32pub fn read_record_stream<R: Read>(
33    lace: &LaceHandle,
34    reader: R,
35    max_bytes: u64,
36) -> Result<RecordStream<R>> {
37    lace.owner.ensure_claimable()?;
38    let chunk_bytes = IO.chunk_bytes();
39    let reservation = IO.reserve(chunk_bytes)?;
40    Ok(RecordStream {
41        reader,
42        decoder: Some(RecordStreamDecoder::new(
43            lace.owner.record_staging.clone(),
44            max_bytes,
45        )),
46        buffer: vec![0; chunk_bytes].into_boxed_slice(),
47        start: 0,
48        end: 0,
49        _reservation: reservation,
50    })
51}
52
53impl<R: Read> RecordStream<R> {
54    fn next_record(&mut self, decoder: &mut RecordStreamDecoder) -> Result<Option<Record>> {
55        loop {
56            if self.start == self.end {
57                let count = loop {
58                    match self.reader.read(&mut self.buffer) {
59                        Ok(count) => break count,
60                        Err(error) if error.kind() == ErrorKind::Interrupted => {}
61                        Err(error) => {
62                            return Err(Error::transport(format!(
63                                "record stream read failed: {error}"
64                            )));
65                        }
66                    }
67                };
68                if count == 0 {
69                    decoder.finish()?;
70                    return Ok(None);
71                }
72                (self.start, self.end) = (0, count);
73            }
74            let (consumed, mut record) = decoder.decode(&self.buffer[self.start..self.end])?;
75            self.start += consumed;
76            while let Some(io) = decoder.take_io() {
77                if let Some(staged) = decoder.complete_io(io.run())? {
78                    record = Some(staged);
79                }
80            }
81            if record.is_some() {
82                return Ok(record);
83            }
84        }
85    }
86}
87
88impl<R: Read> Iterator for RecordStream<R> {
89    type Item = Result<Record>;
90
91    /// The next record; after the end or a failure, `None`. Dropping the
92    /// decoder at either discards a partial record.
93    fn next(&mut self) -> Option<Result<Record>> {
94        let mut decoder = self.decoder.take()?;
95        let result = self.next_record(&mut decoder).transpose();
96        if matches!(result, Some(Ok(_))) {
97            self.decoder = Some(decoder);
98        }
99        result
100    }
101}
102
103#[cfg(test)]
104#[path = "record_stream_tests.rs"]
105mod tests;