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

harmos_storage/query/
reader.rs
