Skip to main content

hermes_core/index/
metadata.rs

1//! Unified index metadata - segments list + vector index state
2//!
3//! This module manages all index-level metadata in a single `metadata.json` file:
4//! - List of committed segments
5//! - Vector index state per field (Flat/Built)
6//! - Trained centroids/codebooks paths
7//!
8//! The workflow is:
9//! 1. During accumulation: segments store Flat vectors, state is Flat
10//! 2. When threshold crossed: train ONCE, update state to Built
11//! 3. On index open: load metadata, skip re-training if already built
12
13use serde::{Deserialize, Serialize};
14use std::collections::HashMap;
15use std::io::Write;
16use std::path::Path;
17
18use crate::dsl::{Schema, VectorIndexType};
19use crate::error::{Error, Result};
20
21/// Metadata file name at index level
22pub const INDEX_META_FILENAME: &str = "metadata.json";
23/// Temp file for atomic writes (write here, then rename to INDEX_META_FILENAME)
24const INDEX_META_TMP_FILENAME: &str = "metadata.json.tmp";
25
26/// Current metadata.json format version written by this build.
27///
28/// `load` refuses metadata stamped with a newer version: serde_json silently
29/// drops fields it does not know about, so loading newer metadata would
30/// misread index state and the next save would destructively rewrite the
31/// unknown fields away.
32pub const INDEX_META_FORMAT_VERSION: u32 = 1;
33
34/// Index-level centroids/codebooks are deliberately bounded before they are
35/// read or decoded. Besides limiting ordinary corruption damage, the matching
36/// bincode limit prevents a tiny forged collection length from requesting an
37/// effectively unbounded allocation.
38pub(crate) const MAX_TRAINED_ARTIFACT_BYTES: usize = 512 * 1024 * 1024;
39
40/// State of vector index for a field
41#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
42pub enum VectorIndexState {
43    /// Accumulating vectors - using Flat (brute-force) search
44    #[default]
45    Flat,
46    /// Index structures built - using ANN search
47    Built {
48        /// Total vector count when training happened
49        vector_count: usize,
50        /// Number of clusters used
51        num_clusters: usize,
52    },
53}
54
55fn default_true() -> bool {
56    true
57}
58
59/// Per-segment metadata stored in index metadata
60/// This allows merge decisions without loading segment files
61#[derive(Debug, Clone, Serialize, Deserialize)]
62pub struct SegmentMetaInfo {
63    /// Number of documents in this segment
64    pub num_docs: u32,
65    /// Parent segment IDs that were merged to produce this segment (empty for fresh segments)
66    pub ancestors: Vec<String>,
67    /// Merge generation: 0 for fresh segments, max(parent generations) + 1 for merged segments
68    pub generation: u32,
69    /// Whether this segment has been reordered via Recursive Graph Bisection (BP).
70    /// Fresh segments and block-copy merges are not reordered. Only segments that have
71    /// been explicitly reordered (via background optimizer or reorder command) are marked true.
72    #[serde(default)]
73    pub reordered: bool,
74    /// Whether the last BP reorder pass ran to natural convergence. False when
75    /// a wall-clock BP budget ended the pass early — the segment is ordered
76    /// better than before, and a later warm-started pass can deepen it.
77    /// Old metadata (field absent) deserializes as converged.
78    #[serde(default = "default_true")]
79    pub bp_converged: bool,
80    /// Number of consecutive budget-exhausted BP rewrites in this segment's
81    /// current reordered lineage. Carried across replacement IDs so the
82    /// optimizer can impose a hard follow-up bound instead of rewriting forever.
83    #[serde(default)]
84    pub bp_unconverged_passes: u32,
85}
86
87/// Per-field vector index metadata
88#[derive(Debug, Clone, Serialize, Deserialize)]
89pub struct FieldVectorMeta {
90    /// Field ID
91    pub field_id: u32,
92    /// Configured index type (target type when built)
93    pub index_type: VectorIndexType,
94    /// Current state
95    pub state: VectorIndexState,
96    /// Path to centroids file (relative to index dir)
97    #[serde(skip_serializing_if = "Option::is_none")]
98    pub centroids_file: Option<String>,
99    /// Path to codebook file (relative to index dir, for ScaNN)
100    #[serde(skip_serializing_if = "Option::is_none")]
101    pub codebook_file: Option<String>,
102}
103
104/// Unified index metadata - single source of truth for index state
105#[derive(Debug, Clone, Serialize, Deserialize)]
106pub struct IndexMetadata {
107    /// Version for compatibility
108    pub version: u32,
109    /// Index schema
110    pub schema: Schema,
111    /// Segment metadata: segment_id -> info (doc count, etc.)
112    /// Using HashMap allows O(1) lookup and stores doc counts for merge decisions
113    #[serde(default)]
114    pub segment_metas: HashMap<String, SegmentMetaInfo>,
115    /// Per-field vector index metadata
116    #[serde(default)]
117    pub vector_fields: HashMap<u32, FieldVectorMeta>,
118    /// Aggregate vector count recorded by all built vector fields.
119    ///
120    /// The per-field `VectorIndexState::Built::vector_count` values are the
121    /// source of truth. This cached aggregate is refreshed whenever a field is
122    /// marked built, rather than being overwritten with whichever field was
123    /// trained last.
124    #[serde(default)]
125    pub total_vectors: usize,
126}
127
128impl IndexMetadata {
129    /// Create new metadata with schema
130    pub fn new(schema: Schema) -> Self {
131        Self {
132            version: INDEX_META_FORMAT_VERSION,
133            schema,
134            segment_metas: HashMap::new(),
135            vector_fields: HashMap::new(),
136            total_vectors: 0,
137        }
138    }
139
140    /// Get segment IDs as a sorted Vec (deterministic ordering)
141    pub fn segment_ids(&self) -> Vec<String> {
142        let mut ids: Vec<String> = self.segment_metas.keys().cloned().collect();
143        ids.sort();
144        ids
145    }
146
147    /// Add a fresh segment (gen=0, no ancestors, not reordered)
148    pub fn add_segment(&mut self, segment_id: String, num_docs: u32) {
149        self.segment_metas.insert(
150            segment_id,
151            SegmentMetaInfo {
152                num_docs,
153                ancestors: Vec::new(),
154                generation: 0,
155                reordered: false,
156                bp_converged: true,
157                bp_unconverged_passes: 0,
158            },
159        );
160    }
161
162    /// Add a merged segment with lineage info
163    pub fn add_merged_segment(
164        &mut self,
165        segment_id: String,
166        num_docs: u32,
167        ancestors: Vec<String>,
168        generation: u32,
169        reordered: bool,
170        bp_converged: bool,
171    ) {
172        self.add_segment_meta(
173            segment_id,
174            SegmentMetaInfo {
175                num_docs,
176                ancestors,
177                generation,
178                reordered,
179                bp_converged,
180                bp_unconverged_passes: 0,
181            },
182        );
183    }
184
185    /// Insert fully constructed lifecycle metadata. Merge/reorder code uses
186    /// this to carry bounded BP lineage; ordinary callers use the safer
187    /// constructors above, which start a fresh lineage.
188    pub(crate) fn add_segment_meta(&mut self, segment_id: String, info: SegmentMetaInfo) {
189        self.segment_metas.insert(segment_id, info);
190    }
191
192    /// Remove a segment
193    pub fn remove_segment(&mut self, segment_id: &str) {
194        self.segment_metas.remove(segment_id);
195    }
196
197    /// Check if segment exists
198    pub fn has_segment(&self, segment_id: &str) -> bool {
199        self.segment_metas.contains_key(segment_id)
200    }
201
202    /// Get segment doc count
203    pub fn segment_doc_count(&self, segment_id: &str) -> Option<u32> {
204        self.segment_metas.get(segment_id).map(|m| m.num_docs)
205    }
206
207    /// Check if a field has been built
208    pub fn is_field_built(&self, field_id: u32) -> bool {
209        self.vector_fields
210            .get(&field_id)
211            .map(|f| matches!(f.state, VectorIndexState::Built { .. }))
212            .unwrap_or(false)
213    }
214
215    /// Get field metadata
216    pub fn get_field_meta(&self, field_id: u32) -> Option<&FieldVectorMeta> {
217        self.vector_fields.get(&field_id)
218    }
219
220    /// Initialize field metadata (called when field is first seen)
221    pub fn init_field(&mut self, field_id: u32, index_type: VectorIndexType) {
222        self.vector_fields
223            .entry(field_id)
224            .or_insert(FieldVectorMeta {
225                field_id,
226                index_type,
227                state: VectorIndexState::Flat,
228                centroids_file: None,
229                codebook_file: None,
230            });
231    }
232
233    /// Mark field as built with trained structures
234    pub fn mark_field_built(
235        &mut self,
236        field_id: u32,
237        vector_count: usize,
238        num_clusters: usize,
239        centroids_file: String,
240        codebook_file: Option<String>,
241    ) {
242        if let Some(field) = self.vector_fields.get_mut(&field_id) {
243            field.state = VectorIndexState::Built {
244                vector_count,
245                num_clusters,
246            };
247            field.centroids_file = Some(centroids_file);
248            field.codebook_file = codebook_file;
249            self.refresh_total_vectors();
250        }
251    }
252
253    /// Refresh the cached aggregate from the authoritative per-field states.
254    ///
255    /// Saturation keeps this infallible metadata helper safe even if it is
256    /// called after loading externally modified metadata with impossible
257    /// counts.
258    pub(crate) fn refresh_total_vectors(&mut self) {
259        self.total_vectors = self
260            .vector_fields
261            .values()
262            .filter_map(|field| match field.state {
263                VectorIndexState::Built { vector_count, .. } => Some(vector_count),
264                VectorIndexState::Flat => None,
265            })
266            .fold(0usize, usize::saturating_add);
267    }
268
269    /// Check if field should be built based on threshold
270    pub fn should_build_field(&self, field_id: u32, threshold: usize) -> bool {
271        // Don't build if already built
272        if self.is_field_built(field_id) {
273            return false;
274        }
275        // Build if we have enough vectors
276        self.total_vectors >= threshold
277    }
278
279    /// Load from directory
280    ///
281    /// If `metadata.json` is missing but `metadata.json.tmp` exists (crash
282    /// between write and rename), recovers from the temp file.
283    pub async fn load<D: crate::directories::Directory>(dir: &D) -> Result<Self> {
284        let path = Path::new(INDEX_META_FILENAME);
285        match dir.open_read(path).await {
286            Ok(slice) => {
287                let bytes = slice.read_bytes().await?;
288                Self::deserialize_versioned(bytes.as_slice())
289            }
290            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
291                // Try recovering from temp file (crash between write and rename)
292                let tmp_path = Path::new(INDEX_META_TMP_FILENAME);
293                let slice = dir.open_read(tmp_path).await?;
294                let bytes = slice.read_bytes().await?;
295                let meta = Self::deserialize_versioned(bytes.as_slice())?;
296                log::warn!("Recovered metadata from temp file (previous crash during save)");
297                Ok(meta)
298            }
299            Err(e) => Err(Error::Io(e)),
300        }
301    }
302
303    /// Deserialize metadata bytes, refusing formats newer than this build.
304    ///
305    /// serde_json silently drops unknown fields, so loading newer metadata
306    /// would misread index state and the next save would destructively
307    /// rewrite the newer fields away. Fail loud instead.
308    fn deserialize_versioned(bytes: &[u8]) -> Result<Self> {
309        let meta: Self =
310            serde_json::from_slice(bytes).map_err(|e| Error::Serialization(e.to_string()))?;
311        if meta.version > INDEX_META_FORMAT_VERSION {
312            return Err(Error::Corruption(format!(
313                "metadata.json format version {} is newer than this build supports (max {}). \
314                 Refusing to load: unknown fields would be silently dropped and destructively \
315                 rewritten on the next save. Open this index with a hermes build that \
316                 understands metadata format version {}",
317                meta.version, INDEX_META_FORMAT_VERSION, meta.version
318            )));
319        }
320        Ok(meta)
321    }
322
323    /// Save to directory (atomic: write temp file, then rename)
324    ///
325    /// Uses write-then-rename so a crash mid-write won't corrupt the
326    /// existing metadata file. On POSIX, rename is atomic.
327    pub async fn save<D: crate::directories::DirectoryWriter>(&self, dir: &D) -> Result<()> {
328        let bytes = self.serialize_to_bytes()?;
329        Self::save_bytes(dir, &bytes).await
330    }
331
332    /// Serialize metadata to bytes (cheap, no I/O).
333    /// Useful when you need to release a lock before doing disk I/O.
334    pub fn serialize_to_bytes(&self) -> Result<Vec<u8>> {
335        serde_json::to_vec_pretty(self).map_err(|e| Error::Serialization(e.to_string()))
336    }
337
338    /// Write pre-serialized metadata bytes to directory (atomic rename + fsync).
339    ///
340    /// The fsync ensures durability: without it, a power failure after rename
341    /// could lose the metadata update on systems with volatile write caches.
342    pub async fn save_bytes<D: crate::directories::DirectoryWriter>(
343        dir: &D,
344        bytes: &[u8],
345    ) -> Result<()> {
346        let tmp_path = Path::new(INDEX_META_TMP_FILENAME);
347        let final_path = Path::new(INDEX_META_FILENAME);
348        // Metadata is tiny, but `DirectoryWriter::write` does not guarantee
349        // the file contents themselves are fsynced. Finish the streaming
350        // writer first (filesystem implementations call `File::sync_all`),
351        // then atomically publish the durable temp file by rename.
352        let mut writer = dir.streaming_writer(tmp_path).await.map_err(Error::Io)?;
353        writer.write_all(bytes).map_err(Error::Io)?;
354        writer.finish().map_err(Error::Io)?;
355        // Rename is the logical commit point: after it succeeds, readers can
356        // observe the new generation and callers must publish the matching
357        // in-memory/tracker state. Directory fsync only strengthens crash
358        // durability. It cannot safely turn an already-visible rename into a
359        // reported pre-commit failure, because cleanup could then delete files
360        // referenced by the metadata now on disk.
361        dir.rename(tmp_path, final_path).await.map_err(Error::Io)?;
362        if let Err(error) = dir.sync().await {
363            log::error!(
364                "[metadata] directory fsync failed after committed rename: {}. \
365                 Continuing with the renamed generation; crash durability is not guaranteed",
366                error,
367            );
368        }
369        Ok(())
370    }
371
372    /// Compatibility loader for callers that only have the persisted field
373    /// map. Invalid/incomplete state is logged and returns `None` rather than a
374    /// partial set.
375    ///
376    /// Index open/build paths use the fallible, schema-aware
377    /// [`Self::try_load_trained_from_fields`] method below.
378    pub async fn load_trained_from_fields<D: crate::directories::Directory>(
379        vector_fields: &HashMap<u32, FieldVectorMeta>,
380        dir: &D,
381    ) -> Option<crate::segment::TrainedVectorStructures> {
382        match Self::load_trained_from_fields_impl(vector_fields, None, dir).await {
383            Ok(trained) => trained,
384            Err(error) => {
385                log::error!("[trained] refusing incomplete/corrupt artifact set: {error}");
386                None
387            }
388        }
389    }
390
391    /// Fallible schema-aware loader used for lifecycle publication.
392    #[cfg_attr(not(feature = "native"), allow(dead_code))]
393    pub(crate) async fn try_load_trained_from_fields<D: crate::directories::Directory>(
394        vector_fields: &HashMap<u32, FieldVectorMeta>,
395        schema: &Schema,
396        dir: &D,
397    ) -> Result<Option<crate::segment::TrainedVectorStructures>> {
398        Self::load_trained_from_fields_impl(vector_fields, Some(schema), dir).await
399    }
400
401    /// Load and validate the complete trained-artifact set described by a
402    /// `vector_fields` snapshot.
403    ///
404    /// This is intentionally all-or-nothing. A `Built` field is a durable
405    /// promise that every artifact required by its configured index exists and
406    /// is compatible with the schema. Returning a partial map would let some
407    /// segment builders publish ANN data while another field was silently
408    /// unusable, and would make the same index behave differently after a
409    /// restart.
410    async fn load_trained_from_fields_impl<D: crate::directories::Directory>(
411        vector_fields: &HashMap<u32, FieldVectorMeta>,
412        schema: Option<&Schema>,
413        dir: &D,
414    ) -> Result<Option<crate::segment::TrainedVectorStructures>> {
415        use std::sync::Arc;
416
417        let mut centroids = rustc_hash::FxHashMap::default();
418        let mut codebooks = rustc_hash::FxHashMap::default();
419        let mut built_fields: Vec<_> = vector_fields
420            .iter()
421            .filter(|(_, meta)| matches!(meta.state, VectorIndexState::Built { .. }))
422            .collect();
423        built_fields.sort_unstable_by_key(|(field_id, _)| **field_id);
424
425        log::debug!(
426            "[trained] loading trained structures, vector_fields={:?}",
427            vector_fields.keys().collect::<Vec<_>>()
428        );
429
430        for (field_id, field_meta) in built_fields {
431            log::debug!(
432                "[trained] field {} state={:?} centroids_file={:?} codebook_file={:?}",
433                field_id,
434                field_meta.state,
435                field_meta.centroids_file,
436                field_meta.codebook_file,
437            );
438            if field_meta.field_id != *field_id {
439                return Err(Error::Corruption(format!(
440                    "trained vector metadata key {field_id} contains field_id {}",
441                    field_meta.field_id
442                )));
443            }
444
445            let schema_config = match schema {
446                None => None,
447                Some(schema) => {
448                    let entry = schema
449                        .get_field_entry(crate::dsl::Field(*field_id))
450                        .ok_or_else(|| {
451                            Error::Corruption(format!(
452                                "trained vector metadata references missing field {field_id}"
453                            ))
454                        })?;
455                    if entry.field_type != crate::dsl::FieldType::DenseVector {
456                        return Err(Error::Corruption(format!(
457                            "trained vector metadata field {field_id} has non-dense schema type {:?}",
458                            entry.field_type
459                        )));
460                    }
461                    let config = entry.dense_vector_config.as_ref().ok_or_else(|| {
462                        Error::Corruption(format!(
463                            "trained vector metadata field {field_id} has no dense-vector configuration"
464                        ))
465                    })?;
466                    if field_meta.index_type != config.index_type {
467                        return Err(Error::Corruption(format!(
468                            "trained vector metadata field {field_id} uses {:?}, schema requires {:?}",
469                            field_meta.index_type, config.index_type
470                        )));
471                    }
472                    Some(config)
473                }
474            };
475            if !matches!(
476                field_meta.index_type,
477                VectorIndexType::IvfRaBitQ | VectorIndexType::ScaNN
478            ) {
479                return Err(Error::Corruption(format!(
480                    "field {field_id} is Built for {:?}, which has no index-level trained artifacts",
481                    field_meta.index_type
482                )));
483            }
484
485            let expected_clusters = match field_meta.state {
486                VectorIndexState::Built { num_clusters, .. } if num_clusters > 0 => num_clusters,
487                VectorIndexState::Built { .. } => {
488                    return Err(Error::Corruption(format!(
489                        "trained vector metadata field {field_id} has zero clusters"
490                    )));
491                }
492                VectorIndexState::Flat => unreachable!("built_fields contains only Built entries"),
493            };
494
495            let centroids_file = field_meta.centroids_file.as_deref().ok_or_else(|| {
496                Error::Corruption(format!(
497                    "trained vector metadata field {field_id} is Built but has no centroids_file"
498                ))
499            })?;
500            let c: crate::structures::CoarseCentroids =
501                load_trained_artifact(dir, *field_id, "centroids", centroids_file).await?;
502            let expected_dim = schema_config.map_or(c.dim, |config| config.dim);
503            let actual_clusters = c.num_clusters as usize;
504            let expected_values = actual_clusters.checked_mul(expected_dim).ok_or_else(|| {
505                Error::Corruption(format!(
506                    "trained centroid dimensions overflow for field {field_id}"
507                ))
508            })?;
509            if actual_clusters == 0
510                || actual_clusters > expected_clusters
511                || c.dim == 0
512                || c.dim != expected_dim
513                || c.centroids.len() != expected_values
514                || c.centroids.iter().any(|value| !value.is_finite())
515            {
516                return Err(Error::Corruption(format!(
517                    "trained centroids for field {field_id} do not match metadata/schema: \
518                     clusters={} (metadata maximum {expected_clusters}), dim={} (expected {}), \
519                     values={} (expected {expected_values})",
520                    c.num_clusters,
521                    c.dim,
522                    expected_dim,
523                    c.centroids.len(),
524                )));
525            }
526            if actual_clusters < expected_clusters {
527                // Older writers persisted the requested cluster count even
528                // though the trainer clamps it to the available sample. This
529                // shape is safe and self-describing in the artifact; accepting
530                // it keeps pre-fix indexes openable. New writers persist the
531                // actual count, so no new mismatch is produced.
532                log::warn!(
533                    "[trained] field {} legacy cluster-count clamp: metadata={}, artifact={}",
534                    field_id,
535                    expected_clusters,
536                    actual_clusters,
537                );
538            }
539            log::debug!(
540                "[trained] field {} loaded centroids ({} clusters)",
541                field_id,
542                c.num_clusters
543            );
544
545            if field_meta.index_type == VectorIndexType::ScaNN {
546                let codebook_file = field_meta.codebook_file.as_deref().ok_or_else(|| {
547                    Error::Corruption(format!(
548                        "trained vector metadata field {field_id} is ScaNN Built but has no codebook_file"
549                    ))
550                })?;
551                let codebook: crate::structures::PQCodebook =
552                    load_trained_artifact(dir, *field_id, "codebook", codebook_file).await?;
553                codebook.validate().map_err(|error| {
554                    Error::Corruption(format!(
555                        "invalid trained codebook for field {field_id}: {error}"
556                    ))
557                })?;
558                if codebook.config.dim != expected_dim {
559                    return Err(Error::Corruption(format!(
560                        "trained codebook for field {field_id} has dimension {}, expected {}",
561                        codebook.config.dim, expected_dim
562                    )));
563                }
564                log::debug!("[trained] field {} loaded codebook", field_id);
565                codebooks.insert(*field_id, Arc::new(codebook));
566            }
567            centroids.insert(*field_id, Arc::new(c));
568        }
569
570        if centroids.is_empty() {
571            Ok(None)
572        } else {
573            Ok(Some(crate::segment::TrainedVectorStructures {
574                centroids,
575                codebooks,
576            }))
577        }
578    }
579}
580
581fn validate_trained_artifact_path(field_id: u32, kind: &str, filename: &str) -> Result<()> {
582    use std::path::Component;
583
584    let path = Path::new(filename);
585    if filename.is_empty()
586        || path.is_absolute()
587        || path.components().any(|component| {
588            matches!(
589                component,
590                Component::ParentDir | Component::RootDir | Component::Prefix(_)
591            )
592        })
593    {
594        return Err(Error::Corruption(format!(
595            "trained {kind} path for field {field_id} is not a safe relative path: '{filename}'"
596        )));
597    }
598    Ok(())
599}
600
601async fn load_trained_artifact<T, D>(
602    dir: &D,
603    field_id: u32,
604    kind: &str,
605    filename: &str,
606) -> Result<T>
607where
608    T: serde::de::DeserializeOwned,
609    D: crate::directories::Directory,
610{
611    validate_trained_artifact_path(field_id, kind, filename)?;
612    let path = Path::new(filename);
613    let file_size = dir.file_size(path).await.map_err(|error| {
614        Error::Corruption(format!(
615            "failed to stat trained {kind} '{filename}' for field {field_id}: {error}"
616        ))
617    })?;
618    validate_trained_artifact_size(field_id, kind, filename, file_size)?;
619    let slice = dir.open_read(path).await.map_err(|error| {
620        Error::Corruption(format!(
621            "failed to open trained {kind} '{filename}' for field {field_id}: {error}"
622        ))
623    })?;
624    validate_trained_artifact_size(field_id, kind, filename, slice.len())?;
625    let bytes = slice.read_bytes().await.map_err(|error| {
626        Error::Corruption(format!(
627            "failed to read trained {kind} '{filename}' for field {field_id}: {error}"
628        ))
629    })?;
630    let (artifact, consumed) = bincode::serde::decode_from_slice::<T, _>(
631        bytes.as_slice(),
632        bincode::config::standard().with_limit::<MAX_TRAINED_ARTIFACT_BYTES>(),
633    )
634    .map_err(|error| {
635        Error::Corruption(format!(
636            "failed to deserialize trained {kind} '{filename}' for field {field_id}: {error}"
637        ))
638    })?;
639    if consumed != bytes.len() {
640        return Err(Error::Corruption(format!(
641            "trained {kind} '{filename}' for field {field_id} has {} trailing bytes",
642            bytes.len() - consumed
643        )));
644    }
645    Ok(artifact)
646}
647
648fn validate_trained_artifact_size(
649    field_id: u32,
650    kind: &str,
651    filename: &str,
652    file_size: u64,
653) -> Result<()> {
654    if file_size > MAX_TRAINED_ARTIFACT_BYTES as u64 {
655        return Err(Error::Corruption(format!(
656            "trained {kind} '{filename}' for field {field_id} is {file_size} bytes, \
657             exceeding the {MAX_TRAINED_ARTIFACT_BYTES}-byte safety limit"
658        )));
659    }
660    Ok(())
661}
662
663#[cfg(test)]
664mod tests {
665    use super::*;
666    use crate::directories::DirectoryWriter;
667
668    #[derive(Clone, Default)]
669    struct SyncFailDirectory(crate::directories::RamDirectory);
670
671    #[async_trait::async_trait]
672    impl crate::directories::Directory for SyncFailDirectory {
673        async fn exists(&self, path: &Path) -> std::io::Result<bool> {
674            self.0.exists(path).await
675        }
676
677        async fn file_size(&self, path: &Path) -> std::io::Result<u64> {
678            self.0.file_size(path).await
679        }
680
681        async fn open_read(&self, path: &Path) -> std::io::Result<crate::directories::FileHandle> {
682            self.0.open_read(path).await
683        }
684
685        async fn read_range(
686            &self,
687            path: &Path,
688            range: std::ops::Range<u64>,
689        ) -> std::io::Result<crate::directories::OwnedBytes> {
690            self.0.read_range(path, range).await
691        }
692
693        async fn list_files(&self, prefix: &Path) -> std::io::Result<Vec<std::path::PathBuf>> {
694            self.0.list_files(prefix).await
695        }
696
697        async fn open_lazy(&self, path: &Path) -> std::io::Result<crate::directories::FileHandle> {
698            self.0.open_lazy(path).await
699        }
700    }
701
702    #[async_trait::async_trait]
703    impl crate::directories::DirectoryWriter for SyncFailDirectory {
704        async fn write(&self, path: &Path, data: &[u8]) -> std::io::Result<()> {
705            self.0.write(path, data).await
706        }
707
708        async fn delete(&self, path: &Path) -> std::io::Result<()> {
709            self.0.delete(path).await
710        }
711
712        async fn rename(&self, from: &Path, to: &Path) -> std::io::Result<()> {
713            self.0.rename(from, to).await
714        }
715
716        async fn sync(&self) -> std::io::Result<()> {
717            Err(std::io::Error::other("injected directory fsync failure"))
718        }
719
720        async fn streaming_writer(
721            &self,
722            path: &Path,
723        ) -> std::io::Result<Box<dyn crate::directories::StreamingWriter>> {
724            self.0.streaming_writer(path).await
725        }
726    }
727
728    fn test_schema() -> Schema {
729        Schema::default()
730    }
731
732    fn dense_schema(index_type: VectorIndexType) -> (Schema, crate::dsl::Field) {
733        let mut builder = crate::dsl::SchemaBuilder::default();
734        let config = match index_type {
735            VectorIndexType::IvfRaBitQ => crate::dsl::DenseVectorConfig::with_ivf(2, Some(1), 1),
736            VectorIndexType::ScaNN => crate::dsl::DenseVectorConfig::with_scann(2, Some(1), 1),
737            other => panic!("unsupported trained test index type: {other:?}"),
738        };
739        let field = builder.add_dense_vector_field_with_config("embedding", true, true, config);
740        (builder.build(), field)
741    }
742
743    fn test_centroids() -> crate::structures::CoarseCentroids {
744        crate::structures::CoarseCentroids {
745            num_clusters: 1,
746            dim: 2,
747            centroids: vec![0.25, 0.75],
748            version: 7,
749            soar_config: None,
750        }
751    }
752
753    async fn write_bincode(
754        directory: &crate::directories::RamDirectory,
755        filename: &str,
756        value: &impl serde::Serialize,
757    ) {
758        let bytes = bincode::serde::encode_to_vec(value, bincode::config::standard()).unwrap();
759        directory.write(Path::new(filename), &bytes).await.unwrap();
760    }
761
762    #[test]
763    fn test_metadata_init() {
764        let mut meta = IndexMetadata::new(test_schema());
765        assert_eq!(meta.total_vectors, 0);
766        assert!(meta.segment_metas.is_empty());
767        assert!(!meta.is_field_built(0));
768
769        meta.init_field(0, VectorIndexType::IvfRaBitQ);
770        assert!(!meta.is_field_built(0));
771        assert!(meta.vector_fields.contains_key(&0));
772    }
773
774    #[tokio::test]
775    async fn load_refuses_metadata_stamped_with_a_newer_format_version() {
776        let directory = crate::directories::RamDirectory::new();
777        let mut metadata = IndexMetadata::new(test_schema());
778        metadata.version = 2;
779        metadata.save(&directory).await.unwrap();
780
781        let error = IndexMetadata::load(&directory)
782            .await
783            .expect_err("metadata from a newer format version must be refused, not silently pruned")
784            .to_string();
785        assert!(error.contains("version 2"), "{error}");
786        assert!(error.contains("newer"), "{error}");
787    }
788
789    #[tokio::test]
790    async fn tmp_recovery_refuses_metadata_stamped_with_a_newer_format_version() {
791        let directory = crate::directories::RamDirectory::new();
792        let mut metadata = IndexMetadata::new(test_schema());
793        metadata.version = 2;
794        let bytes = metadata.serialize_to_bytes().unwrap();
795        // Simulate a crash between write and rename: only the temp file exists.
796        directory
797            .write(Path::new(INDEX_META_TMP_FILENAME), &bytes)
798            .await
799            .unwrap();
800
801        let error = IndexMetadata::load(&directory)
802            .await
803            .expect_err("temp-file recovery must apply the same version gate")
804            .to_string();
805        assert!(error.contains("version 2"), "{error}");
806    }
807
808    #[tokio::test]
809    async fn save_treats_post_rename_sync_failure_as_committed() {
810        let directory = SyncFailDirectory::default();
811        let mut metadata = IndexMetadata::new(test_schema());
812        metadata.add_segment("committed".to_string(), 7);
813
814        metadata.save(&directory).await.unwrap();
815
816        let loaded = IndexMetadata::load(&directory).await.unwrap();
817        assert_eq!(loaded.segment_doc_count("committed"), Some(7));
818    }
819
820    #[tokio::test]
821    async fn trained_artifacts_load_only_when_the_complete_built_set_is_valid() {
822        let mut builder = crate::dsl::SchemaBuilder::default();
823        let config = crate::dsl::DenseVectorConfig::with_ivf(2, Some(1), 1);
824        let first = builder.add_dense_vector_field_with_config(
825            "first_embedding",
826            true,
827            true,
828            config.clone(),
829        );
830        let second =
831            builder.add_dense_vector_field_with_config("second_embedding", true, true, config);
832        let schema = builder.build();
833        let directory = crate::directories::RamDirectory::new();
834        let mut metadata = IndexMetadata::new(schema.clone());
835        metadata.init_field(first.0, VectorIndexType::IvfRaBitQ);
836        metadata.init_field(second.0, VectorIndexType::IvfRaBitQ);
837        metadata.mark_field_built(first.0, 10, 1, "field_0_centroids.bin".into(), None);
838        metadata.mark_field_built(second.0, 10, 1, "field_1_centroids.bin".into(), None);
839        write_bincode(&directory, "field_0_centroids.bin", &test_centroids()).await;
840
841        let error = IndexMetadata::try_load_trained_from_fields(
842            &metadata.vector_fields,
843            &schema,
844            &directory,
845        )
846        .await
847        .err()
848        .expect("missing artifact must fail the complete load")
849        .to_string();
850        assert!(error.contains("field_1_centroids.bin"), "{error}");
851        assert!(error.contains("field 1"), "{error}");
852        assert!(
853            IndexMetadata::load_trained_from_fields(&metadata.vector_fields, &directory)
854                .await
855                .is_none(),
856            "the compatibility API must also fail closed instead of returning the valid subset"
857        );
858    }
859
860    #[tokio::test]
861    async fn index_open_fails_closed_when_built_artifact_is_missing() {
862        let (schema, field) = dense_schema(VectorIndexType::IvfRaBitQ);
863        let directory = crate::directories::RamDirectory::new();
864        let mut metadata = IndexMetadata::new(schema);
865        metadata.init_field(field.0, VectorIndexType::IvfRaBitQ);
866        metadata.mark_field_built(field.0, 10, 1, "missing_centroids.bin".into(), None);
867        metadata.save(&directory).await.unwrap();
868
869        let error = match crate::index::Index::open(directory, crate::index::IndexConfig::default())
870            .await
871        {
872            Ok(_) => panic!("Index::open accepted a Built field with no artifact"),
873            Err(error) => error.to_string(),
874        };
875        assert!(error.contains("missing_centroids.bin"), "{error}");
876    }
877
878    #[tokio::test]
879    async fn scann_built_state_requires_a_codebook() {
880        let (schema, field) = dense_schema(VectorIndexType::ScaNN);
881        let directory = crate::directories::RamDirectory::new();
882        let mut metadata = IndexMetadata::new(schema.clone());
883        metadata.init_field(field.0, VectorIndexType::ScaNN);
884        metadata.mark_field_built(field.0, 10, 1, "field_0_centroids.bin".into(), None);
885        write_bincode(&directory, "field_0_centroids.bin", &test_centroids()).await;
886
887        let error = IndexMetadata::try_load_trained_from_fields(
888            &metadata.vector_fields,
889            &schema,
890            &directory,
891        )
892        .await
893        .err()
894        .expect("ScaNN Built state without a codebook must fail")
895        .to_string();
896        assert!(error.contains("has no codebook_file"), "{error}");
897    }
898
899    #[tokio::test]
900    async fn legacy_requested_cluster_count_accepts_a_clamped_artifact() {
901        let mut builder = crate::dsl::SchemaBuilder::default();
902        let field = builder.add_dense_vector_field_with_config(
903            "embedding",
904            true,
905            true,
906            crate::dsl::DenseVectorConfig::with_ivf(2, Some(4), 1),
907        );
908        let schema = builder.build();
909        let directory = crate::directories::RamDirectory::new();
910        let mut metadata = IndexMetadata::new(schema.clone());
911        metadata.init_field(field.0, VectorIndexType::IvfRaBitQ);
912        metadata.mark_field_built(field.0, 1, 4, "field_0_centroids.bin".into(), None);
913        write_bincode(&directory, "field_0_centroids.bin", &test_centroids()).await;
914
915        let trained = IndexMetadata::try_load_trained_from_fields(
916            &metadata.vector_fields,
917            &schema,
918            &directory,
919        )
920        .await
921        .unwrap()
922        .unwrap();
923        assert_eq!(trained.centroids[&field.0].num_clusters, 1);
924    }
925
926    #[tokio::test]
927    async fn trained_artifact_loader_rejects_trailing_data() {
928        let (schema, field) = dense_schema(VectorIndexType::IvfRaBitQ);
929        let directory = crate::directories::RamDirectory::new();
930        let mut metadata = IndexMetadata::new(schema.clone());
931        metadata.init_field(field.0, VectorIndexType::IvfRaBitQ);
932        metadata.mark_field_built(field.0, 10, 1, "field_0_centroids.bin".into(), None);
933        let mut bytes =
934            bincode::serde::encode_to_vec(test_centroids(), bincode::config::standard()).unwrap();
935        bytes.extend_from_slice(&[0xaa, 0xbb]);
936        directory
937            .write(Path::new("field_0_centroids.bin"), &bytes)
938            .await
939            .unwrap();
940
941        let error = IndexMetadata::try_load_trained_from_fields(
942            &metadata.vector_fields,
943            &schema,
944            &directory,
945        )
946        .await
947        .err()
948        .expect("trailing artifact bytes must fail validation")
949        .to_string();
950        assert!(error.contains("trailing bytes"), "{error}");
951    }
952
953    #[test]
954    fn trained_artifact_size_limit_rejects_before_reading() {
955        let error = validate_trained_artifact_size(
956            3,
957            "centroids",
958            "field_3_centroids.bin",
959            MAX_TRAINED_ARTIFACT_BYTES as u64 + 1,
960        )
961        .unwrap_err()
962        .to_string();
963        assert!(error.contains("exceeding"), "{error}");
964        assert!(error.contains("field 3"), "{error}");
965    }
966
967    #[tokio::test]
968    async fn trained_artifact_decode_limit_rejects_forged_collection_length() {
969        let (schema, field) = dense_schema(VectorIndexType::IvfRaBitQ);
970        let directory = crate::directories::RamDirectory::new();
971        let mut metadata = IndexMetadata::new(schema.clone());
972        metadata.init_field(field.0, VectorIndexType::IvfRaBitQ);
973        metadata.mark_field_built(field.0, 10, 1, "field_0_centroids.bin".into(), None);
974
975        // CoarseCentroids begins with num_clusters=1, dim=2, then the Vec
976        // length. Bincode's standard varint marker 253 introduces a u64; this
977        // tiny payload claims an impossible f32 vector and must hit the decode
978        // limit before any large allocation is attempted.
979        let mut bytes = vec![1, 2, 253];
980        bytes.extend_from_slice(&u64::MAX.to_le_bytes());
981        directory
982            .write(Path::new("field_0_centroids.bin"), &bytes)
983            .await
984            .unwrap();
985
986        let error = IndexMetadata::try_load_trained_from_fields(
987            &metadata.vector_fields,
988            &schema,
989            &directory,
990        )
991        .await
992        .err()
993        .expect("forged collection length must fail the bounded decoder")
994        .to_string();
995        assert!(error.contains("failed to deserialize"), "{error}");
996    }
997
998    #[test]
999    fn test_metadata_segments() {
1000        let mut meta = IndexMetadata::new(test_schema());
1001        meta.add_segment("abc123".to_string(), 50);
1002        meta.add_segment("def456".to_string(), 100);
1003        assert_eq!(meta.segment_metas.len(), 2);
1004        assert_eq!(meta.segment_doc_count("abc123"), Some(50));
1005        assert_eq!(meta.segment_doc_count("def456"), Some(100));
1006
1007        // Overwrites existing
1008        meta.add_segment("abc123".to_string(), 75);
1009        assert_eq!(meta.segment_metas.len(), 2);
1010        assert_eq!(meta.segment_doc_count("abc123"), Some(75));
1011
1012        meta.remove_segment("abc123");
1013        assert_eq!(meta.segment_metas.len(), 1);
1014        assert!(meta.has_segment("def456"));
1015        assert!(!meta.has_segment("abc123"));
1016    }
1017
1018    #[test]
1019    fn test_mark_field_built() {
1020        let mut meta = IndexMetadata::new(test_schema());
1021        meta.init_field(0, VectorIndexType::IvfRaBitQ);
1022        meta.total_vectors = 10000;
1023
1024        assert!(!meta.is_field_built(0));
1025
1026        meta.mark_field_built(0, 10000, 256, "field_0_centroids.bin".to_string(), None);
1027
1028        assert!(meta.is_field_built(0));
1029        let field = meta.get_field_meta(0).unwrap();
1030        assert_eq!(
1031            field.centroids_file.as_deref(),
1032            Some("field_0_centroids.bin")
1033        );
1034    }
1035
1036    #[test]
1037    fn total_vectors_is_aggregate_of_built_field_counts() {
1038        let mut meta = IndexMetadata::new(test_schema());
1039        meta.init_field(7, VectorIndexType::IvfRaBitQ);
1040        meta.init_field(3, VectorIndexType::ScaNN);
1041
1042        // Build in reverse field-id order to ensure the result is not tied to
1043        // HashMap or training iteration order.
1044        meta.mark_field_built(7, 400, 20, "field_7_centroids.bin".to_string(), None);
1045        assert_eq!(meta.total_vectors, 400);
1046        meta.mark_field_built(
1047            3,
1048            250,
1049            15,
1050            "field_3_centroids.bin".to_string(),
1051            Some("field_3_codebook.bin".to_string()),
1052        );
1053        assert_eq!(meta.total_vectors, 650);
1054
1055        // Rebuilding a field replaces its contribution; it does not add a
1056        // duplicate training snapshot.
1057        meta.mark_field_built(7, 425, 20, "field_7_centroids.bin".to_string(), None);
1058        assert_eq!(meta.total_vectors, 675);
1059    }
1060
1061    #[test]
1062    fn test_should_build_field() {
1063        let mut meta = IndexMetadata::new(test_schema());
1064        meta.init_field(0, VectorIndexType::IvfRaBitQ);
1065
1066        // Below threshold
1067        meta.total_vectors = 500;
1068        assert!(!meta.should_build_field(0, 1000));
1069
1070        // Above threshold
1071        meta.total_vectors = 1500;
1072        assert!(meta.should_build_field(0, 1000));
1073
1074        // Already built - should not build again
1075        meta.mark_field_built(0, 1500, 256, "centroids.bin".to_string(), None);
1076        assert!(!meta.should_build_field(0, 1000));
1077    }
1078
1079    #[test]
1080    fn test_serialization() {
1081        let mut meta = IndexMetadata::new(test_schema());
1082        meta.add_segment("seg1".to_string(), 100);
1083        meta.init_field(0, VectorIndexType::IvfRaBitQ);
1084        meta.total_vectors = 5000;
1085
1086        let json = serde_json::to_string_pretty(&meta).unwrap();
1087        let loaded: IndexMetadata = serde_json::from_str(&json).unwrap();
1088
1089        assert_eq!(loaded.segment_ids().len(), meta.segment_ids().len());
1090        assert_eq!(loaded.segment_doc_count("seg1"), Some(100));
1091        assert_eq!(loaded.total_vectors, meta.total_vectors);
1092        assert!(loaded.vector_fields.contains_key(&0));
1093    }
1094
1095    #[test]
1096    fn old_metadata_defaults_the_bp_retry_counter() {
1097        let mut meta = IndexMetadata::new(test_schema());
1098        meta.add_segment("legacy".to_string(), 10);
1099        let mut json = serde_json::to_value(&meta).unwrap();
1100        json["segment_metas"]["legacy"]
1101            .as_object_mut()
1102            .unwrap()
1103            .remove("bp_unconverged_passes");
1104
1105        let loaded: IndexMetadata = serde_json::from_value(json).unwrap();
1106        assert_eq!(loaded.segment_metas["legacy"].bp_unconverged_passes, 0);
1107    }
1108
1109    #[test]
1110    fn test_merged_segment_lineage() {
1111        let mut meta = IndexMetadata::new(test_schema());
1112        meta.add_segment("a".to_string(), 50);
1113        meta.add_segment("b".to_string(), 75);
1114
1115        // Fresh segments: gen=0, no ancestors
1116        assert_eq!(meta.segment_metas["a"].generation, 0);
1117        assert!(meta.segment_metas["a"].ancestors.is_empty());
1118
1119        // Merge a+b → c
1120        meta.add_merged_segment(
1121            "c".to_string(),
1122            125,
1123            vec!["a".to_string(), "b".to_string()],
1124            1,
1125            false,
1126            true,
1127        );
1128        assert_eq!(meta.segment_metas["c"].generation, 1);
1129        assert_eq!(meta.segment_metas["c"].ancestors, vec!["a", "b"]);
1130        assert_eq!(meta.segment_doc_count("c"), Some(125));
1131
1132        // Merge c+d → e (gen should be 2)
1133        meta.add_segment("d".to_string(), 30);
1134        meta.add_merged_segment(
1135            "e".to_string(),
1136            155,
1137            vec!["c".to_string(), "d".to_string()],
1138            2,
1139            false,
1140            true,
1141        );
1142        assert_eq!(meta.segment_metas["e"].generation, 2);
1143    }
1144}