haematite-cli 0.21.7

//! V3: the next page of the embedding job's walk, over two shards.

use std::error::Error;
use std::num::NonZeroUsize;

use haematite::db::DatabaseConfig;
use haematite::store::NodeCacheBudget;
use serde_json::value::RawValue;

use super::{Pending, Waiting};
use crate::serve::protocol::{Change, CollectionWrite, DocumentWrite, Metric, Mutation, Version};
use crate::serve::storage::embed_job::{Made, Page, Reach, field_key};
use crate::serve::storage::embedded::{self, Declaration, ModelName, VectorKind};
use crate::serve::storage::{
    Captured, Failure, Storage, branch_name, document_key, reads, revision_key, version,
};

type TestResult = Result<(), Box<dyn Error>>;

const FIELD: &str = "body_vec";

fn create() -> Result<(tempfile::TempDir, Storage), Box<dyn Error>> {
    create_with(2)
}

fn create_with(shard_count: usize) -> Result<(tempfile::TempDir, Storage), Box<dyn Error>> {
    let directory = tempfile::tempdir()?;
    let mut storage = Storage::create(
        DatabaseConfig {
            data_dir: directory.path().join("database"),
            shard_count,
            distributed: None,
            executor_threads: Some(1),
            node_cache_budget: Some(NodeCacheBudget::Unlimited),
            retention: haematite::Retention::reclaim(0),
        },
        NodeCacheBudget::Unlimited,
        fixture_long_wait_interval(),
    )?;
    mutate(
        &mut storage,
        Change::Collection(CollectionWrite {
            collection: "notes".to_owned(),
            schema: None,
        }),
    )?;
    Ok((directory, storage))
}

fn head(storage: &Storage) -> Result<Version, Failure> {
    storage
        .refs
        .get(&branch_name("main"))
        .map(version)
        .ok_or_else(|| Failure::new("branch_missing", "main"))
}

fn mutate(storage: &mut Storage, change: Change) -> Result<Version, Failure> {
    let at = head(storage)?;
    let receipt = storage.mutate(&Mutation {
        request_id: format!("pending-{}-{}", at.generation, at.sequence),
        branch: "main".to_owned(),
        expected: at,
        change,
    })?;
    Ok(receipt.version)
}

/// Writes record `id` as the JSON `document`.
fn write(storage: &mut Storage, id: &str, document: &str) -> Result<Version, Box<dyn Error>> {
    let change = Change::Put(DocumentWrite {
        collection: "notes".to_owned(),
        document_id: id.to_owned(),
        document: RawValue::from_string(document.to_owned())?,
    });
    Ok(mutate(storage, change)?)
}

fn put(storage: &mut Storage, id: &str, body: &str) -> Result<Version, Box<dyn Error>> {
    write(storage, id, &format!(r#"{{"body":"{body}"}}"#))
}

fn model() -> ModelName {
    ModelName {
        name: "stand-in".to_owned(),
        revision: "1".to_owned(),
        manifest_sha256: "abc".to_owned(),
    }
}

fn declare(storage: &mut Storage) -> Result<(), Box<dyn Error>> {
    let declaration = Declaration {
        kind: VectorKind::Dense,
        source: "body".to_owned(),
        model: model(),
        dimensions: NonZeroUsize::new(2).ok_or("two is not zero")?,
        distance: Metric::Cosine,
    };
    let bytes = embedded::encode_declaration(&declaration)?;
    let key = field_key("notes", FIELD);
    storage.commit_namespace(&branch_name("main"), &[(key, Some(bytes))])?;
    Ok(())
}

fn pages(records: usize) -> Result<NonZeroUsize, Box<dyn Error>> {
    Ok(NonZeroUsize::new(records).ok_or("no records")?)
}

const INPUT_BYTES: usize = 64;

fn pending(storage: &Storage, records: usize) -> Result<Option<Pending>, Box<dyn Error>> {
    let input = pages(INPUT_BYTES)?;
    Ok(storage.embed_pending("main", "notes", FIELD, pages(records)?, input)?)
}

/// Commits `pending` as the job would: a vector for each text, a removal for
/// each record gone, and the page's reach.
fn commit(storage: &mut Storage, pending: &Pending, records: usize) -> TestResult {
    let mut made = Vec::new();
    let mut removed = Vec::new();
    for waiting in &pending.waiting {
        match waiting {
            Waiting::Text { id, made_from, .. } => made.push(Made {
                id: id.clone(),
                made_from: *made_from,
                vector: vec![1.0, 0.0],
            }),
            Waiting::Removed { id } => removed.push(id.clone()),
        }
    }
    let page = Page {
        collection: "notes".to_owned(),
        field: FIELD.to_owned(),
        model: model(),
        made,
        removed,
        cleared: Vec::new(),
        windowed: Vec::new(),
        toward: pending.toward,
        reach: pending.reach.clone(),
    };
    storage.embed_commit("main", &page, pages(records)?)?;
    Ok(())
}

/// Runs the job to the end of its walk, a page of `records` at a time, and
/// gives every page it took.
fn run(storage: &mut Storage, records: usize) -> Result<Vec<Pending>, Box<dyn Error>> {
    let mut taken = Vec::new();
    for _page in 0..100 {
        let Some(next) = pending(storage, records)? else {
            return Ok(taken);
        };
        commit(storage, &next, records)?;
        taken.push(next);
    }
    Err("the walk did not end in 100 pages".into())
}

/// Runs one walk to its `End` page, a page of `records` at a time, and gives
/// every page it took; a walk begun after it, toward a later head, is not
/// taken.
fn walk(storage: &mut Storage, records: usize) -> Result<Vec<Pending>, Box<dyn Error>> {
    let mut taken = Vec::new();
    for _page in 0..100 {
        let next = pending(storage, records)?.ok_or("the walk ended with no End page")?;
        commit(storage, &next, records)?;
        let ended = next.reach == Reach::End;
        taken.push(next);
        if ended {
            return Ok(taken);
        }
    }
    Err("the walk did not end in 100 pages".into())
}

fn ids(pending: &Pending) -> Vec<String> {
    let mut ids: Vec<String> = pending
        .waiting
        .iter()
        .map(|waiting| match waiting {
            Waiting::Text { id, .. } | Waiting::Removed { id } => id.clone(),
        })
        .chain(pending.no_text.iter().cloned())
        .chain(pending.too_long.iter().cloned())
        .chain(pending.empty_input.iter().cloned())
        .collect();
    ids.sort();
    ids
}

fn stored(storage: &Storage, id: &str) -> Result<bool, Box<dyn Error>> {
    let name = embedded::companion_name("main");
    let record = storage.refs.get(&name).ok_or("no companion")?;
    let at = Captured {
        branch: name.clone(),
        version: version(record),
        shards: record.shards.clone(),
    };
    let key = embedded::vector_key("notes", FIELD, id);
    Ok(reads::get_bytes(&storage.reader, &at, &key)?.is_some())
}

/// Removes record `id`'s keys directly: the delete verb refuses a collection
/// with a declared field until V6.
fn remove(storage: &mut Storage, id: &str) -> TestResult {
    let removed = [
        (document_key("notes", id), None),
        (revision_key("notes", id), None),
    ];
    storage.commit_namespace(&branch_name("main"), &removed)?;
    Ok(())
}

/// Fails while the first walk does not start from the empty tree: it lists
/// every record, in id order, with its text and the revision it is at.
#[test]
fn the_first_walk_lists_every_record_with_its_text() -> TestResult {
    let (_directory, mut storage) = create()?;
    for (id, body) in [("n-1", "a"), ("n-2", "b"), ("n-3", "c")] {
        put(&mut storage, id, body)?;
    }
    declare(&mut storage)?;
    let first = pending(&storage, 64)?.ok_or("nothing pending")?;
    assert_eq!(first.toward, head(&storage)?);
    assert_eq!(first.reach, Reach::End);
    let texts: Vec<(&str, &str)> = first
        .waiting
        .iter()
        .filter_map(|waiting| match waiting {
            Waiting::Text { id, text, .. } => Some((id.as_str(), text.as_str())),
            Waiting::Removed { .. } => None,
        })
        .collect();
    assert_eq!(texts, [("n-1", "a"), ("n-2", "b"), ("n-3", "c")]);
    Ok(())
}

/// Fails while the job can claim `End` with rows left, or a page's `Upto` can
/// pass a record another shard still holds: seven records over two shards,
/// two a page, are taken in one id order, each once, every page but the last
/// reaching `Upto` its last id and the last claiming `End`; then nothing is
/// pending and every record has its vector.
#[test]
fn the_job_claims_end_only_when_no_shard_has_more() -> TestResult {
    let (_directory, mut storage) = create()?;
    for index in 1..=7 {
        put(&mut storage, &format!("n-{index}"), "text")?;
    }
    declare(&mut storage)?;
    let taken = run(&mut storage, 2)?;
    let (last, middle) = taken.split_last().ok_or("no pages")?;
    assert_eq!(last.reach, Reach::End);
    let mut seen: Vec<String> = Vec::new();
    for page in middle {
        let Reach::Upto(after) = &page.reach else {
            return Err("a middle page claimed the end".into());
        };
        let page_ids = ids(page);
        assert!(
            page_ids.iter().all(|id| id <= after),
            "{page_ids:?} past {after}"
        );
        if let Some(first) = page_ids.first() {
            assert!(
                seen.iter().all(|id| id < first),
                "{first} came after a later id"
            );
        }
        seen.extend(page_ids);
        assert!(seen.iter().all(|id| id <= after));
    }
    seen.extend(ids(last));
    let all: Vec<String> = (1..=7).map(|index| format!("n-{index}")).collect();
    assert_eq!(seen, all, "each record once, in one id order");
    for id in &all {
        assert!(stored(&storage, id)?, "{id} has no vector");
    }
    assert_eq!(pending(&storage, 2)?, None);
    Ok(())
}

/// Fails while a finished walk is walked again: after it, only what changed
/// since is listed.
#[test]
fn after_a_walk_only_later_changes_are_listed() -> TestResult {
    let (_directory, mut storage) = create()?;
    for index in 1..=3 {
        put(&mut storage, &format!("n-{index}"), "text")?;
    }
    declare(&mut storage)?;
    run(&mut storage, 64)?;
    put(&mut storage, "n-2", "again")?;
    put(&mut storage, "n-9", "new")?;
    let next = pending(&storage, 64)?.ok_or("nothing pending")?;
    assert_eq!(ids(&next), ["n-2", "n-9"]);
    Ok(())
}

/// Fails while a removed record is not listed as removed.
#[test]
fn a_removed_record_is_listed_as_removed() -> TestResult {
    let (_directory, mut storage) = create()?;
    put(&mut storage, "n-1", "text")?;
    put(&mut storage, "n-2", "text")?;
    declare(&mut storage)?;
    run(&mut storage, 64)?;
    remove(&mut storage, "n-1")?;
    let next = pending(&storage, 64)?.ok_or("nothing pending")?;
    let gone = Waiting::Removed {
        id: "n-1".to_owned(),
    };
    assert_eq!(next.waiting, [gone]);
    run(&mut storage, 64)?;
    assert!(!stored(&storage, "n-1")?);
    Ok(())
}

/// Fails while text that cannot be embedded is sent, or cut: a number, a
/// missing field and text past the input bound are counted, not sent, and
/// the walk still ends.
#[test]
fn text_that_cannot_be_embedded_is_counted_not_sent() -> TestResult {
    let (_directory, mut storage) = create()?;
    write(&mut storage, "n-1", r#"{"body":7}"#)?;
    write(&mut storage, "n-2", r#"{"title":"no body"}"#)?;
    put(&mut storage, "n-3", &"x".repeat(INPUT_BYTES + 1))?;
    put(&mut storage, "n-4", "fits")?;
    declare(&mut storage)?;
    let first = pending(&storage, 64)?.ok_or("nothing pending")?;
    assert_eq!(first.no_text, ["n-1", "n-2"]);
    assert_eq!(first.too_long, ["n-3"]);
    assert_eq!(first.waiting.len(), 1);
    run(&mut storage, 64)?;
    assert!(!stored(&storage, "n-3")?);
    assert_eq!(pending(&storage, 64)?, None);
    Ok(())
}

/// EMB-E1: a text empty or only whitespace is one the worker refuses. The
/// walk names it as passed, never lists it to send, and moves past it.
#[test]
fn an_empty_or_whitespace_text_is_passed_not_listed() -> TestResult {
    let (_directory, mut storage) = create()?;
    put(&mut storage, "n-1", "")?;
    put(&mut storage, "n-2", "fits")?;
    put(&mut storage, "n-3", "  ")?;
    put(&mut storage, "n-4", "fits too")?;
    put(&mut storage, "n-5", r"\n\t")?;
    declare(&mut storage)?;
    let first = pending(&storage, 64)?.ok_or("nothing pending")?;
    assert_eq!(first.empty_input, ["n-1", "n-3", "n-5"]);
    let listed: Vec<&str> = first
        .waiting
        .iter()
        .map(|waiting| match waiting {
            Waiting::Text { id, .. } | Waiting::Removed { id } => id.as_str(),
        })
        .collect();
    assert_eq!(listed, ["n-2", "n-4"]);
    run(&mut storage, 64)?;
    assert!(!stored(&storage, "n-1")?);
    assert!(stored(&storage, "n-2")?);
    assert_eq!(pending(&storage, 64)?, None);
    Ok(())
}

/// Fails while a walk follows the head: a write during a walk is not in it,
/// and is the next walk's.
#[test]
fn a_walk_keeps_its_target_when_the_head_moves() -> TestResult {
    let (_directory, mut storage) = create()?;
    for index in 1..=5 {
        put(&mut storage, &format!("n-{index}"), "text")?;
    }
    declare(&mut storage)?;
    let target = head(&storage)?;
    let first = pending(&storage, 2)?.ok_or("nothing pending")?;
    commit(&mut storage, &first, 2)?;
    put(&mut storage, "n-8", "later")?;
    let taken = walk(&mut storage, 2)?;
    assert!(taken.iter().all(|page| page.toward == target));
    assert!(
        taken
            .iter()
            .all(|page| !ids(page).contains(&"n-8".to_owned()))
    );
    let next = pending(&storage, 2)?.ok_or("the later write was lost")?;
    assert_eq!(ids(&next), ["n-8"]);
    Ok(())
}

/// Fails while a shard that stopped with more lets the page claim `End`: on
/// one shard, a page of two from five records ends at the shard's own stop,
/// so every page but the last reaches `Upto`, and all five are taken.
#[test]
fn one_shard_with_more_never_claims_end() -> TestResult {
    let (_directory, mut storage) = create_with(1)?;
    for index in 1..=5 {
        put(&mut storage, &format!("n-{index}"), "text")?;
    }
    declare(&mut storage)?;
    let first = pending(&storage, 2)?.ok_or("nothing pending")?;
    assert_eq!(first.reach, Reach::Upto("n-2".to_owned()));
    let mut seen = ids(&first);
    commit(&mut storage, &first, 2)?;
    for page in run(&mut storage, 2)? {
        seen.extend(ids(&page));
    }
    let all: Vec<String> = (1..=5).map(|index| format!("n-{index}")).collect();
    assert_eq!(seen, all);
    Ok(())
}

/// `JT1b`. Fails while a shard holding exactly a page's changes keeps the walk
/// open: n-1, n-2 and n-3 are all in shard 0 of two (blake3 of the revision
/// key), and a page of three that reads them all claims `End`, so nothing is
/// pending after it.
#[test]
fn a_page_that_reads_the_last_record_claims_end() -> TestResult {
    let (_directory, mut storage) = create()?;
    for index in 1..=3 {
        put(&mut storage, &format!("n-{index}"), "text")?;
    }
    declare(&mut storage)?;
    let first = pending(&storage, 3)?.ok_or("nothing pending")?;
    assert_eq!(ids(&first), ["n-1", "n-2", "n-3"]);
    assert_eq!(first.reach, Reach::End);
    commit(&mut storage, &first, 3)?;
    assert_eq!(pending(&storage, 3)?, None);
    Ok(())
}

/// `JT1b`. Guards the page past a full shard: n-1, n-2 and n-3 in shard 0 and
/// n-4 in shard 1, a page of three reaches `Upto` n-3 and the next takes n-4
/// and claims `End`.
#[test]
fn a_full_shard_beside_another_reaches_its_last_then_ends() -> TestResult {
    let (_directory, mut storage) = create()?;
    for index in 1..=4 {
        put(&mut storage, &format!("n-{index}"), "text")?;
    }
    declare(&mut storage)?;
    let first = pending(&storage, 3)?.ok_or("nothing pending")?;
    assert_eq!(ids(&first), ["n-1", "n-2", "n-3"]);
    assert_eq!(first.reach, Reach::Upto("n-3".to_owned()));
    commit(&mut storage, &first, 3)?;
    let last = pending(&storage, 3)?.ok_or("n-4 was skipped")?;
    assert_eq!(ids(&last), ["n-4"]);
    assert_eq!(last.reach, Reach::End);
    Ok(())
}

/// `JT1b`. Guards the change past the page: n-1, n-2, n-3 and n-8 are all in
/// shard 0, one more than a page of three. The page keeps three and reaches
/// `Upto` n-3, skipping nothing; the next starts after n-3, takes n-8 alone
/// and claims `End`.
#[test]
fn a_shard_one_past_the_page_is_cut_to_the_page_and_its_last_comes_next() -> TestResult {
    let (_directory, mut storage) = create()?;
    for id in ["n-1", "n-2", "n-3", "n-8"] {
        put(&mut storage, id, "text")?;
    }
    declare(&mut storage)?;
    let first = pending(&storage, 3)?.ok_or("nothing pending")?;
    assert_eq!(ids(&first), ["n-1", "n-2", "n-3"]);
    assert_eq!(first.reach, Reach::Upto("n-3".to_owned()));
    commit(&mut storage, &first, 3)?;
    let last = pending(&storage, 3)?.ok_or("n-8 was skipped")?;
    assert_eq!(ids(&last), ["n-8"]);
    assert_eq!(last.reach, Reach::End);
    Ok(())
}

/// `JT1b`. Fails while a page setting with no room for the change past it is
/// taken: the walk refuses it by name instead of overflowing its budget.
#[test]
fn a_page_setting_with_no_room_past_it_is_refused_by_name() -> TestResult {
    let (_directory, mut storage) = create()?;
    put(&mut storage, "n-1", "text")?;
    declare(&mut storage)?;
    let refused = storage
        .embed_pending(
            "main",
            "notes",
            FIELD,
            pages(usize::MAX)?,
            pages(INPUT_BYTES)?,
        )
        .err()
        .ok_or("a page of usize::MAX was read")?;
    assert_eq!(refused.code, "embed_page_records_overflow");
    Ok(())
}

/// STOR-003 MD4-b: a walk's target that is 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; any other refusal is itself.
#[test]
fn every_unreadable_target_is_embed_walk_released_naming_why() {
    let at = Version {
        generation: 1,
        sequence: 4,
    };
    for code in [
        "version_missing",
        "version_released",
        "version_generation_ended",
    ] {
        let mapped = super::walk_unreadable("main", at, Failure::new(code, "why"));
        assert_eq!(mapped.code, "embed_walk_released", "{code}");
        let named = format!("({code}: why)");
        assert!(mapped.message.contains(&named), "{named}");
    }
    let other = super::walk_unreadable("main", at, Failure::new("generation_mismatch", "why"));
    assert_eq!(other.code, "generation_mismatch");
}

#[test]
fn held_pending_reads_do_not_take_foreground_storage() -> TestResult {
    let (directory, mut storage) = create()?;
    declare(&mut storage)?;
    put(&mut storage, "one", "held text")?;
    let pending_reader = storage.pending_reader();
    let storage = std::sync::Mutex::new(storage);
    let admission = crate::serve::embedding::Admission::new(std::time::Duration::from_secs(1))?;
    let writer = admission.foreground(&storage, false)?;
    let (page, declaration) = pending_reader
        .read("main", "notes", FIELD, pages(10)?, pages(INPUT_BYTES)?)?
        .ok_or("pending page missing")?;
    assert_eq!(declaration.source, "body");
    assert_eq!(page.waiting.len(), 1);
    assert!(
        matches!(&page.waiting[0], Waiting::Text { id, text, .. } if id == "one" && text == "held text")
    );
    drop(writer);
    drop(storage);
    drop(directory);
    Ok(())
}

const fn fixture_long_wait_interval() -> std::time::Duration {
    const LONG_WAIT_INTERVAL: std::time::Duration = std::time::Duration::from_secs(1);
    LONG_WAIT_INTERVAL
}