harmos_runtime/streams/
reader.rs

1/* ***
2 *
3 * Harmos
4 *
5 * SPDX-License-Identifier: Apache-2.0
6 *
7 * ***
8 */
9
10/* === Modules === */
11
12/* === Re-Exports === */
13
14/* === Imports === */
15
16// Standard Library Imports
17use std::any::Any;
18use std::sync::Arc;
19use std::sync::atomic::{AtomicBool, Ordering};
20// Third Party Imports
21
22// Local Imports
23use crate::{state::At, state::Position, storage::Cursor};
24
25/* === Constants === */
26
27/* === Macros === */
28
29/* === Type Aliases === */
30
31/* === Structs === */
32
33#[derive(Clone, Copy, Debug)]
34pub struct ReadLimits {
35    pub rows: usize,
36    pub row_bytes: usize,
37    pub batch_bytes: usize,
38    pub batch_rows: usize,
39    pub segments: usize,
40    pub segment_bytes: usize,
41    pub decoded_bytes: usize,
42    pub pinned_bytes: usize,
43}
44
45#[derive(Clone, Copy, Debug)]
46pub struct PageLimits {
47    pub max_rows: usize,
48    pub max_encoded_bytes: usize,
49}
50
51#[derive(Debug)]
52pub struct RowPage<S> {
53    pub rows: Vec<At<Arc<S>>>,
54    pub encoded_bytes: usize,
55    pub complete: bool,
56}
57
58struct Ready<S> {
59    row: At<Arc<S>>,
60    encoded_bytes: usize,
61}
62
63struct Owned {
64    source: Box<dyn RowSource>,
65    _permit: tokio::sync::OwnedSemaphorePermit,
66}
67
68pub struct RowReader<S> {
69    owned: Option<Owned>,
70    cut: ReadCut,
71    limits: ReadLimits,
72    returned: usize,
73    done: bool,
74    partition: Option<super::Partition<S>>,
75    pending: Option<Ready<S>>,
76}
77
78/* === Enums === */
79
80#[derive(Clone, Copy, Debug, PartialEq, Eq)]
81pub enum ReadCut {
82    Window { delivered: Cursor },
83    Series { epoch: u64, delivered: Cursor, ordinal: u64 },
84}
85
86#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
87pub enum ReadError {
88    #[error("the active row reader budget is exhausted")]
89    Busy,
90
91    #[error("row read exceeds its {0} limit")]
92    Limit(&'static str),
93
94    #[error("the attached series adapter does not support bounded reads")]
95    Unsupported,
96
97    #[error("bounded row decoding does not support {0}")]
98    EncodingUnsupported(&'static str),
99
100    #[error("no series adapter is attached")]
101    Unserved,
102
103    #[error("row history is retained only above {floor:?}")]
104    Retained { floor: Position },
105
106    #[error("selected stream capture is incomplete or uncertified")]
107    IncompleteCapture,
108
109    #[error("row source is unreadable: {0}")]
110    Unreadable(String),
111
112    #[error("row reader is closed")]
113    Closed,
114}
115
116/* === Traits === */
117
118pub trait RowSource: Send {
119    fn cut(&self) -> ReadCut;
120
121    #[allow(clippy::type_complexity)]
122    fn next(&mut self, cancelled: &AtomicBool) -> Result<Option<At<Arc<dyn Any + Send + Sync>>>, ReadError>;
123}
124
125/* === Implementations === */
126
127impl Default for ReadLimits {
128    fn default() -> Self {
129        Self {
130            rows: 100_000,
131            row_bytes: 1024 * 1024,
132            batch_bytes: 8 * 1024 * 1024,
133            batch_rows: 65_536,
134            segments: 256,
135            segment_bytes: 64 * 1024 * 1024,
136            decoded_bytes: 32 * 1024 * 1024,
137            pinned_bytes: 16 * 1024 * 1024,
138        }
139    }
140}
141
142impl Default for PageLimits {
143    fn default() -> Self {
144        Self {
145            max_rows: 256,
146            max_encoded_bytes: 256 * 1024,
147        }
148    }
149}
150
151impl<S: serde::Serialize + Send + Sync + 'static> RowReader<S> {
152    pub fn cut(&self) -> ReadCut {
153        self.cut
154    }
155
156    pub async fn next_page(&mut self, page_limits: PageLimits) -> Result<RowPage<S>, ReadError> {
157        if page_limits.max_rows == 0 || page_limits.max_encoded_bytes == 0 {
158            return Err(ReadError::Limit("positive page limits"));
159        }
160        let Some(mut owned) = self.owned.take() else {
161            return if self.done {
162                Ok(RowPage {
163                    rows: Vec::new(),
164                    encoded_bytes: 0,
165                    complete: true,
166                })
167            } else {
168                Err(ReadError::Closed)
169            };
170        };
171        let partition = self.partition.take();
172        let mut pending = self.pending.take();
173        let returned = self.returned;
174        let limits = self.limits;
175        struct Cancel(Arc<AtomicBool>);
176        impl Drop for Cancel {
177            fn drop(&mut self) {
178                self.0.store(true, Ordering::Release);
179            }
180        }
181        let cancelled = Arc::new(AtomicBool::new(false));
182        let _cancel = Cancel(cancelled.clone());
183        let result = tokio::task::spawn_blocking(move || {
184            let mut page = RowPage {
185                rows: Vec::new(),
186                encoded_bytes: 0usize,
187                complete: false,
188            };
189            loop {
190                if cancelled.load(Ordering::Acquire) {
191                    return Err(ReadError::Closed);
192                }
193                let next = match pending.take() {
194                    Some(ready) => Some(ready),
195                    None => loop {
196                        if cancelled.load(Ordering::Acquire) {
197                            return Err(ReadError::Closed);
198                        }
199                        let Some(row) = owned.source.next(&cancelled)? else {
200                            break None;
201                        };
202                        let typed = super::row_value::<S>(row.value.as_ref())
203                            .ok_or_else(|| ReadError::Unreadable("source row type mismatch".into()))?;
204                        if partition.as_ref().is_some_and(|p| !(p.matches)(typed)) {
205                            continue;
206                        }
207                        // Validate The Borrowed Original Before Copying Its Row.
208                        let encoded_bytes = encoded_size(typed, limits.row_bytes)?;
209                        let row_value = if let Some(managed) = row.value.downcast_ref::<super::ManagedPublication>() {
210                            Arc::new(managed.row().clone()) as Arc<dyn std::any::Any + Send + Sync>
211                        } else {
212                            row.value
213                        };
214                        let value = row_value.downcast::<S>().map_err(|_| ReadError::Unreadable("source row type mismatch".into()))?;
215                        break Some(Ready {
216                            row: At { value, position: row.position },
217                            encoded_bytes,
218                        });
219                    },
220                };
221                let Some(ready) = next else {
222                    page.complete = true;
223                    break;
224                };
225                if returned + page.rows.len() == limits.rows {
226                    return Err(ReadError::Limit("rows"));
227                }
228                let bytes = page.encoded_bytes.checked_add(ready.encoded_bytes);
229                if page.rows.len() == page_limits.max_rows || bytes.is_none_or(|bytes| bytes > page_limits.max_encoded_bytes) {
230                    if page.rows.is_empty() {
231                        return Err(ReadError::Limit("page row bytes"));
232                    }
233                    // One row of lookahead makes terminal completion exact,
234                    // even at a full-page boundary. Acquisition reserves this
235                    // slot from pinned_bytes; its encoding is measured once.
236                    pending = Some(ready);
237                    break;
238                }
239                page.encoded_bytes = bytes.expect("within the checked page budget");
240                page.rows.push(ready.row);
241            }
242            Ok((owned, partition, pending, page))
243        })
244        .await
245        .map_err(|_| ReadError::Closed)?;
246        let (owned, partition, pending, page) = result?;
247        self.partition = partition;
248        self.returned += page.rows.len();
249        self.done = page.complete;
250        if !page.complete {
251            self.owned = Some(owned);
252            self.pending = pending;
253        }
254        Ok(page)
255    }
256}
257
258impl<S: Clone + serde::Serialize + Send + Sync + 'static> super::Query<S> {
259    pub async fn reader(self, limits: ReadLimits) -> Result<RowReader<S>, ReadError> {
260        if [
261            limits.rows,
262            limits.row_bytes,
263            limits.batch_bytes,
264            limits.batch_rows,
265            limits.segments,
266            limits.segment_bytes,
267            limits.decoded_bytes,
268            limits.pinned_bytes,
269        ]
270        .contains(&0)
271        {
272            return Err(ReadError::Limit("positive limits"));
273        }
274        let reserved = limits.row_bytes.checked_add(std::mem::size_of::<Ready<S>>()).ok_or(ReadError::Limit("pending row reservation"))?;
275        let mut source_limits = limits;
276        source_limits.pinned_bytes = limits.pinned_bytes.checked_sub(reserved).ok_or(ReadError::Limit("pending row reservation"))?;
277        let permit = self.tier.reader_permit()?;
278        let end = match self.as_of {
279            Some(cut) => match self.end {
280                std::ops::Bound::Unbounded => std::ops::Bound::Included(cut),
281                std::ops::Bound::Included(end) if end > cut => std::ops::Bound::Included(cut),
282                std::ops::Bound::Excluded(end) if end > cut => std::ops::Bound::Included(cut),
283                end => end,
284            },
285            None => self.end,
286        };
287        let label_bytes = self
288            .partition
289            .as_ref()
290            .map_or(Some(0usize), |partition| partition.labels.iter().try_fold(0usize, |total, label| total.checked_add(label.len() + 32)));
291        if label_bytes.is_none_or(|bytes| bytes > limits.pinned_bytes) {
292            return Err(ReadError::Limit("track bytes"));
293        }
294        let track = self.partition.as_ref().map(|p| p.labels.clone()).unwrap_or_default();
295        let opened = self.tier.reader(self.definition, track, self.start, end, source_limits, permit).await?;
296        let source = match opened.source {
297            Err(ReadError::Unserved) => self.window.reader::<S>(self.start, end, source_limits)?,
298            result => result?,
299        };
300        Ok(RowReader {
301            cut: source.cut(),
302            owned: Some(Owned { source, _permit: opened.permit }),
303            limits,
304            returned: 0,
305            done: false,
306            partition: self.partition,
307            pending: None,
308        })
309    }
310}
311
312/* === Public Functions === */
313
314pub(crate) fn encoded_size<S: serde::Serialize>(row: &S, limit: usize) -> Result<usize, ReadError> {
315    struct Count {
316        size: usize,
317        limit: usize,
318    }
319    impl std::io::Write for Count {
320        fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
321            self.size = self
322                .size
323                .checked_add(bytes.len())
324                .filter(|size| *size <= self.limit)
325                .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::FileTooLarge, "row byte limit"))?;
326            Ok(bytes.len())
327        }
328        fn flush(&mut self) -> std::io::Result<()> {
329            Ok(())
330        }
331    }
332    let mut count = Count { size: 0, limit };
333    serde_json::to_writer(&mut count, row).map_err(|error| {
334        if error.io_error_kind() == Some(std::io::ErrorKind::FileTooLarge) {
335            ReadError::Limit("row bytes")
336        } else {
337            ReadError::Unreadable(error.to_string())
338        }
339    })?;
340    Ok(count.size)
341}
342
343/* === Private Functions === */
344
345/* === Tests === */
harmos_runtime/streams/reader.rs

harmos_runtime/streams/
reader.rs
