1use 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
11pub struct RecordStream<R> {
14 reader: R,
15 decoder: Option<RecordStreamDecoder>,
17 buffer: Box<[u8]>,
18 start: usize,
19 end: usize,
20 _reservation: IoReservation,
21}
22
23pub 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 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;