haematite-cli 0.21.0

use std::num::{NonZeroU32, NonZeroUsize};
use std::sync::mpsc::{Receiver, channel};
use std::time::Duration;

use haematite::db::DatabaseConfig;
use haematite::store::NodeCacheBudget;
use serde_json::{Value, json, value::RawValue};

use super::{ResponseFailure, Service, fits_serialized};
use crate::serve::protocol::{
    BranchPage, Change, CollectionPage, CollectionWrite, DocumentRead, DocumentWrite, HistoryPage,
    Mutation, Operation, Selection, Version,
};
use crate::serve::storage::{Receipt, Storage};

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

#[test]
fn serialized_result_check_stops_at_the_exact_byte_limit() -> Result<(), Box<dyn std::error::Error>>
{
    let value = json!({"raw_json": "{\"large\":18446744073709551616001}"});
    let exact = serde_json::to_vec(&value)?.len();
    let exact_limit = NonZeroUsize::new(exact).ok_or("empty serialized value")?;
    let short_limit = NonZeroUsize::new(exact - 1).ok_or("serialized value too short")?;

    assert!(fits_serialized(&value, exact_limit)?);
    assert!(!fits_serialized(&value, short_limit)?);
    Ok(())
}

const PAGE_LIMIT: NonZeroU32 = NonZeroU32::MIN.saturating_add(9);

fn main_at(version: Option<Version>) -> Selection {
    Selection {
        branch: "main".to_owned(),
        version,
    }
}

fn put_one(id: &str, document: &str, expected: Version) -> Result<Mutation, serde_json::Error> {
    Ok(Mutation {
        request_id: id.to_owned(),
        branch: "main".to_owned(),
        expected,
        change: Change::Put(DocumentWrite {
            collection: "records".to_owned(),
            document_id: "one".to_owned(),
            document: RawValue::from_string(document.to_owned())?,
        }),
    })
}

fn get_one(service: &Service) -> Result<Value, ResponseFailure> {
    let read = Operation::Get(DocumentRead {
        selection: main_at(None),
        collection: "records".to_owned(),
        document_id: "one".to_owned(),
    });
    service.dispatch(read)
}

/// Every read the test thread makes while the one writer is paused after its
/// store commit and before its publication. None may wait on the writer.
fn read_while_paused(service: &Service, paused: &Receiver<()>, published: Version) -> TestResult {
    paused
        .recv_timeout(Duration::from_secs(10))
        .map_err(|error| format!("paused signal not received within 10 s: {error}"))?;
    let expected = serde_json::to_value(published)?;
    for _ in 0..10 {
        let read = get_one(service)?;
        assert_eq!(read["selection"]["version"], expected);
        assert_eq!(read["document"]["raw_json"], r#"{"n":1}"#);
    }
    let collections = Operation::Collections(CollectionPage {
        selection: main_at(None),
        after: None,
        limit: PAGE_LIMIT,
    });
    let collections = service.dispatch(collections)?;
    assert_eq!(collections["selection"]["version"], expected);
    let history = Operation::History(HistoryPage {
        branch: "main".to_owned(),
        after: None,
        limit: PAGE_LIMIT,
    });
    let history = service.dispatch(history)?;
    let versions = history["entries"]
        .as_array()
        .ok_or("history has no entry array")?;
    let last = versions.last().ok_or("history has no entries")?;
    assert_eq!(last["version"], expected);
    let branches = Operation::Branches(BranchPage {
        after: None,
        limit: PAGE_LIMIT,
    });
    let branches = service.dispatch(branches)?;
    let main = branches["entries"]
        .as_array()
        .ok_or("branches.list has no entry array")?
        .iter()
        .find(|entry| entry["branch"] == "main")
        .ok_or("branches.list has no main")?;
    assert_eq!(main["version"], expected);
    let second = put_one("second", r#"{"n":3}"#, published)?;
    let refusal = service
        .dispatch(Operation::Mutate(second))
        .err()
        .ok_or("second mutation was admitted beside the paused writer")?;
    assert_eq!(refusal.code, "busy");
    Ok(())
}

#[test]
fn reads_answer_from_the_confirmed_catalogue_while_the_writer_is_paused() -> TestResult {
    let dir = tempfile::tempdir()?;
    let storage = Storage::create(
        DatabaseConfig {
            data_dir: dir.path().join("database"),
            shard_count: 2,
            distributed: None,
            executor_threads: Some(2),
            node_cache_budget: Some(NodeCacheBudget::Unlimited),
            retention: haematite::Retention::KeepAll,
        },
        NodeCacheBudget::Unlimited,
        fixture_long_wait_interval(),
    )?;
    let service = Service::new(storage, PAGE_LIMIT)?;
    let initial = service.catalogue.capture(&main_at(None))?.version;
    let declare = Operation::Mutate(Mutation {
        request_id: "declare".to_owned(),
        branch: "main".to_owned(),
        expected: initial,
        change: Change::Collection(CollectionWrite {
            collection: "records".to_owned(),
            schema: None,
        }),
    });
    service.dispatch(declare)?;
    let declared = service.catalogue.capture(&main_at(None))?.version;
    let first = Operation::Mutate(put_one("first", r#"{"n":1}"#, declared)?);
    let first = service.dispatch(first)?;
    let v1 = serde_json::from_value::<Receipt>(first)?.version;
    let write = Operation::Mutate(put_one("write", r#"{"n":2}"#, v1)?);

    let written = std::thread::scope(|scope| -> Result<Value, Box<dyn std::error::Error>> {
        let (paused_signal, paused) = channel();
        let (release, release_signal) = channel();
        let service = &service;
        let writer = scope.spawn(move || {
            Storage::arm_publish_pause(paused_signal, release_signal);
            service.dispatch(write)
        });
        let observed = read_while_paused(service, &paused, v1);
        // Release the writer only once every read above has returned.
        let released = release.send(());
        let written = writer.join().map_err(|_| "writer thread panicked")?;
        observed?;
        released?;
        written.map_err(Into::into)
    })?;

    let next = Version {
        generation: v1.generation,
        sequence: v1.sequence + 1,
    };
    assert_eq!(serde_json::from_value::<Receipt>(written)?.version, next);
    let read = get_one(&service)?;
    assert_eq!(read["selection"]["version"], serde_json::to_value(next)?);
    assert_eq!(read["document"]["raw_json"], r#"{"n":2}"#);
    Ok(())
}

#[test]
fn nonzero_service_byte_limit_readback() -> TestResult {
    for (given, expected) in [
        (None, usize::MAX),
        (Some(1), 1),
        (Some(37), 37),
        (Some(usize::MAX), usize::MAX),
    ] {
        let given = given
            .map(|value| NonZeroUsize::new(value).ok_or("zero limit"))
            .transpose()?;
        assert_eq!(
            super::response_limit(std::hint::black_box(given))?.get(),
            expected
        );
    }
    Ok(())
}

#[test]
fn nonzero_query_running_limit_readback() -> TestResult {
    for (given, expected) in [
        (None, 1),
        (Some(1), 1),
        (Some(3), 3),
        (Some(usize::MAX), usize::MAX),
    ] {
        let settings = given
            .map(|value| -> Result<_, Box<dyn std::error::Error>> {
                Ok(crate::serve::query::Settings {
                    time_ms: std::num::NonZeroU64::MIN,
                    memory_bytes: std::num::NonZeroU64::MIN,
                    outputs_per_record: std::num::NonZeroU64::MIN,
                    scan_records: std::num::NonZeroU64::MIN.saturating_add(1),
                    max_running: NonZeroUsize::new(value).ok_or("zero running limit")?,
                    worker: std::path::PathBuf::from("unused-worker"),
                })
            })
            .transpose()?;
        let limit = super::running_limit(std::hint::black_box(settings.as_ref()))?;
        assert_eq!(limit.get(), expected);
        if expected <= 3 {
            let slots = crate::serve::query::admission::Slots::new(limit);
            let held = (0..expected)
                .map(|_| slots.try_acquire())
                .collect::<Result<Vec<_>, _>>()?;
            assert_eq!(
                slots.try_acquire().err().ok_or("extra slot admitted")?.code,
                "query_busy"
            );
            drop(held);
            assert!(slots.try_acquire().is_ok());
        }
    }
    Ok(())
}

fn nonzero_service_fixture() -> Result<(Service, tempfile::TempDir), Box<dyn std::error::Error>> {
    use crate::serve::protocol::{Cardinality, Edge, Node, Relation, VectorWrite};
    let directory = tempfile::tempdir()?;
    let mut storage = Storage::create(
        DatabaseConfig {
            data_dir: directory.path().join("database"),
            shard_count: 1,
            distributed: None,
            executor_threads: Some(1),
            node_cache_budget: Some(NodeCacheBudget::Unlimited),
            retention: haematite::Retention::KeepAll,
        },
        NodeCacheBudget::Unlimited,
        fixture_long_wait_interval(),
    )?;
    let mut changes = vec![Change::Collection(CollectionWrite {
        collection: "records".to_owned(),
        schema: None,
    })];
    for id in ["a", "b", "c"] {
        changes.push(Change::Put(DocumentWrite {
            collection: "records".to_owned(),
            document_id: id.to_owned(),
            document: RawValue::from_string("{}".to_owned())?,
        }));
    }
    for relation in ["links", "other"] {
        changes.push(Change::RelationPut(Relation {
            relation: relation.to_owned(),
            from_collection: "records".to_owned(),
            to_collection: "records".to_owned(),
            cardinality: Cardinality::ManyToMany,
        }));
    }
    for id in ["b", "c"] {
        changes.push(Change::VectorPut(VectorWrite {
            collection: "records".to_owned(),
            document_id: id.to_owned(),
            field: "embedding".to_owned(),
            vector: vec![1.0, 2.0],
        }));
        changes.push(Change::EdgePut(Edge {
            from: Node {
                collection: "records".to_owned(),
                document_id: "a".to_owned(),
            },
            relation: "links".to_owned(),
            to: Node {
                collection: "records".to_owned(),
                document_id: id.to_owned(),
            },
            properties: json!({}),
        }));
    }
    for (index, change) in changes.into_iter().enumerate() {
        let expected = storage.capture(&main_at(None))?.version;
        storage.mutate(&Mutation {
            request_id: format!("nonzero/{index}"),
            branch: "main".to_owned(),
            expected,
            change,
        })?;
    }
    Ok((Service::new(storage, PAGE_LIMIT)?, directory))
}

fn nonzero_read_operation(kind: usize) -> Result<Operation, Box<dyn std::error::Error>> {
    use crate::serve::protocol::{EdgePage, GraphPage, Node, RelationsRead, VectorPage};
    Ok(match kind {
        0 => Operation::VectorList(VectorPage {
            selection: main_at(None),
            collection: "records".to_owned(),
            field: "embedding".to_owned(),
            after: None,
            limit: PAGE_LIMIT,
        }),
        1 => Operation::Relations(RelationsRead {
            selection: main_at(None),
        }),
        2 => Operation::EdgesOut(EdgePage {
            selection: main_at(None),
            node: Node {
                collection: "records".to_owned(),
                document_id: "a".to_owned(),
            },
            relation: None,
            after: None,
            limit: PAGE_LIMIT,
        }),
        3 => Operation::GraphExport(GraphPage {
            selection: main_at(None),
            after: None,
            limit: PAGE_LIMIT,
            relations: None,
            collections: None,
            members: None,
            records: None,
            unrelated: None,
        }),
        _ => return Err("unknown read kind".into()),
    })
}

fn nonzero_bounded_service_read(kind: usize, key: &str) -> TestResult {
    let (service, directory) = nonzero_service_fixture()?;
    let access = crate::serve::access::request(&crate::serve::access::Standalone)?;
    let all = service.dispatch_read(&access, nonzero_read_operation(kind)?, None)?;
    let entries = all[key].as_array().ok_or("entries absent")?;
    assert_eq!(entries.len(), 2);
    let exact = if kind == 1 {
        serde_json::to_vec(entries)?.len()
    } else {
        serde_json::to_vec(&entries[0])?.len()
    };
    for given in [None, Some(usize::MAX)] {
        let given = given
            .map(|value| NonZeroUsize::new(value).ok_or("zero budget"))
            .transpose()?;
        let read = service.dispatch_read(
            &access,
            nonzero_read_operation(kind)?,
            std::hint::black_box(given),
        )?;
        assert_eq!(read[key], all[key]);
    }
    let refused = service
        .dispatch_read(
            &access,
            nonzero_read_operation(kind)?,
            Some(NonZeroUsize::MIN),
        )
        .err()
        .ok_or("minimum did not bound the read")?;
    assert_eq!(refused.code, "response_limit");
    let limit = NonZeroUsize::new(exact).ok_or("empty serialized entry")?;
    let read = service.dispatch_read(&access, nonzero_read_operation(kind)?, Some(limit))?;
    if kind == 1 {
        assert_eq!(read[key], all[key]);
        let short = NonZeroUsize::new(exact - 1).ok_or("empty short budget")?;
        let refused = service
            .dispatch_read(&access, nonzero_read_operation(kind)?, Some(short))
            .err()
            .ok_or("short aggregate budget accepted")?;
        assert_eq!(refused.code, "response_limit");
    } else {
        assert_eq!(read[key], json!([entries[0].clone()]));
        assert!(!read["next_after"].is_null());
    }
    drop(service);
    drop(directory);
    Ok(())
}

#[test]
fn nonzero_service_vector_read_stops() -> TestResult {
    nonzero_bounded_service_read(0, "vectors")
}

#[test]
fn nonzero_service_relations_read_stops() -> TestResult {
    nonzero_bounded_service_read(1, "relations")
}

#[test]
fn nonzero_service_edge_read_stops() -> TestResult {
    nonzero_bounded_service_read(2, "edges")
}

#[test]
fn nonzero_service_graph_read_stops() -> TestResult {
    nonzero_bounded_service_read(3, "edges")
}

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