haematite-cli 0.35.0

//! STOR-050 R1 and R2 over the served connection: a stage's lineage is a row
//! in the stage's own commit, history answers it, a merge carries one row
//! listing the source's stages, and a line keeps every version until a
//! discard names it, after which a reclaim pass leaves none of its nodes.

use std::collections::{BTreeSet, HashSet};
use std::error::Error;
use std::io;
use std::num::{NonZeroU32, NonZeroUsize};
use std::os::unix::net::UnixStream;
use std::path::Path;
use std::sync::Arc;
use std::thread::JoinHandle;
use std::time::Duration;

use haematite::db::vacuum::CANONICAL_REFS_DIR;
use haematite::db::{DatabaseConfig, ReadOnlyDatabase};
use haematite::store::NodeCacheBudget;
use haematite::{DiskStore, Hash, Node};
use serde_json::{Value, json};

use super::{Limits, connection};
use crate::serve::{frame, service::Service, storage::Storage};

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

const SHARDS: usize = 2;

struct Served {
    client: Option<UnixStream>,
    worker: Option<JoinHandle<io::Result<()>>>,
    service: Arc<Service>,
    directory: tempfile::TempDir,
    max_frame: NonZeroU32,
    next_id: u64,
}

impl Served {
    fn new() -> Result<Self, Box<dyn Error>> {
        let directory = tempfile::tempdir()?;
        let storage = Storage::create(
            DatabaseConfig {
                data_dir: directory.path().join("database"),
                shard_count: SHARDS,
                distributed: None,
                executor_threads: Some(2),
                node_cache_budget: Some(NodeCacheBudget::Unlimited),
                retention: haematite::Retention::reclaim(0),
            },
            NodeCacheBudget::Unlimited,
            Duration::from_secs(1),
        )?;
        let service = Arc::new(Service::new(
            storage,
            NonZeroU32::new(64).ok_or("zero page")?,
        )?);
        let max_frame = NonZeroU32::new(1 << 20).ok_or("zero frame")?;
        let limits = Limits {
            max_frame,
            max_clients: NonZeroUsize::new(1).ok_or("zero clients")?,
            frame_timeout: Duration::from_secs(600),
        };
        let (client, peer) = UnixStream::pair()?;
        let serving = Arc::clone(&service);
        let worker = std::thread::Builder::new()
            .name("served-lines".to_owned())
            .spawn(move || connection(peer, &serving, limits))?;
        let mut fixture = Self {
            client: Some(client),
            worker: Some(worker),
            service,
            directory,
            max_frame,
            next_id: 0,
        };
        fixture.hello()?;
        Ok(fixture)
    }

    fn send(&mut self, operation: &Value) -> Result<Value, Box<dyn Error>> {
        self.next_id += 1;
        let id = format!("lines-{}", self.next_id);
        let request = json!({"id": id, "operation": operation});
        let client = self.client.as_mut().ok_or("closed fixture")?;
        frame::write(client, &serde_json::to_vec(&request)?, self.max_frame)?;
        let bytes = frame::read(client, self.max_frame)?.ok_or("connection closed")?;
        let answer: Value = serde_json::from_slice(&bytes)?;
        assert_eq!(answer["id"], id.as_str());
        Ok(answer)
    }

    /// The answer's result; an error answer fails the test with its text.
    fn ok(&mut self, operation: &Value) -> Result<Value, Box<dyn Error>> {
        let answer = self.send(operation)?;
        if let Some(error) = answer.get("error") {
            return Err(format!("{} refused: {error}", operation["method"]).into());
        }
        answer
            .get("result")
            .cloned()
            .ok_or_else(|| "missing result".into())
    }

    /// The refusal's code; a result fails the test.
    fn refused(&mut self, operation: &Value) -> Result<Value, Box<dyn Error>> {
        let answer = self.send(operation)?;
        answer
            .get("error")
            .cloned()
            .ok_or_else(|| format!("{} was not refused: {answer}", operation["method"]).into())
    }

    fn hello(&mut self) -> TestResult {
        self.ok(&json!({"method":"hello", "params":{
            "protocol": crate::serve::protocol::PROTOCOL,
            "capabilities": [crate::serve::protocol::RECEIPT_CAPABILITY],
            "max_frame": self.max_frame,
        }}))?;
        Ok(())
    }

    fn head(&mut self, branch: &str) -> Result<Value, Box<dyn Error>> {
        let answer = self.ok(&json!({"method":"collections.list", "params":{
            "selection":{"branch":branch, "version":null}, "after":null, "limit":1,
        }}))?;
        Ok(answer["selection"]["version"].clone())
    }

    fn apply_params(
        &mut self,
        id: &str,
        branch: &str,
        change: Value,
        stage: Option<Value>,
    ) -> Result<Value, Box<dyn Error>> {
        let expected = self.head(branch)?;
        let mut params = json!({"request_id": id, "branch": branch, "expected": expected});
        params["change"] = change;
        if let Some(stage) = stage {
            params["stage"] = stage;
        }
        Ok(json!({"method":"mutation.apply", "params": params}))
    }

    fn apply(
        &mut self,
        id: &str,
        branch: &str,
        change: Value,
        stage: Option<Value>,
    ) -> Result<Value, Box<dyn Error>> {
        let operation = self.apply_params(id, branch, change, stage)?;
        self.ok(&operation)
    }

    fn history(&mut self, branch: &str) -> Result<Vec<Value>, Box<dyn Error>> {
        let answer = self.ok(&json!({"method":"history.list", "params":{
            "branch": branch, "after": null, "limit": 64,
        }}))?;
        assert_eq!(answer["next_after"], Value::Null, "one page holds it");
        Ok(answer["entries"]
            .as_array()
            .ok_or("history entries")?
            .clone())
    }

    fn documents(&mut self, branch: &str) -> Result<Value, Box<dyn Error>> {
        self.ok(&json!({"method":"documents.list", "params":{
            "selection":{"branch":branch, "version":null}, "collection":"stages",
            "after":null, "limit":64,
        }}))
    }

    fn data_dir(&self) -> std::path::PathBuf {
        self.directory.path().join("database")
    }

    fn finish(mut self) -> TestResult {
        drop(self.client.take());
        let worker = self.worker.take().ok_or("no worker")?;
        worker.join().map_err(|_| "connection thread panicked")??;
        Ok(())
    }
}

fn collection() -> Value {
    json!({"kind":"collection.put", "value":{"collection":"stages", "schema":null}})
}

fn put(document_id: &str, rows: u64) -> Value {
    json!({"kind":"documents.put_many", "value":{"collection":"stages", "documents":[
        {"document_id": document_id, "document": {"rows": rows, "by": document_id}},
    ]}})
}

fn fork(name: &str, from: &Value, line: bool) -> Value {
    if line {
        json!({"kind":"branch.fork", "value":{"name":name, "from":from, "line":true}})
    } else {
        json!({"kind":"branch.fork", "value":{"name":name, "from":from}})
    }
}

fn discard(name: &str) -> Value {
    json!({"kind":"branch.discard", "value":{"name":name}})
}

fn stage(line: &str, name: &str, inputs: &[(&str, &Value)]) -> Value {
    let inputs: Vec<Value> = inputs
        .iter()
        .map(|(branch, version)| json!({"branch": branch, "version": version}))
        .collect();
    json!({"line": line, "run": "run-1", "stage": name, "subject": "ward-7", "inputs": inputs})
}

fn entry_at<'a>(entries: &'a [Value], version: &Value) -> Result<&'a Value, Box<dyn Error>> {
    entries
        .iter()
        .find(|entry| &entry["version"] == version)
        .ok_or_else(|| format!("no history entry at {version}").into())
}

/// R1 red: a stage object of two inputs is written with its change, and
/// history answers it at that version; the fork's own version has none.
#[test]
fn a_stage_of_two_inputs_commits_with_its_change_and_history_answers_it() -> TestResult {
    let mut served = Served::new()?;
    let made = served.apply("collection", "main", collection(), None)?;
    let main = made["version"].clone();
    let forked = served.apply("fork", "main", fork("run-1", &main, false), None)?;
    let anchor = forked["version"].clone();
    let sent = stage("run-1", "extract", &[("main", &main), ("run-1", &anchor)]);
    let written = served.apply("extract", "run-1", put("extract", 3), Some(sent.clone()))?;
    let at = written["version"].clone();
    let entries = served.history("run-1")?;
    assert_eq!(entry_at(&entries, &at)?["stage"], sent, "the stage as sent");
    assert_eq!(entry_at(&entries, &anchor)?["stage"], Value::Null);
    assert_eq!(
        served.documents("run-1")?["entries"]
            .as_array()
            .map(Vec::len),
        Some(1),
        "the stage's data is there with it"
    );
    served.finish()
}

/// A stage member that is no name, or a stage beside branch control, is
/// refused by name and nothing is written.
#[test]
fn a_stage_member_that_is_no_name_is_refused_by_name() -> TestResult {
    let mut served = Served::new()?;
    let made = served.apply("collection", "main", collection(), None)?;
    let main = made["version"].clone();
    let mut empty = stage("main", "extract", &[("main", &main)]);
    empty["inputs"][0]["branch"] = json!("");
    let operation = served.apply_params("empty", "main", put("a", 1), Some(empty))?;
    let error = served.refused(&operation)?;
    assert_eq!(error["code"], "invalid_identifier", "{error}");
    assert_eq!(error["path"], "/stage/inputs/0/branch", "{error}");
    assert_eq!(error["outcome"], "not_applied", "{error}");
    let mut free = stage("main", "extract", &[("main", &main)]);
    free["subject"] = json!({"ward": 7});
    let operation = served.apply_params("free", "main", put("a", 1), Some(free))?;
    assert_eq!(served.refused(&operation)?["code"], "invalid_params");
    let beside = stage("main", "extract", &[("main", &main)]);
    let operation =
        served.apply_params("beside", "main", fork("other", &main, true), Some(beside))?;
    let error = served.refused(&operation)?;
    assert_eq!(error["code"], "invalid_operation", "{error}");
    assert_eq!(error["path"], "/stage", "{error}");
    assert_eq!(served.head("main")?, main, "nothing was written");
    served.finish()
}

/// Every stage row (`s\0`) at a version's roots, as raw keys.
fn stage_rows(data_dir: &Path, entry: &Value) -> Result<usize, Box<dyn Error>> {
    let reader = ReadOnlyDatabase::open(data_dir)?;
    let mut rows = 0;
    for (shard, root) in roots(entry)? {
        let page = reader.range_page_at(
            shard,
            root,
            b"s\0",
            b"s\x01",
            NonZeroUsize::new(1024).ok_or("zero")?,
        )?;
        rows += page.entries.len();
    }
    Ok(rows)
}

/// R1 red: a merge of a branch with three stage commits onto its parent
/// leaves one stage row there, naming the source and its version and
/// listing the three; after the source is discarded history still says so.
#[test]
fn a_merge_carries_one_stage_row_listing_the_sources_stages() -> TestResult {
    let mut served = Served::new()?;
    let made = served.apply("collection", "main", collection(), None)?;
    let main = made["version"].clone();
    let forked = served.apply("fork", "main", fork("run-2", &main, false), None)?;
    let mut previous = forked["version"].clone();
    let mut sent = Vec::new();
    for name in ["extract", "score", "rank"] {
        let stage = stage("run-2", name, &[("run-2", &previous)]);
        let written = served.apply(name, "run-2", put(name, 1), Some(stage.clone()))?;
        previous = written["version"].clone();
        sent.push(json!({"branch":"run-2", "version": previous, "stage": stage}));
    }
    let merge = json!({"kind":"branch.merge", "value":{"from":"run-2", "reviewed": previous}});
    let merged = served.apply("merge", "main", merge, None)?;
    let at = merged["version"].clone();
    let entries = served.history("main")?;
    let entry = entry_at(&entries, &at)?;
    let expected = json!({"from":{"branch":"run-2", "version": previous}, "stages": sent});
    assert_eq!(entry["merged"], expected, "one row lists the three");
    assert_eq!(entry["stage"], Value::Null);
    assert_eq!(
        stage_rows(&served.data_dir(), entry)?,
        1,
        "exactly one stage row on the target"
    );
    served.apply("discard", "run-2", discard("run-2"), None)?;
    let entries = served.history("main")?;
    assert_eq!(
        entry_at(&entries, &at)?["merged"],
        expected,
        "the lineage outlives its source"
    );
    served.finish()
}

fn hash(text: &str) -> Result<Hash, Box<dyn Error>> {
    let mut bytes = [0_u8; 32];
    if text.len() != 64 {
        return Err(format!("not a hash: {text}").into());
    }
    for (index, byte) in bytes.iter_mut().enumerate() {
        let pair = text.get(index * 2..index * 2 + 2).ok_or("short hash")?;
        *byte = u8::from_str_radix(pair, 16)?;
    }
    Ok(Hash::from_bytes(bytes))
}

fn roots(entry: &Value) -> Result<Vec<(usize, Hash)>, Box<dyn Error>> {
    entry["roots"]
        .as_array()
        .ok_or("roots")?
        .iter()
        .map(|root| {
            let shard = usize::try_from(root["shard"].as_u64().ok_or("shard")?)?;
            Ok((shard, hash(root["root"].as_str().ok_or("root")?)?))
        })
        .collect()
}

fn store(data_dir: &Path, shard: usize) -> Result<DiskStore, Box<dyn Error>> {
    Ok(DiskStore::new(
        data_dir.join(format!("shard-{shard}")).join("store"),
        NodeCacheBudget::Unlimited,
    )?)
}

/// Every node reached from `roots`, by shard, read from disk.
fn reached(
    data_dir: &Path,
    roots: &[(usize, Hash)],
) -> Result<HashSet<(usize, Hash)>, Box<dyn Error>> {
    let mut seen = HashSet::new();
    for shard in 0..SHARDS {
        let store = store(data_dir, shard)?;
        let mut stack: Vec<Hash> = roots
            .iter()
            .filter(|(on, _)| *on == shard)
            .map(|(_, root)| *root)
            .collect();
        while let Some(node) = stack.pop() {
            if !seen.insert((shard, node)) {
                continue;
            }
            let read = store.get(&node)?.ok_or_else(|| format!("{node} missing"))?;
            if let Node::Internal(internal) = read.as_ref() {
                stack.extend(internal.children().iter().map(|(_, child)| *child));
            }
        }
    }
    Ok(seen)
}

fn on_disk(data_dir: &Path, shard: usize, node: &Hash) -> Result<bool, Box<dyn Error>> {
    Ok(store(data_dir, shard)?.get(node)?.is_some())
}

/// The service namespace's head roots, read from its record on disk.
fn service_roots(data_dir: &Path) -> Result<Vec<(usize, Hash)>, Box<dyn Error>> {
    let record = haematite::branch::RefStore::read_record(
        data_dir.join(CANONICAL_REFS_DIR),
        haematite::branch::history::SERVICE_BRANCH,
    )?
    .ok_or("no service record")?;
    Ok(record
        .shards
        .iter()
        .map(|shard| (shard.shard_id, shard.head))
        .collect())
}

/// R2 and the line end to end: fork a line, three stage applies, history
/// with the three stage objects, a fork of the line at stage 1's version,
/// both discarded, a reclaim pass to rest, and every node only the two
/// lines reached is gone (counted), while main reads as before.
#[test]
fn a_line_keeps_its_stages_forks_at_stage_one_and_leaves_no_node_once_discarded() -> TestResult {
    let mut served = Served::new()?;
    let made = served.apply("collection", "main", collection(), None)?;
    let main = made["version"].clone();
    served.apply("seed", "main", put("seed", 0), None)?;
    let forked = served.apply("fork-line", "main", fork("line-a", &main, true), None)?;
    let anchor = forked["version"].clone();
    let mut previous = anchor;
    let mut versions = Vec::new();
    let mut sent = Vec::new();
    for name in ["extract", "score", "rank"] {
        let stage = stage("line-a", name, &[("line-a", &previous)]);
        let written = served.apply(name, "line-a", put(name, 7), Some(stage.clone()))?;
        previous = written["version"].clone();
        versions.push(previous.clone());
        sent.push(stage);
    }
    let entries = served.history("line-a")?;
    assert_eq!(entries.len(), 4, "the fork and three stages");
    for (version, stage) in versions.iter().zip(&sent) {
        assert_eq!(&entry_at(&entries, version)?["stage"], stage);
    }
    let stage_one = versions.first().ok_or("no stage")?.clone();
    let rerun = served.apply(
        "fork-rerun",
        "line-a",
        fork("line-b", &stage_one, true),
        None,
    )?;
    let rerun_at = rerun["version"].clone();
    let again = stage("line-b", "score", &[("line-a", &stage_one)]);
    served.apply("score-again", "line-b", put("score", 8), Some(again))?;
    let data_dir = served.data_dir();
    let line_roots: Vec<(usize, Hash)> = served
        .history("line-a")?
        .iter()
        .chain(served.history("line-b")?.iter())
        .map(roots)
        .collect::<Result<Vec<_>, _>>()?
        .into_iter()
        .flatten()
        .collect();
    // Every version main keeps, and the service namespace, keep their nodes.
    let main_entries = served.history("main")?;
    let documents_before = served.documents("main")?;
    let mut kept_roots: Vec<(usize, Hash)> = main_entries
        .iter()
        .map(roots)
        .collect::<Result<Vec<_>, _>>()?
        .into_iter()
        .flatten()
        .collect();
    kept_roots.extend(service_roots(&data_dir)?);
    let kept = reached(&data_dir, &kept_roots)?;
    let only_lines: BTreeSet<(usize, String)> = reached(&data_dir, &line_roots)?
        .into_iter()
        .filter(|node| !kept.contains(node))
        .map(|(shard, node)| (shard, node.to_string()))
        .collect();
    assert!(!only_lines.is_empty(), "the lines wrote nodes of their own");
    let first = served.apply("discard-a", "line-a", discard("line-a"), None)?;
    assert_eq!(first["versions_released"], 4, "{first}");
    let second = served.apply("discard-b", "line-b", discard("line-b"), None)?;
    assert_eq!(second["versions_released"], 2, "{second}");
    served.service.activate_reclaim()?;
    served.service.quiesce_reclaim()?;
    let mut left = 0_usize;
    for (shard, node) in &only_lines {
        if on_disk(&data_dir, *shard, &hash(node)?)? {
            left += 1;
        }
    }
    eprintln!(
        "STOR-050 R2 pruning: {} nodes only the two lines reached, {left} left after the \
         discards and a reclaim pass to rest",
        only_lines.len()
    );
    assert_eq!(
        left,
        0,
        "of {} nodes only the lines reached",
        only_lines.len()
    );
    assert_eq!(
        served.documents("main")?,
        documents_before,
        "main as before"
    );
    assert_eq!(
        served.history("main")?,
        main_entries,
        "main's history as before"
    );
    let gone = served.refused(&json!({"method":"collections.list", "params":{
        "selection":{"branch":"line-a", "version": stage_one}, "after":null, "limit":1,
    }}))?;
    assert_eq!(gone["code"], "version_released", "{gone}");
    assert_eq!(rerun_at["sequence"], "0");
    served.finish()
}