searchi/crates/searchi-service/tests/budget_google_cpu.rs

//! Actual Google CPU inference during background indexing; not a scale benchmark.
use anyhow::Result;
use searchi_core::SearchIndex;
use searchi_semantic::{GemmaEmbedder, TextWindow};
use searchi_service::{
    EmbeddingProvider, HybridSearch, InferenceBudget, InferenceRole, PreparedGraphs, SearchOptions,
};
use std::{
    fs,
    sync::{
        atomic::{AtomicBool, AtomicUsize, Ordering},
        mpsc, Arc,
    },
    time::Duration,
};

struct ObservedGoogle {
    inner: GemmaEmbedder,
    background: bool,
    in_flight: Arc<AtomicBool>,
    overlap: Arc<AtomicBool>,
    entered: mpsc::Sender<()>,
    calls: Arc<AtomicUsize>,
}
impl EmbeddingProvider for ObservedGoogle {
    fn fingerprint(&self) -> &str {
        EmbeddingProvider::fingerprint(&self.inner)
    }
    fn dimensions(&self) -> usize {
        EmbeddingProvider::dimensions(&self.inner)
    }
    fn split_windows(&mut self, text: &str, title: &str) -> Result<Vec<TextWindow>> {
        self.inner.split_windows(text, title)
    }
    fn embed_documents(&mut self, text: &[&str], title: &str) -> Result<Vec<Vec<f32>>> {
        assert!(self.background);
        self.calls.fetch_add(1, Ordering::Relaxed);
        self.in_flight.store(true, Ordering::Release);
        let _ = self.entered.send(());
        let result = self.inner.embed_documents(text, title);
        self.in_flight.store(false, Ordering::Release);
        result
    }
    fn embed_query(&mut self, text: &str) -> Result<Vec<f32>> {
        if self.in_flight.load(Ordering::Acquire) {
            self.overlap.store(true, Ordering::Release);
        }
        self.inner.embed_query(text)
    }
    fn embed_query_debug(&mut self, text: &str) -> Result<Vec<f32>> {
        if self.in_flight.load(Ordering::Acquire) {
            self.overlap.store(true, Ordering::Release);
        }
        self.inner.embed_query_debug(text)
    }
    fn query_diagnostics(&self) -> std::collections::BTreeMap<String, u64> {
        self.inner.query_diagnostics()
    }
    fn worker_info(&self) -> Option<searchi_semantic::WorkerInfo> {
        Some(self.inner.info.clone())
    }
}
#[test]
#[ignore = "Requires verified real Google worker, two CPU cores; CPU only"]
fn foreground_google_query_runs_while_background_encoder_is_admitted() -> Result<()> {
    assert_eq!(std::env::var("SEARCHI_FORCE_CPU")?, "1");
    let model = std::env::var_os("SEARCHI_GEMMA_MODEL_DIR").expect("Verified model required");
    let worker = std::env::var_os("SEARCHI_GEMMA_WORKER").expect("Normal worker required");
    let budget = InferenceBudget::new(2);
    assert_eq!(
        budget.threads(),
        2,
        "This concurrent test requires two allocated CPU cores"
    );
    let temp = tempfile::tempdir()?;
    let root = temp.path().join("sources");
    fs::create_dir(&root)?;
    fs::write(
        root.join("garden.txt"),
        "Tomato plants need regular watering and plenty of sunlight.",
    )?;
    fs::write(
        root.join("supplier.txt"),
        "The supplier contract ends with sixty days written notice.",
    )?;
    let db = temp.path().join("index.db");
    let mut lexical = SearchIndex::open(&db)?;
    lexical.index_folder(&root)?;
    let graphs = PreparedGraphs::default();
    let mut initial = HybridSearch::open(&db, &model, &worker, 1)?;
    initial.share_prepared_graphs(&graphs)?;
    initial.index_embeddings()?;
    drop(initial);
    let in_flight = Arc::new(AtomicBool::new(false));
    let overlap = Arc::new(AtomicBool::new(false));
    let calls = Arc::new(AtomicUsize::new(0));
    let (entered, events) = mpsc::channel();
    let make = |background| -> Result<HybridSearch> {
        let role = if background {
            InferenceRole::Background
        } else {
            InferenceRole::Foreground
        };
        let threads = budget.worker_threads(role);
        let startup = budget.acquire(role, threads)?;
        let inner = GemmaEmbedder::load(&model, &worker, threads)?;
        assert_eq!(inner.info.device, "cpu");
        drop(startup);
        let mut service = HybridSearch::with_provider(
            &db,
            Box::new(ObservedGoogle {
                inner,
                background,
                in_flight: in_flight.clone(),
                overlap: overlap.clone(),
                entered: entered.clone(),
                calls: calls.clone(),
            }),
        )?
        .with_budget(&budget, role, threads);
        service.share_prepared_graphs(&graphs)?;
        Ok(service)
    };
    let mut reader = make(false)?;
    reader.warmup()?;
    let mut indexer = make(true)?;
    fs::write(
        root.join("maintenance.txt"),
        "Technical maintenance notes for equipment inventory scheduling and routine inspections. "
            .repeat(800),
    )?;
    lexical.index_folder(&root)?;
    let (report, sample) = std::thread::scope(|scope| -> Result<_> {
        let indexing = scope.spawn(move || indexer.index_embeddings());
        let mut sample = None;
        for query in [
            "How should we water tomato plants?",
            "What do tomato plants need for growth?",
            "How can we keep tomatoes healthy?",
        ] {
            events.recv_timeout(Duration::from_secs(90))?;
            let pending = budget.pending_foreground();
            let response = reader.search_measured(query, 10, &SearchOptions::default(), true)?;
            drop(pending);
            assert_eq!(response.results[0].filename, "garden.txt");
            let trace = response.metrics.trace.as_ref().expect("Detailed trace");
            assert!(trace
                .iter()
                .any(|span| span.counters.get("query_cache_hit") == Some(&0)));
            let wait = trace
                .iter()
                .find(|span| span.name == "semantic.inference.admission_wait")
                .expect("Actual admission span")
                .duration_ms;
            sample = Some((response.metrics.total_ms, wait));
            if overlap.load(Ordering::Acquire) {
                break;
            }
        }
        let report = indexing.join().expect("Index thread panic")?;
        Ok((report, sample))
    })?;
    assert!(
        overlap.load(Ordering::Acquire),
        "Must observe genuine query inference while a background call remains active"
    );
    assert_eq!(report.reused_sources, 2);
    assert_eq!(report.encoded_sources, 1);
    assert_eq!(
        report.document_inference_calls,
        calls.load(Ordering::Relaxed)
    );
    assert!(report.document_inference_calls > 1);
    let (query_ms, admission_ms) = sample.unwrap();
    eprintln!("Concurrent Google CPU fixture: total_budget=2 query_threads=1 indexing_threads=1 query_ms={query_ms:.3} admission_ms={admission_ms:.3} document_calls={}; not scale/p95/UI evidence", report.document_inference_calls);
    Ok(())
}