Skip to main content

holographic_memory/core/
admin.rs

1// Copyright 2024-2026 WritersLogic Contributors
2// SPDX-License-Identifier: Apache-2.0
3
4use anyhow::{anyhow, bail, ensure, Context, Result};
5use crc32fast::Hasher;
6use memmap2::Mmap;
7use serde::Serialize;
8use serde_json::Value;
9use std::collections::BTreeSet;
10use std::fs::File;
11use std::path::Path;
12use std::time::{SystemTime, UNIX_EPOCH};
13
14use super::storage::{
15    is_log_end, FormatManifest, FORMAT_MAGIC, FORMAT_MANIFEST, HEADER_SIZE, LEGACY_SEGMENT_SIZE,
16    MAX_RAW_FRAME_SIZE, SEGMENT_SIZE, STORAGE_FORMAT_VERSION,
17};
18use super::store_lock::{StoreLock, LOCK_FILE};
19use super::{HmsConfig, HmsCore};
20
21const ARENA_DIRECTORY: &str = "vectors_data.bin";
22
23#[derive(Debug, Clone, Serialize)]
24#[serde(rename_all = "camelCase")]
25pub struct StoreInspection {
26    pub store_path: String,
27    pub format_version: u32,
28    pub segment_count: usize,
29    pub frame_count: usize,
30    pub logical_bytes: u64,
31    pub allocated_bytes: u64,
32    pub checksums_verified: bool,
33    /// Manifests whose checksum sidecar was verified (`verify` only).
34    #[serde(default, skip_serializing_if = "Vec::is_empty")]
35    pub manifests_verified: Vec<String>,
36    /// Document cross-check against the opened store (`verify` only).
37    #[serde(skip_serializing_if = "Option::is_none")]
38    pub documents: Option<DocumentCheck>,
39}
40
41#[derive(Debug, Clone, Serialize)]
42#[serde(rename_all = "camelCase")]
43pub struct DocumentCheck {
44    /// False when the store was not opened; `reason` says why.
45    pub checked: bool,
46    #[serde(skip_serializing_if = "Option::is_none")]
47    pub reason: Option<String>,
48    /// Documents the transaction log says are live; absent when the frames are encrypted.
49    #[serde(skip_serializing_if = "Option::is_none")]
50    pub expected: Option<usize>,
51    pub present: usize,
52    pub missing: Vec<String>,
53}
54
55/// Inspect storage framing without opening or changing the HMS engine.
56pub fn inspect_store(path: impl AsRef<Path>, verify_checksums: bool) -> Result<StoreInspection> {
57    let store_path = path.as_ref();
58    let arena_path = store_path.join(ARENA_DIRECTORY);
59    let manifest_path = arena_path.join(FORMAT_MANIFEST);
60    let manifest: FormatManifest = serde_json::from_slice(
61        &std::fs::read(&manifest_path)
62            .with_context(|| format!("missing storage manifest {}", manifest_path.display()))?,
63    )
64    .with_context(|| format!("invalid storage manifest {}", manifest_path.display()))?;
65
66    if manifest.magic != FORMAT_MAGIC {
67        bail!("unknown storage magic {}", manifest.magic);
68    }
69    if manifest.version > STORAGE_FORMAT_VERSION {
70        bail!(
71            "storage format {} is newer than supported format {}",
72            manifest.version,
73            STORAGE_FORMAT_VERSION
74        );
75    }
76    let expected = if manifest.version == 3 {
77        SEGMENT_SIZE
78    } else {
79        LEGACY_SEGMENT_SIZE
80    };
81    if manifest.segment_size != expected {
82        bail!(
83            "storage segment size mismatch: manifest={}, runtime={}",
84            manifest.segment_size,
85            expected
86        );
87    }
88
89    let arena_path = super::durable_file::arena_generation(&arena_path)?;
90    let mut segment_count = 0usize;
91    let mut frame_count = 0usize;
92    let mut logical_bytes = 0u64;
93    let mut allocated_bytes = 0u64;
94    loop {
95        let segment_path = arena_path.join(format!("seg_{segment_count}.bin"));
96        if !segment_path.exists() {
97            break;
98        }
99        let file = File::open(&segment_path)?;
100        let metadata = file.metadata()?;
101        let final_segment = !arena_path
102            .join(format!("seg_{}.bin", segment_count + 1))
103            .exists();
104        let valid_size = if manifest.version >= 4 && final_segment {
105            (SEGMENT_SIZE as u64..=manifest.segment_size as u64).contains(&metadata.len())
106                && metadata.len() % SEGMENT_SIZE as u64 == 0
107        } else {
108            metadata.len() == manifest.segment_size as u64
109        };
110        if !valid_size {
111            bail!(
112                "segment {} has size {}, expected {}",
113                segment_path.display(),
114                metadata.len(),
115                manifest.segment_size
116            );
117        }
118        allocated_bytes = allocated_bytes
119            .checked_add(metadata.len())
120            .context("arena allocated byte count overflow")?;
121        let mmap = unsafe { Mmap::map(&file)? };
122        let (frames, bytes) = inspect_segment(&mmap, &segment_path, verify_checksums)?;
123        frame_count += frames;
124        logical_bytes += bytes as u64;
125        segment_count += 1;
126    }
127    if segment_count == 0 {
128        bail!("store contains no arena segments");
129    }
130    for entry in std::fs::read_dir(&arena_path)? {
131        let name = entry?.file_name();
132        let name = name.to_string_lossy();
133        if let Some(id) = name
134            .strip_prefix("seg_")
135            .and_then(|value| value.strip_suffix(".bin"))
136            .and_then(|value| value.parse::<usize>().ok())
137        {
138            if id >= segment_count {
139                bail!("arena segments are not contiguous (unexpected seg_{id}.bin)");
140            }
141        }
142    }
143
144    Ok(StoreInspection {
145        store_path: store_path.display().to_string(),
146        format_version: manifest.version,
147        segment_count,
148        frame_count,
149        logical_bytes,
150        allocated_bytes,
151        checksums_verified: verify_checksums,
152        manifests_verified: Vec::new(),
153        documents: None,
154    })
155}
156
157/// `inspect_store(path, true)` plus the manifest checksums and a document cross-check: the
158/// store is opened through the engine, which authenticates and replays every frame, and
159/// every document the log says is live must resolve there. Corruption fails the verify
160/// instead of being served as a shorter store.
161pub fn verify_store(
162    path: impl AsRef<Path>,
163    encryption_passphrase_env: Option<String>,
164) -> Result<StoreInspection> {
165    let store = path.as_ref();
166    let mut inspection = inspect_store(store, true)?;
167    let arena_path = store.join(ARENA_DIRECTORY);
168    inspection.manifests_verified = super::durable_file::verify_manifests(store, &arena_path)?;
169    let schema_path = store.join("store.json");
170    let schema: Value = serde_json::from_slice(
171        &std::fs::read(&schema_path)
172            .with_context(|| format!("missing store schema {}", schema_path.display()))?,
173    )
174    .with_context(|| format!("invalid store schema {}", schema_path.display()))?;
175    let dimensions = u32::try_from(
176        schema["dimensions"]
177            .as_u64()
178            .context("store.json lacks dimensions")?,
179    )?;
180    let mut config = HmsConfig::default();
181    if let Some(name) = schema["text_encoder"].as_str() {
182        config.text_encoder = name
183            .parse()
184            .map_err(|e| anyhow!("store.json text encoder {name}: {e}"))?;
185    }
186    config.meaning.enabled = schema["meaning_enabled"].as_bool().unwrap_or(false);
187    config.embedding_space = serde_json::from_value(schema["embedding_space"].clone())
188        .context("store.json embedding_space")?;
189    let expected = if schema["encrypted"].as_bool().unwrap_or(false) {
190        let Some(env) = encryption_passphrase_env else {
191            inspection.documents = Some(DocumentCheck {
192                checked: false,
193                reason: Some(
194                    "encrypted store: pass --encryption-passphrase-env to open it and \
195                     authenticate every frame"
196                        .into(),
197                ),
198                expected: None,
199                present: 0,
200                missing: Vec::new(),
201            });
202            return Ok(inspection);
203        };
204        config.security.encryption_enabled = true;
205        config.security.encryption_passphrase_env = Some(env);
206        None
207    } else {
208        Some(live_document_ids(&arena_path)?)
209    };
210    let core = HmsCore::new(dimensions, Some(store.display().to_string()), Some(config))
211        .context("opening the store for the document cross-check")?;
212    inspection.documents = Some(match expected {
213        Some(ids) => {
214            let missing = missing_documents(&ids, &core);
215            DocumentCheck {
216                checked: true,
217                reason: None,
218                expected: Some(ids.len()),
219                present: ids.len() - missing.len(),
220                missing,
221            }
222        }
223        None => DocumentCheck {
224            checked: true,
225            reason: Some(
226                "every frame authenticated on open; the presence cross-check needs plaintext \
227                 frames"
228                    .into(),
229            ),
230            expected: None,
231            present: 0,
232            missing: Vec::new(),
233        },
234    });
235    Ok(inspection)
236}
237
238/// Document ids the plaintext transaction log leaves live: every `Document` id minus later
239/// `DeleteDocument`s. Frame bounds and checksums were validated by `inspect_store` first.
240fn live_document_ids(arena_path: &Path) -> Result<BTreeSet<String>> {
241    let generation = super::durable_file::arena_generation(arena_path)?;
242    let mut live = BTreeSet::new();
243    for segment in 0usize.. {
244        let segment_path = generation.join(format!("seg_{segment}.bin"));
245        if !segment_path.exists() {
246            break;
247        }
248        let mmap = unsafe { Mmap::map(&File::open(&segment_path)?)? };
249        let mut offset = 0usize;
250        while offset + HEADER_SIZE <= mmap.len() {
251            let header = &mmap[offset..offset + HEADER_SIZE];
252            if is_log_end(header) {
253                break;
254            }
255            let raw_len = u32::from_le_bytes(header[4..8].try_into()?) as usize;
256            let comp_len = u32::from_le_bytes(header[8..12].try_into()?) as usize;
257            let frame_end = offset + HEADER_SIZE + comp_len;
258            ensure!(
259                frame_end <= mmap.len(),
260                "truncated frame in {} at byte {offset}",
261                segment_path.display()
262            );
263            let payload = &mmap[offset + HEADER_SIZE..frame_end];
264            let decoded = if comp_len < raw_len {
265                lz4_flex::decompress(payload, raw_len)?
266            } else {
267                payload.to_vec()
268            };
269            let json = super::engine::journal::payload_json(&decoded)
270                .with_context(|| {
271                    format!(
272                        "frame at byte {offset} of {} is not a transaction frame",
273                        segment_path.display()
274                    )
275                })?
276                .with_context(|| {
277                    format!(
278                        "frame at byte {offset} of {} is not a transaction",
279                        segment_path.display()
280                    )
281                })?;
282            let mutations: Vec<Value> = serde_json::from_slice(&json).with_context(|| {
283                format!(
284                    "frame at byte {offset} of {} is not a transaction",
285                    segment_path.display()
286                )
287            })?;
288            for mutation in &mutations {
289                let id_of = |variant: &str| {
290                    mutation
291                        .get(variant)
292                        .and_then(|m| m.get("id"))
293                        .and_then(Value::as_str)
294                };
295                if let Some(id) = id_of("Document") {
296                    live.insert(id.to_string());
297                } else if let Some(id) = id_of("DeleteDocument") {
298                    live.remove(id);
299                }
300            }
301            offset = frame_end;
302        }
303    }
304    Ok(live)
305}
306
307fn missing_documents(expected: &BTreeSet<String>, core: &HmsCore) -> Vec<String> {
308    expected
309        .iter()
310        .filter(|id| !core.has_document(id))
311        .cloned()
312        .collect()
313}
314
315pub(crate) fn inspect_segment(
316    data: &[u8],
317    path: &Path,
318    verify_checksums: bool,
319) -> Result<(usize, usize)> {
320    let mut offset = 0usize;
321    let mut frames = 0usize;
322    while offset + HEADER_SIZE <= data.len() {
323        let expected_crc = u32::from_le_bytes(data[offset..offset + 4].try_into()?);
324        let raw_len = u32::from_le_bytes(data[offset + 4..offset + 8].try_into()?) as usize;
325        let comp_len = u32::from_le_bytes(data[offset + 8..offset + 12].try_into()?) as usize;
326        if raw_len == 0 && comp_len == 0 {
327            if is_log_end(&data[offset..offset + HEADER_SIZE]) {
328                return Ok((frames, offset));
329            }
330            bail!(
331                "invalid end-of-log header in {} at byte {}",
332                path.display(),
333                offset
334            );
335        }
336        if raw_len == 0 || raw_len > MAX_RAW_FRAME_SIZE || comp_len == 0 || comp_len > data.len() {
337            bail!(
338                "invalid frame header in {} at byte {}",
339                path.display(),
340                offset
341            );
342        }
343        let frame_end = offset
344            .checked_add(HEADER_SIZE + comp_len)
345            .ok_or_else(|| anyhow!("frame length overflow in {}", path.display()))?;
346        if frame_end > data.len() {
347            bail!("truncated frame in {} at byte {}", path.display(), offset);
348        }
349        if verify_checksums {
350            let payload = &data[offset + HEADER_SIZE..frame_end];
351            let decoded = if comp_len < raw_len {
352                lz4_flex::decompress(payload, raw_len).with_context(|| {
353                    format!("invalid LZ4 frame in {} at byte {}", path.display(), offset)
354                })?
355            } else {
356                payload.to_vec()
357            };
358            let mut hasher = Hasher::new();
359            hasher.update(&decoded);
360            if hasher.finalize() != expected_crc {
361                bail!("checksum mismatch in {} at byte {}", path.display(), offset);
362            }
363        }
364        frames += 1;
365        offset = frame_end;
366    }
367    Ok((frames, offset))
368}
369
370/// Create a verified, point-in-time copy of a current-format store.
371/// The source is exclusively locked for the full copy and verification. The
372/// destination appears only after verification succeeds.
373pub fn migrate_store(
374    source: impl AsRef<Path>,
375    destination: impl AsRef<Path>,
376) -> Result<StoreInspection> {
377    let source = source.as_ref();
378    let destination = destination.as_ref();
379    if destination.exists() {
380        bail!(
381            "migration destination already exists: {}",
382            destination.display()
383        );
384    }
385    let source_canonical = source.canonicalize()?;
386    let parent = destination.parent().unwrap_or_else(|| Path::new("."));
387    std::fs::create_dir_all(parent)?;
388    let parent_canonical = parent.canonicalize()?;
389    if parent_canonical.starts_with(&source_canonical) {
390        bail!("migration destination cannot be inside the source store");
391    }
392
393    let _source_lock = StoreLock::acquire(source)?;
394    // Verify after acquiring the lock so the exact snapshot being copied is
395    // the snapshot that passed integrity validation.
396    inspect_store(source, true)?;
397    let nonce = SystemTime::now().duration_since(UNIX_EPOCH)?.as_nanos();
398    let temp = parent.join(format!(".hms-migrate-{}-{nonce}", std::process::id()));
399    let result = (|| -> Result<StoreInspection> {
400        copy_tree(source, &temp)?;
401        let inspection = inspect_store(&temp, true)?;
402        std::fs::rename(&temp, destination).with_context(|| {
403            format!(
404                "failed to atomically publish migration {} -> {}",
405                temp.display(),
406                destination.display()
407            )
408        })?;
409        Ok(StoreInspection {
410            store_path: destination.display().to_string(),
411            ..inspection
412        })
413    })();
414    if result.is_err() && temp.exists() {
415        let _ = std::fs::remove_dir_all(&temp);
416    }
417    result
418}
419
420fn copy_tree(source: &Path, destination: &Path) -> Result<()> {
421    std::fs::create_dir(destination)?;
422    for entry in std::fs::read_dir(source)? {
423        let entry = entry?;
424        let file_type = entry.file_type()?;
425        if file_type.is_symlink() {
426            bail!("migration refuses symbolic link {}", entry.path().display());
427        }
428        if entry.file_name() == LOCK_FILE {
429            continue;
430        }
431        let target = destination.join(entry.file_name());
432        if file_type.is_dir() {
433            copy_tree(&entry.path(), &target)?;
434        } else if file_type.is_file() {
435            std::fs::copy(entry.path(), &target)?;
436            std::fs::set_permissions(&target, entry.metadata()?.permissions())?;
437        } else {
438            bail!("migration refuses special file {}", entry.path().display());
439        }
440    }
441    Ok(())
442}
443
444/// Re-encode original JSONL source documents into a new, atomically published store.
445pub fn reencode_documents(
446    input: impl AsRef<Path>,
447    destination: impl AsRef<Path>,
448    dimensions: u32,
449    config: super::HmsConfig,
450) -> Result<StoreInspection> {
451    use std::io::{BufRead, BufReader, Read};
452    let destination = destination.as_ref();
453    if destination.exists() {
454        bail!("re-encoding destination already exists");
455    }
456    let parent = destination
457        .parent()
458        .filter(|p| !p.as_os_str().is_empty())
459        .unwrap_or_else(|| Path::new("."));
460    std::fs::create_dir_all(parent)?;
461    let temporary = tempfile::Builder::new()
462        .prefix(".hms-reencode-")
463        .tempdir_in(parent)?;
464    let mut source = BufReader::new(File::open(input.as_ref())?);
465    {
466        let hms = super::HmsCore::new(
467            dimensions,
468            Some(temporary.path().display().to_string()),
469            Some(config),
470        )?;
471        let mut line_number = 0;
472        loop {
473            let mut line = String::new();
474            let bytes = source
475                .by_ref()
476                .take(16 * 1024 * 1024 + 1)
477                .read_line(&mut line)?;
478            if bytes == 0 {
479                break;
480            }
481            line_number += 1;
482            if bytes > 16 * 1024 * 1024 {
483                bail!("JSONL line {line_number} exceeds 16 MiB");
484            }
485            if line.trim().is_empty() {
486                continue;
487            }
488            let document: super::documents::DocumentInput = serde_json::from_str(&line)
489                .with_context(|| format!("invalid source document on line {line_number}"))?;
490            hms.memorize_document(document)
491                .with_context(|| format!("cannot re-encode line {line_number}"))?;
492        }
493        hms.flush()?;
494    }
495    let inspection = inspect_store(temporary.path(), true)?;
496    super::durable_file::sync_directory(temporary.path())?;
497    std::fs::rename(temporary.path(), destination)?;
498    super::durable_file::sync_directory(parent)?;
499    Ok(StoreInspection {
500        store_path: destination.display().to_string(),
501        ..inspection
502    })
503}
504
505#[cfg(test)]
506mod tests {
507    use super::*;
508    use crate::core::documents::DocumentInput;
509    use crate::core::entangled::EntangledHVec;
510    use crate::core::HmsCore;
511
512    /// `verify` cross-checks the opened store against the log: a healthy store reports every
513    /// live document present, and a logged document the store cannot resolve is named.
514    #[test]
515    fn verify_cross_checks_logged_documents() -> Result<()> {
516        let root = tempfile::tempdir()?;
517        let store = root.path().join("store");
518        {
519            let hms = HmsCore::new(256, Some(store.display().to_string()), None)?;
520            for n in 0..3 {
521                hms.memorize_document(DocumentInput {
522                    id: format!("doc{n}"),
523                    text: format!("document {n} about harbors and ledgers"),
524                    source_uri: None,
525                    version: None,
526                    metadata: None,
527                    chunk_words: None,
528                    overlap_words: None,
529                    store_text: None,
530                    embeddings: None,
531                    title_weight: None,
532                })?;
533            }
534            assert!(hms.delete_document("doc1")?);
535            hms.flush()?;
536        }
537        let verified = verify_store(&store, None)?;
538        assert!(verified
539            .manifests_verified
540            .iter()
541            .any(|m| m == "store.json"));
542        assert!(verified
543            .manifests_verified
544            .iter()
545            .any(|m| m == "format.json"));
546        let documents = verified
547            .documents
548            .expect("plaintext stores are cross-checked");
549        assert_eq!(
550            (documents.checked, documents.expected, documents.present),
551            (true, Some(2), 2)
552        );
553        assert!(documents.missing.is_empty());
554
555        let core = HmsCore::new(256, Some(store.display().to_string()), None)?;
556        let expected: BTreeSet<String> = ["doc0", "doc2", "doc9"]
557            .into_iter()
558            .map(String::from)
559            .collect();
560        assert_eq!(
561            missing_documents(&expected, &core),
562            vec!["doc9".to_string()]
563        );
564        Ok(())
565    }
566
567    #[test]
568    fn inspect_verify_and_migrate_roundtrip() -> Result<()> {
569        let root = tempfile::tempdir()?;
570        let source = root.path().join("source");
571        {
572            let hms = HmsCore::new(256, Some(source.display().to_string()), None)?;
573            hms.memorize("one".into(), EntangledHVec::new_deterministic(256, 1))?;
574        }
575
576        let inspected = inspect_store(&source, false)?;
577        assert_eq!(inspected.frame_count, 1);
578        assert!(!inspected.checksums_verified);
579        let verified = inspect_store(&source, true)?;
580        assert!(verified.checksums_verified);
581        let destination = root.path().join("destination");
582        let migrated = migrate_store(&source, &destination)?;
583        assert_eq!(migrated.frame_count, 1);
584        let reopened = HmsCore::new(256, Some(destination.display().to_string()), None)?;
585        assert_eq!(reopened.index_status().vector_count, 1);
586        Ok(())
587    }
588
589    #[test]
590    fn verification_rejects_corruption() -> Result<()> {
591        use std::fs::OpenOptions;
592        use std::io::{Seek, SeekFrom, Write};
593        let root = tempfile::tempdir()?;
594        let source = root.path().join("source");
595        {
596            let hms = HmsCore::new(256, Some(source.display().to_string()), None)?;
597            hms.memorize("one".into(), EntangledHVec::new_deterministic(256, 1))?;
598        }
599        let segment = source.join(ARENA_DIRECTORY).join("seg_0.bin");
600        let mut file = OpenOptions::new().write(true).open(segment)?;
601        file.seek(SeekFrom::Start(HEADER_SIZE as u64))?;
602        file.write_all(&[0xff])?;
603        file.sync_all()?;
604        assert!(inspect_store(&source, true).is_err());
605        Ok(())
606    }
607}