haematite_cli/serve/storage/

embed_pending.rs

//! V3 (ruled 5 Oct 08:15, 08:27 and 08:31): the next page of the embedding
//! job's walk for one declared field of one branch.
//!
//! A walk is the difference over the field's collection's revision keys
//! (`m\0{collection}\0`) from the position's `through` (the empty tree before
//! the first walk ends) to its target, `toward`: the head when the walk
//! begins, kept while the head moves. Its cost is in proportion to the
//! changes, never to the store: the engine's diff skips equal subtrees, and
//! each change costs one read of its document at `toward`.
//!
//! The shards are walked together, in one id order. Each gives at most one
//! change more than a page after the walk's `after` (JT1b): the engine stops a
//! shard the moment it holds its budget, so a budget of exactly a page would
//! stop a shard holding exactly a page's changes as if it had more, and the
//! walk that has read its last record would not say so. One that has more
//! stops at its `next_after`, and the page keeps only the changes at or
//! before the smallest such key, so no record another shard still holds can
//! fall before the page's `Upto`; a shard that gave the change past the page
//! is cut back to a page, which then reaches `Upto` its last change. The
//! page claims `End` only when no shard has more and all it holds fits in
//! one page.
//!
//! A shard's stop can be a key that is no change, or no record at all (an
//! internal node's separator, settled by the engine's `settle_boundary`), so
//! `Upto` is always the last change the page holds. A page with no change and
//! a stop goes on from that stop within the same call: the engine settles a
//! stop only strictly past the key it was given (`diff_page.rs` `settle`) and
//! refuses by name when it cannot, so this always moves and never spins.
//!
//! EMB-E2 (j): while a field's rewalk marker is in its companion, its open
//! walk runs from the empty tree to `through`, and lists only the texts with
//! no vector and no windows there: the ones passed as longer than the model.
//! Nothing else is named, so no record is counted twice.

use std::num::NonZeroUsize;
use std::sync::Arc;

use haem_frames::embedding::{ItemCode, input_refusal};
use haematite::db::ObservedChange;
use haematite::tree::{DiffBudget, DiffInterval, Hash, empty_root_hash};
use serde_json::Value;

use super::embed_job::Reach;
use super::embedded::{self, Declaration};
use super::{
    Captured, Catalogue, DocumentRevision, Failure, Held, Storage, document_key, identifier, reads,
    revision_key,
};
use crate::serve::protocol::{Selection, Version};

/// Most tree nodes one shard's diff page may load.
const DIFF_NODES: usize = 4096;

/// One change the job must make a vector for, or remove.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Waiting {
    /// The record's text at `toward`, and the revision it is at.
    Text {
        id: String,
        made_from: Version,
        text: String,
    },
    Removed {
        id: String,
    },
}

/// The next page of a walk.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Pending {
    pub toward: Version,
    pub waiting: Vec<Waiting>,
    /// Records with no text in the source field, or text that is not a
    /// string: no vector is made (Q10).
    pub no_text: Vec<String>,
    /// Records whose text is longer than a worker takes: never cut, no
    /// vector is made (Q10).
    pub too_long: Vec<String>,
    /// EMB-E1: records whose text is empty or only whitespace, which the
    /// worker refuses as `embed_input_empty`: passed with no vector, never
    /// sent, so the walk goes on past them.
    pub empty_input: Vec<String>,
    pub reach: Reach,
}

impl Storage {
    pub(in crate::serve) fn pending_reader(&self) -> PendingReader {
        PendingReader {
            reader: Arc::clone(&self.reader),
            catalogue: self.catalogue(),
        }
    }

    /// The next page of `branch`'s walk for `collection.field`, at most
    /// `page_records` changes; `None` when the position is at the head and
    /// no walk is open.
    ///
    /// # Errors
    /// `branch_missing`; `embed_field_not_declared`; `embed_walk_released`
    /// when a version the walk needs is no longer kept;
    /// `embed_record_id_corrupt`; whatever a read or the diff answers.
    pub fn embed_pending(
        &self,
        branch: &str,
        collection: &str,
        field: &str,
        page_records: NonZeroUsize,
        input_bytes: NonZeroUsize,
    ) -> Result<Option<Pending>, Failure> {
        self.pending_reader()
            .read(branch, collection, field, page_records, input_bytes)
            .map(|page| page.map(|(pending, _)| pending))
    }
}

/// Immutable authority for a pending page; no writer admission is held.
pub(in crate::serve) struct PendingReader {
    reader: Arc<haematite::db::ReadOnlyDatabase>,
    catalogue: Catalogue,
}

pub(in crate::serve) struct PendingPage {
    pub pending: Pending,
    pub declaration: Declaration,
    pub position: Option<embedded::Position>,
}

impl PendingReader {
    pub fn read(
        &self,
        branch: &str,
        collection: &str,
        field: &str,
        page_records: NonZeroUsize,
        input_bytes: NonZeroUsize,
    ) -> Result<Option<(Pending, Declaration)>, Failure> {
        self.read_page(branch, collection, field, page_records, input_bytes)
            .map(|page| page.map(|page| (page.pending, page.declaration)))
    }

    pub fn read_page(
        &self,
        branch: &str,
        collection: &str,
        field: &str,
        page_records: NonZeroUsize,
        input_bytes: NonZeroUsize,
    ) -> Result<Option<PendingPage>, Failure> {
        identifier(branch)?;
        identifier(collection)?;
        identifier(field)?;
        let head = self.catalogue.capture_held(&Selection {
            branch: branch.to_owned(),
            version: None,
        })?;
        let declaration = self.declared(&head, collection, field)?;
        let companion = self.catalogue.capture_companion_held(branch)?;
        let position = self.position(companion.as_ref(), collection, field)?;
        let (from_branch, through) = position
            .as_ref()
            .map_or((branch, None), |at| (at.branch.as_str(), at.through));
        let open = position.as_ref().and_then(|at| at.toward.as_ref());
        let toward = open.map_or(head.version, |open| open.version);
        if open.is_none() && from_branch == branch && through == Some(head.version) {
            return Ok(None);
        }
        // A rewalk is an open walk toward the position's own `through`.
        let rewalk = open.is_some_and(|open| through == Some(open.version))
            && self.rewalking(companion.as_ref(), collection, field)?;
        let from_hold = through
            .filter(|_| !rewalk)
            .map(|through| self.kept(from_branch, through))
            .transpose()?;
        let from: Vec<Hash> = from_hold.as_ref().map_or_else(
            || vec![empty_root_hash(); self.reader.shard_count()],
            |held| held.shards.iter().map(|shard| shard.head).collect(),
        );
        let to = self.kept(branch, toward)?;
        let after = open.map(|open| revision_key(collection, &open.after));
        let (changes, more) =
            self.walk_page(collection, &from, &to, after.as_deref(), page_records)?;
        let mut pending = Pending {
            toward,
            waiting: Vec::new(),
            no_text: Vec::new(),
            too_long: Vec::new(),
            empty_input: Vec::new(),
            reach: Reach::End,
        };
        let sort = Sort {
            to: &to,
            declaration: &declaration,
            collection,
            field,
            limit: u64::try_from(input_bytes.get()).unwrap_or(u64::MAX),
            rewalk: companion.as_deref().filter(|_| rewalk),
        };
        for change in &changes {
            self.sort(&mut pending, change, &sort)?;
        }
        if more {
            let last = changes.last().ok_or_else(|| {
                Failure::new(
                    "embed_walk_stalled",
                    "a page with more after it holds no change",
                )
            })?;
            pending.reach = Reach::Upto(record_id(collection, last.key())?);
        }
        Ok(Some(PendingPage {
            pending,
            declaration,
            position,
        }))
    }

    fn declared(
        &self,
        at: &Captured,
        collection: &str,
        field: &str,
    ) -> Result<Declaration, Failure> {
        let undeclared = || {
            Failure::new(
                "embed_field_not_declared",
                format!("{collection}.{field} is not declared embedded"),
            )
        };
        let bytes = reads::get_bytes(
            &self.reader,
            at,
            &super::embed_job::field_key(collection, field),
        )?
        .ok_or_else(undeclared)?;
        match embedded::read_field(&bytes)? {
            embedded::Field::Embedded(declaration) => Ok(declaration),
            embedded::Field::Supplied { .. } => Err(undeclared()),
        }
    }

    /// Puts one change of the walk in its place on `pending`. EMB-E1: each
    /// text is checked as the worker checks a request, so the walk never
    /// lists one the worker would refuse. EMB-E2 (j): a rewalk lists only a
    /// text with no vector and no windows, and names nothing else.
    ///
    /// # Errors
    /// `embed_record_id_corrupt`; whatever a read answers.
    fn sort(
        &self,
        pending: &mut Pending,
        change: &ObservedChange,
        sort: &Sort<'_>,
    ) -> Result<(), Failure> {
        let id = record_id(sort.collection, change.key())?;
        let revision = match change {
            ObservedChange::Removed { .. } => {
                pending.waiting.push(Waiting::Removed { id });
                return Ok(());
            }
            ObservedChange::Added { value, .. }
            | ObservedChange::Modified {
                new_value: value, ..
            } => value,
        };
        if let Some(companion) = sort.rewalk
            && self.has_vectors(companion, sort, &id)?
        {
            return Ok(());
        }
        let made_from = serde_json::from_slice::<DocumentRevision>(revision)
            .map_err(Failure::storage)?
            .written_at;
        let named = sort.rewalk.is_none();
        match text(
            &self.reader,
            sort.to,
            sort.declaration,
            sort.collection,
            &id,
        )? {
            None if named => pending.no_text.push(id),
            Some(text) => match input_refusal(&text, sort.limit) {
                None => pending.waiting.push(Waiting::Text {
                    id,
                    made_from,
                    text,
                }),
                Some(ItemCode::EmbedInputEmpty) if named => pending.empty_input.push(id),
                Some(_) if named => pending.too_long.push(id),
                Some(_) => {}
            },
            None => {}
        }
        Ok(())
    }

    /// EMB-E2 (j): whether `collection.field`'s rewalk marker is in its
    /// companion.
    fn rewalking(
        &self,
        companion: Option<&Held>,
        collection: &str,
        field: &str,
    ) -> Result<bool, Failure> {
        let Some(companion) = companion else {
            return Ok(false);
        };
        let key = embedded::rewalk_key(collection, field);
        Ok(reads::get_bytes(&self.reader, companion, &key)?.is_some())
    }

    /// EMB-E2 (j): whether `id` has a vector or windows in `companion`: two
    /// point reads, the vector key and its first window's.
    fn has_vectors(
        &self,
        companion: &Captured,
        sort: &Sort<'_>,
        id: &str,
    ) -> Result<bool, Failure> {
        let vector = embedded::vector_key(sort.collection, sort.field, id);
        let window = embedded::window_key(sort.collection, sort.field, id, 0);
        Ok(
            reads::get_bytes(&self.reader, companion, &vector)?.is_some()
                || reads::get_bytes(&self.reader, companion, &window)?.is_some(),
        )
    }

    fn position(
        &self,
        companion: Option<&Held>,
        collection: &str,
        field: &str,
    ) -> Result<Option<embedded::Position>, Failure> {
        let Some(companion) = companion else {
            return Ok(None);
        };
        reads::get_bytes(
            &self.reader,
            companion,
            &embedded::position_key(collection, field),
        )?
        .map(|bytes| embedded::read_position(&bytes))
        .transpose()
    }

    /// One page of the difference, in key order across every shard, and
    /// whether the walk has more after it; never empty when it has.
    fn walk_page(
        &self,
        collection: &str,
        from: &[Hash],
        to: &Captured,
        after: Option<&[u8]>,
        page_records: NonZeroUsize,
    ) -> Result<(Vec<ObservedChange>, bool), Failure> {
        let mut after = after.map(<[u8]>::to_vec);
        loop {
            let (changes, stop) =
                self.walk_once(collection, from, to, after.as_deref(), page_records)?;
            match stop {
                Some(stop) if changes.is_empty() => after = Some(stop),
                stop => return Ok((changes, stop.is_some())),
            }
        }
    }

    /// One page of the difference after `after`, and the key it stops at
    /// when the walk has more after it.
    fn walk_once(
        &self,
        collection: &str,
        from: &[Hash],
        to: &Captured,
        after: Option<&[u8]>,
        page_records: NonZeroUsize,
    ) -> Result<(Vec<ObservedChange>, Option<Vec<u8>>), Failure> {
        let lower = format!("m\0{collection}\0").into_bytes();
        let upper = format!("m\0{collection}\u{1}").into_bytes();
        let budget = DiffBudget {
            // One past the page, so a shard holding exactly a page finishes
            // with no stop; one that gives the change past it is cut below.
            max_changes: page_records.checked_add(1).ok_or_else(|| {
                Failure::new(
                    "embed_page_records_overflow",
                    "the page setting leaves no room to look one change past it",
                )
            })?,
            max_nodes: NonZeroUsize::new(DIFF_NODES)
                .ok_or_else(|| Failure::new("diff_budget", "node budget must be nonzero"))?,
        };
        let mut changes = Vec::new();
        // The smallest key at which a shard stopped with more to give.
        let mut bound: Option<Vec<u8>> = None;
        for shard in &to.shards {
            let from_root = from.get(shard.shard_id).copied().ok_or_else(|| {
                Failure::new(
                    "incomplete_roots",
                    "the walk's start does not name every shard",
                )
            })?;
            let page = self
                .reader
                .changes_page_in_range(
                    shard.shard_id,
                    from_root,
                    shard.head,
                    DiffInterval {
                        lower: Some(&lower),
                        upper: Some(&upper),
                    },
                    after,
                    budget,
                )
                .map_err(Failure::observer)?;
            if let Some(next) = page.next_after
                && bound.as_ref().is_none_or(|bound| next < *bound)
            {
                bound = Some(next);
            }
            changes.extend(page.entries);
        }
        if let Some(bound) = &bound {
            changes.retain(|change| change.key() <= bound.as_slice());
        }
        changes.sort_by(|left, right| left.key().cmp(right.key()));
        if changes.len() > page_records.get() {
            changes.truncate(page_records.get());
            let last = changes.last().map(|change| change.key().to_vec());
            return Ok((changes, last));
        }
        Ok((changes, bound))
    }

    /// The roots of `branch` at `at`: its head, or, through the catalogue's
    /// view, a version its record keeps or a kept history row (STOR-003 MD4).
    fn kept(&self, branch: &str, at: Version) -> Result<Held, Failure> {
        self.catalogue
            .capture_held(&Selection {
                branch: branch.to_owned(),
                version: Some(at),
            })
            .map_err(|failure| walk_unreadable(branch, at, failure))
    }
}

/// What a capture of the walk's target answered, as the walk sees it: a
/// version missing, released or of an ended generation cannot be read, and
/// each is `embed_walk_released` with the capture's own code kept in the
/// message (STOR-003 MD4-b); anything else is answered as itself.
fn walk_unreadable(branch: &str, at: Version, failure: Failure) -> Failure {
    const UNREADABLE: [&str; 3] = [
        "version_missing",
        "version_released",
        "version_generation_ended",
    ];
    if UNREADABLE.contains(&failure.code.as_str()) {
        Failure::new(
            "embed_walk_released",
            format!(
                "{branch} at {}.{} is no longer kept, so the walk cannot be read ({}: {})",
                at.generation, at.sequence, failure.code, failure.message
            ),
        )
    } else {
        failure
    }
}

/// The record id of a revision key in `collection`.
/// What a page's changes are sorted against.
struct Sort<'a> {
    to: &'a Captured,
    declaration: &'a Declaration,
    collection: &'a str,
    field: &'a str,
    /// The most bytes of text a worker takes.
    limit: u64,
    /// EMB-E2 (j): the companion, when the walk is a rewalk.
    rewalk: Option<&'a Captured>,
}

fn record_id(collection: &str, key: &[u8]) -> Result<String, Failure> {
    let corrupt = || Failure::new("embed_record_id_corrupt", "a revision key names no record");
    let prefix = format!("m\0{collection}\0");
    let id = key.strip_prefix(prefix.as_bytes()).ok_or_else(corrupt)?;
    String::from_utf8(id.to_vec()).map_err(|_not_utf8| corrupt())
}

/// The declared source field's text in a record at `at`, if it has one.
fn text(
    reader: &haematite::db::ReadOnlyDatabase,
    at: &Captured,
    declaration: &Declaration,
    collection: &str,
    id: &str,
) -> Result<Option<String>, Failure> {
    let Some(bytes) = reads::get_bytes(reader, at, &document_key(collection, id))? else {
        return Ok(None);
    };
    let document: Value = serde_json::from_slice(&bytes).map_err(Failure::storage)?;
    Ok(document
        .get(&declaration.source)
        .and_then(Value::as_str)
        .map(str::to_owned))
}

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