haematite-cli 0.21.8

//! The one place a query gets records from: documents at one captured moment,
//! in (collection, id) order, handed on only when the request's access allows.
//!
//! A source of documents reads each record's body. Each captured shard holds
//! at most one candidate; a shard's next candidate is read on the call after
//! its last one was taken, never before, so the source never reads a record it
//! was not asked for.
//!
//! A source of origins reads each document's revision entry (`m\0`) and no
//! body, in the same (collection, id) order. Each shard reads ahead in
//! batches: as many entries as the page may still take, at most
//! `ORIGINS_BATCH`, and at least one. So a page takes no more than its limit,
//! and reads at most one batch per shard that it does not take.
//!
//! A source takes at most its limit of candidates, counting those the request
//! may not read. At the limit, with candidates left, it stops and names the
//! last place it took, given or not, so the next page starts just after it: a
//! long run of records the caller cannot read is crossed page by page, and
//! the place sits only in a sealed continuation.

#[cfg(test)]
use std::cell::Cell;
use std::collections::VecDeque;
use std::num::NonZeroUsize;

use haematite::db::{ObserverError, ReadOnlyDatabase};
use haematite::{Hash, TreeError};
use haematite_query::Refusal;
use haematite_query::handle::Error;
use haematite_query::page::{Origin, Record, RecordSource};
use serde_json::value::RawValue;

use crate::serve::access::Access;
use crate::serve::storage::{Captured, DocumentRevision, document_key, revision_key};

const DOCUMENTS: &[u8] = b"d\0";
const REVISIONS: &[u8] = b"m\0";

/// The most revision entries one shard reads at once for a source of origins.
pub const ORIGINS_BATCH: NonZeroUsize = NonZeroUsize::MIN.saturating_add(63);

/// What a source reads for each record.
#[derive(Clone, Copy, PartialEq, Eq)]
enum Kind {
    /// The document's body and its revision: a record for the worker.
    Documents,
    /// The revision entry alone: an origin, with no body.
    Origins,
}

impl Kind {
    const fn prefix(self) -> &'static [u8] {
        match self {
            Self::Documents => DOCUMENTS,
            Self::Origins => REVISIONS,
        }
    }

    fn key(self, collection: &str, id: &str) -> Vec<u8> {
        match self {
            Self::Documents => document_key(collection, id),
            Self::Origins => revision_key(collection, id),
        }
    }
}

/// Which documents a query reads.
pub enum Scope {
    /// Every document of every collection.
    All,
    /// The documents of one collection.
    Collection(String),
}

/// A document key: the place a source starts from, inclusive unless `after`.
pub struct Place {
    pub collection: String,
    pub id: String,
    /// Start just after this document: it was passed at a page's limit.
    pub after: bool,
}

type RawRow = (Vec<u8>, Vec<u8>);

/// One captured shard's place in the merge.
struct Shard {
    id: usize,
    root: Hash,
    /// The first key this shard has not yet been read from.
    from: Vec<u8>,
    /// Entries read and not yet taken, in key order; the next read is made
    /// once they are all taken.
    ahead: VecDeque<RawRow>,
    /// A read came back short: nothing is left past `from`.
    ended: bool,
}

/// Records at a captured moment that the request's one `Access` may read.
pub struct AuthorizedSource<'a> {
    kind: Kind,
    access: &'a Access,
    reader: &'a ReadOnlyDatabase,
    at: &'a Captured,
    to: Vec<u8>,
    shards: Vec<Shard>,
    scan_limit: u64,
    /// Every physical candidate taken, including those the request may not read.
    scanned: u64,
    /// The key of the last candidate taken.
    last: Option<Vec<u8>>,
    /// Where the source stopped at its limit with candidates left.
    stopped: Option<Origin>,
    #[cfg(test)]
    reads: Cell<u64>,
    #[cfg(test)]
    decoded: Cell<u64>,
}

impl<'a> AuthorizedSource<'a> {
    /// Logical candidates read, including records denied by request access.
    pub(in crate::serve) const fn scanned(&self) -> u64 {
        self.scanned
    }

    pub(super) fn has_candidates(&mut self) -> Result<bool, Error> {
        self.fill()?;
        Ok(self.shards.iter().any(|shard| !shard.ahead.is_empty()))
    }

    pub(super) fn last_place(&self) -> Result<Option<Origin>, Error> {
        let Some(last) = self.last.as_deref() else {
            return Ok(None);
        };
        let (collection, id) = parts(self.kind, last)?;
        Ok(Some(Origin {
            collection: collection.to_owned(),
            id: id.to_owned(),
            version: RawValue::from_string("null".to_owned())
                .map_err(|error| read_failed(error.to_string()))?,
            score: None,
        }))
    }

    /// A source over `scope` at `at`, from `start` (inclusive) when given.
    /// Nothing is read until the first `next_record`.
    ///
    /// # Errors
    ///
    /// `query_protocol` when `start` lies outside `scope`.
    pub fn new(
        access: &'a Access,
        reader: &'a ReadOnlyDatabase,
        at: &'a Captured,
        scope: &Scope,
        start: Option<&Place>,
        scan_limit: u64,
    ) -> Result<Self, Error> {
        let span = Span { scope, start };
        Self::open(Kind::Documents, access, reader, at, &span, scan_limit)
    }

    /// As [`Self::new`], giving each record's origin only: its revision entry
    /// is read, its body never. Each record's value is `null`.
    ///
    /// # Errors
    ///
    /// As [`Self::new`].
    pub fn origins(
        access: &'a Access,
        reader: &'a ReadOnlyDatabase,
        at: &'a Captured,
        scope: &Scope,
        start: Option<&Place>,
        scan_limit: u64,
    ) -> Result<Self, Error> {
        let span = Span { scope, start };
        Self::open(Kind::Origins, access, reader, at, &span, scan_limit)
    }

    fn open(
        kind: Kind,
        access: &'a Access,
        reader: &'a ReadOnlyDatabase,
        at: &'a Captured,
        span: &Span<'_>,
        scan_limit: u64,
    ) -> Result<Self, Error> {
        let prefix = match span.scope {
            Scope::All => kind.prefix().to_vec(),
            // A collection's keys are its documents' keys with an empty id.
            Scope::Collection(collection) => kind.key(collection, ""),
        };
        let mut to = prefix.clone();
        let last = to
            .last_mut()
            .ok_or_else(|| protocol("empty scope prefix"))?;
        *last = last
            .checked_add(1)
            .ok_or_else(|| protocol("unbounded scope prefix"))?;
        let from = match span.start {
            None => prefix,
            Some(place) => {
                let mut key = kind.key(&place.collection, &place.id);
                if !key.starts_with(&prefix) || key >= to {
                    return Err(protocol("the start place lies outside the query's scope"));
                }
                if place.after {
                    key.push(0); // As `take` does: ids never contain NUL.
                }
                key
            }
        };
        let shards = at
            .shards
            .iter()
            .map(|shard| Shard {
                id: shard.shard_id,
                root: shard.head,
                from: from.clone(),
                ahead: VecDeque::new(),
                ended: false,
            })
            .collect();
        Ok(Self {
            kind,
            access,
            reader,
            at,
            to,
            shards,
            scan_limit,
            scanned: 0,
            last: None,
            stopped: None,
            #[cfg(test)]
            reads: Cell::new(0),
            #[cfg(test)]
            decoded: Cell::new(0),
        })
    }

    /// How many entries a shard reads at once: one document, or as many
    /// origins as this page may still take, at most `ORIGINS_BATCH`, and at
    /// least one so the source can tell its end from its limit.
    fn batch(&self) -> NonZeroUsize {
        match self.kind {
            Kind::Documents => NonZeroUsize::MIN,
            Kind::Origins => {
                let left = self.scan_limit.saturating_sub(self.scanned);
                let left = usize::try_from(left).unwrap_or(usize::MAX);
                NonZeroUsize::new(left.min(ORIGINS_BATCH.get())).unwrap_or(NonZeroUsize::MIN)
            }
        }
    }

    /// Reads the next batch for every shard whose entries were all taken.
    fn fill(&mut self) -> Result<(), Error> {
        let batch = self.batch();
        for shard in &mut self.shards {
            if shard.ended || !shard.ahead.is_empty() {
                continue;
            }
            #[cfg(test)]
            {
                self.reads.set(self.reads.get() + 1);
            }
            let entries = self
                .reader
                .range_page_at(shard.id, shard.root, &shard.from, &self.to, batch)
                .map_err(read_failure)?
                .entries;
            shard.ended = entries.len() < batch.get();
            if let Some((key, _)) = entries.last() {
                shard.from.clone_from(key);
                shard.from.push(0); // Exclusive of the key read; ids never contain NUL.
            }
            shard.ahead.extend(entries);
        }
        Ok(())
    }

    /// Takes the least entry across the shards; `None` when none is left.
    fn take(&mut self) -> Option<RawRow> {
        let shard = self
            .shards
            .iter_mut()
            .filter(|shard| !shard.ahead.is_empty())
            .min_by(|left, right| left.ahead.front().cmp(&right.ahead.front()))?;
        shard.ahead.pop_front()
    }

    /// At the limit: the end when no candidate is left, else a stop after the
    /// last place taken. A limit of none moves nowhere and is refused.
    fn stop(&mut self) -> Result<Option<Record>, Error> {
        if self.shards.iter().all(|shard| shard.ahead.is_empty()) {
            return Ok(None);
        }
        let last = self.last.as_deref().ok_or_else(|| {
            Error::new(
                Refusal::QueryScanLimit,
                "a page that may read no records moves nowhere",
            )
        })?;
        let (collection, id) = parts(self.kind, last)?;
        let version = RawValue::from_string("null".to_owned())
            .map_err(|error| read_failed(format!("a passed place: {error}")))?;
        self.stopped = Some(Origin {
            collection: collection.to_owned(),
            id: id.to_owned(),
            version,
            score: None,
        });
        Ok(None)
    }

    /// The live value at `key` at the captured root of the shard that owns it.
    fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Error> {
        let shard_id = self.reader.shard_for(key);
        let shard = self
            .at
            .shards
            .iter()
            .find(|shard| shard.shard_id == shard_id)
            .ok_or_else(|| {
                Error::new(
                    Refusal::QueryMomentGone,
                    format!("the captured moment does not name shard {shard_id}"),
                )
            })?;
        #[cfg(test)]
        {
            self.reads.set(self.reads.get() + 1);
        }
        self.reader
            .get_at(shard_id, shard.head, key)
            .map_err(read_failure)
    }

    /// The record for a document the request may read: its body and the
    /// version its own revision was written at.
    fn record(&self, collection: &str, id: &str, body: &[u8]) -> Result<Record, Error> {
        #[cfg(test)]
        {
            self.decoded.set(self.decoded.get() + 1);
        }
        let value: Box<RawValue> = serde_json::from_slice(body).map_err(|error| {
            read_failed(format!("document {collection}/{id} is not JSON: {error}"))
        })?;
        let revision = self
            .get(&revision_key(collection, id))?
            .ok_or_else(|| read_failed(format!("document {collection}/{id} has no revision")))?;
        Ok(Record {
            value,
            origin: origin(collection, id, &revision)?,
        })
    }

    /// The record for a revision entry the request may read: its origin, and
    /// a `null` value, for no body is read.
    fn origin_record(collection: &str, id: &str, revision: &[u8]) -> Result<Record, Error> {
        let value = RawValue::from_string("null".to_owned())
            .map_err(|error| read_failed(format!("an origin's value: {error}")))?;
        Ok(Record {
            value,
            origin: origin(collection, id, revision)?,
        })
    }

    /// The record of a document a search found, carrying its `score`, when
    /// the request may read it. A vector whose document is gone at the
    /// captured moment is a read failure, never skipped.
    pub fn found(&self, collection: &str, id: &str, score: f64) -> Result<Option<Record>, Error> {
        if !self.access.can_read_document(collection, id) {
            return Ok(None);
        }
        let body = self
            .get(&document_key(collection, id))?
            .ok_or_else(|| read_failed(format!("{collection}/{id} has a vector but no record")))?;
        let mut record = self.record(collection, id, &body)?;
        record.origin.score = Some(score);
        Ok(Some(record))
    }
}

impl RecordSource for AuthorizedSource<'_> {
    fn next_record(&mut self) -> Result<Option<Record>, Error> {
        loop {
            if self.scanned >= self.scan_limit {
                // Entries already read say whether any are left; a shard is
                // read again only when none are.
                if self.shards.iter().all(|shard| shard.ahead.is_empty()) {
                    self.fill()?;
                }
                return self.stop();
            }
            self.fill()?;
            let Some((key, body)) = self.take() else {
                return Ok(None);
            };
            self.scanned += 1;
            let (collection, id) = parts(self.kind, &key)?;
            if !self.access.can_read_document(collection, id) {
                self.last = Some(key);
                continue;
            }
            let record = match self.kind {
                Kind::Documents => self.record(collection, id, &body)?,
                Kind::Origins => Self::origin_record(collection, id, &body)?,
            };
            self.last = Some(key);
            return Ok(Some(record));
        }
    }

    fn passed(&self) -> Option<Origin> {
        self.stopped.clone()
    }
}

/// The start of a source's range.
struct Span<'a> {
    scope: &'a Scope,
    start: Option<&'a Place>,
}

/// The origin of `collection`/`id`, from its revision entry: the version it
/// was written at.
fn origin(collection: &str, id: &str, revision: &[u8]) -> Result<Origin, Error> {
    let revision: DocumentRevision = serde_json::from_slice(revision).map_err(|error| {
        read_failed(format!(
            "document {collection}/{id} revision is unreadable: {error}"
        ))
    })?;
    let version = serde_json::value::to_raw_value(&revision.written_at)
        .map_err(|error| read_failed(format!("version did not serialise: {error}")))?;
    Ok(Origin {
        collection: collection.to_owned(),
        id: id.to_owned(),
        version,
        score: None,
    })
}

/// A document or revision key's collection and id.
fn parts(kind: Kind, key: &[u8]) -> Result<(&str, &str), Error> {
    key.strip_prefix(kind.prefix())
        .and_then(|rest| std::str::from_utf8(rest).ok())
        .and_then(|rest| rest.split_once('\0'))
        .ok_or_else(|| read_failed("a stored document key does not name a collection and id"))
}

/// A node the captured moment needs is gone: named, never read around.
fn read_failure(error: ObserverError) -> Error {
    match error {
        ObserverError::Tree(TreeError::MissingNode { hash }) => Error::new(
            Refusal::QueryMomentGone,
            format!("the captured moment no longer holds node {hash}"),
        ),
        ObserverError::RootUnavailable { shard_id, root } => Error::new(
            Refusal::QueryMomentGone,
            format!("shard {shard_id} cannot read the captured root {root}"),
        ),
        other => read_failed(format!("record read failed: {other}")),
    }
}

fn protocol(message: &str) -> Error {
    Error::new(Refusal::QueryProtocol, message)
}

/// A store read that failed for a cause other than a moment that is gone.
fn read_failed(message: impl Into<String>) -> Error {
    Error::new(Refusal::QueryReadFailed, message)
}

#[cfg(test)]
#[path = "source_tests.rs"]
mod tests;