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::Arc;
21
22// Third Party Imports
23use arrow_array::RecordBatch;
24use tokio::sync::OwnedSemaphorePermit;
25
26// Local Imports
27use super::execute::Plan;
28use super::object_schemas::NormalizedObjects;
29use super::{QueryCompletion, QueryOutput, Source};
30use crate::storage::arrow::{ArrowBatch, ArrowSchema, RetainedMemory, resident_bytes, retain_native};
31use crate::storage::backend::ScanPage;
32use crate::storage::capacity::{Held, Pool, Refusal, Reservation};
33use crate::storage::config::SCAN_ROWS;
34use crate::storage::manager::{ENTRY_OVERHEAD, Inner};
35use crate::storage::memory::{MemoryBank, Shortfall};
36use crate::storage::selection::Entry;
37use crate::storage::snapshot::Snapshot;
38use crate::storage::{FileReference, ObjectEntry, SchemaKey, StorageManager, StorageSchema, StoredRecord, StoredRow};
39use crate::{Error, MemorySize, Result};
40use harmos_domain::model::storage::ArrowIpc;
41use harmos_domain::operations::storage::Query;
42
43/* === Constants === */
44
45// Two Protobuf Keys (Five Bytes Each) And Two Length Varints (Ten Bytes Each),
46// Less The Postcard Variant And Length Bytes They Replace.
47const PROTOBUF_FRAMING: usize = 30;
48
49/* === Macros === */
50
51/* === Type Aliases === */
52
53/* === Structs === */
54
55pub struct QueryReader {
56    inner: Arc<Inner>,
57    held: Held,
58    query: Query,
59    snapshot: Option<Snapshot>,
60    active: Option<OwnedSemaphorePermit>,
61    _description: Reservation,
62    schema: Option<(StorageSchema, RetainedMemory)>,
63    objects: Option<NormalizedObjects>,
64    plan: Option<(Plan, RetainedMemory)>,
65    after: u64,
66    rows: NonZeroUsize,
67    source_done: bool,
68    pending: Option<RecordBatch>,
69    pending_offset: usize,
70    delivered: u64,
71    exhausted: bool,
72    failed: bool,
73    inflight: bool,
74}
75
76/* === Enums === */
77
78#[derive(Clone, Debug)]
79pub enum QueryArrow {
80    Schema(ArrowSchema),
81    Batch(ArrowBatch),
82}
83enum Rows {
84    Objects(Vec<ObjectEntry>),
85    Records(Vec<StoredRecord>),
86    Stream(Vec<StoredRow>),
87    Schemas(Vec<StorageSchema>),
88    Files(Vec<FileReference>),
89}
90
91/* === Traits === */
92
93/* === Implementations === */
94
95impl QueryArrow {
96    pub fn output(&self) -> QueryOutput {
97        match self {
98            Self::Schema(schema) => QueryOutput::Schema(ArrowIpc::new(schema.ipc_bytes().to_vec())),
99            Self::Batch(batch) => QueryOutput::Batch(ArrowIpc::new(batch.ipc_bytes().to_vec())),
100        }
101    }
102}
103
104impl TryFrom<&QueryOutput> for QueryArrow {
105    type Error = Error;
106
107    fn try_from(output: &QueryOutput) -> Result<Self> {
108        Ok(match output {
109            QueryOutput::Schema(ipc) => Self::Schema(ArrowSchema::from_ipc(ipc.bytes())?),
110            QueryOutput::Batch(ipc) => Self::Batch(ArrowBatch::from_ipc(ipc.bytes())?),
111        })
112    }
113}
114
115impl MemorySize for QueryArrow {
116    fn memory_size(&self) -> usize {
117        std::mem::size_of::<Self>()
118            + match self {
119                Self::Schema(value) => value.memory_size().saturating_sub(std::mem::size_of_val(value)),
120                Self::Batch(value) => value.memory_size().saturating_sub(std::mem::size_of_val(value)),
121            }
122    }
123}
124impl<C> StorageManager<C> {
125    pub async fn query(&self, query: Query) -> Result<QueryReader> {
126        query.validate()?;
127        self.inner.primary.history().check(&query.at.history)?;
128
129        // Admit A Query Slot Before Charging Its Description. Queued Callers
130        // Hold No Read Bytes That An Active Reader Might Need To Finish.
131        let active = query_slot(&self.inner).await?;
132        let held = Held::default();
133        let description_bytes = query.memory_size().saturating_add(ENTRY_OVERHEAD);
134        let description = Reservation::wait(&self.inner, Pool::Read, &held, description_bytes, None).await.map_err(Refusal::into_error)?;
135
136        // Pin The Requested Image For The Whole Invocation.
137        let snapshot = self.snapshot(&query.at).await?;
138        let after = match &query.source {
139            Source::Stream(selection) => selection.after,
140            _ => 0,
141        };
142        Ok(QueryReader {
143            inner: Arc::clone(&self.inner),
144            held,
145            query,
146            snapshot: Some(snapshot),
147            active: Some(active),
148            _description: description,
149            schema: None,
150            objects: None,
151            plan: None,
152            after,
153            rows: SCAN_ROWS,
154            source_done: false,
155            pending: None,
156            pending_offset: 0,
157            delivered: 0,
158            exhausted: false,
159            failed: false,
160            inflight: false,
161        })
162    }
163}
164
165impl QueryReader {
166    pub async fn next(&mut self) -> Result<Option<QueryArrow>> {
167        if self.failed || self.inflight {
168            self.fail();
169            return Err(Error::NotReady { domain: "failed or interrupted storage query" });
170        }
171        self.inflight = true;
172        let result = self.advance().await;
173        self.inflight = false;
174        if result.is_err() {
175            self.fail();
176        }
177        result
178    }
179
180    pub fn completion(&self) -> Result<QueryCompletion> {
181        if !self.exhausted || self.failed || self.inflight {
182            return Err(Error::NotReady { domain: "storage query completion" });
183        }
184        Ok(QueryCompletion { at: self.query.at.clone(), rows: self.delivered })
185    }
186
187    fn fail(&mut self) {
188        self.failed = true;
189        self.pending = None;
190        self.plan = None;
191        self.schema = None;
192        self.objects = None;
193        self.release();
194    }
195
196    // Free The Active Slot And Snapshot Once No More Source Input Is Needed.
197    fn release(&mut self) {
198        self.snapshot = None;
199        self.active = None;
200    }
201
202    async fn advance(&mut self) -> Result<Option<QueryArrow>> {
203        if self.exhausted {
204            return Ok(None);
205        }
206        if self.plan.is_none() {
207            let schema = self.prepare().await?;
208            return Ok(Some(QueryArrow::Schema(schema)));
209        }
210        loop {
211            if let Some(batch) = self.deliver().await? {
212                return Ok(Some(QueryArrow::Batch(batch)));
213            }
214            let (plan, _) = self.plan.as_ref().expect("prepared plan");
215            if plan.is_satisfied() {
216                self.source_done = true;
217            }
218            if !self.source_done {
219                self.next_page().await?;
220                continue;
221            }
222            self.release();
223            if !self.finish().await? {
224                self.exhausted = true;
225                return Ok(None);
226            }
227        }
228    }
229
230    // Resolve Source Schemas, Validate The Plan And Fix The Output Schema.
231    async fn prepare(&mut self) -> Result<ArrowSchema> {
232        let snapshot = self.snapshot.as_ref().expect("prepared query owns its snapshot");
233        let mut transforms = std::mem::take(&mut self.query.transforms);
234        match &self.query.source {
235            Source::Objects(selection) => {
236                let objects = NormalizedObjects::resolve(&self.inner, &self.held, snapshot, selection, &mut transforms).await?;
237                self.objects = Some(objects);
238            }
239            Source::Records(selection) => {
240                let schema = schema_at(&self.inner, &self.held, snapshot, &selection.schema).await?;
241                self.schema = Some(schema);
242            }
243            Source::Stream(selection) => {
244                let schema = schema_at(&self.inner, &self.held, snapshot, &selection.schema).await?;
245                self.schema = Some(schema);
246            }
247            Source::Schemas(_) | Source::Files(_) => {}
248        }
249
250        // Derive The Source Schema From An Empty Conversion.
251        let mut credit = 0usize;
252        loop {
253            let bank = admit(&self.inner, &self.held, credit).await?;
254            let attempt = (|| {
255                let empty = self.convert(None, &bank)?;
256                let plan = Plan::new(empty.schema(), transforms.clone(), self.query.slice)?;
257                let schema = ArrowSchema::try_new(plan.output_schema())?;
258                let retained =
259                    bank.reserve(plan.memory_size().saturating_add(schema.memory_size()).saturating_add(transforms.memory_size()))?;
260                Ok::<_, Error>((plan, schema, retained))
261            })();
262            match attempt {
263                Ok((plan, schema, retained)) => {
264                    let output = schema.retain(retained.clone());
265                    self.plan = Some((plan, retained));
266                    return Ok(output);
267                }
268                Err(error) => credit = retry(&bank, error, None)?,
269            }
270        }
271    }
272
273    // Load, Convert And Process One Source Page As One Reloadable Attempt.
274    async fn next_page(&mut self) -> Result<()> {
275        let mut credit = 0usize;
276        loop {
277            let bank = match MemoryBank::admit(&self.inner, &self.held, credit, None).await {
278                Ok(bank) => bank,
279                Err(Refusal::Exceeded(_)) if self.rows.get() > 1 => {
280                    self.rows = NonZeroUsize::new(self.rows.get() / 2).unwrap();
281                    credit = 0;
282                    continue;
283                }
284                Err(refusal) => return Err(refusal.into_error()),
285            };
286            match self.page_attempt(&bank).await {
287                Ok(()) => return Ok(()),
288                Err(error) => {
289                    let mut rows = self.rows;
290                    credit = retry(&bank, error, Some(&mut rows))?;
291                    self.rows = rows;
292                }
293            }
294        }
295    }
296
297    async fn page_attempt(&mut self, bank: &MemoryBank) -> Result<()> {
298        // Charge The Raw Page Before Backend I/O.
299        let scan_bytes = self.inner.settings.scan_bytes();
300        let page = ScanPage { after: self.after, rows: self.rows, bytes: scan_bytes };
301        let mut raw = bank.scope();
302        raw.claim(self.rows.get().saturating_mul(ENTRY_OVERHEAD).saturating_add(scan_bytes.get()))?;
303        let snapshot = self.snapshot.as_ref().expect("running query owns its snapshot");
304        let (rows, bytes) = match &self.query.source {
305            Source::Objects(selection) => scan(snapshot, selection, page, Rows::Objects).await?,
306            Source::Records(selection) => scan(snapshot, selection, page, Rows::Records).await?,
307            Source::Stream(selection) => scan(snapshot, selection, page, Rows::Stream).await?,
308            Source::Schemas(selection) => scan(snapshot, selection, page, Rows::Schemas).await?,
309            Source::Files(selection) => scan(snapshot, selection, page, Rows::Files).await?,
310        };
311        // A Single Oversized Entry Is Charged At Its Actual Size.
312        raw.claim(bytes.saturating_sub(scan_bytes.get()))?;
313        let Some(last) = rows.last() else {
314            self.source_done = true;
315            return Ok(());
316        };
317
318        // Convert And Run The Page; State Changes Only After Every Claim Succeeded.
319        let batch = self.convert(Some(&rows), bank)?;
320        let retained = bank.reserve(resident_bytes(&batch)?)?;
321        let batch = retain_native(batch, retained)?;
322        drop(rows);
323        let (plan, _) = self.plan.as_mut().expect("prepared plan");
324        let output = plan.page(&batch, bank, &mut || Ok(()))?;
325        self.pending = output.filter(|output| output.num_rows() != 0);
326        self.pending_offset = 0;
327        self.after = last;
328        Ok(())
329    }
330
331    // Produce The Global Output Once; `false` When Nothing Remains.
332    async fn finish(&mut self) -> Result<bool> {
333        let Some((plan, _)) = self.plan.as_mut() else {
334            return Ok(false);
335        };
336        let mut credit = plan.finish_bytes()?;
337        loop {
338            let bank = admit(&self.inner, &self.held, credit).await?;
339            let (plan, _) = self.plan.as_mut().expect("prepared plan");
340            match plan.finish(&bank, &mut || Ok(())) {
341                Ok(Some(output)) if output.num_rows() != 0 => {
342                    self.pending = Some(output);
343                    self.pending_offset = 0;
344                    return Ok(true);
345                }
346                Ok(_) => return Ok(false),
347                Err(error) => credit = retry(&bank, error, None)?,
348            }
349        }
350    }
351
352    // Split Pending Output Into Response Batches Within Both Ceilings.
353    async fn deliver(&mut self) -> Result<Option<ArrowBatch>> {
354        let Some(batch) = &self.pending else {
355            return Ok(None);
356        };
357        if self.pending_offset >= batch.num_rows() {
358            self.pending = None;
359            self.pending_offset = 0;
360            return Ok(None);
361        }
362        let settings = self.inner.settings.response;
363        let mut rows = settings.batch_rows.get().min(batch.num_rows() - self.pending_offset);
364        loop {
365            let slice = batch.slice(self.pending_offset, rows);
366            let allowance = resident_bytes(&slice)?.saturating_mul(4).saturating_add(settings.batch_bytes.get());
367            let mut reservation =
368                Reservation::wait(&self.inner, Pool::Read, &self.held, allowance, None).await.map_err(Refusal::into_error)?;
369            let output = ArrowBatch::try_new(slice)?;
370
371            // The Wire Output Adds Its Envelope Around The Arrow Bytes Field.
372            let envelope = QueryOutput::Batch(ArrowIpc::new(output.ipc_bytes().to_vec()));
373            // Protobuf Frames This One-Variant Bytes Envelope With Two Keys And Two
374            // Length Prefixes Beyond The Postcard Bytes; Bound Both Encodings.
375            let encoded = crate::Message::encoded_len(&envelope).saturating_add(PROTOBUF_FRAMING);
376            drop(envelope);
377            if encoded <= settings.batch_bytes.get() {
378                self.pending_offset += rows;
379                self.delivered += rows as u64;
380                reservation.shrink_to(output.encoding_memory_size()?);
381                return output.retain(reservation.retain()).map(Some);
382            }
383            if rows == 1 {
384                return Err(Error::ResourceExhausted { resource: "oversized query row" });
385            }
386            rows /= 2;
387        }
388    }
389
390    fn convert(&self, rows: Option<&Rows>, bank: &MemoryBank) -> Result<RecordBatch> {
391        let schema = self.schema.as_ref().map(|(schema, _)| schema);
392        match (&self.query.source, rows) {
393            (Source::Objects(_), rows) => {
394                let objects = self.objects.as_ref().expect("resolved object schemas");
395                match rows {
396                    Some(Rows::Objects(rows)) => objects.map(rows, bank),
397                    _ => Ok(objects.empty()),
398                }
399            }
400            (Source::Records(_), Some(Rows::Records(rows))) => super::source::records(schema.expect("record schema"), rows, bank),
401            (Source::Records(_), _) => super::source::records(schema.expect("record schema"), &[], bank),
402            (Source::Stream(_), Some(Rows::Stream(rows))) => super::source::stream(schema.expect("stream schema"), rows, bank),
403            (Source::Stream(_), _) => super::source::stream(schema.expect("stream schema"), &[], bank),
404            (Source::Schemas(_), Some(Rows::Schemas(rows))) => super::source::schemas(rows, bank),
405            (Source::Schemas(_), _) => super::source::schemas(&[], bank),
406            (Source::Files(_), Some(Rows::Files(rows))) => super::source::files(rows, bank),
407            (Source::Files(_), _) => super::source::files(&[], bank),
408        }
409    }
410}
411
412impl Rows {
413    fn last(&self) -> Option<u64> {
414        let position = match self {
415            Self::Objects(rows) => rows.last().map(Entry::position),
416            Self::Records(rows) => rows.last().map(Entry::position),
417            Self::Stream(rows) => rows.last().map(Entry::position),
418            Self::Schemas(rows) => rows.last().map(Entry::position),
419            Self::Files(rows) => rows.last().map(Entry::position),
420        };
421        position.map(|position| position.sequence)
422    }
423}
424
425/* === Public Functions === */
426
427pub(crate) async fn schema_at(
428    inner: &Arc<Inner>, held: &Held, snapshot: &Snapshot, key: &SchemaKey,
429) -> Result<(StorageSchema, RetainedMemory)> {
430    find_schema(inner, held, snapshot, key)
431        .await?
432        .ok_or(Error::InvalidInput { field: "storage schema", reason: "schema is unavailable at requested cursor" })
433}
434
435pub(crate) async fn find_schema(
436    inner: &Arc<Inner>, held: &Held, snapshot: &Snapshot, key: &SchemaKey,
437) -> Result<Option<(StorageSchema, RetainedMemory)>> {
438    let selection = super::SchemaSelection { schema: Some(key.clone()) };
439    let scan_bytes = inner.settings.scan_bytes();
440    let mut credit = 0usize;
441    loop {
442        let bank = admit(inner, held, credit).await?;
443        let attempt = async {
444            let mut raw = bank.scope();
445            raw.claim(scan_bytes.get().saturating_add(ENTRY_OVERHEAD))?;
446            let page = ScanPage { after: 0, rows: NonZeroUsize::MIN, bytes: scan_bytes };
447            let (mut schemas, _) = snapshot.scan::<StorageSchema>(&selection, page).await?;
448            let Some(schema) = schemas.pop() else {
449                return Ok::<_, Error>(None);
450            };
451            let retained = bank.reserve(schema.memory_size().saturating_add(ENTRY_OVERHEAD))?;
452            Ok(Some((schema, retained)))
453        };
454        match attempt.await {
455            Ok(resolved) => return Ok(resolved),
456            Err(error) => credit = retry(&bank, error, None)?,
457        }
458    }
459}
460
461pub(crate) async fn query_slot(inner: &Arc<Inner>) -> Result<OwnedSemaphorePermit> {
462    if let Ok(active) = Arc::clone(&inner.query_active).try_acquire_owned() {
463        return Ok(active);
464    }
465    let _pending = Arc::clone(&inner.query_pending)
466        .try_acquire_owned()
467        .map_err(|_| Error::ResourceExhausted { resource: "pending storage queries" })?;
468    Arc::clone(&inner.query_active).acquire_owned().await.map_err(|_| Error::Closing { domain: "storage queries" })
469}
470
471pub(crate) async fn admit(inner: &Arc<Inner>, held: &Held, credit: usize) -> Result<MemoryBank> {
472    MemoryBank::admit(inner, held, credit, None).await.map_err(Refusal::into_error)
473}
474
475pub(crate) fn retry(bank: &MemoryBank, error: Error, rows: Option<&mut NonZeroUsize>) -> Result<usize> {
476    match bank.shortfall() {
477        Some(Shortfall::Busy) => Ok(bank.requested()),
478        Some(Shortfall::Exceeded) => match rows {
479            Some(rows) if rows.get() > 1 => {
480                *rows = NonZeroUsize::new(rows.get() / 2).unwrap();
481                Ok(0)
482            }
483            _ => Err(Error::ResourceExhausted { resource: "storage read memory" }),
484        },
485        None => Err(error),
486    }
487}
488
489/* === Private Functions === */
490
491async fn scan<T: Entry>(snapshot: &Snapshot, selection: &T::Selection, page: ScanPage, wrap: fn(Vec<T>) -> Rows) -> Result<(Rows, usize)> {
492    let (entries, bytes) = snapshot.scan::<T>(selection, page).await?;
493    Ok((wrap(entries), bytes))
494}
495
496/* === Tests === */
harmos_storage/query/reader.rs

harmos_storage/query/
reader.rs
