haematite-cli 0.21.3

//! Measures durable fixture writes and complete served queries before speed changes.
//! Each pass creates one service; it reports observations without timing assertions.

#![forbid(unsafe_code)]

#[cfg(unix)]
#[path = "probe_query_ms/fixture.rs"]
mod fixture;
#[cfg(unix)]
#[path = "probe_query_ms/options.rs"]
mod options;
#[cfg(unix)]
#[path = "probe_query_ms/served.rs"]
mod served;

#[cfg(unix)]
use std::collections::BTreeMap;
#[cfg(unix)]
use std::error::Error;
#[cfg(unix)]
use std::io::{self, BufWriter, Write};
#[cfg(unix)]
use std::path::Path;
#[cfg(unix)]
use std::time::{Duration, Instant};

#[cfg(unix)]
use haematite_cli::client::{Client, Reply, refused};
#[cfg(unix)]
use haematite_cli::serve::protocol::{QueryPage, QueryRequest, Selection, Version};
#[cfg(unix)]
use haematite_cli::serve::query::Rows;
#[cfg(unix)]
use serde_json::Value;

#[cfg(unix)]
type ProbeResult<T> = Result<T, Box<dyn Error>>;

#[cfg(unix)]
fn main() -> ProbeResult<()> {
    let options = options::Options::read()?;
    let binaries = options::Binaries::find()?;
    let commit = options::commit()?;
    let machine = options::machine()?;
    let mut directory = options::Directory::create(options.dir.as_deref(), options.keep)?;
    let result = pass(&options, &binaries, &directory.path, &commit, &machine);
    let cleanup = directory.finish();
    match (result, cleanup) {
        (Ok(()), Ok(())) => Ok(()),
        (Err(error), Ok(())) | (Ok(()), Err(error)) => Err(error),
        (Err(error), Err(cleanup)) => Err(format!("{error}; directory cleanup: {cleanup}").into()),
    }
}

#[cfg(not(unix))]
fn main() -> std::process::ExitCode {
    eprintln!("probe_query_ms requires Unix sockets");
    std::process::ExitCode::FAILURE
}

#[cfg(unix)]
fn pass(
    options: &options::Options,
    binaries: &options::Binaries,
    directory: &Path,
    commit: &str,
    machine: &str,
) -> ProbeResult<()> {
    let mut table = Table::new()?;
    header(&mut table, options, binaries, directory, commit, machine)?;
    let config = served::configure(directory, binaries, options)?;
    let started = Instant::now();
    let mut running = served::Running::start(&binaries.serve, &config)?;
    table.measured(
        "serve startup",
        started.elapsed(),
        1,
        "waits on stderr ready signal",
    )?;
    table.metadata("serve child pid", &running.pid()?.to_string())?;
    table.metadata(
        "serve software version (owned registry entry)",
        running.version()?,
    )?;
    let result = measure(options, directory, &mut table);
    let stopped = running.stop();
    match (&result, &stopped) {
        (Ok(()), Ok(())) => table.metadata("serve cleanup", "owned child stopped by pid and reaped; owned registry publication removed; permanent name-lock anchor retained")?,
        (_, Err(error)) => table.metadata("serve cleanup failed", &error.to_string())?,
        (Err(_), Ok(())) => {}
    }
    table.flush()?;
    match (result, stopped) {
        (Ok(()), Ok(())) => Ok(()),
        (Err(error), Ok(())) | (Ok(()), Err(error)) => Err(error),
        (Err(error), Err(stop)) => Err(format!("{error}; serve cleanup: {stop}").into()),
    }
}

#[cfg(unix)]
fn header(
    table: &mut Table,
    options: &options::Options,
    binaries: &options::Binaries,
    directory: &Path,
    commit: &str,
    machine: &str,
) -> ProbeResult<()> {
    table.metadata("commit (checkout HEAD read at runtime)", commit)?;
    table.metadata("machine", machine)?;
    table.metadata(
        "fixture",
        &format!(
            "N={} D={} R={} patients={} wards=20 text_seed={} vector_seed={}",
            options.records,
            options.dims,
            options.repeat,
            options.records / 10,
            fixture::TEXT_SEED,
            fixture::VECTOR_SEED
        ),
    )?;
    table.metadata(
        "haematite binary used",
        &binaries.serve.display().to_string(),
    )?;
    table.metadata(
        "query_worker binary used",
        &binaries.query.display().to_string(),
    )?;
    table.metadata("embedding worker", "not used: caller-supplied seeded vectors and record questions need no model or haematite-embed-stand-in")?;
    table.metadata(
        "directory",
        &format!(
            "{}; {}",
            directory.display(),
            if options.keep {
                "kept after stopping serve"
            } else {
                "removed after stopping serve"
            }
        ),
    )?;
    table.metadata("timing scope", "query.run through all pages, including wire, worker startup and JSON decoding; fresh client hello excluded; no client result or plan cache; one serve shares OS/store caches")?;
    table.metadata("limits", "32 MiB frames; 256 MiB query memory; 120 s query budget; 300 s socket I/O bound; scan covers fixture")?;
    Ok(())
}

#[cfg(unix)]
fn measure(options: &options::Options, directory: &Path, table: &mut Table) -> ProbeResult<()> {
    let socket = directory.join("serve.sock");
    let client = Client::connect(&socket, served::IO_BOUND)?;
    table.metadata(
        "serve hello (reported protocol version and limits)",
        client.hello().get(),
    )?;
    table.metadata(
        "hello software version",
        "not reported by hello; software version comes from the owned child registry entry",
    )?;
    let mut writer = served::Writer::new(client)?;
    fixture::build(&mut writer, options, directory, table)?;
    let head = writer.head;
    drop(writer);
    let note = fixture::note_id(0);
    let queries = [
        "notes | limit(100)".to_owned(),
        "notes | 'select(.status == \"open\")' | limit(100)".to_owned(),
        "notes | rg(\"troponin\") | limit(100)".to_owned(),
        "notes | rg(\"TROPONIN\"; \"-i\") | count".to_owned(),
        "notes | sort(.written; desc) | limit(100)".to_owned(),
        "notes | group_by(.ward) | count".to_owned(),
        "match (n:notes)-[:about]->(p:patients) | limit(100)".to_owned(),
        "match (n:notes)-[:about]->(p:patients)-[:in]->(w:wards) | limit(100)".to_owned(),
        "{ notes { body patient { id } } } | limit(100)".to_owned(),
        format!("notes | similar(.body_vec; record(\"{note}\"); cosine) | limit(10)"),
        format!(
            "notes | similar(.body_vec; record(\"{note}\"); cosine) | limit(10) | '{{id: $id, score: $score}}'"
        ),
        format!("notes | rg(\"troponin\") | rerank(.body_vec; record(\"{note}\"); cosine)"),
    ];
    for query in queries {
        time_query(&socket, &query, head, options.repeat, table)?;
    }
    Ok(())
}

#[cfg(unix)]
#[derive(Debug, PartialEq, Eq)]
struct Count {
    rows: u64,
    aggregate: Option<u64>,
}

#[cfg(unix)]
enum QueryOutcome {
    Complete(Count),
    Refused { reason: String, rows: u64 },
}

#[cfg(unix)]
fn time_query(
    socket: &Path,
    query: &str,
    head: Version,
    repeat: usize,
    table: &mut Table,
) -> ProbeResult<()> {
    let mut durations = Vec::new();
    durations.try_reserve_exact(repeat)?;
    let mut previous: Option<Count> = None;
    for sample in 0..repeat {
        let mut client = Client::connect(socket, served::IO_BOUND)?;
        let started = Instant::now();
        let count = match query_once(&mut client, query, head) {
            Ok(QueryOutcome::Complete(count)) => count,
            Ok(QueryOutcome::Refused { reason, rows }) => {
                table.metadata(
                    query,
                    &format!(
                        "{reason}; attempt {}/{repeat}; {rows} rows received before refusal; no timing claimed",
                        sample + 1,
                    ),
                )?;
                return Ok(());
            }
            Err(error) => {
                table.metadata(query, &format!("query failed: {error}; no timing claimed"))?;
                return Err(error);
            }
        };
        durations.push(started.elapsed());
        if previous.as_ref().is_some_and(|previous| previous != &count) {
            return Err(format!("query counts changed across repetitions: {query}").into());
        }
        previous = Some(count);
    }
    let count = previous.ok_or("query has no samples")?;
    let (minimum, median) = timings(&mut durations)?;
    let detail = count.aggregate.map_or_else(
        || {
            if query.contains("rg(") {
                format!(
                    "complete; {repeat} repetitions; returned hits={}",
                    count.rows
                )
            } else {
                format!("complete; {repeat} repetitions")
            }
        },
        |count| format!("complete; {repeat} repetitions; count={count}"),
    );
    table.row(
        query,
        &[
            format!("{minimum:.3}"),
            format!("{median:.3}"),
            count.rows.to_string(),
            String::new(),
            String::new(),
            detail,
        ],
    )
}

#[cfg(unix)]
fn query_once(client: &mut Client, query: &str, head: Version) -> ProbeResult<QueryOutcome> {
    let mut request = QueryRequest {
        selection: Selection {
            branch: "main".to_owned(),
            version: Some(head),
        },
        query: query.to_owned(),
        parameters: BTreeMap::new(),
        page: QueryPage {
            rows: u64::from(client.max_page().get()).try_into()?,
            bytes: u64::from(client.max_frame().get().saturating_sub(1024)).try_into()?,
        },
        continuation: None,
        rows: Rows::Values,
    };
    let mut count = Count {
        rows: 0,
        aggregate: None,
    };
    loop {
        let raw = serde_json::value::to_raw_value(&request)?;
        let answer = match client.call("query.run", Some(&raw))? {
            Reply::Answered(answer) => answer,
            Reply::Refused(refusal) => {
                return Ok(QueryOutcome::Refused {
                    reason: refused(&refusal),
                    rows: count.rows,
                });
            }
        };
        let answer: Value = serde_json::from_str(answer.get())?;
        let version: Version = serde_json::from_value(
            answer
                .get("selection")
                .and_then(|selection| selection.get("version"))
                .ok_or("query reply has no version")?
                .clone(),
        )?;
        if version != head {
            return Err("query answered at another version".into());
        }
        let rows = answer
            .get("rows")
            .and_then(Value::as_array)
            .ok_or("query reply has no rows")?;
        count.rows = count
            .rows
            .checked_add(u64::try_from(rows.len())?)
            .ok_or("query row count exceeds u64")?;
        if query.ends_with("| count") {
            for row in rows {
                count.aggregate = Some(
                    row.get("value")
                        .and_then(Value::as_u64)
                        .ok_or("count reply is not an unsigned integer")?,
                );
            }
        }
        match answer.get("next").ok_or("query reply has no next")? {
            Value::Null => break,
            Value::String(next) if !next.is_empty() => {
                if request.continuation.as_deref() == Some(next.as_str()) {
                    return Err("query repeated its continuation without advancing".into());
                }
                request.continuation = Some(next.clone());
            }
            _ => return Err("query continuation is neither null nor nonempty text".into()),
        }
    }
    if query.ends_with("| count") && count.rows != 1 {
        return Err("count reply did not return exactly one aggregate row".into());
    }
    Ok(QueryOutcome::Complete(count))
}

#[cfg(unix)]
fn timings(durations: &mut [Duration]) -> ProbeResult<(f64, f64)> {
    durations.sort_unstable();
    let minimum = durations
        .first()
        .ok_or("query has no minimum")?
        .as_secs_f64()
        * 1000.0;
    let midpoint = durations.len() / 2;
    let upper = durations
        .get(midpoint)
        .ok_or("query has no median")?
        .as_secs_f64()
        * 1000.0;
    let median = if durations.len().is_multiple_of(2) {
        let lower = durations
            .get(midpoint - 1)
            .ok_or("query has no lower median")?
            .as_secs_f64()
            * 1000.0;
        lower.midpoint(upper)
    } else {
        upper
    };
    Ok((minimum, median))
}

#[cfg(unix)]
struct Table {
    output: BufWriter<io::Stdout>,
}

#[cfg(unix)]
impl Table {
    fn new() -> ProbeResult<Self> {
        let mut table = Self {
            output: BufWriter::new(io::stdout()),
        };
        writeln!(
            table.output,
            "| Measurement | min ms | median ms | rows/items | items/s | bytes | result |"
        )?;
        writeln!(
            table.output,
            "| --- | ---: | ---: | ---: | ---: | ---: | --- |"
        )?;
        Ok(table)
    }

    fn row(&mut self, label: &str, cells: &[String; 6]) -> ProbeResult<()> {
        write!(self.output, "| {} ", cell(label))?;
        for value in cells {
            write!(self.output, "| {} ", cell(value))?;
        }
        writeln!(self.output, "|")?;
        self.flush()
    }

    fn metadata(&mut self, label: &str, detail: &str) -> ProbeResult<()> {
        self.row(
            label,
            &[
                String::new(),
                String::new(),
                String::new(),
                String::new(),
                String::new(),
                detail.to_owned(),
            ],
        )
    }

    fn measured(
        &mut self,
        label: &str,
        elapsed: Duration,
        count: u32,
        detail: &str,
    ) -> ProbeResult<()> {
        let seconds = elapsed.as_secs_f64();
        if seconds <= 0.0 {
            return Err("measurement clock returned zero elapsed time".into());
        }
        let millis = format!("{:.3}", seconds * 1000.0);
        self.row(
            label,
            &[
                millis.clone(),
                millis,
                count.to_string(),
                format!("{:.3}", f64::from(count) / seconds),
                String::new(),
                detail.to_owned(),
            ],
        )
    }

    fn size(
        &mut self,
        label: &str,
        elapsed: Duration,
        bytes: u64,
        detail: &str,
    ) -> ProbeResult<()> {
        self.row(
            label,
            &[
                format!("{:.3}", elapsed.as_secs_f64() * 1000.0),
                String::new(),
                String::new(),
                String::new(),
                bytes.to_string(),
                detail.to_owned(),
            ],
        )
    }

    fn flush(&mut self) -> ProbeResult<()> {
        Ok(self.output.flush()?)
    }
}

#[cfg(unix)]
fn cell(text: &str) -> String {
    text.replace('&', "&amp;")
        .replace('|', "&#124;")
        .replace('<', "&lt;")
        .replace('>', "&gt;")
        .replace('\r', " ")
        .replace('\n', "<br>")
}