1use 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 #[serde(default, skip_serializing_if = "Vec::is_empty")]
35 pub manifests_verified: Vec<String>,
36 #[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 pub checked: bool,
46 #[serde(skip_serializing_if = "Option::is_none")]
47 pub reason: Option<String>,
48 #[serde(skip_serializing_if = "Option::is_none")]
50 pub expected: Option<usize>,
51 pub present: usize,
52 pub missing: Vec<String>,
53}
54
55pub 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
157pub 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
238fn 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
370pub 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 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
444pub 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 #[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}