searchi/crates/searchi-core/examples/performance.rs

//! CPU-only synthetic retrieval benchmark. This does not measure filesystem
//! scanning, extraction, embedding inference, or desktop rendering performance.
use anyhow::{bail, Context, Result};
use rusqlite::{params, Connection};
use searchi_core::SearchIndex;
use serde::Serialize;
use std::{env, fs, path::PathBuf, time::Instant};

#[derive(Serialize)]
struct Measurements {
    query: String,
    samples: usize,
    p50_ms: f64,
    first_query_ms: f64,
    first_query_target_met: bool,
    p95_ms: f64,
    target_p95_ms: f64,
    target_met: bool,
}

#[derive(Serialize)]
struct Report {
    background_prepare_seconds: f64,
    reopened_background_prepare_seconds: f64,
    background_representative_queries: Vec<String>,
    prepared_frequency_count: usize,
    first_foreground_distinct_names: DistinctMeasurements,
    reopened_reader_first_names: DistinctMeasurements,
    fixture: &'static str,
    files: usize,
    documents: usize,
    seed_seconds: f64,
    migration_seconds: f64,
    fixture_reused: bool,
    database_bytes: u64,
    database_path: String,
    architecture: &'static str,
    os: &'static str,
    cpu_only: bool,
    queries: Vec<Measurements>,
}

#[derive(Serialize)]
struct DistinctMeasurements {
    description: &'static str,
    includes_reader_open: bool,
    queries: Vec<String>,
    samples_ms: Vec<f64>,
    p50_ms: f64,
    p95_ms: f64,
    first_sample_ms: f64,
    first_sample_target_met: bool,
    maximum_ms: f64,
    samples_exceeding_target: usize,
    target_p95_ms: f64,
    target_met: bool,
}

fn distinct_names(
    index: &SearchIndex,
    database: &std::path::Path,
    files: usize,
    reopen: bool,
) -> Result<DistinctMeasurements> {
    let count = 20.min(files / 2);
    let mut queries = Vec::with_capacity(count);
    let mut samples_ms = Vec::with_capacity(count);
    for sample in 0..count {
        // Two disjoint deterministic sets spread across the entire corpus.
        // Neither set repeats the four-case exact filename at files-1.
        let number = sample * (files - 2) / count + usize::from(reopen);
        let query = format!("invoice_{number:07}.txt");
        let start = Instant::now();
        let results = if reopen {
            SearchIndex::open(database)?.search(&query, 20)?
        } else {
            index.search(&query, 20)?
        };
        let elapsed = start.elapsed().as_secs_f64() * 1000.0;
        let first = results
            .first()
            .context("distinct exact query returned no results")?;
        if first.filename != query || first.match_kind != "exact_filename" {
            bail!("distinct query lost exact filename priority");
        }
        queries.push(query);
        samples_ms.push(elapsed);
    }
    let mut sorted = samples_ms.clone();
    sorted.sort_by(f64::total_cmp);
    let percentile = |fraction: f64| {
        sorted[((sorted.len() as f64 * fraction).ceil() as usize).saturating_sub(1)]
    };
    let p95_ms = percentile(0.95);
    eprintln!(
        "{} distinct first filename queries (reopened={reopen}): p95 {p95_ms:.3} ms",
        count
    );
    let first_sample_ms = samples_ms[0];
    let samples_exceeding_target = samples_ms.iter().filter(|sample| **sample >= 50.0).count();
    Ok(DistinctMeasurements {
        description: if reopen {
            "each never-repeated filename on a newly opened reader, persisted statistics, no prewarm; includes open+search"
        } else {
            "distinct foreground filenames; no benchmark query warmups; production background representative queries may overlap"
        },
        includes_reader_open: reopen,
        queries,
        samples_ms,
        p50_ms: percentile(0.5),
        p95_ms,
        first_sample_ms,
        first_sample_target_met: first_sample_ms < 50.0,
        maximum_ms: sorted[sorted.len() - 1],
        samples_exceeding_target,
        target_p95_ms: 50.0,
        target_met: p95_ms < 50.0,
    })
}

fn main() -> Result<()> {
    let arguments = env::args().skip(1).collect::<Vec<_>>();
    let mut files = 10_000usize;
    let mut documents = 1_000usize;
    let mut samples = 20usize;
    let mut database = PathBuf::from("target/benchmarks/synthetic.db");
    let mut reuse = false;
    for pair in arguments.chunks(2) {
        if pair.len() != 2 {
            bail!("usage: performance --files N --documents N --samples N --db NEW_PATH");
        }
        match pair[0].as_str() {
            "--files" => files = pair[1].parse()?,
            "--documents" => documents = pair[1].parse()?,
            "--samples" => samples = pair[1].parse()?,
            "--db" => database = PathBuf::from(&pair[1]),
            "--reuse" => reuse = pair[1].parse()?,
            option => bail!("unknown argument {option}"),
        }
    }
    if !(2..=2_000_000).contains(&files)
        || documents == 0
        || documents > files
        || !(5..=1000).contains(&samples)
    {
        bail!("files must be2..2000000, documents1..files, samples5..1000");
    }
    if database.exists() && !reuse {
        bail!(
            "benchmark database must be a new path: {}",
            database.display()
        );
    }
    if reuse && !database.exists() {
        bail!("reuse requires an existing synthetic fixture");
    }
    if let Some(parent) = database.parent() {
        fs::create_dir_all(parent)?;
    }
    // Seed through the actual versioned schema and its normal indexing triggers.
    if reuse {
        let connection = Connection::open(&database)?;
        let (actual_files, actual_documents, foreign_roots): (usize,usize,usize) = connection.query_row(
            "SELECT count(*),count(NULLIF(content,'')),sum(root!='synthetic://benchmark') FROM files",[],
            |row|Ok((row.get(0)?,row.get(1)?,row.get(2)?)))?;
        if (actual_files, actual_documents, foreign_roots) != (files, documents, 0) {
            bail!("existing fixture does not match configured synthetic corpus");
        }
    }
    let migration_started = Instant::now();
    let index = SearchIndex::open(&database)?;
    drop(index);
    let migration_seconds = migration_started.elapsed().as_secs_f64();
    let start = Instant::now();
    let mut connection = Connection::open(&database)?;
    connection.execute_batch("PRAGMA synchronous=NORMAL;")?;
    if !reuse {
        for batch_start in (0..files).step_by(5000) {
            let transaction = connection.transaction()?;
            {
                let mut statement = transaction.prepare(
                    "INSERT INTO files(path,root,filename,content,size,mtime,extension)
                VALUES (?1,'synthetic://benchmark',?2,?3,?4,0,'txt')",
                )?;
                let mut passage = transaction.prepare(
                    "INSERT INTO passages(text,file_path,kind,number) VALUES (?1,?2,'paragraph',1)",
                )?;
                for number in batch_start..(batch_start + 5000).min(files) {
                    let filename = format!("invoice_{number:07}.txt");
                    let path = format!(
                        "synthetic://benchmark/folder_{:04}/{filename}",
                        number % 1000
                    );
                    let content = if number < documents {
                        format!("Supplier termination cancellation clause budget analysis marker{number:07}")
                    } else {
                        String::new()
                    };
                    statement.execute(params![path, filename, content, content.len() as i64])?;
                    if number < documents {
                        passage.execute(params![content, path])?;
                    }
                }
            }
            transaction.commit()?;
            if batch_start % 100_000 == 0 {
                eprintln!(
                    "seeded {} / {files} synthetic file records",
                    (batch_start + 5000).min(files)
                );
            }
        }
    }
    connection.execute_batch("PRAGMA wal_checkpoint(TRUNCATE);")?;
    drop(connection);
    let seed_seconds = if reuse {
        0.0
    } else {
        start.elapsed().as_secs_f64()
    };
    let database_bytes = fs::metadata(&database)?.len();
    let mut index = SearchIndex::open(&database)?;
    let connection = Connection::open(&database)?;
    let mut statement=connection.prepare("SELECT filename FROM files WHERE id=(SELECT min(id) FROM files) OR id=(SELECT max(id) FROM files)")?;
    let mut background_representative_queries = statement
        .query_map([], |row| row.get::<_, String>(0))?
        .collect::<rusqlite::Result<Vec<_>>>()?;
    let words = background_representative_queries
        .iter()
        .flat_map(|name| {
            name.split(|character: char| !character.is_alphabetic())
                .filter(|word| word.chars().count() >= 3)
                .take(3)
        })
        .map(str::to_owned)
        .collect::<std::collections::BTreeSet<_>>();
    background_representative_queries.extend(words);
    let text: String = connection.query_row(
        "SELECT substr(content,1,256) FROM files WHERE content!='' ORDER BY id LIMIT 1",
        [],
        |row| row.get(0),
    )?;
    background_representative_queries.extend(
        text.split_whitespace()
            .filter(|term| term.chars().all(char::is_alphanumeric))
            .take(3)
            .map(str::to_owned),
    );
    drop(statement);
    drop(connection);
    let preparation = Instant::now();
    let prepared_frequency_count = index.prewarm()?;
    let background_prepare_seconds = preparation.elapsed().as_secs_f64();
    // Measure user-like first inputs before the old repeated-query cases.
    let first_foreground_distinct_names = distinct_names(&index, &database, files, false)?;
    let reopened_reader_first_names = distinct_names(&index, &database, files, true)?;
    let exact = format!("invoice_{:07}.txt", files - 1);
    let unique_content = format!("marker{:07}", documents - 1);
    let cases = [
        (exact, 50.0),
        ("invoice_".into(), 50.0),
        (unique_content, 150.0),
        ("cancellation".into(), 150.0),
    ];
    let mut queries = Vec::new();
    for (query, target) in cases {
        let first_started = Instant::now();
        index.search(&query, 20)?;
        let first_query_ms = first_started.elapsed().as_secs_f64() * 1000.0;
        for _ in 0..2 {
            index.search(&query, 20)?;
        }
        let mut durations = Vec::with_capacity(samples);
        for _ in 0..samples {
            let started = Instant::now();
            let results = index
                .search(&query, 20)
                .with_context(|| format!("benchmark query {query}"))?;
            if results.is_empty() {
                bail!("synthetic query {query} unexpectedly returned no results");
            }
            std::hint::black_box(results);
            durations.push(started.elapsed().as_secs_f64() * 1000.0);
        }
        durations.sort_by(f64::total_cmp);
        let percentile = |fraction: f64| {
            durations[((durations.len() as f64 * fraction).ceil() as usize).saturating_sub(1)]
        };
        let p95_ms = percentile(0.95);
        eprintln!("query {query:?}: p95 {p95_ms:.3} ms; target < {target:.0} ms");
        queries.push(Measurements {
            query,
            samples,
            p50_ms: percentile(0.5),
            first_query_ms,
            first_query_target_met: first_query_ms < target,
            p95_ms,
            target_p95_ms: target,
            target_met: p95_ms < target,
        });
    }
    // Measure restart preparation after foreground samples, so a second
    // prewarm cannot secretly improve the first-input measurements above.
    let mut restarted = SearchIndex::open(&database)?;
    let restart_start = Instant::now();
    restarted.prewarm()?;
    let reopened_background_prepare_seconds = restart_start.elapsed().as_secs_f64();
    let report = Report {
        background_prepare_seconds,
        reopened_background_prepare_seconds,
        background_representative_queries,
        prepared_frequency_count,
        first_foreground_distinct_names,
        reopened_reader_first_names,
        fixture:
            "synthetic SQLite metadata and text; no filesystem extraction or semantic inference",
        files,
        documents,
        seed_seconds,
        migration_seconds,
        fixture_reused: reuse,
        database_bytes,
        database_path: database.display().to_string(),
        architecture: env::consts::ARCH,
        os: env::consts::OS,
        cpu_only: true,
        queries,
    };
    println!("{}", serde_json::to_string_pretty(&report)?);
    Ok(())
}