laurus 0.10.0

//! Integration tests for crash-safe atomic writes of the production HNSW
//! index files (issue #784; segmented layout since #882).
//!
//! Every on-disk artifact — per-segment `.hnsw` files (written by
//! `HnswIndexWriter::write`), the `segments.json` manifest (#879), and the
//! deletion bitmap — is written to a `.tmp` file and atomically
//! `rename_file`d into place. A crash between writing a temp file and the
//! rename therefore leaves the previously committed state intact (the
//! orphaned temp is ignored and swept on reopen), and a successful commit
//! leaves no temp file behind.

use async_trait::async_trait;
use std::any::Any;
use std::collections::HashMap;
use std::collections::HashSet;
use std::io::Write;
use std::sync::Arc;

use laurus::lexical::LexicalIndexConfig;
use laurus::storage::Storage;
use laurus::storage::memory::{MemoryStorage, MemoryStorageConfig};
use laurus::vector::Vector;
use laurus::vector::core::distance::DistanceMetric;
use laurus::vector::core::field::HnswOption;
use laurus::vector::store::config::VectorFieldConfig;
use laurus::vector::store::request::{
    QueryVector, VectorScoreMode, VectorSearchParams, VectorSearchRequest,
};
use laurus::vector::{FieldOption, VectorIndexConfig, VectorSearchQuery};
use laurus::{DataValue, Document};
use laurus::{EmbedInput, EmbedInputType, Embedder};
use laurus::{LaurusError, Result};

const DIM: usize = 16;
const N: u64 = 50;
const STEP: f32 = 0.01;
/// First sealed segment of the (default, #882) segmented layout, under the
/// "vec" field's own sub-namespace (Issue #948:
/// `MultiFieldVectorIndex` gives every vector field its own
/// `PrefixedStorage` directory rather than sharing the storage root).
const HNSW_FILE: &str = "vec/segment_000000.hnsw";
/// An orphaned staging file from a simulated crashed segment write.
const HNSW_TMP: &str = "vec/segment_000001.hnsw.tmp";

#[derive(Debug)]
struct MockEmbedder {
    dimension: usize,
}

#[async_trait]
impl Embedder for MockEmbedder {
    async fn embed(&self, input: &EmbedInput<'_>) -> Result<Vector> {
        match input {
            EmbedInput::Text(_) => Ok(Vector::new(vec![0.0; self.dimension])),
            _ => Err(LaurusError::invalid_argument("text only")),
        }
    }
    fn supported_input_types(&self) -> Vec<EmbedInputType> {
        vec![EmbedInputType::Text]
    }
    fn name(&self) -> &str {
        "mock"
    }
    fn as_any(&self) -> &dyn Any {
        self
    }
}

fn doc_vec(i: u64) -> Vec<f32> {
    let theta = i as f32 * STEP;
    let mut v = vec![0.0; DIM];
    v[0] = theta.cos();
    v[1] = theta.sin();
    v
}

fn query_vec() -> Vec<f32> {
    let mut v = vec![0.0; DIM];
    v[0] = 1.0;
    v
}

fn hnsw() -> FieldOption {
    FieldOption::Hnsw(HnswOption {
        dimension: DIM,
        distance: DistanceMetric::Cosine,
        m: 16,
        ef_construction: 100,
        default_ef_search: None,
        base_weight: 1.0,
        quantizer: Default::default(),
        rerank_storage: None,
        embedder: None,
        pq_codebook_path: None,
    })
}

fn make_config() -> VectorIndexConfig {
    let mut field_configs = HashMap::new();
    field_configs.insert(
        "vec".to_string(),
        VectorFieldConfig {
            vector: Some(hnsw()),
            lexical: None,
        },
    );
    VectorIndexConfig {
        fields: field_configs,
        embedder: Arc::new(MockEmbedder { dimension: DIM }),
        default_fields: vec!["vec".to_string()],
        metadata: HashMap::new(),
        deletion_config: laurus::DeletionConfig {
            auto_compaction: false,
            ..Default::default()
        },
        shard_id: 0,
        metadata_config: LexicalIndexConfig::default(),
    }
}

async fn build_committed(storage: Arc<dyn Storage>) -> laurus::vector::VectorStore {
    let store = laurus::vector::VectorStore::new(storage, make_config()).unwrap();
    for id in 0..N {
        let doc = Document::builder()
            .add_field("vec", DataValue::Vector(doc_vec(id)))
            .build();
        store.upsert_document_by_internal_id(id, doc).await.unwrap();
    }
    store.commit().await.unwrap();
    store
}

fn request() -> VectorSearchRequest {
    VectorSearchRequest {
        query: VectorSearchQuery::Vectors(vec![QueryVector {
            vector: Vector::new(query_vec()),
            weight: 1.0,
            fields: Some(vec!["vec".into()]),
        }]),
        params: VectorSearchParams {
            limit: 10,
            score_mode: VectorScoreMode::WeightedSum,
            fields: None,
            allowed_ids: None,
            ..Default::default()
        },
    }
}

fn hit_ids(results: &laurus::vector::VectorSearchResults) -> HashSet<u64> {
    results.hits.iter().map(|h| h.doc_id).collect()
}

#[tokio::test(flavor = "multi_thread")]
async fn commit_leaves_no_temp_file() {
    let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new(MemoryStorageConfig::default()));
    let store = build_committed(storage.clone()).await;

    // After a successful commit the sealed segment and the manifest are in
    // place, and every temp file has been renamed away.
    assert!(storage.file_exists(HNSW_FILE), "sealed segment must exist");
    assert!(
        storage.file_exists("vec/segments.json"),
        "manifest must exist"
    );
    let temps: Vec<String> = storage
        .list_files()
        .unwrap()
        .into_iter()
        .filter(|f| f.ends_with(".tmp"))
        .collect();
    assert!(
        temps.is_empty(),
        "a successful commit must not leave any .tmp behind, got {temps:?}"
    );
    assert_eq!(hit_ids(&store.search(request()).unwrap()).len(), 10);
}

#[tokio::test(flavor = "multi_thread")]
async fn orphaned_temp_from_crashed_write_is_ignored() {
    let storage: Arc<dyn Storage> = Arc::new(MemoryStorage::new(MemoryStorageConfig::default()));
    let store = build_committed(storage.clone()).await;
    let before = hit_ids(&store.search(request()).unwrap());
    assert_eq!(before.len(), 10);
    drop(store);

    // Simulate a crash *during* a later segment write: the temp file was
    // created but the atomic rename never happened. The committed segment
    // must be untouched.
    {
        let mut out = storage.create_output(HNSW_TMP).unwrap();
        out.write_all(b"partially-written-garbage-from-a-crashed-commit")
            .unwrap();
        out.close().unwrap();
    }
    assert!(
        storage.file_exists(HNSW_FILE),
        "committed segment still present"
    );

    // Reopening reads the valid committed segment, ignores the orphaned
    // temp (the #879 sweep removes it), and returns the same results.
    let reopened = laurus::vector::VectorStore::new(storage.clone(), make_config()).unwrap();
    let after = hit_ids(&reopened.search(request()).unwrap());
    assert_eq!(
        after, before,
        "the committed index must survive an orphaned temp file from a crashed write"
    );
    assert!(
        !storage.file_exists(HNSW_TMP),
        "the orphaned staging file must be swept on reopen (#879)"
    );
}