searchi/crates/searchi-service/src/prepared.rs

//! Process-local handoff of immutable graphs prepared by background indexing.
use anyhow::{ensure, Result};
use std::{
    collections::VecDeque,
    path::{Path, PathBuf},
    sync::{Arc, Mutex},
};
use usearch::Index;

/// Share between services for the same database. No model or query data is shared.
#[derive(Clone, Default)]
pub struct PreparedGraphs(Arc<Mutex<State>>);
#[derive(Default)]
struct State {
    identity: Option<(PathBuf, usize)>,
    entries: VecDeque<Entry>,
}
#[derive(Clone)]
struct Entry {
    generation: i64,
    fingerprint: String,
    checksum: String,
    graph: Arc<Index>,
}
impl PreparedGraphs {
    pub(crate) fn bind(&self, database: &Path, dimensions: usize) -> Result<()> {
        let identity = (database.canonicalize()?, dimensions);
        let mut state = self
            .0
            .lock()
            .map_err(|_| anyhow::anyhow!("Prepared graph registry unavailable"))?;
        if let Some(existing) = &state.identity {
            ensure!(
                *existing == identity,
                "Prepared graphs belong to another database or dimension"
            );
        } else {
            state.identity = Some(identity);
        }
        Ok(())
    }
    pub(crate) fn get(
        &self,
        generation: i64,
        fingerprint: &str,
        checksum: &str,
    ) -> Result<Option<Arc<Index>>> {
        let state = self
            .0
            .lock()
            .map_err(|_| anyhow::anyhow!("Prepared graph registry unavailable"))?;
        Ok(state
            .entries
            .iter()
            .find(|entry| {
                entry.generation == generation
                    && entry.fingerprint == fingerprint
                    && entry.checksum == checksum
            })
            .map(|entry| entry.graph.clone()))
    }
    pub(crate) fn publish(
        &self,
        generation: i64,
        fingerprint: String,
        checksum: String,
        graph: Arc<Index>,
    ) -> Result<Publication> {
        let mut state = self
            .0
            .lock()
            .map_err(|_| anyhow::anyhow!("Prepared graph registry unavailable"))?;
        state.entries.push_back(Entry {
            generation,
            fingerprint,
            checksum,
            graph: graph.clone(),
        });
        while state.entries.len() > 2 {
            state.entries.pop_front();
        }
        Ok(Publication {
            registry: self.clone(),
            graph,
            committed: false,
        })
    }
}
/// Failed SQL publication removes only this publisher's entry.
pub(crate) struct Publication {
    registry: PreparedGraphs,
    graph: Arc<Index>,
    committed: bool,
}
impl Publication {
    pub(crate) fn commit(mut self) {
        self.committed = true;
    }
}
impl Drop for Publication {
    fn drop(&mut self) {
        if !self.committed {
            if let Ok(mut state) = self.registry.0.lock() {
                state
                    .entries
                    .retain(|entry| !Arc::ptr_eq(&entry.graph, &self.graph));
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    fn graph() -> Arc<Index> {
        Arc::new(
            Index::new(&usearch::IndexOptions {
                dimensions: 768,
                ..Default::default()
            })
            .unwrap(),
        )
    }
    #[test]
    fn identity_checksum_and_rollback_are_enforced() -> Result<()> {
        let directory = tempfile::tempdir()?;
        let database = directory.path().join("index.db");
        std::fs::write(&database, [])?;
        let other = directory.path().join("other.db");
        std::fs::write(&other, [])?;
        let registry = PreparedGraphs::default();
        registry.bind(&database, 768)?;
        assert!(registry.bind(&other, 768).is_err());
        assert!(registry.bind(&database, 256).is_err());
        let active = graph();
        registry
            .publish(1, "policy".into(), "one".into(), active.clone())?
            .commit();
        assert!(registry.get(1, "policy", "wrong")?.is_none());
        assert!(registry.get(1, "wrong", "one")?.is_none());
        {
            let _pending = registry.publish(2, "policy".into(), "two".into(), graph())?;
            assert!(registry.get(2, "policy", "two")?.is_some());
        }
        assert!(registry.get(2, "policy", "two")?.is_none());
        assert!(Arc::ptr_eq(
            &active,
            &registry.get(1, "policy", "one")?.unwrap()
        ));
        registry
            .publish(2, "policy".into(), "two".into(), graph())?
            .commit();
        registry
            .publish(3, "policy".into(), "three".into(), graph())?
            .commit();
        assert_eq!(registry.0.lock().unwrap().entries.len(), 2);
        assert!(registry.get(1, "policy", "one")?.is_none());
        assert_eq!(active.dimensions(), 768); // Eviction never invalidates a retained reader.
        Ok(())
    }
}