lace_tokio/
record_stream.rs

1//! Reading a Lace-010 record stream (`.lrs`) from an async byte reader.
2
3use lace_sans_io::runtime_support::io_budget::IoReservation;
4use lace_sans_io::runtime_support::record_sink::SinkIoDone;
5use lace_sans_io::runtime_support::record_stream::RecordStreamDecoder;
6use tokio::io::{AsyncRead, AsyncReadExt};
7use tokio::task::JoinHandle;
8
9use crate::io::IO;
10use crate::{Error, LaceHandle, Record, Result};
11
12/// The validated records of one record stream, in stream order; see
13/// [`read_record_stream`].
14pub struct RecordStream<R> {
15    reader: R,
16    /// `None` once the stream has ended or failed.
17    decoder: Option<RecordStreamDecoder>,
18    buffer: Box<[u8]>,
19    start: usize,
20    end: usize,
21    /// Staging IO on the blocking pool, kept across a cancelled
22    /// [`next`](RecordStream::next).
23    io_task: Option<JoinHandle<SinkIoDone>>,
24    _reservation: IoReservation,
25}
26
27/// Read the records of a record stream from `reader`, one at a time.
28///
29/// Each record is validated as it arrives and received like an ILTP
30/// response: into `lace`'s record staging, written on the blocking pool, or
31/// into memory when its store has none. A record whose canonical length
32/// exceeds what remains of `max_bytes` fails at its start, before any of its
33/// bytes are held. An invalid, oversized or truncated record, or a failed
34/// read, ends the stream with one error and discards its partial bytes. The
35/// reader's chunk buffer is reserved from the IO budget until the stream
36/// drops.
37pub fn read_record_stream<R: AsyncRead + Unpin>(
38    lace: &LaceHandle,
39    reader: R,
40    max_bytes: u64,
41) -> Result<RecordStream<R>> {
42    lace.owner.claim()?;
43    let chunk_bytes = IO.chunk_bytes();
44    let reservation = IO.reserve(chunk_bytes)?;
45    Ok(RecordStream {
46        reader,
47        decoder: Some(RecordStreamDecoder::new(
48            lace.owner.record_staging.clone(),
49            max_bytes,
50        )),
51        buffer: vec![0; chunk_bytes].into_boxed_slice(),
52        start: 0,
53        end: 0,
54        io_task: None,
55        _reservation: reservation,
56    })
57}
58
59impl<R: AsyncRead + Unpin> RecordStream<R> {
60    /// The next record; after the end or a failure, `None`. Dropping the
61    /// decoder at either discards a partial record. Cancelling this future
62    /// loses no input.
63    pub async fn next(&mut self) -> Option<Result<Record>> {
64        self.decoder.as_ref()?;
65        let result = self.next_record().await.transpose();
66        if !matches!(result, Some(Ok(_))) {
67            self.decoder = None;
68        }
69        result
70    }
71
72    async fn next_record(&mut self) -> Result<Option<Record>> {
73        loop {
74            if let Some(record) = self.settle_io().await? {
75                return Ok(Some(record));
76            }
77            let decoder = self.decoder.as_mut().expect("a live record stream");
78            if self.start == self.end {
79                let count = self.reader.read(&mut self.buffer).await.map_err(|error| {
80                    Error::transport(format!("record stream read failed: {error}"))
81                })?;
82                if count == 0 {
83                    decoder.finish()?;
84                    return Ok(None);
85                }
86                (self.start, self.end) = (0, count);
87            }
88            let (consumed, record) = decoder.decode(&self.buffer[self.start..self.end])?;
89            self.start += consumed;
90            if record.is_some() {
91                return Ok(record);
92            }
93        }
94    }
95
96    /// Run the decoder's pending staging IO, if any, on the blocking pool;
97    /// the record it completed.
98    async fn settle_io(&mut self) -> Result<Option<Record>> {
99        loop {
100            let decoder = self.decoder.as_mut().expect("a live record stream");
101            let task = match &mut self.io_task {
102                Some(task) => task,
103                task => {
104                    let Some(io) = decoder.take_io() else {
105                        return Ok(None);
106                    };
107                    task.insert(tokio::task::spawn_blocking(move || io.run()))
108                }
109            };
110            let done = task.await;
111            self.io_task = None;
112            let done = done.map_err(|error| {
113                Error::integration(format!("record staging task failed: {error}"))
114            })?;
115            let decoder = self.decoder.as_mut().expect("a live record stream");
116            if let Some(record) = decoder.complete_io(done)? {
117                return Ok(Some(record));
118            }
119        }
120    }
121}
122
123#[cfg(test)]
124#[path = "record_stream_tests.rs"]
125mod tests;
lace_tokio/record_stream.rs

lace_tokio/
record_stream.rs
