nodedb_vector/collection/

checkpoint.rs

// SPDX-License-Identifier: Apache-2.0

//! Checkpoint serialization and deserialization for `VectorCollection`.
//!
//! ## On-disk framing
//!
//! **Plaintext** checkpoints are raw MessagePack bytes (existing format).
//! The first 4 bytes are never `SEGV`, so detection is unambiguous.
//!
//! **Encrypted** checkpoints use the following layout:
//!
//! ```text
//! [SEGV (4B)] [version_u16_le (2B)] [cipher_alg_u8 (1B)] [kid_u8 (1B)]
//! [epoch (4B)] [reserved (4B)] [AES-256-GCM ciphertext of msgpack payload]
//! ```
//!
//! The first 16 bytes form a `SegmentPreamble` (reusing the existing preamble
//! layout with a distinct `SEGV` magic). These 16 bytes are included as AAD,
//! preventing preamble-swap attacks. The nonce is `(epoch, lsn=0)` — epoch
//! provides per-write uniqueness even without an LSN.

use std::collections::HashMap;

use nodedb_types::{Surrogate, VectorQuantization};
use serde::{Deserialize, Serialize};

use crate::collection::payload_index::PayloadIndexSetSnapshot;
use crate::collection::segment::{DEFAULT_SEAL_THRESHOLD, SealedSegment};
use crate::collection::tier::StorageTier;
use crate::distance::DistanceMetric;
use crate::error::VectorError;
use crate::flat::FlatIndex;
use crate::hnsw::{HnswIndex, HnswParams};
use crate::quantize::pq::PqCodec;
use crate::quantize::sq8::Sq8Codec;

use super::lifecycle::VectorCollection;

/// Magic bytes identifying an encrypted vector checkpoint. Shared with
/// `nodedb-spatial`'s SEGV checkpoint format.
const SEGV_MAGIC: [u8; 4] = *b"SEGV";

/// Encrypt `plaintext` into the SEGV envelope from [`nodedb_wal::crypto`].
fn encrypt_checkpoint(
    key: &nodedb_wal::crypto::WalEncryptionKey,
    plaintext: &[u8],
) -> Result<Vec<u8>, VectorError> {
    nodedb_wal::crypto::encrypt_segment_envelope(key, &SEGV_MAGIC, plaintext).map_err(|e| {
        VectorError::CheckpointEncryptionError {
            detail: e.to_string(),
        }
    })
}

/// Decrypt an encrypted checkpoint blob (starting at byte 0, which is `SEGV`).
fn decrypt_checkpoint(
    key: &nodedb_wal::crypto::WalEncryptionKey,
    blob: &[u8],
) -> Result<Vec<u8>, VectorError> {
    nodedb_wal::crypto::decrypt_segment_envelope(key, &SEGV_MAGIC, blob).map_err(|e| {
        VectorError::CheckpointEncryptionError {
            detail: e.to_string(),
        }
    })
}

#[derive(Serialize, Deserialize, zerompk::ToMessagePack, zerompk::FromMessagePack)]
pub(crate) struct CollectionSnapshot {
    pub dim: usize,
    pub params_m: usize,
    pub params_m0: usize,
    pub params_ef_construction: usize,
    pub params_metric: u8,
    pub next_id: u32,
    pub growing_base_id: u32,
    pub growing_vectors: Vec<Vec<f32>>,
    pub growing_deleted: Vec<bool>,
    pub sealed_segments: Vec<SealedSnapshot>,
    pub building_segments: Vec<BuildingSnapshot>,
    /// `(global_vector_id, surrogate_u32)` pairs.
    #[serde(default)]
    pub surrogate_map: Vec<(u32, u32)>,
    /// `(document_surrogate_u32, [global_vector_ids])` pairs.
    #[serde(default)]
    pub multi_doc_map: Vec<(u32, Vec<u32>)>,
    /// Quantization mode for the collection-level codec-dispatch index.
    /// Serialised as a u8 matching `VectorQuantization` discriminants.
    /// 0 = None (default, backward-compatible).
    #[serde(default)]
    pub quantization_tag: u8,
    /// Serialised `PayloadIndexSetSnapshot` (msgpack bytes).
    /// Empty vec = no payload indexes (default, backward-compatible).
    #[serde(default)]
    pub payload_index_bytes: Vec<u8>,
}

#[derive(Serialize, Deserialize, zerompk::ToMessagePack, zerompk::FromMessagePack)]
pub(crate) struct SealedSnapshot {
    pub base_id: u32,
    pub hnsw_bytes: Vec<u8>,
    #[serde(default)]
    pub pq_bytes: Option<Vec<u8>>,
    #[serde(default)]
    pub pq_codes: Option<Vec<u8>>,
    /// Serialized [`Sq8Codec`] bytes (magic + version + msgpack).
    /// Present when SQ8 quantization is active and PQ is absent.
    /// `None` means no SQ8 quantization for this segment.
    #[serde(default)]
    pub sq8_bytes: Option<Vec<u8>>,
    /// Pre-quantized SQ8 codes for all vectors in this segment.
    /// Layout: `[v0_d0, v0_d1, ..., v1_d0, ...]` (dim bytes per vector).
    /// `None` when SQ8 is not configured.
    #[serde(default)]
    pub sq8_codes: Option<Vec<u8>>,
}

#[derive(Serialize, Deserialize, zerompk::ToMessagePack, zerompk::FromMessagePack)]
pub(crate) struct BuildingSnapshot {
    pub base_id: u32,
    pub vectors: Vec<Vec<f32>>,
    #[serde(default)]
    pub deleted: Vec<bool>,
}

impl VectorCollection {
    /// Serialize all segments for checkpointing.
    ///
    /// When `kek` is `Some`, the MessagePack payload is wrapped in an
    /// AES-256-GCM encrypted envelope with a `SEGV` preamble. When `None`,
    /// raw MessagePack bytes are returned (existing plaintext format).
    ///
    /// Returns an empty `Vec` on serialization failure (callers treat this as a
    /// skip signal, consistent with the pre-existing error handling).
    pub fn checkpoint_to_bytes(
        &self,
        kek: Option<&nodedb_wal::crypto::WalEncryptionKey>,
    ) -> Vec<u8> {
        let snapshot = CollectionSnapshot {
            dim: self.dim,
            params_m: self.params.m,
            params_m0: self.params.m0,
            params_ef_construction: self.params.ef_construction,
            params_metric: self.params.metric as u8,
            next_id: self.next_id,
            growing_base_id: self.growing_base_id,
            growing_vectors: (0..self.growing.len() as u32)
                .filter_map(|i| self.growing.get_vector_raw(i).map(|v| v.to_vec()))
                .collect(),
            growing_deleted: (0..self.growing.len() as u32)
                .map(|i| self.growing.is_deleted(i))
                .collect(),
            sealed_segments: self
                .sealed
                .iter()
                .map(|s| {
                    let (pq_bytes, pq_codes) = match &s.pq {
                        Some((codec, codes)) => (codec.to_bytes().ok(), Some(codes.clone())),
                        None => (None, None),
                    };
                    // Only serialize SQ8 when PQ is absent — a segment never carries both.
                    let (sq8_bytes, sq8_codes) = if pq_bytes.is_none() {
                        match &s.sq8 {
                            Some((codec, codes)) => (Some(codec.to_bytes()), Some(codes.clone())),
                            None => (None, None),
                        }
                    } else {
                        (None, None)
                    };
                    SealedSnapshot {
                        base_id: s.base_id,
                        hnsw_bytes: s.index.checkpoint_to_bytes(),
                        pq_bytes,
                        pq_codes,
                        sq8_bytes,
                        sq8_codes,
                    }
                })
                .collect(),
            building_segments: self
                .building
                .iter()
                .map(|b| BuildingSnapshot {
                    base_id: b.base_id,
                    vectors: (0..b.flat.len() as u32)
                        .filter_map(|i| b.flat.get_vector_raw(i).map(|v| v.to_vec()))
                        .collect(),
                    deleted: (0..b.flat.len() as u32)
                        .map(|i| b.flat.is_deleted(i))
                        .collect(),
                })
                .collect(),
            surrogate_map: self
                .surrogate_map
                .iter()
                .map(|(&k, s)| (k, s.as_u32()))
                .collect(),
            multi_doc_map: self
                .multi_doc_map
                .iter()
                .map(|(k, v)| (k.as_u32(), v.clone()))
                .collect(),
            quantization_tag: quantization_to_tag(self.quantization),
            payload_index_bytes: {
                let snap = self.payload.to_snapshot();
                match zerompk::to_msgpack_vec(&snap) {
                    Ok(bytes) => bytes,
                    Err(e) => {
                        tracing::warn!(
                            error = %e,
                            "vector payload index snapshot serialization failed"
                        );
                        return Vec::new();
                    }
                }
            },
        };
        let msgpack = match zerompk::to_msgpack_vec(&snapshot) {
            Ok(bytes) => bytes,
            Err(e) => {
                tracing::warn!(error = %e, "vector collection checkpoint serialization failed");
                return Vec::new();
            }
        };

        if let Some(key) = kek {
            match encrypt_checkpoint(key, &msgpack) {
                Ok(encrypted) => encrypted,
                Err(e) => {
                    tracing::warn!(error = %e, "vector collection checkpoint encryption failed");
                    Vec::new()
                }
            }
        } else {
            msgpack
        }
    }

    /// Restore a collection from checkpoint bytes.
    ///
    /// `kek` controls the expected framing:
    /// - `None` → the file must be plaintext MessagePack (starting with bytes
    ///   that are NOT `SEGV`). If the file starts with `SEGV` and no key is
    ///   provided, returns `Err(CheckpointEncryptedNoKey)`.
    /// - `Some(key)` → encryption is **required**. If the file starts with
    ///   `SEGV`, it is decrypted with `key`. If the file is plaintext, returns
    ///   `Err(CheckpointPlaintextKeyRequired)` — refuse to silently load
    ///   unencrypted data when the operator has enabled at-rest encryption.
    pub fn from_checkpoint(
        bytes: &[u8],
        kek: Option<&nodedb_wal::crypto::WalEncryptionKey>,
    ) -> Result<Option<Self>, VectorError> {
        let is_encrypted = bytes.len() >= 4 && bytes[0..4] == SEGV_MAGIC;

        let msgpack: Vec<u8>;
        let msgpack_ref: &[u8];

        if is_encrypted {
            if let Some(key) = kek {
                msgpack = decrypt_checkpoint(key, bytes)?;
                msgpack_ref = &msgpack;
            } else {
                return Err(VectorError::CheckpointEncryptedNoKey);
            }
        } else if kek.is_some() {
            return Err(VectorError::CheckpointPlaintextKeyRequired);
        } else {
            msgpack_ref = bytes;
        }

        let snap: CollectionSnapshot = match zerompk::from_msgpack(msgpack_ref) {
            Ok(s) => s,
            Err(_) => return Ok(None),
        };
        let metric = match snap.params_metric {
            0 => DistanceMetric::L2,
            1 => DistanceMetric::Cosine,
            2 => DistanceMetric::InnerProduct,
            3 => DistanceMetric::Manhattan,
            4 => DistanceMetric::Chebyshev,
            5 => DistanceMetric::Hamming,
            6 => DistanceMetric::Jaccard,
            7 => DistanceMetric::Pearson,
            _ => DistanceMetric::Cosine,
        };
        let params = HnswParams {
            m: snap.params_m,
            m0: snap.params_m0,
            ef_construction: snap.params_ef_construction,
            metric,
        };

        let mut growing = FlatIndex::new(snap.dim, metric);
        for (i, v) in snap.growing_vectors.iter().enumerate() {
            let deleted = snap.growing_deleted.get(i).copied().unwrap_or(false);
            if deleted {
                growing.insert_tombstoned(v.clone());
            } else {
                growing.insert(v.clone());
            }
        }

        // no-governor: cold restore path; segment count bounded by collection config
        let mut sealed = Vec::with_capacity(snap.sealed_segments.len());
        for ss in &snap.sealed_segments {
            if let Some(index) = HnswIndex::from_checkpoint(&ss.hnsw_bytes).ok().flatten() {
                let pq = match (&ss.pq_bytes, &ss.pq_codes) {
                    (Some(bytes), Some(codes)) => PqCodec::from_bytes(bytes)
                        .ok()
                        .map(|codec| (codec, codes.clone())),
                    _ => None,
                };
                // Restore SQ8 from persisted bytes — never recompute on load.
                // A segment never carries both PQ and SQ8.
                let sq8 = if pq.is_some() {
                    None
                } else {
                    match (&ss.sq8_bytes, &ss.sq8_codes) {
                        (Some(codec_bytes), Some(codes)) => Sq8Codec::from_bytes(codec_bytes)
                            .ok()
                            .map(|codec| (codec, codes.clone())),
                        _ => None,
                    }
                };
                sealed.push(SealedSegment {
                    index,
                    base_id: ss.base_id,
                    sq8,
                    pq,
                    tier: StorageTier::L0Ram,
                    mmap_vectors: None,
                });
            }
        }

        for bs in &snap.building_segments {
            let mut index = HnswIndex::new(snap.dim, params.clone());
            for v in &bs.vectors {
                index
                    .insert(v.clone())
                    .expect("dimension guaranteed by checkpoint");
            }
            // Replay building-segment tombstones onto the HNSW index.
            for (i, &dead) in bs.deleted.iter().enumerate() {
                if dead {
                    index.delete(i as u32);
                }
            }
            let sq8 = VectorCollection::build_sq8_for_index(&index);
            sealed.push(SealedSegment {
                index,
                base_id: bs.base_id,
                sq8,
                pq: None,
                tier: StorageTier::L0Ram,
                mmap_vectors: None,
            });
        }

        let next_segment_id = (sealed.len() + 1) as u32;

        let index_config = crate::index_config::IndexConfig {
            hnsw: params.clone(),
            ..crate::index_config::IndexConfig::default()
        };
        Ok(Some(Self {
            growing,
            growing_base_id: snap.growing_base_id,
            sealed,
            building: Vec::new(),
            params,
            next_id: snap.next_id,
            next_segment_id,
            dim: snap.dim,
            data_dir: None,
            ram_budget_bytes: 0,
            mmap_fallback_count: 0,
            mmap_segment_count: 0,
            surrogate_map: snap
                .surrogate_map
                .iter()
                .map(|&(k, s)| (k, Surrogate::new(s)))
                .collect(),
            surrogate_to_local: snap
                .surrogate_map
                .iter()
                .map(|&(k, s)| (Surrogate::new(s), k))
                .collect(),
            multi_doc_map: snap
                .multi_doc_map
                .into_iter()
                .map(|(k, v)| (Surrogate::new(k), v))
                .collect::<HashMap<_, _>>(),
            seal_threshold: DEFAULT_SEAL_THRESHOLD,
            index_config,
            codec_dispatch: None,
            quantization: quantization_from_tag(snap.quantization_tag),
            payload: if snap.payload_index_bytes.is_empty() {
                super::payload_index::PayloadIndexSet::default()
            } else {
                zerompk::from_msgpack::<PayloadIndexSetSnapshot>(&snap.payload_index_bytes)
                    .map(super::payload_index::PayloadIndexSet::from_snapshot)
                    .unwrap_or_default()
            },
            arena_index: None,
        }))
    }
}

/// Encode a `VectorQuantization` to a u8 tag for storage.
fn quantization_to_tag(q: VectorQuantization) -> u8 {
    match q {
        VectorQuantization::None => 0,
        VectorQuantization::Sq8 => 1,
        VectorQuantization::Pq => 2,
        VectorQuantization::RaBitQ => 3,
        VectorQuantization::Bbq => 4,
        VectorQuantization::Binary => 5,
        VectorQuantization::Ternary => 6,
        VectorQuantization::Opq => 7,
        _ => 0,
    }
}

/// Decode a u8 tag back to `VectorQuantization`.
fn quantization_from_tag(tag: u8) -> VectorQuantization {
    match tag {
        0 => VectorQuantization::None,
        1 => VectorQuantization::Sq8,
        2 => VectorQuantization::Pq,
        3 => VectorQuantization::RaBitQ,
        4 => VectorQuantization::Bbq,
        5 => VectorQuantization::Binary,
        6 => VectorQuantization::Ternary,
        7 => VectorQuantization::Opq,
        _ => VectorQuantization::None,
    }
}

#[cfg(test)]
mod tests {
    use crate::collection::lifecycle::VectorCollection;
    use crate::distance::DistanceMetric;
    use crate::hnsw::HnswParams;

    /// SQ8 calibration data must survive a checkpoint round-trip without
    /// being recomputed. Verifies that the O(N*dim) rebuild-on-restart bug
    /// is eliminated: `sq8` on the restored sealed segment is `Some` and
    /// its `inv_scales` match the original exactly.
    #[test]
    fn checkpoint_roundtrip_preserves_sq8() {
        use crate::collection::lifecycle::VectorCollection;
        use crate::hnsw::{HnswIndex, HnswParams};

        let params = HnswParams {
            metric: crate::distance::DistanceMetric::L2,
            ..HnswParams::default()
        };
        // 1024 vectors of dim=8 — enough to pass the ≥1000 threshold in
        // `build_sq8_for_index`, so sq8 will be Some after complete_build.
        // Uses plain HNSW (default IndexType::Hnsw) so SQ8 is selected.
        let mut coll = VectorCollection::with_seal_threshold(8, params, 1024);
        for i in 0..1024u32 {
            let mut v = vec![0.0f32; 8];
            for (d, slot) in v.iter_mut().enumerate() {
                *slot = ((i as f32) * 0.01 + (d as f32) * 0.1).sin();
            }
            coll.insert(v);
        }
        let req = coll.seal("sq8_test").expect("seal produced request");
        let mut idx = HnswIndex::new(req.dim, req.params.clone());
        for v in &req.vectors {
            idx.insert(v.clone()).unwrap();
        }
        coll.complete_build(req.segment_id, idx);

        let sealed = coll.sealed_segments();
        assert!(!sealed.is_empty(), "expected at least one sealed segment");
        let orig_sq8 = sealed[0]
            .sq8
            .as_ref()
            .expect("sq8 must be Some after complete_build with ≥1000 vectors");
        let orig_dim = orig_sq8.0.dim();
        // Capture serialized form as ground truth.
        let orig_bytes = orig_sq8.0.to_bytes();

        let checkpoint = coll.checkpoint_to_bytes(None);
        let restored = VectorCollection::from_checkpoint(&checkpoint, None)
            .unwrap()
            .unwrap();

        let restored_sealed = restored.sealed_segments();
        assert!(!restored_sealed.is_empty());
        let restored_sq8 = restored_sealed[0]
            .sq8
            .as_ref()
            .expect("sq8 must be Some after restoring checkpoint — never recomputed");

        assert_eq!(restored_sq8.0.dim(), orig_dim, "dim mismatch after restore");
        // Byte-level equality guarantees calibration data is persisted, not recomputed.
        assert_eq!(
            restored_sq8.0.to_bytes(),
            orig_bytes,
            "sq8 codec bytes differ — calibration data was recomputed rather than persisted"
        );
    }

    #[test]
    fn checkpoint_roundtrip() {
        let mut coll = VectorCollection::new(
            3,
            HnswParams {
                metric: DistanceMetric::L2,
                ..HnswParams::default()
            },
        );
        for i in 0..50u32 {
            coll.insert(vec![i as f32, 0.0, 0.0]);
        }
        let bytes = coll.checkpoint_to_bytes(None);
        let restored = VectorCollection::from_checkpoint(&bytes, None)
            .unwrap()
            .unwrap();
        assert_eq!(restored.len(), 50);
        assert_eq!(restored.dim(), 3);

        let results = restored.search(&[25.0, 0.0, 0.0], 1, 64);
        assert_eq!(results[0].id, 25);
    }

    /// Payload bitmap indexes registered on a vector-primary collection
    /// must survive a checkpoint round-trip — otherwise `WHERE` filters
    /// would silently return zero rows after a node restart.
    #[test]
    fn checkpoint_roundtrip_preserves_payload_bitmap() {
        use crate::collection::PayloadIndexKind;
        use crate::collection::payload_index::FilterPredicate;
        use nodedb_types::Value;
        use std::collections::HashMap;

        let mut coll = VectorCollection::new(
            3,
            HnswParams {
                metric: DistanceMetric::L2,
                ..HnswParams::default()
            },
        );
        coll.payload
            .add_index("category".to_string(), PayloadIndexKind::Equality);
        for i in 0u32..10 {
            let node_id = coll.insert(vec![i as f32, 0.0, 0.0]);
            let mut fields = HashMap::new();
            let cat = if i % 2 == 0 { "A" } else { "B" };
            fields.insert("category".to_string(), Value::String(cat.to_string()));
            coll.payload.insert_row(node_id, &fields);
        }

        let bytes = coll.checkpoint_to_bytes(None);
        let restored = VectorCollection::from_checkpoint(&bytes, None)
            .unwrap()
            .unwrap();

        let pred = FilterPredicate::Eq {
            field: "category".to_string(),
            value: Value::String("A".to_string()),
        };
        let bm = restored
            .payload
            .pre_filter(&pred)
            .expect("payload index 'category' must be present after restore");
        assert_eq!(
            bm.len(),
            5,
            "5 rows of category=A must survive checkpoint round-trip"
        );
    }
}