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 { max_rows: 256, max_encoded_bytes: 256 * 1024 }
145    }
146}
147
148impl<S: serde::Serialize + Send + Sync + 'static> RowReader<S> {
149    pub fn cut(&self) -> ReadCut {
150        self.cut
151    }
152
153    pub async fn next_page(&mut self, page_limits: PageLimits) -> Result<RowPage<S>, ReadError> {
154        if page_limits.max_rows == 0 || page_limits.max_encoded_bytes == 0 {
155            return Err(ReadError::Limit("positive page limits"));
156        }
157        let Some(mut owned) = self.owned.take() else {
158            return if self.done { Ok(RowPage { rows: Vec::new(), encoded_bytes: 0, complete: true }) } else { Err(ReadError::Closed) };
159        };
160        let partition = self.partition.take();
161        let mut pending = self.pending.take();
162        let returned = self.returned;
163        let limits = self.limits;
164        struct Cancel(Arc<AtomicBool>);
165        impl Drop for Cancel {
166            fn drop(&mut self) {
167                self.0.store(true, Ordering::Release);
168            }
169        }
170        let cancelled = Arc::new(AtomicBool::new(false));
171        let _cancel = Cancel(cancelled.clone());
172        let result = tokio::task::spawn_blocking(move || {
173            let mut page = RowPage { rows: Vec::new(), encoded_bytes: 0usize, complete: false };
174            loop {
175                if cancelled.load(Ordering::Acquire) {
176                    return Err(ReadError::Closed);
177                }
178                let next = match pending.take() {
179                    Some(ready) => Some(ready),
180                    None => loop {
181                        if cancelled.load(Ordering::Acquire) {
182                            return Err(ReadError::Closed);
183                        }
184                        let Some(row) = owned.source.next(&cancelled)? else {
185                            break None;
186                        };
187                        let typed = super::row_value::<S>(row.value.as_ref())
188                            .ok_or_else(|| ReadError::Unreadable("source row type mismatch".into()))?;
189                        if partition.as_ref().is_some_and(|p| !(p.matches)(typed)) {
190                            continue;
191                        }
192                        // Validate the borrowed original before allocating a legacy raw copy.
193                        let encoded_bytes = encoded_size(typed, limits.row_bytes)?;
194                        let row_value = if let Some(managed) = row.value.downcast_ref::<super::ManagedPublication>() {
195                            Arc::new(managed.row().clone()) as Arc<dyn std::any::Any + Send + Sync>
196                        } else {
197                            row.value
198                        };
199                        let value = row_value.downcast::<S>().map_err(|_| ReadError::Unreadable("source row type mismatch".into()))?;
200                        break Some(Ready { row: At { value, position: row.position }, encoded_bytes });
201                    },
202                };
203                let Some(ready) = next else {
204                    page.complete = true;
205                    break;
206                };
207                if returned + page.rows.len() == limits.rows {
208                    return Err(ReadError::Limit("rows"));
209                }
210                let bytes = page.encoded_bytes.checked_add(ready.encoded_bytes);
211                if page.rows.len() == page_limits.max_rows || bytes.is_none_or(|bytes| bytes > page_limits.max_encoded_bytes) {
212                    if page.rows.is_empty() {
213                        return Err(ReadError::Limit("page row bytes"));
214                    }
215                    // One row of lookahead makes terminal completion exact,
216                    // even at a full-page boundary. Acquisition reserves this
217                    // slot from pinned_bytes; its encoding is measured once.
218                    pending = Some(ready);
219                    break;
220                }
221                page.encoded_bytes = bytes.expect("within the checked page budget");
222                page.rows.push(ready.row);
223            }
224            Ok((owned, partition, pending, page))
225        })
226        .await
227        .map_err(|_| ReadError::Closed)?;
228        let (owned, partition, pending, page) = result?;
229        self.partition = partition;
230        self.returned += page.rows.len();
231        self.done = page.complete;
232        if !page.complete {
233            self.owned = Some(owned);
234            self.pending = pending;
235        }
236        Ok(page)
237    }
238}
239
240impl<S: Clone + serde::Serialize + Send + Sync + 'static> super::Query<S> {
241    pub async fn reader(self, limits: ReadLimits) -> Result<RowReader<S>, ReadError> {
242        if [
243            limits.rows,
244            limits.row_bytes,
245            limits.batch_bytes,
246            limits.batch_rows,
247            limits.segments,
248            limits.segment_bytes,
249            limits.decoded_bytes,
250            limits.pinned_bytes,
251        ]
252        .contains(&0)
253        {
254            return Err(ReadError::Limit("positive limits"));
255        }
256        let reserved = limits.row_bytes.checked_add(std::mem::size_of::<Ready<S>>()).ok_or(ReadError::Limit("pending row reservation"))?;
257        let mut source_limits = limits;
258        source_limits.pinned_bytes = limits.pinned_bytes.checked_sub(reserved).ok_or(ReadError::Limit("pending row reservation"))?;
259        let permit = self.tier.reader_permit()?;
260        let end = match self.as_of {
261            Some(cut) => match self.end {
262                std::ops::Bound::Unbounded => std::ops::Bound::Included(cut),
263                std::ops::Bound::Included(end) if end > cut => std::ops::Bound::Included(cut),
264                std::ops::Bound::Excluded(end) if end > cut => std::ops::Bound::Included(cut),
265                end => end,
266            },
267            None => self.end,
268        };
269        let label_bytes = self
270            .partition
271            .as_ref()
272            .map_or(Some(0usize), |partition| partition.labels.iter().try_fold(0usize, |total, label| total.checked_add(label.len() + 32)));
273        if label_bytes.is_none_or(|bytes| bytes > limits.pinned_bytes) {
274            return Err(ReadError::Limit("track bytes"));
275        }
276        let track = self.partition.as_ref().map(|p| p.labels.clone()).unwrap_or_default();
277        let opened = self.tier.reader(self.definition, track, self.start, end, source_limits, permit).await?;
278        let source = match opened.source {
279            Err(ReadError::Unserved) => self.window.reader::<S>(self.start, end, source_limits)?,
280            result => result?,
281        };
282        Ok(RowReader {
283            cut: source.cut(),
284            owned: Some(Owned { source, _permit: opened.permit }),
285            limits,
286            returned: 0,
287            done: false,
288            partition: self.partition,
289            pending: None,
290        })
291    }
292}
293
294/* === Public Functions === */
295
296pub(crate) fn encoded_size<S: serde::Serialize>(row: &S, limit: usize) -> Result<usize, ReadError> {
297    struct Count {
298        size: usize,
299        limit: usize,
300    }
301    impl std::io::Write for Count {
302        fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
303            self.size = self
304                .size
305                .checked_add(bytes.len())
306                .filter(|size| *size <= self.limit)
307                .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::FileTooLarge, "row byte limit"))?;
308            Ok(bytes.len())
309        }
310        fn flush(&mut self) -> std::io::Result<()> {
311            Ok(())
312        }
313    }
314    let mut count = Count { size: 0, limit };
315    serde_json::to_writer(&mut count, row).map_err(|error| {
316        if error.io_error_kind() == Some(std::io::ErrorKind::FileTooLarge) {
317            ReadError::Limit("row bytes")
318        } else {
319            ReadError::Unreadable(error.to_string())
320        }
321    })?;
322    Ok(count.size)
323}
324
325/* === Private Functions === */
326
327/* === Tests === */
harmos_runtime/streams/reader.rs

harmos_runtime/streams/
reader.rs
