haematite_cli/serve/storage/

embed_job.rs

//! V2 and V2b (ruled 5 Oct 08:15 and 08:27): one page of the embedding job's
//! work, committed into the branch's companion namespace with the job's
//! position, in one commit.
//!
//! A page names at most the page setting's records (default 64, several
//! worker calls), so one commit is bounded by the setting, never by the
//! store. The branch itself never moves, so no client's `expected` is ever
//! refused because the job wrote (Q1).
//!
//! The job walks the difference from the position's `through` to a target
//! version, `toward`, in key order, a page at a time. A page in the middle of
//! the walk moves its `after`; the page that ends it makes `toward` the new
//! `through`. A crash between pages resumes after the last committed `after`,
//! so no record is embedded twice and none is lost. Whether a page ends its
//! walk is the job's claim: storage cannot tell a last page from one with
//! records still after it, so V3 proves the job claims the end only when the
//! difference has no more rows. Within one walk `after` never goes back, and
//! a new walk cannot begin while one is open.
//!
//! A page's text was read at its walk's `toward`. A vector is written only
//! while its record's revision at the branch's head is still the one its
//! text was read at. A record rewritten or removed after `toward` is skipped:
//! the next walk lists it again, so nothing is written as current for text it
//! was not made from. A revision at or before `toward` that is not the one
//! given means the page was not read there, and the whole page is refused.
//!
//! The whole page is also refused, with nothing written, for a vector whose
//! length differs from the declaration's `dimensions` or one made by a model
//! other than the declared one.
//!
//! EMB-E2: a record longer than the model reads may come back as its
//! windows. They are written in the same commit as the rest of the page,
//! every one or none, under `vw\0` keys with their spans, and the record's
//! `v\0` vector of any earlier text is deleted. Any record a page touches
//! loses the windows of its earlier text past those it now has: window 0
//! names how many there are, so that is one read, never a scan.
//!
//! Owed elsewhere: a discarded branch's companion is removed in the same
//! finishing sequence as the discard (the discard box), and a caller's
//! `vector.put` to a field declared embedded is refused (V6).

use std::num::NonZeroUsize;

use haematite::branch::{CommitDurability, CommitRequest};

use super::embedded::{self, Declaration, Field, ModelName, Position, Provenance, Toward};
use super::{
    BranchKind, Captured, DocumentRevision, Failure, Storage, branch_name, create_branch_with_kind,
    current_timestamp, identifier, open_branch, reads, revision_key, version,
};
use crate::serve::protocol::Version;

/// Records one commit may name when no setting says otherwise.
pub const DEFAULT_PAGE_RECORDS: usize = 64;

/// One record's vector, and the revision its text was read at.
#[derive(Debug, Clone, PartialEq)]
pub struct Made {
    pub id: String,
    pub made_from: Version,
    pub vector: Vec<f32>,
}

/// EMB-E2: one window of a record's text, its bytes `[start, end)`, and its
/// vector.
#[derive(Debug, Clone, PartialEq)]
pub struct Window {
    pub start: u64,
    pub end: u64,
    pub vector: Vec<f32>,
}

/// EMB-E2: a record's windows in text order, and the revision its text was
/// read at.
#[derive(Debug, Clone, PartialEq)]
pub struct MadeWindows {
    pub id: String,
    pub made_from: Version,
    pub windows: Vec<Window>,
}

/// One page of the job's work on one declared field of one branch.
#[derive(Debug, Clone, PartialEq)]
pub struct Page {
    pub collection: String,
    pub field: String,
    /// The model the worker's Ready named for these vectors.
    pub model: ModelName,
    pub made: Vec<Made>,
    /// EMB-E2: records longer than the model reads, each embedded whole as
    /// its windows.
    pub windowed: Vec<MadeWindows>,
    /// Records removed from the branch between the walk's start and
    /// `toward`.
    pub removed: Vec<String>,
    /// RED4: records still there whose text at `toward` makes no vector,
    /// too long in bytes or for the model: a vector made from an earlier
    /// text is deleted, so none is kept for text the record no longer has.
    pub cleared: Vec<String>,
    /// The walk's target: the branch version the page's text was read at.
    pub toward: Version,
    /// How far through the walk this page reaches.
    pub reach: Reach,
}

/// How far a page reaches in its walk.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Reach {
    /// The walk goes on after this record id; every record up to and
    /// including it is finished.
    Upto(String),
    /// The job claims the walk has no records after this page's.
    End,
}

/// What one page's commit did.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Committed {
    pub written: usize,
    pub removed: usize,
    /// Changed after `toward`: no vector written, and listed again.
    pub skipped: Vec<String>,
    pub position: Position,
}

const fn order(at: Version) -> (u64, u64) {
    (at.generation, at.sequence)
}

/// The branch key of a field's `w\0` value, as vectors.rs lays it out.
#[must_use]
pub fn field_key(collection: &str, field: &str) -> Vec<u8> {
    format!("w\0{collection}\0{field}").into_bytes()
}

/// Whose versions the kept `through` is, that `through`, and the last record
/// the open walk finished, after checking the page belongs where the
/// position is.
///
/// A position may name another branch, as a fork's inherited one names its
/// parent (F-1): its `through` is that branch's version, and stays labelled
/// so until the walk ends (V2c). An open walk in this companion is always
/// this branch's own, since a companion is inherited only with no walk open
/// (F-2).
fn walk(
    previous: Option<&Position>,
    branch: &str,
    page: &Page,
) -> Result<(String, Option<Version>, Option<String>), Failure> {
    let Some(previous) = previous else {
        return Ok((branch.to_owned(), None, None));
    };
    if let Some(open) = &previous.toward {
        if open.version != page.toward {
            return Err(Failure::new(
                "embed_walk_open",
                "a new walk cannot begin while one is open",
            ));
        }
        if let Reach::Upto(after) = &page.reach
            && after.as_bytes() < open.after.as_bytes()
        {
            return Err(Failure::new(
                "embed_walk_backwards",
                "a page cannot reach back past the walk's last record",
            ));
        }
        return Ok((
            previous.branch.clone(),
            previous.through,
            Some(open.after.clone()),
        ));
    }
    if previous.branch == branch
        && previous
            .through
            .is_some_and(|through| order(page.toward) <= order(through))
    {
        return Err(Failure::new(
            "embed_position_backwards",
            "a walk must reach past the job's position",
        ));
    }
    Ok((previous.branch.clone(), previous.through, None))
}

/// One page's companion writes, as they are gathered.
struct Writes<'page> {
    /// The companion's head before this page, `None` before its first.
    at: Option<Captured>,
    page: &'page Page,
    pending: Vec<(Vec<u8>, Option<Vec<u8>>)>,
    dimensions: usize,
}

impl Writes<'_> {
    /// # Errors
    /// `embed_vector_length` for a vector of another length than declared.
    fn length(&self, id: &str, vector: &[f32]) -> Result<(), Failure> {
        if vector.len() == self.dimensions {
            return Ok(());
        }
        Err(Failure::new(
            "embed_vector_length",
            format!(
                "record {id} has {} members; the field is declared with {}",
                vector.len(),
                self.dimensions
            ),
        ))
    }
}

fn stale(id: &str) -> Failure {
    Failure::new(
        "embed_page_stale_text",
        format!("record {id} was not read at the page's version"),
    )
}

fn page_size(branch: &str, page: &Page, page_records: NonZeroUsize) -> Result<usize, Failure> {
    identifier(branch)?;
    identifier(&page.collection)?;
    identifier(&page.field)?;
    let named = page
        .made
        .len()
        .saturating_add(page.windowed.len())
        .saturating_add(page.removed.len())
        .saturating_add(page.cleared.len());
    if named > page_records.get() {
        return Err(Failure::new(
            "embed_page_too_large",
            format!("a page names {named} records; the page setting is {page_records}"),
        ));
    }
    Ok(named)
}

fn page_position(branch: &str, page: &Page, owner: String, through: Option<Version>) -> Position {
    match &page.reach {
        Reach::Upto(after) => Position {
            branch: owner,
            through,
            toward: Some(Toward {
                version: page.toward,
                after: after.clone(),
            }),
        },
        Reach::End => Position {
            branch: branch.to_owned(),
            through: Some(page.toward),
            toward: None,
        },
    }
}

fn validate_page_walk(page: &Page, done: Option<&str>) -> Result<(), Failure> {
    let ids = page
        .made
        .iter()
        .map(|made| &made.id)
        .chain(page.windowed.iter().map(|made| &made.id))
        .chain(&page.removed)
        .chain(&page.cleared);
    for id in ids {
        let past = done.is_some_and(|done| id.as_bytes() <= done.as_bytes());
        let beyond = match &page.reach {
            Reach::Upto(after) => id.as_bytes() > after.as_bytes(),
            Reach::End => false,
        };
        if past || beyond {
            return Err(Failure::new(
                "embed_page_outside_walk",
                format!("record {id} is not between the walk's last record and this page's"),
            ));
        }
    }
    Ok(())
}

impl Storage {
    /// Commits `page` into `branch`'s companion: its vectors, its removals
    /// and the new position, in one commit. The companion is made on the
    /// first page.
    ///
    /// # Errors
    /// `embed_page_too_large` past `page_records`; `branch_missing`;
    /// `embed_position_ahead` for `toward` past the head;
    /// `embed_field_not_declared`; `embed_model_mismatch`;
    /// `embed_walk_open`; `embed_position_backwards`; `embed_walk_backwards`;
    /// `embed_page_outside_walk`; `embed_vector_length`;
    /// `embed_page_stale_text`; whatever a read or the commit answers.
    pub fn embed_commit(
        &mut self,
        branch: &str,
        page: &Page,
        page_records: NonZeroUsize,
    ) -> Result<Committed, Failure> {
        let named = page_size(branch, page, page_records)?;
        let record = self
            .refs
            .get(&branch_name(branch))
            .ok_or_else(|| Failure::new("branch_missing", branch))?;
        let head = Captured {
            branch: branch.to_owned(),
            version: version(record),
            shards: record.shards.clone(),
        };
        if order(page.toward) > order(head.version) {
            return Err(Failure::new(
                "embed_position_ahead",
                "a page cannot complete a version the branch has not reached",
            ));
        }
        let declaration = self.declared(&head, &page.collection, &page.field)?;
        let tag = embedded::model_tag(&page.model)?;
        if tag != embedded::model_tag(&declaration.model)? {
            return Err(Failure::new(
                "embed_model_mismatch",
                format!(
                    "the vectors are {}'s; the field is declared for {}",
                    page.model.name, declaration.model.name
                ),
            ));
        }
        let companion = embedded::companion_name(branch);
        let position_key = embedded::position_key(&page.collection, &page.field);
        let previous = self.position(&companion, &position_key)?;
        let (owner, through, done) = walk(previous.as_ref(), branch, page)?;
        let position = page_position(branch, page, owner, through);
        validate_page_walk(page, done.as_deref())?;
        let mut writes = Writes {
            at: self.companion_at(&companion),
            page,
            pending: Vec::with_capacity(named.saturating_add(1)),
            dimensions: declaration.dimensions.get(),
        };
        let (mut written, mut skipped) = (0_usize, Vec::new());
        for made in &page.made {
            if self.write_made(&head, &mut writes, made, tag)? {
                written = written.saturating_add(1);
            } else {
                skipped.push(made.id.clone());
            }
        }
        for made in &page.windowed {
            if self.write_windows(&head, &mut writes, made, tag)? {
                written = written.saturating_add(1);
            } else {
                skipped.push(made.id.clone());
            }
        }
        for id in &page.removed {
            if let Some(now) = self.revision(&head, &page.collection, id)?
                && order(now) <= order(page.toward)
            {
                return Err(stale(id));
            }
            let key = embedded::vector_key(&page.collection, &page.field, id);
            writes.pending.push((key, None));
            self.clear_windows(&mut writes, id, 0)?;
        }
        // RED4: whatever the record holds now, no vector of its text at
        // `toward` exists: one changed since is listed by the next walk.
        for id in &page.cleared {
            let key = embedded::vector_key(&page.collection, &page.field, id);
            writes.pending.push((key, None));
            self.clear_windows(&mut writes, id, 0)?;
        }
        self.end_rewalk(&mut writes)?;
        let mut writes = writes.pending;
        writes.push((position_key, Some(embedded::encode_position(&position)?)));
        if self.refs.get(&companion).is_none() {
            create_branch_with_kind(
                &companion,
                self.writer.base_roots().map_err(Failure::database)?,
                BranchKind::Namespace,
                &mut self.refs,
                &self.registry,
                current_timestamp(),
            )
            .map_err(Failure::branch_commit)?;
        }
        self.commit_namespace(&companion, &writes)?;
        Ok(Committed {
            written,
            removed: page.removed.len(),
            skipped,
            position,
        })
    }

    /// Writes `made`'s one vector when its text is current, clearing any
    /// windows an earlier walk wrote for it (EMB-E2); `false` when skipped.
    ///
    /// # Errors
    /// `embed_vector_length`; `embed_page_stale_text`; whatever a read answers.
    fn write_made(
        &self,
        head: &Captured,
        writes: &mut Writes<'_>,
        made: &Made,
        tag: embedded::ModelTag,
    ) -> Result<bool, Failure> {
        writes.length(&made.id, &made.vector)?;
        if !self.current(head, writes.page, &made.id, made.made_from)? {
            return Ok(false);
        }
        let provenance = Provenance {
            made_from: made.made_from,
            tag,
        };
        let value = embedded::encode_embedded(&provenance, &made.vector);
        let key = embedded::vector_key(&writes.page.collection, &writes.page.field, &made.id);
        writes.pending.push((key, Some(value)));
        self.clear_windows(writes, &made.id, 0)?;
        Ok(true)
    }

    /// EMB-E2 (j): a page that ends its field's walk removes the rewalk
    /// marker, if there is one, in the same commit.
    ///
    /// # Errors
    /// Whatever a read answers.
    fn end_rewalk(&self, writes: &mut Writes<'_>) -> Result<(), Failure> {
        let (Reach::End, Some(at)) = (&writes.page.reach, &writes.at) else {
            return Ok(());
        };
        let key = embedded::rewalk_key(&writes.page.collection, &writes.page.field);
        if reads::get_bytes(&self.reader, at, &key)?.is_some() {
            writes.pending.push((key, None));
        }
        Ok(())
    }

    /// Whether `id`'s text at `head` is still the one read at `made_from`:
    /// `false` when it changed after `toward`, and the next walk lists it.
    ///
    /// # Errors
    /// `embed_page_stale_text` for text not read at the page's version.
    fn current(
        &self,
        head: &Captured,
        page: &Page,
        id: &str,
        made_from: Version,
    ) -> Result<bool, Failure> {
        if order(made_from) > order(page.toward) {
            return Err(stale(id));
        }
        match self.revision(head, &page.collection, id)? {
            Some(now) if now == made_from => Ok(true),
            Some(now) if order(now) <= order(page.toward) => Err(stale(id)),
            _ => Ok(false),
        }
    }

    /// EMB-E2: `made`'s windows, each with its span and the window count,
    /// replacing the record's single vector and any windows of an earlier
    /// text past these; `false` when its text changed after `toward`.
    ///
    /// # Errors
    /// `embed_vector_length`; `embed_windows_invalid` for fewer than two
    /// windows, more than [`embedded::WINDOWS_MOST`], or spans that do not
    /// cover the text in order; as [`Self::current`].
    fn write_windows(
        &self,
        head: &Captured,
        writes: &mut Writes<'_>,
        made: &MadeWindows,
        tag: embedded::ModelTag,
    ) -> Result<bool, Failure> {
        let page = writes.page;
        let invalid = || {
            Failure::new(
                "embed_windows_invalid",
                format!(
                    "record {}'s windows do not cover its text in order",
                    made.id
                ),
            )
        };
        let count = u32::try_from(made.windows.len()).map_err(|_| invalid())?;
        if !(2..=embedded::WINDOWS_MOST).contains(&count) {
            return Err(invalid());
        }
        let mut previous: Option<&Window> = None;
        for window in &made.windows {
            writes.length(&made.id, &window.vector)?;
            let follows = previous.map_or(window.start == 0, |before| {
                (before.start..=before.end).contains(&window.start)
                    && window.start != before.start
                    && window.end > before.end
            });
            if !follows || window.end <= window.start {
                return Err(invalid());
            }
            previous = Some(window);
        }
        if !self.current(head, page, &made.id, made.made_from)? {
            return Ok(false);
        }
        let provenance = Provenance {
            made_from: made.made_from,
            tag,
        };
        let key = embedded::vector_key(&page.collection, &page.field, &made.id);
        writes.pending.push((key, None));
        for (index, window) in (0..count).zip(&made.windows) {
            let span = embedded::Span {
                count,
                start: window.start,
                end: window.end,
            };
            let key = embedded::window_key(&page.collection, &page.field, &made.id, index);
            let value = embedded::encode_window(span, &provenance, &window.vector);
            writes.pending.push((key, Some(value)));
        }
        self.clear_windows(writes, &made.id, count)?;
        Ok(true)
    }

    /// EMB-E2: deletes `id`'s windows from `keep` on, as the companion holds
    /// them before this page: window 0 names their count, so this is one
    /// read, never a scan.
    ///
    /// # Errors
    /// `vector_corrupt` for a stored window 0 that names no count; whatever
    /// the read answers.
    fn clear_windows(&self, writes: &mut Writes<'_>, id: &str, keep: u32) -> Result<(), Failure> {
        let Some(at) = &writes.at else {
            return Ok(());
        };
        let page = writes.page;
        let first = embedded::window_key(&page.collection, &page.field, id, 0);
        let Some(bytes) = reads::get_bytes(&self.reader, at, &first)? else {
            return Ok(());
        };
        let count = embedded::window_span(&bytes)?.count;
        for index in keep..count {
            let key = embedded::window_key(&page.collection, &page.field, id, index);
            writes.pending.push((key, None));
        }
        Ok(())
    }

    /// The companion's head, `None` before its first page.
    fn companion_at(&self, companion: &str) -> Option<Captured> {
        self.refs.get(companion).map(|record| Captured {
            branch: companion.to_owned(),
            version: version(record),
            shards: record.shards.clone(),
        })
    }

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

    /// A record's revision at `head`, or `None` when it does not exist.
    fn revision(
        &self,
        head: &Captured,
        collection: &str,
        id: &str,
    ) -> Result<Option<Version>, Failure> {
        reads::get_bytes(&self.reader, head, &revision_key(collection, id))?
            .map(|bytes| {
                serde_json::from_slice::<DocumentRevision>(&bytes)
                    .map(|row| row.written_at)
                    .map_err(Failure::storage)
            })
            .transpose()
    }

    /// The job's position in `companion`, or `None` before its first page.
    ///
    /// # Errors
    /// Propagates companion read failures or an invalid stored embedding position.
    pub fn position(&self, companion: &str, key: &[u8]) -> Result<Option<Position>, Failure> {
        let Some(record) = self.refs.get(companion) else {
            return Ok(None);
        };
        let at = Captured {
            branch: companion.to_owned(),
            version: version(record),
            shards: record.shards.clone(),
        };
        reads::get_bytes(&self.reader, &at, key)?
            .map(|bytes| embedded::read_position(&bytes))
            .transpose()
    }

    /// Writes `declaration` as `collection.field`'s `w\0` value on `branch`.
    /// It stands in for the declaring change (V6) and is to be removed when
    /// that change is on main (ruled 08:49).
    ///
    /// # Errors
    /// Propagates declaration encoding or namespace commit failures.
    #[cfg(test)]
    pub fn declare_for_tests(
        &mut self,
        branch: &str,
        collection: &str,
        field: &str,
        declaration: &Declaration,
    ) -> Result<(), Failure> {
        let bytes = embedded::encode_declaration(declaration)?;
        let key = field_key(collection, field);
        self.commit_namespace(&branch_name(branch), &[(key, Some(bytes))])
    }

    /// Writes `position` for `collection.field` into `owner`'s companion,
    /// making the companion if it has none. It stands in for a fork's
    /// inheriting, which does not exist yet (ruled 08:55, F2).
    ///
    /// # Errors
    /// Propagates position encoding, companion creation or namespace commit failures.
    #[cfg(test)]
    pub fn position_for_tests(
        &mut self,
        owner: &str,
        collection: &str,
        field: &str,
        position: &Position,
    ) -> Result<(), Failure> {
        let companion = embedded::companion_name(owner);
        if self.refs.get(&companion).is_none() {
            create_branch_with_kind(
                &companion,
                self.writer.base_roots().map_err(Failure::database)?,
                BranchKind::Namespace,
                &mut self.refs,
                &self.registry,
                current_timestamp(),
            )
            .map_err(Failure::branch_commit)?;
        }
        let key = embedded::position_key(collection, field);
        let bytes = embedded::encode_position(position)?;
        self.commit_namespace(&companion, &[(key, Some(bytes))])
    }

    /// Writes `writes` into the namespace or branch `name` in one durable
    /// commit, and publishes its record.
    pub(super) fn commit_namespace(
        &mut self,
        name: &str,
        writes: &[(Vec<u8>, Option<Vec<u8>>)],
    ) -> Result<(), Failure> {
        let at = self
            .refs
            .get(name)
            .map(version)
            .ok_or_else(|| Failure::new("reference_missing", name))?;
        let handle =
            open_branch(name, &self.refs, &self.registry).map_err(Failure::branch_commit)?;
        for (key, value) in writes {
            let shard = self.writer.shard_for(key);
            value
                .as_ref()
                .map_or_else(
                    || handle.delete(shard, key),
                    |value| handle.put(shard, key, value),
                )
                .map_err(Failure::branch)?;
        }
        self.writer
            .commit(
                &handle,
                &self.registry,
                CommitRequest {
                    durability: CommitDurability::Durable {
                        refs: &mut self.refs,
                    },
                    extra_parents: &[],
                    timestamp: current_timestamp(),
                },
            )
            .map_err(|error| {
                let installed = self.refs.get(name).is_some_and(|now| version(now) != at);
                Failure::new(
                    if installed {
                        "embed_durability_unknown"
                    } else {
                        "embed_write_refused"
                    },
                    error.to_string(),
                )
            })?;
        self.publish_record(name)
    }
}

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