harmos_storage/query/
reader.rs

1/* ***
2 *
3 * Harmos
4 *
5 * SPDX-License-Identifier: Apache-2.0
6 *
7 * ***
8 */
9
10// Storage Query Result Reader
11
12/* === Modules === */
13
14/* === Re-Exports === */
15
16/* === Imports === */
17
18// Standard Library Imports
19use std::num::NonZeroUsize;
20use std::sync::{
21    Arc,
22    atomic::{AtomicBool, Ordering},
23};
24
25// Third Party Imports
26use arrow_array::RecordBatch;
27use tokio::sync::OwnedSemaphorePermit;
28
29// Local Imports
30use super::execute::Plan;
31use super::object_schemas::NormalizedObjects;
32use super::transform::{ProjectExecution, Transform, project::Project};
33use super::{QueryCompletion, QueryOutput, Source};
34use crate::storage::arrow::{ArrowBatch, ArrowSchema};
35use crate::storage::backend::ScanPage;
36use crate::storage::config::SCAN_ROWS;
37use crate::storage::manager::Inner;
38use crate::storage::selection::Entry;
39use crate::storage::snapshot::Snapshot;
40use crate::storage::{FileReference, ObjectEntry, SchemaKey, StorageManager, StorageSchema, StoredRecord, StoredRow};
41use crate::{Error, MemorySize, Result};
42use harmos_domain::model::storage::ArrowIpc;
43
44/* === Constants === */
45
46// Two Protobuf Keys (Five Bytes Each) And Two Length Varints (Ten Bytes Each),
47// Less The Postcard Variant And Length Bytes They Replace.
48const PROTOBUF_FRAMING: usize = 30;
49
50/* === Macros === */
51
52/* === Type Aliases === */
53
54/* === Structs === */
55
56#[derive(Clone, Debug)]
57pub struct Query {
58    pub at: harmos_domain::model::storage::StorageCursor,
59    pub source: Source,
60    pub transforms: Vec<harmos_domain::model::storage::query::transform::Transform>,
61    pub slice: harmos_domain::model::storage::query::RowSlice,
62}
63
64pub struct QueryReader {
65    context: crate::execution::Context,
66    inner: Arc<Inner>,
67    query: Query,
68    snapshot: Option<Snapshot>,
69    active: Option<OwnedSemaphorePermit>,
70    schema: Option<StorageSchema>,
71    objects: Option<NormalizedObjects>,
72    plan: Option<Plan>,
73    projection: Option<(Project, Arc<arrow_schema::Schema>)>,
74    after: u64,
75    rows: NonZeroUsize,
76    source_done: bool,
77    pending: Option<RecordBatch>,
78    pending_offset: usize,
79    delivered: u64,
80    exhausted: bool,
81    failed: bool,
82    inflight: bool,
83}
84
85/* === Enums === */
86
87#[derive(Clone, Debug)]
88pub enum QueryArrow {
89    Schema(ArrowSchema),
90    Batch(ArrowBatch),
91}
92struct CancelWork(Arc<AtomicBool>);
93
94impl Drop for CancelWork {
95    fn drop(&mut self) {
96        self.0.store(true, Ordering::Release);
97    }
98}
99
100enum Rows {
101    Objects(Vec<ObjectEntry>),
102    Records(Vec<StoredRecord>),
103    Stream(Vec<StoredRow>),
104    Schemas(Vec<StorageSchema>),
105    Files(Vec<FileReference>),
106}
107
108/* === Traits === */
109
110/* === Implementations === */
111
112impl Query {
113    pub fn validate(&self) -> harmos_domain::Result<()> {
114        self.source.validate()?;
115        if self.transforms.len() > 64 {
116            return Err(harmos_domain::model::storage::query::transform::invalid("too many query transforms"));
117        }
118        Ok(())
119    }
120}
121
122impl MemorySize for Query {
123    fn memory_size(&self) -> usize {
124        std::mem::size_of::<Self>()
125            .saturating_add(self.at.memory_size().saturating_sub(std::mem::size_of_val(&self.at)))
126            .saturating_add(self.source.memory_size().saturating_sub(std::mem::size_of_val(&self.source)))
127            .saturating_add(self.transforms.memory_size().saturating_sub(std::mem::size_of_val(&self.transforms)))
128            .saturating_add(self.slice.memory_size().saturating_sub(std::mem::size_of_val(&self.slice)))
129    }
130}
131
132impl QueryArrow {
133    /// Bounds The Transport Value Before IPC Conversion.
134    pub fn output_memory_size(&self) -> usize {
135        let capacity = match self {
136            Self::Schema(schema) => schema.ipc_capacity(),
137            Self::Batch(batch) => batch.ipc_capacity(),
138        };
139        std::mem::size_of::<QueryOutput>().saturating_add(capacity)
140    }
141
142    #[tracing::instrument(name = "storage.query.reader.output", level = "trace", skip_all)]
143    pub fn output(self) -> QueryOutput {
144        match self {
145            Self::Schema(schema) => QueryOutput::Schema(ArrowIpc::from(schema)),
146            Self::Batch(batch) => QueryOutput::Batch(ArrowIpc::from(batch)),
147        }
148    }
149}
150
151impl TryFrom<QueryOutput> for QueryArrow {
152    type Error = Error;
153
154    fn try_from(output: QueryOutput) -> Result<Self> {
155        Ok(match output {
156            QueryOutput::Schema(ipc) => Self::Schema(ArrowSchema::from_ipc(ipc.into_bytes())?),
157            QueryOutput::Batch(ipc) => Self::Batch(ArrowBatch::from_ipc(ipc.into_bytes())?),
158        })
159    }
160}
161
162impl MemorySize for QueryArrow {
163    fn memory_size(&self) -> usize {
164        std::mem::size_of::<Self>()
165            + match self {
166                Self::Schema(value) => value.memory_size().saturating_sub(std::mem::size_of_val(value)),
167                Self::Batch(value) => value.memory_size().saturating_sub(std::mem::size_of_val(value)),
168            }
169    }
170}
171impl<C> StorageManager<C> {
172    #[tracing::instrument(name = "storage.manager.query", level = "trace", skip_all, fields(outcome = "cancelled"))]
173    pub async fn query(&self, query: Query) -> Result<QueryReader> {
174        let result = async {
175            query.validate()?;
176            self.inner.primary.history().check(&query.at.history)?;
177
178            let active = query_slot(&self.inner).await?;
179
180            // Pin The Requested Image For The Whole Invocation.
181            let snapshot = self.snapshot(&query.at).await?;
182            let after = match &query.source {
183                Source::Stream(selection) => selection.after,
184                _ => 0,
185            };
186            Ok(QueryReader {
187                context: crate::execution::Context::current(),
188                inner: Arc::clone(&self.inner),
189                query,
190                snapshot: Some(snapshot),
191                active: Some(active),
192                schema: None,
193                objects: None,
194                plan: None,
195                projection: None,
196                after,
197                rows: SCAN_ROWS,
198                source_done: false,
199                pending: None,
200                pending_offset: 0,
201                delivered: 0,
202                exhausted: false,
203                failed: false,
204                inflight: false,
205            })
206        }
207        .await;
208        tracing::Span::current().record("outcome", if result.is_ok() { "completed" } else { "failed" });
209        result
210    }
211}
212
213impl QueryReader {
214    pub async fn next(&mut self) -> Result<Option<QueryArrow>> {
215        let context = self.context.clone();
216        let span = context.sync(|| tracing::trace_span!("storage.reader.page"));
217        let future = async {
218            if self.failed || self.inflight {
219                self.fail();
220                return Err(Error::NotReady {
221                    domain: "failed or interrupted storage query",
222                });
223            }
224            self.inflight = true;
225            let result = self.advance().await;
226            self.inflight = false;
227            if result.is_err() {
228                self.fail();
229            }
230            result
231        };
232        context.scope(tracing::Instrument::instrument(future, span)).await
233    }
234
235    #[tracing::instrument(name = "storage.query.reader.completion", level = "trace", skip_all)]
236    pub fn completion(&self) -> Result<QueryCompletion> {
237        if !self.exhausted || self.failed || self.inflight {
238            return Err(Error::NotReady { domain: "storage query completion" });
239        }
240        Ok(QueryCompletion {
241            at: self.query.at.clone(),
242            rows: self.delivered,
243        })
244    }
245
246    fn fail(&mut self) {
247        self.failed = true;
248        self.pending = None;
249        self.plan = None;
250        self.schema = None;
251        self.objects = None;
252        self.projection = None;
253        self.release();
254    }
255
256    // Release Admission And Any Remaining Snapshot At Completion Or Failure.
257    fn release(&mut self) {
258        self.snapshot = None;
259        self.active = None;
260    }
261
262    #[tracing::instrument(name = "storage.query.reader.advance", level = "trace", skip_all)]
263    async fn advance(&mut self) -> Result<Option<QueryArrow>> {
264        if self.exhausted {
265            return Ok(None);
266        }
267        if self.plan.is_none() {
268            let schema = self.prepare().await?;
269            return Ok(Some(QueryArrow::Schema(schema)));
270        }
271        loop {
272            if let Some(batch) = self.deliver().await? {
273                return Ok(Some(QueryArrow::Batch(batch)));
274            }
275            let plan = self.plan.as_ref().expect("prepared plan");
276            if plan.is_satisfied() {
277                self.source_done = true;
278            }
279            if !self.source_done {
280                self.next_page().await?;
281                continue;
282            }
283            self.snapshot = None;
284            if !self.finish().await? {
285                self.release();
286                self.exhausted = true;
287                return Ok(None);
288            }
289        }
290    }
291
292    // Resolve Source Schemas, Validate The Plan And Fix The Output Schema.
293    #[tracing::instrument(name = "storage.query.reader.prepare", level = "trace", skip_all)]
294    async fn prepare(&mut self) -> Result<ArrowSchema> {
295        let snapshot = self.snapshot.as_ref().expect("prepared query owns its snapshot");
296        let mut transforms = std::mem::take(&mut self.query.transforms);
297        if transforms.len() > super::execute::MAX_TRANSFORMS {
298            return Err(super::transform::invalid("too many query transforms"));
299        }
300        match &self.query.source {
301            Source::Objects(selection) => {
302                let objects = NormalizedObjects::resolve(&self.inner, snapshot, selection, &mut transforms).await?;
303                self.objects = Some(objects);
304            }
305            Source::Records(selection) => {
306                let schema = schema_at(&self.inner, snapshot, &selection.schema).await?;
307                self.schema = Some(schema);
308            }
309            Source::Stream(selection) => {
310                let schema = schema_at(&self.inner, snapshot, &selection.schema).await?;
311                self.schema = Some(schema);
312            }
313            Source::Schemas(_) | Source::Files(_) => {}
314        }
315
316        let mut empty = self.convert(None)?;
317        crate::storage::arrow::validate_schema(&empty.schema())?;
318        // Payload-Only Projections Decode Their Selected Columns Directly.
319        if matches!(self.query.source, Source::Records(_) | Source::Stream(_))
320            && let Some(Transform::Project(project)) = transforms.first()
321            && project.fields.iter().all(|field| field.field.fields.first().is_some_and(|name| name == "payload"))
322        {
323            empty = project.apply(&empty, &mut || Ok(()))?;
324            let Transform::Project(project) = transforms.remove(0) else { unreachable!() };
325            self.projection = Some((project, empty.schema()));
326        }
327        let plan = Plan::new(empty.schema(), transforms, self.query.slice)?;
328        let schema = ArrowSchema::try_new(plan.output_schema())?;
329        self.plan = Some(plan);
330        Ok(schema)
331    }
332
333    // Load, Convert And Process One Source Page As One Reloadable Attempt.
334    #[tracing::instrument(name = "storage.query.reader.next_page", level = "trace", skip_all, fields(after = self.after, rows = self.rows.get()))]
335    async fn next_page(&mut self) -> Result<()> {
336        self.page_attempt().await
337    }
338
339    #[tracing::instrument(name = "storage.query.reader.page_attempt", level = "trace", skip_all)]
340    async fn page_attempt(&mut self) -> Result<()> {
341        let scan_bytes = self.inner.settings.scan_bytes();
342        let page = ScanPage {
343            after: self.after,
344            rows: self.rows,
345            bytes: scan_bytes,
346        };
347
348        let snapshot = self.snapshot.as_ref().expect("running query owns its snapshot");
349
350        let (rows, _) = match &self.query.source {
351            Source::Objects(selection) => scan(snapshot, selection, page, Rows::Objects).await?,
352            Source::Records(selection) => scan(snapshot, selection, page, Rows::Records).await?,
353            Source::Stream(selection) => scan(snapshot, selection, page, Rows::Stream).await?,
354            Source::Schemas(selection) => scan(snapshot, selection, page, Rows::Schemas).await?,
355            Source::Files(selection) => scan(snapshot, selection, page, Rows::Files).await?,
356        };
357
358        let Some(last) = rows.last() else {
359            self.source_done = true;
360            return Ok(());
361        };
362
363        // Advance Only After The Complete Page Converts.
364        let batch = self.convert(Some(&rows))?;
365        drop(rows);
366        let output = self.work(Some(batch)).await?;
367        self.pending = output.filter(|output| output.num_rows() != 0);
368        self.pending_offset = 0;
369        self.after = last;
370        Ok(())
371    }
372
373    // Produce The Global Output Once; `false` When Nothing Remains.
374    #[tracing::instrument(name = "storage.query.reader.finish", level = "trace", skip_all)]
375    async fn finish(&mut self) -> Result<bool> {
376        if self.plan.is_none() {
377            return Ok(false);
378        }
379        match self.work(None).await? {
380            Some(output) if output.num_rows() != 0 => {
381                self.pending = Some(output);
382                self.pending_offset = 0;
383                Ok(true)
384            }
385            _ => Ok(false),
386        }
387    }
388
389    // A Dropped Caller Cancels Blocking Work Before Its Next Bounded Step.
390    async fn work(&mut self, batch: Option<RecordBatch>) -> Result<Option<RecordBatch>> {
391        let mut plan = self.plan.take().expect("prepared query plan");
392        let cancelled = Arc::new(AtomicBool::new(false));
393        let guard = CancelWork(Arc::clone(&cancelled));
394        let context = crate::execution::Context::current();
395        let result = tokio::task::spawn_blocking(move || {
396            context.blocking(|| {
397                let mut check = || {
398                    if cancelled.load(Ordering::Acquire) { Err(Error::Closing { domain: "cancelled storage query" }) } else { Ok(()) }
399                };
400                let output = match batch {
401                    Some(batch) => plan.page(&batch, &mut check),
402                    None => plan.finish(&mut check),
403                }?;
404                Ok::<_, Error>((plan, output))
405            })
406        })
407        .await
408        .map_err(|error| Error::backend("query worker", crate::BackendFailure::Permanent, error))??;
409        drop(guard);
410        self.plan = Some(result.0);
411        Ok(result.1)
412    }
413
414    // Split Pending Output Into Response Batches Within Both Ceilings.
415    #[tracing::instrument(name = "storage.query.reader.deliver", level = "trace", skip_all)]
416    async fn deliver(&mut self) -> Result<Option<ArrowBatch>> {
417        let Some(batch) = &self.pending else {
418            return Ok(None);
419        };
420        if self.pending_offset >= batch.num_rows() {
421            self.pending = None;
422            self.pending_offset = 0;
423            return Ok(None);
424        }
425        let settings = self.inner.settings.response;
426        let mut rows = settings.batch_rows.get().min(batch.num_rows() - self.pending_offset);
427        loop {
428            let slice = batch.slice(self.pending_offset, rows);
429            let output = ArrowBatch::try_new(slice)?;
430
431            // The Wire Output Adds Its Envelope Around The Arrow Bytes Field.
432            let envelope = QueryOutput::Batch(ArrowIpc::new(output.ipc_bytes().to_vec()));
433            // Protobuf Frames This One-Variant Bytes Envelope With Two Keys And Two
434            // Length Prefixes Beyond The Postcard Bytes; Bound Both Encodings.
435            let encoded = crate::Message::encoded_size(&envelope).saturating_add(PROTOBUF_FRAMING);
436            drop(envelope);
437            if encoded <= settings.batch_bytes.get() {
438                self.pending_offset += rows;
439                self.delivered += rows as u64;
440                return Ok(Some(output));
441            }
442            if rows == 1 {
443                return Err(Error::ResourceExhausted { resource: "oversized query row" });
444            }
445            rows /= 2;
446        }
447    }
448
449    #[tracing::instrument(name = "storage.query.reader.convert", level = "trace", skip_all)]
450    fn convert(&self, rows: Option<&Rows>) -> Result<RecordBatch> {
451        let schema = self.schema.as_ref();
452        match (&self.query.source, rows) {
453            (Source::Objects(_), rows) => {
454                let objects = self.objects.as_ref().expect("resolved object schemas");
455                match rows {
456                    Some(Rows::Objects(rows)) => objects.map(rows),
457                    _ => Ok(objects.empty()),
458                }
459            }
460            (Source::Records(_), Some(Rows::Records(rows))) => match &self.projection {
461                Some((project, output)) => super::source::projected(schema.expect("record schema"), rows.iter(), project, output.clone()),
462                None => super::source::records(schema.expect("record schema"), rows),
463            },
464            (Source::Records(_), _) => super::source::records(schema.expect("record schema"), &[]),
465            (Source::Stream(_), Some(Rows::Stream(rows))) => match &self.projection {
466                Some((project, output)) => {
467                    let records = rows.iter().map(|row| &row.record);
468                    super::source::projected(schema.expect("stream schema"), records, project, output.clone())
469                }
470                None => super::source::stream(schema.expect("stream schema"), rows),
471            },
472            (Source::Stream(_), _) => super::source::stream(schema.expect("stream schema"), &[]),
473            (Source::Schemas(_), Some(Rows::Schemas(rows))) => super::source::schemas(rows),
474            (Source::Schemas(_), _) => super::source::schemas(&[]),
475            (Source::Files(_), Some(Rows::Files(rows))) => super::source::files(rows),
476            (Source::Files(_), _) => super::source::files(&[]),
477        }
478    }
479}
480
481impl Rows {
482    fn last(&self) -> Option<u64> {
483        let position = match self {
484            Self::Objects(rows) => rows.last().map(Entry::position),
485            Self::Records(rows) => rows.last().map(Entry::position),
486            Self::Stream(rows) => rows.last().map(Entry::position),
487            Self::Schemas(rows) => rows.last().map(Entry::position),
488            Self::Files(rows) => rows.last().map(Entry::position),
489        };
490        position.map(|position| position.sequence)
491    }
492}
493
494/* === Public Functions === */
495
496#[tracing::instrument(name = "storage.query.reader.schema_at", level = "trace", skip_all)]
497pub(crate) async fn schema_at(inner: &Arc<Inner>, snapshot: &Snapshot, key: &SchemaKey) -> Result<StorageSchema> {
498    find_schema(inner, snapshot, key).await?.ok_or(Error::InvalidInput {
499        field: "storage schema",
500        reason: "schema is unavailable at requested cursor",
501    })
502}
503
504#[tracing::instrument(name = "storage.query.reader.find_schema", level = "trace", skip_all)]
505pub(crate) async fn find_schema(inner: &Arc<Inner>, snapshot: &Snapshot, key: &SchemaKey) -> Result<Option<StorageSchema>> {
506    let selection = super::SchemaSelection { schema: Some(key.clone()) };
507    let page = ScanPage {
508        after: 0,
509        rows: NonZeroUsize::MIN,
510        bytes: inner.settings.scan_bytes(),
511    };
512    let (mut schemas, _) = snapshot.scan::<StorageSchema>(&selection, page).await?;
513    Ok(schemas.pop())
514}
515
516#[tracing::instrument(name = "storage.query.reader.query_slot", level = "trace", skip_all)]
517pub(crate) async fn query_slot(inner: &Arc<Inner>) -> Result<OwnedSemaphorePermit> {
518    if let Ok(active) = Arc::clone(&inner.query_active).try_acquire_owned() {
519        return Ok(active);
520    }
521    let _pending = Arc::clone(&inner.query_pending)
522        .try_acquire_owned()
523        .map_err(|_| Error::ResourceExhausted { resource: "pending storage queries" })?;
524    Arc::clone(&inner.query_active).acquire_owned().await.map_err(|_| Error::Closing { domain: "storage queries" })
525}
526
527/* === Private Functions === */
528
529#[tracing::instrument(name = "storage.query.reader.scan", level = "trace", skip_all)]
530async fn scan<T: Entry>(snapshot: &Snapshot, selection: &T::Selection, page: ScanPage, wrap: fn(Vec<T>) -> Rows) -> Result<(Rows, usize)> {
531    let (entries, bytes) = snapshot.scan::<T>(selection, page).await?;
532    Ok((wrap(entries), bytes))
533}
534
535/* === Tests === */
harmos_storage/query/reader.rs

harmos_storage/query/
reader.rs
