haematite-cli 0.21.5

//! SPAN-MEASURE: the one-flush span of haematite's largest production-shaped commit.
//!
//! A measurement, not a gate: it prints and asserts nothing. The span is what
//! `Catalogue::prepare_advance` returns for one served write: the data nodes,
//! the history rows the advance spills and the catalogue nodes for the head
//! move, in the one plan the format-4 commit writer would publish. The writes
//! are haematite's own (`prepare_change`, the count rows and the receipt), and the
//! frame is encoded by `CommitFrame::encode` with unbounded limits.
//!
//! The largest commit is one import of `max_page × 64` documents (the import
//! bound in mutations.rs) sized to fill `max_frame`, written when the kept
//! window is full and no kept version has a row, so all K rows spill. The
//! control is one small document written after the collection is declared,
//! when nothing spills. `HAEM_SPAN_MAX_PAGE`, `HAEM_SPAN_MAX_FRAME`,
//! `HAEM_SPAN_DOC_BYTES`, `HAEM_SPAN_KEPT` and `HAEM_SPAN_SHARDS` override the
//! shape; every line prints the shape it used.

use std::collections::BTreeMap;
use std::error::Error;
use std::fs;
use std::num::{NonZeroU32, NonZeroUsize};
use std::path::Path;

use haematite::branch::KeptLimit;
use haematite::branch::catalogue::{
    Advance, Catalogue, Entry, Expected, NodeBudget, Prepared, ShardChanges, Write,
};
use haematite::branch::history::{self, SERVICE_BRANCH};
use haematite::db::{DEFAULT_INTERNAL_TARGET_BYTES, DEFAULT_LEAF_TARGET_BYTES, DatabaseConfig};
use haematite::store::{Disk, NodeCacheBudget};
use haematite::tree::{Hash as NodeHash, LeafNode, Node, Policy};
use haematite::{CommitFrame, CommitLimits, CommitNode};
use serde_json::value::RawValue;

use super::mutations::PreparedWrite;
use super::{Receipt, Storage, branch_name, counts, receipt_key, version};
use crate::serve::protocol::{
    Change, CollectionWrite, DocumentImport, DocumentWrite, ImportedDocument, Mutation, Selection,
    Version,
};

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

/// Today's limits, from `fence/commit_writer_support.rs` (`LIMITS`).
const LIMIT_NODES: u64 = 256;
const LIMIT_FRAME_BYTES: u64 = 65_536;
const LIMIT_NODE_BYTES: u64 = 65_536;

struct Shape {
    max_page: NonZeroU32,
    max_frame: usize,
    doc_bytes: Option<usize>,
    kept: usize,
    shards: usize,
}

fn setting(name: &str, default: usize) -> TestResult<usize> {
    std::env::var(name).map_or(Ok(default), |text| Ok(text.parse()?))
}

impl Shape {
    fn from_env() -> TestResult<Self> {
        let doc_bytes = setting("HAEM_SPAN_DOC_BYTES", 0)?;
        Ok(Self {
            max_page: NonZeroU32::new(u32::try_from(setting("HAEM_SPAN_MAX_PAGE", 64)?)?)
                .ok_or("HAEM_SPAN_MAX_PAGE must be nonzero")?,
            max_frame: setting("HAEM_SPAN_MAX_FRAME", 1_048_576)?,
            doc_bytes: (doc_bytes > 0).then_some(doc_bytes),
            kept: setting("HAEM_SPAN_KEPT", 64)?,
            shards: setting("HAEM_SPAN_SHARDS", 1)?,
        })
    }

    /// The import bound in mutations.rs: `max_page × 64` documents.
    fn documents(&self) -> TestResult<usize> {
        Ok(usize::try_from(self.max_page.get())?
            .checked_mul(64)
            .ok_or("document count overflow")?)
    }

    /// Text bytes per document so the import request fills `max_frame`.
    fn text_bytes(&self) -> TestResult<usize> {
        if let Some(bytes) = self.doc_bytes {
            return Ok(bytes);
        }
        // Per document: the id, the JSON members around it and the text quotes.
        let per_document = self
            .max_frame
            .saturating_sub(4_096)
            .checked_div(self.documents()?)
            .ok_or("zero documents")?;
        Ok(per_document.saturating_sub(64).max(1))
    }
}

#[derive(Default)]
struct Measured {
    nodes: usize,
    per_shard: BTreeMap<usize, usize>,
    plain_bytes: u64,
    file_bytes: u64,
    largest_plain: u64,
    largest_file: u64,
    records: usize,
    receipt_bytes: usize,
    frame_bytes: u64,
}

const fn policy() -> Policy {
    Policy::v2(DEFAULT_LEAF_TARGET_BYTES, DEFAULT_INTERNAL_TARGET_BYTES)
}

const fn budget() -> NodeBudget {
    NodeBudget {
        nodes: NonZeroUsize::MAX,
        uncompressed_bytes: NonZeroUsize::MAX,
    }
}

/// Store handles beside the live writer, one per shard, for reading its trees.
/// This bypasses the writer lock, which is safe only because the store keeps
/// everything and no vacuum runs. A shard the writer has not yet created holds
/// no nodes, so its handle opens in `scratch` and the live layout is untouched.
fn stores(data_dir: &Path, scratch: &Path, shards: usize) -> TestResult<Vec<Disk>> {
    (0..shards)
        .map(|shard| {
            let live = data_dir.join(format!("shard-{shard}"));
            let parent = if live.is_dir() {
                live
            } else {
                let unused = scratch.join(format!("shard-{shard}"));
                fs::create_dir_all(&unused)?;
                unused
            };
            Ok(Disk::new(parent.join("store"), NodeCacheBudget::Unlimited)?)
        })
        .collect()
}

/// A catalogue holding main and the service branch as the ref store has them.
fn catalogue(storage: &Storage, stores: &mut [Disk], kept: KeptLimit) -> TestResult<Catalogue> {
    let first = stores.first_mut().ok_or("no shard")?;
    let empty = first.put(&Node::Leaf(LeafNode::new(Vec::new())?))?;
    let base = Catalogue::at::<haematite::store::Error>(empty, policy())?;
    let writes = [branch_name("main"), SERVICE_BRANCH.to_owned()]
        .into_iter()
        .map(|name| {
            Ok(Write::Put {
                entry: Entry {
                    record: storage.refs.get(&name).ok_or("record missing")?.clone(),
                    kept: storage.refs.kept(&name).to_vec(),
                },
                limit: kept,
                expected: Expected::Absent,
            })
        })
        .collect::<TestResult<Vec<_>>>()?;
    let plan = base.prepare(stores, writes, budget())?;
    for node in &plan.nodes {
        stores
            .get_mut(node.shard)
            .ok_or("node shard missing")?
            .put(&node.node)?;
    }
    Ok(Catalogue::at::<haematite::store::Error>(
        plan.root,
        policy(),
    )?)
}

/// The writes `Storage::mutate` would commit for `request`, grouped by shard.
fn writes(storage: &Storage, request: &Mutation) -> TestResult<(Vec<ShardChanges>, Vec<u8>)> {
    let at = storage.capture(&Selection {
        branch: request.branch.clone(),
        version: None,
    })?;
    let next = Version {
        generation: at.version.generation,
        sequence: at.version.sequence.checked_add(1).ok_or("sequence")?,
    };
    let mut writes: Vec<PreparedWrite> = storage.prepare_change(&at, &request.change, next)?;
    counts::append(&storage.reader, &at, &request.change, &mut writes)?;
    let receipt = serde_json::to_vec(&Receipt {
        request_id: request.request_id.clone(),
        fingerprint: NodeHash::of(&serde_json::to_vec(request)?).to_string(),
        branch: request.branch.clone(),
        version: next,
        retention: None,
    })?;
    writes.push((
        receipt_key(&request.branch, next.generation, &request.request_id),
        Some(receipt.clone()),
    ));
    let mut shards = BTreeMap::<usize, Vec<PreparedWrite>>::new();
    for (key, value) in writes {
        shards
            .entry(storage.writer.shard_for(&key))
            .or_default()
            .push((key, value));
    }
    let changes = shards
        .into_iter()
        .map(|(shard, writes)| ShardChanges { shard, writes })
        .collect();
    Ok((changes, receipt))
}

/// The plan the commit writer would publish for `request`, and its sizes.
fn measure(
    storage: &Storage,
    data_dir: &Path,
    shape: &Shape,
    request: &Mutation,
) -> TestResult<Measured> {
    let kept = KeptLimit::new(shape.kept)?;
    let scratch = tempfile::tempdir()?;
    let mut handles = stores(data_dir, scratch.path(), shape.shards)?;
    let catalogue = catalogue(storage, &mut handles, kept)?;
    let (changes, receipt) = writes(storage, request)?;
    let name = branch_name(&request.branch);
    let record = storage.refs.get(&name).ok_or("main missing")?;
    let plan = catalogue.prepare_advance(
        &handles,
        Advance {
            name: name.clone(),
            created: record.created,
            sequence: record.seq,
            timestamp: record.timestamp.saturating_add(1),
            parents: Vec::new(),
            changes,
            kept_limit: kept,
            service_kept_limit: kept,
            history_identity: history::identity(&name),
        },
        budget(),
        |key| storage.writer.shard_for(key),
    )?;
    sizes(&plan, receipt)
}

/// Node and frame sizes; file bytes come from publishing into a scratch store.
fn sizes(plan: &Prepared, receipt: Vec<u8>) -> TestResult<Measured> {
    let scratch = tempfile::tempdir()?;
    let mut measured = Measured {
        nodes: plan.nodes.len(),
        records: plan.records.len(),
        receipt_bytes: receipt.len(),
        ..Measured::default()
    };
    let mut listed = Vec::new();
    for node in &plan.nodes {
        *measured.per_shard.entry(node.shard).or_default() += 1;
        let uncompressed = u64::try_from(node.node.serialise_and_hash().0.len())?;
        let directory = scratch.path().join(format!("shard-{}", node.shard));
        if !directory.exists() {
            fs::create_dir(&directory)?;
        }
        let store_dir = directory.join("store");
        let mut store = Disk::new(&store_dir, NodeCacheBudget::Unlimited)?;
        let hash: NodeHash = store.put(&node.node)?;
        let file = fs::metadata(hash.node_file_in(&store_dir))?.len();
        measured.plain_bytes = measured.plain_bytes.saturating_add(uncompressed);
        measured.file_bytes = measured.file_bytes.saturating_add(file);
        measured.largest_plain = measured.largest_plain.max(uncompressed);
        measured.largest_file = measured.largest_file.max(file);
        listed.push(CommitNode {
            shard: u64::try_from(node.shard)?,
            hash: hash.into_bytes(),
            class: 0,
        });
    }
    listed.sort_unstable_by_key(|node| (node.shard, node.hash));
    listed.dedup_by_key(|node| (node.shard, node.hash));
    measured.frame_bytes = u64::try_from(frame(plan, receipt, listed)?.len())?;
    Ok(measured)
}

fn frame(plan: &Prepared, receipt: Vec<u8>, nodes: Vec<CommitNode>) -> TestResult<Vec<u8>> {
    let mut refs_blob = u32::try_from(plan.records.len())?.to_le_bytes().to_vec();
    for record in &plan.records {
        refs_blob.extend_from_slice(&u32::try_from(record.len())?.to_le_bytes());
        refs_blob.extend_from_slice(record);
    }
    CommitFrame {
        flags: 0,
        uuid: [0; 16],
        sequence: 1,
        timestamp: 0,
        base_catalog_hash: plan.base.into_bytes(),
        new_catalog_hash: plan.root.into_bytes(),
        refs_blob,
        receipt,
        nodes,
    }
    .encode(CommitLimits {
        max_frame_bytes: u64::MAX,
        max_nodes: u64::MAX,
        max_uncompressed_node_bytes: u64::MAX,
    })
    .map_err(|refusal| format!("frame encode refused: {refusal:?}").into())
}

fn report(kind: &str, shape: &Shape, documents: usize, text: usize, measured: &Measured) {
    println!(
        "SPAN kind={kind} shards={} kept={} max_page={} max_frame={} documents={documents} \
         text_bytes={text} nodes={} per_shard={:?} plain_bytes={} file_bytes={} \
         largest_plain={} largest_file={} records={} receipt_bytes={} frame_bytes={} \
         limit_nodes={LIMIT_NODES} limit_frame={LIMIT_FRAME_BYTES} limit_node={LIMIT_NODE_BYTES} \
         over_nodes={} over_frame={} over_node={} over_file={}",
        shape.shards,
        shape.kept,
        shape.max_page,
        shape.max_frame,
        measured.nodes,
        measured.per_shard,
        measured.plain_bytes,
        measured.file_bytes,
        measured.largest_plain,
        measured.largest_file,
        measured.records,
        measured.receipt_bytes,
        measured.frame_bytes,
        u64::try_from(measured.nodes).unwrap_or(u64::MAX) > LIMIT_NODES,
        measured.frame_bytes > LIMIT_FRAME_BYTES,
        measured.largest_plain > LIMIT_NODE_BYTES,
        measured.largest_file > LIMIT_NODE_BYTES,
    );
}

fn put(storage: &mut Storage, id: &str, text: &str) -> TestResult {
    let expected = version(
        storage
            .refs
            .get(&branch_name("main"))
            .ok_or("main missing")?,
    );
    storage.mutate(&Mutation {
        request_id: format!("warm-{id}"),
        branch: "main".to_owned(),
        expected,
        change: Change::Put(DocumentWrite {
            collection: "records".to_owned(),
            document_id: id.to_owned(),
            document: RawValue::from_string(serde_json::json!({ "text": text }).to_string())?,
        }),
    })?;
    Ok(())
}

fn import(storage: &Storage, documents: usize, text: usize) -> TestResult<Mutation> {
    let body = "x".repeat(text);
    let documents = (0..documents)
        .map(|index| {
            Ok(ImportedDocument {
                document_id: format!("doc-{index:06}"),
                document: RawValue::from_string(serde_json::json!({ "text": body }).to_string())?,
            })
        })
        .collect::<TestResult<Vec<_>>>()?;
    Ok(Mutation {
        request_id: "span-largest".to_owned(),
        branch: "main".to_owned(),
        expected: version(
            storage
                .refs
                .get(&branch_name("main"))
                .ok_or("main missing")?,
        ),
        change: Change::PutMany(DocumentImport {
            collection: "records".to_owned(),
            documents,
        }),
    })
}

#[test]
fn span_measure_largest_production_commit_against_today_limits() -> TestResult {
    let shape = Shape::from_env()?;
    let directory = tempfile::tempdir()?;
    let data_dir = directory.path().join("database");
    let mut storage = Storage::create_kept(
        DatabaseConfig {
            data_dir: data_dir.clone(),
            shard_count: shape.shards,
            distributed: None,
            executor_threads: Some(1),
            node_cache_budget: Some(NodeCacheBudget::Unlimited),
            retention: haematite::Retention::KeepAll,
        },
        NodeCacheBudget::Unlimited,
        KeptLimit::new(shape.kept)?,
        fixture_long_wait_interval(),
    )?;
    let expected = version(
        storage
            .refs
            .get(&branch_name("main"))
            .ok_or("main missing")?,
    );
    storage.mutate(&Mutation {
        request_id: "declare".to_owned(),
        branch: "main".to_owned(),
        expected,
        change: Change::Collection(CollectionWrite {
            collection: "records".to_owned(),
            schema: None,
        }),
    })?;
    // The control: one small document, nothing spilled.
    let control = Mutation {
        request_id: "span-control".to_owned(),
        branch: "main".to_owned(),
        expected: version(
            storage
                .refs
                .get(&branch_name("main"))
                .ok_or("main missing")?,
        ),
        change: Change::Put(DocumentWrite {
            collection: "records".to_owned(),
            document_id: "control".to_owned(),
            document: RawValue::from_string(r#"{"text":"c"}"#.to_owned())?,
        }),
    };
    let measured = measure(&storage, &data_dir, &shape, &control)?;
    report("control", &shape, 1, 1, &measured);
    // K writes fill the kept window, so the next advance spills every row.
    for index in 0..shape.kept {
        put(&mut storage, &format!("warm-{index:04}"), "w")?;
    }
    let documents = shape.documents()?;
    let text = shape.text_bytes()?;
    let largest = import(&storage, documents, text)?;
    let request_bytes = serde_json::to_vec(&largest)?.len();
    let measured = measure(&storage, &data_dir, &shape, &largest)?;
    println!("SPAN request_bytes={request_bytes}");
    report("largest", &shape, documents, text, &measured);
    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
}