1#![cfg(not(target_arch = "wasm32"))]
49
50use crate::storage::compression::{self, Codec, CompressionStats};
51use crate::{Document, RagError, Result};
52use serde::{Deserialize, Serialize};
53use std::fs::{self, File};
54use std::io::{Read, Write};
55use std::path::{Component, Path, PathBuf};
56use std::sync::atomic::{AtomicU64, Ordering};
57use std::time::{SystemTime, UNIX_EPOCH};
58
59const STORAGE_VERSION: u32 = 2;
60const DATA_EXTENSION: &str = "data";
61const META_EXTENSION: &str = "meta";
62const TMP_EXTENSION: &str = "tmp";
63static TMP_FILE_COUNTER: AtomicU64 = AtomicU64::new(0);
64
65#[derive(Debug, Clone, Serialize, Deserialize)]
70pub struct StorageMetadata {
71 pub version: u32,
73 pub created_at: u64,
75 pub updated_at: u64,
77 pub item_type: String,
79 pub compression: Codec,
81 pub original_size: usize,
83 pub compressed_size: usize,
85}
86
87impl StorageMetadata {
88 fn new(
90 item_type: String,
91 compression: Codec,
92 original_size: usize,
93 compressed_size: usize,
94 ) -> Self {
95 let now = SystemTime::now()
96 .duration_since(UNIX_EPOCH)
97 .unwrap()
98 .as_secs();
99
100 Self {
101 version: STORAGE_VERSION,
102 created_at: now,
103 updated_at: now,
104 item_type,
105 compression,
106 original_size,
107 compressed_size,
108 }
109 }
110
111 fn touch(&mut self) {
113 self.updated_at = SystemTime::now()
114 .duration_since(UNIX_EPOCH)
115 .unwrap()
116 .as_secs();
117 }
118}
119
120#[derive(Debug)]
135pub struct FileStorage {
136 base_path: PathBuf,
137 codec: Codec,
138}
139
140impl FileStorage {
141 pub fn new(base_path: impl AsRef<Path>) -> Result<Self> {
164 Self::with_codec(base_path, Codec::None)
165 }
166
167 pub fn with_codec(base_path: impl AsRef<Path>, codec: Codec) -> Result<Self> {
190 let base_path = base_path.as_ref().to_path_buf();
191
192 if !base_path.exists() {
194 fs::create_dir_all(&base_path).map_err(|e| {
195 RagError::StorageError(format!("Failed to create storage directory: {}", e))
196 })?;
197 }
198
199 if !base_path.is_dir() {
201 return Err(RagError::StorageError(format!(
202 "Storage path is not a directory: {}",
203 base_path.display()
204 )));
205 }
206
207 Ok(Self { base_path, codec })
208 }
209
210 pub fn save_document(&self, id: &str, document: &Document) -> Result<CompressionStats> {
244 Self::validate_item_name(id)?;
245
246 let serialized = serde_json::to_vec(document)
248 .map_err(|e| RagError::StorageError(format!("JSON serialization failed: {}", e)))?;
249
250 let (compressed, stats) = compression::compress_with(&serialized, self.codec)
252 .map_err(|e| RagError::StorageError(format!("Compression failed: {}", e)))?;
253
254 let metadata = if self.exists(id) {
256 let mut meta = self.get_metadata(id)?;
257 meta.touch();
258 meta.original_size = stats.original_size;
259 meta.compressed_size = stats.compressed_size;
260 meta.compression = stats.codec;
261 meta
262 } else {
263 StorageMetadata::new(
264 "document".to_string(),
265 stats.codec,
266 stats.original_size,
267 stats.compressed_size,
268 )
269 };
270
271 let data_path = self.item_path(id);
273 self.write_atomic(&data_path, &compressed)?;
274
275 let meta_path = self.metadata_path(id);
277 let meta_bytes = serde_json::to_vec(&metadata)
278 .map_err(|e| RagError::StorageError(format!("metadata serialize failed: {}", e)))?;
279 self.write_atomic(&meta_path, &meta_bytes)?;
280
281 Ok(stats)
282 }
283
284 pub fn load_document(&self, id: &str) -> Result<Document> {
310 Self::validate_item_name(id)?;
311
312 if !self.exists(id) {
314 return Err(RagError::StorageError(format!(
315 "Document not found: {}",
316 id
317 )));
318 }
319
320 let metadata = self.get_metadata(id)?;
322
323 if metadata.version != STORAGE_VERSION {
325 return Err(RagError::StorageError(format!(
326 "Incompatible storage version: expected {}, got {}",
327 STORAGE_VERSION, metadata.version
328 )));
329 }
330
331 let data_path = self.item_path(id);
333 let mut file = File::open(&data_path)?;
334 let mut compressed = Vec::new();
335 file.read_to_end(&mut compressed)?;
336
337 if compressed.len() != metadata.compressed_size {
339 return Err(RagError::StorageError(format!(
340 "Data corruption detected: size mismatch for {}",
341 id
342 )));
343 }
344
345 let decompressed = compression::decompress(&compressed)
347 .map_err(|e| RagError::StorageError(format!("Decompression failed: {}", e)))?;
348
349 let document: Document = serde_json::from_slice(&decompressed)
351 .map_err(|e| RagError::StorageError(format!("JSON deserialization failed: {}", e)))?;
352
353 Ok(document)
354 }
355
356 pub fn save_flat_index(
371 &self,
372 name: &str,
373 index: &FlatIndexWrapper,
374 ) -> Result<CompressionStats> {
375 Self::validate_item_name(name)?;
376 self.save_with_metadata(name, index, "flat_index")
377 }
378
379 pub fn load_flat_index(&self, name: &str) -> Result<FlatIndexWrapper> {
393 Self::validate_item_name(name)?;
394 self.load_with_metadata(name)
395 }
396
397 pub fn save_hnsw_index(
412 &self,
413 name: &str,
414 index: &HNSWIndexWrapper,
415 ) -> Result<CompressionStats> {
416 Self::validate_item_name(name)?;
417 self.save_with_metadata(name, index, "hnsw_index")
418 }
419
420 pub fn load_hnsw_index(&self, name: &str) -> Result<HNSWIndexWrapper> {
434 Self::validate_item_name(name)?;
435 self.load_with_metadata(name)
436 }
437
438 pub fn delete(&self, name: &str) -> Result<()> {
465 Self::validate_item_name(name)?;
466
467 let data_path = self.item_path(name);
468 let meta_path = self.metadata_path(name);
469
470 if data_path.exists() {
472 fs::remove_file(&data_path).map_err(|e| {
473 RagError::StorageError(format!("Failed to delete data file: {}", e))
474 })?;
475 }
476
477 if meta_path.exists() {
479 fs::remove_file(&meta_path).map_err(|e| {
480 RagError::StorageError(format!("Failed to delete metadata file: {}", e))
481 })?;
482 }
483
484 Ok(())
485 }
486
487 pub fn list(&self) -> Result<Vec<String>> {
513 let entries = fs::read_dir(&self.base_path).map_err(|e| {
514 RagError::StorageError(format!("Failed to read storage directory: {}", e))
515 })?;
516
517 let mut names = std::collections::HashSet::new();
518
519 for entry in entries {
520 let entry = entry.map_err(|e| {
521 RagError::StorageError(format!("Failed to read directory entry: {}", e))
522 })?;
523
524 let path = entry.path();
525 if path.is_file() {
526 if let Some(ext) = path.extension() {
527 if ext == DATA_EXTENSION || ext == META_EXTENSION {
528 if let Some(stem) = path.file_stem() {
529 if let Some(name) = stem.to_str() {
530 names.insert(name.to_string());
531 }
532 }
533 }
534 }
535 }
536 }
537
538 let mut result: Vec<String> = names.into_iter().collect();
539 result.sort();
540 Ok(result)
541 }
542
543 pub fn get_metadata(&self, name: &str) -> Result<StorageMetadata> {
569 Self::validate_item_name(name)?;
570
571 let meta_path = self.metadata_path(name);
572
573 if !meta_path.exists() {
574 return Err(RagError::StorageError(format!(
575 "Metadata not found for item: {}",
576 name
577 )));
578 }
579
580 let mut file = File::open(&meta_path)?;
581 let mut contents = Vec::new();
582 file.read_to_end(&mut contents)?;
583
584 let metadata: StorageMetadata = serde_json::from_slice(&contents)
586 .or_else(|_| bincode::deserialize::<StorageMetadata>(&contents))
587 .map_err(|e| RagError::StorageError(format!("metadata deserialize failed: {}", e)))?;
588 Ok(metadata)
589 }
590
591 pub fn total_size(&self) -> Result<u64> {
615 let entries = fs::read_dir(&self.base_path).map_err(|e| {
616 RagError::StorageError(format!("Failed to read storage directory: {}", e))
617 })?;
618
619 let mut total = 0u64;
620
621 for entry in entries {
622 let entry = entry.map_err(|e| {
623 RagError::StorageError(format!("Failed to read directory entry: {}", e))
624 })?;
625
626 let metadata = entry.metadata()?;
627 if metadata.is_file() {
628 total += metadata.len();
629 }
630 }
631
632 Ok(total)
633 }
634
635 pub fn clear(&self) -> Result<()> {
658 let entries = fs::read_dir(&self.base_path).map_err(|e| {
659 RagError::StorageError(format!("Failed to read storage directory: {}", e))
660 })?;
661
662 for entry in entries {
663 let entry = entry.map_err(|e| {
664 RagError::StorageError(format!("Failed to read directory entry: {}", e))
665 })?;
666
667 let path = entry.path();
668 if path.is_file() {
669 fs::remove_file(&path)
670 .map_err(|e| RagError::StorageError(format!("Failed to delete file: {}", e)))?;
671 }
672 }
673
674 Ok(())
675 }
676
677 pub fn exists(&self, name: &str) -> bool {
700 if Self::is_invalid_item_name(name) {
701 return false;
702 }
703 self.item_path(name).exists() && self.metadata_path(name).exists()
704 }
705
706 fn is_invalid_item_name(name: &str) -> bool {
709 if name.is_empty() {
710 return true;
711 }
712
713 if name.contains('\0') {
715 return true;
716 }
717
718 if name.contains('/') || name.contains('\\') {
721 return true;
722 }
723
724 let path = Path::new(name);
725 if path.is_absolute() {
726 return true;
727 }
728
729 let mut components = path.components();
731 match components.next() {
732 Some(Component::Normal(_)) => {
733 if components.next().is_some() {
734 return true;
735 }
736 }
737 _ => return true,
738 }
739
740 let base_name = name.split('.').next().unwrap_or(name);
745 let stem_upper = base_name.to_ascii_uppercase();
746 let is_reserved = matches!(
747 stem_upper.as_str(),
748 "CON"
749 | "PRN"
750 | "AUX"
751 | "NUL"
752 | "COM1"
753 | "COM2"
754 | "COM3"
755 | "COM4"
756 | "COM5"
757 | "COM6"
758 | "COM7"
759 | "COM8"
760 | "COM9"
761 | "LPT1"
762 | "LPT2"
763 | "LPT3"
764 | "LPT4"
765 | "LPT5"
766 | "LPT6"
767 | "LPT7"
768 | "LPT8"
769 | "LPT9"
770 );
771 if is_reserved {
772 return true;
773 }
774
775 false
776 }
777
778 fn validate_item_name(name: &str) -> Result<()> {
779 if Self::is_invalid_item_name(name) {
780 return Err(RagError::StorageError(format!(
781 "Invalid item name: '{}'. Names must be a single path segment",
782 name
783 )));
784 }
785 Ok(())
786 }
787
788 fn item_path(&self, name: &str) -> PathBuf {
790 self.base_path.join(format!("{}.{}", name, DATA_EXTENSION))
791 }
792
793 fn metadata_path(&self, name: &str) -> PathBuf {
795 self.base_path.join(format!("{}.{}", name, META_EXTENSION))
796 }
797
798 fn write_atomic(&self, path: &Path, data: &[u8]) -> Result<()> {
816 let filename = path.file_name().and_then(|f| f.to_str()).unwrap_or("item");
818 let counter = TMP_FILE_COUNTER.fetch_add(1, Ordering::Relaxed);
819 let tmp_path = path.with_file_name(format!(
820 "{}.{}.{}.{}",
821 filename,
822 std::process::id(),
823 counter,
824 TMP_EXTENSION
825 ));
826
827 {
829 let mut file = File::create(&tmp_path)?;
830 file.write_all(data)?;
831 file.sync_all()?; }
833
834 fs::rename(&tmp_path, path).map_err(|e| {
836 let _ = fs::remove_file(&tmp_path);
838 RagError::IoError(e)
839 })?;
840
841 Ok(())
842 }
843
844 fn save_with_metadata<T: Serialize>(
848 &self,
849 name: &str,
850 item: &T,
851 item_type: &str,
852 ) -> Result<CompressionStats> {
853 let serialized = serde_json::to_vec(item)
855 .map_err(|e| RagError::StorageError(format!("JSON serialization failed: {}", e)))?;
856
857 let (compressed, stats) = compression::compress_with(&serialized, self.codec)
859 .map_err(|e| RagError::StorageError(format!("Compression failed: {}", e)))?;
860
861 let metadata = if self.exists(name) {
863 let mut meta = self.get_metadata(name)?;
864 meta.touch();
865 meta.original_size = stats.original_size;
866 meta.compressed_size = stats.compressed_size;
867 meta.compression = stats.codec;
868 meta
869 } else {
870 StorageMetadata::new(
871 item_type.to_string(),
872 stats.codec,
873 stats.original_size,
874 stats.compressed_size,
875 )
876 };
877
878 let data_path = self.item_path(name);
880 self.write_atomic(&data_path, &compressed)?;
881
882 let meta_path = self.metadata_path(name);
884 let meta_bytes = serde_json::to_vec(&metadata)
885 .map_err(|e| RagError::StorageError(format!("metadata serialize failed: {}", e)))?;
886 self.write_atomic(&meta_path, &meta_bytes)?;
887
888 Ok(stats)
889 }
890
891 fn load_with_metadata<T: for<'de> Deserialize<'de>>(&self, name: &str) -> Result<T> {
895 if !self.exists(name) {
897 return Err(RagError::StorageError(format!("Item not found: {}", name)));
898 }
899
900 let metadata = self.get_metadata(name)?;
902
903 if metadata.version != STORAGE_VERSION {
905 return Err(RagError::StorageError(format!(
906 "Incompatible storage version: expected {}, got {}",
907 STORAGE_VERSION, metadata.version
908 )));
909 }
910
911 let data_path = self.item_path(name);
913 let mut file = File::open(&data_path)?;
914 let mut compressed = Vec::new();
915 file.read_to_end(&mut compressed)?;
916
917 if compressed.len() != metadata.compressed_size {
919 return Err(RagError::StorageError(format!(
920 "Data corruption detected: size mismatch for {}",
921 name
922 )));
923 }
924
925 let decompressed = compression::decompress(&compressed)
927 .map_err(|e| RagError::StorageError(format!("Decompression failed: {}", e)))?;
928
929 let item: T = serde_json::from_slice::<T>(&decompressed)
931 .map_err(|e| RagError::StorageError(format!("JSON deserialization failed: {}", e)))?;
932
933 Ok(item)
934 }
935}
936
937#[derive(Debug, Clone, Serialize, Deserialize)]
942pub struct FlatIndexWrapper {
943 pub embedding_dim: usize,
944 pub documents: Vec<Document>,
945}
946
947impl FlatIndexWrapper {
948 pub fn from_index(index: &crate::index::FlatIndex) -> Self {
950 Self {
951 embedding_dim: index.embedding_dim(),
952 documents: index.get_all_documents(),
953 }
954 }
955
956 pub fn to_index(&self) -> Result<crate::index::FlatIndex> {
958 let mut index = crate::index::FlatIndex::new(self.embedding_dim);
959 index.add_batch(self.documents.clone())?;
960 Ok(index)
961 }
962}
963
964#[derive(Debug, Clone, Serialize, Deserialize)]
969pub struct HNSWIndexWrapper {
970 pub embedding_dim: usize,
971 pub documents: Vec<Document>,
972 pub config: HNSWConfigWrapper,
973}
974
975#[derive(Debug, Clone, Serialize, Deserialize)]
977pub struct HNSWConfigWrapper {
978 #[serde(default)]
981 pub metric: crate::index::hnsw::DistanceMetric,
982 #[serde(default)]
985 pub storage: crate::index::hnsw::Storage,
986 #[serde(default = "default_rerank_candidates")]
987 pub rerank_candidates: usize,
988 pub m: usize,
989 pub m0: usize,
990 pub ef_construction: usize,
991 pub ef_search: usize,
992 pub ml: f32,
993 #[serde(default = "default_use_heuristic")]
994 pub use_heuristic: bool,
995 #[serde(default)]
996 pub extend_candidates: bool,
997 #[serde(default = "default_keep_pruned")]
998 pub keep_pruned_connections: bool,
999 #[serde(default)]
1007 pub seed: Option<u64>,
1008}
1009
1010fn default_use_heuristic() -> bool {
1011 true
1012}
1013fn default_keep_pruned() -> bool {
1014 true
1015}
1016
1017impl From<&crate::index::HNSWConfig> for HNSWConfigWrapper {
1018 fn from(config: &crate::index::HNSWConfig) -> Self {
1019 Self {
1020 m: config.m,
1021 m0: config.m0,
1022 ef_construction: config.ef_construction,
1023 ef_search: config.ef_search,
1024 ml: config.ml,
1025 use_heuristic: config.use_heuristic,
1026 extend_candidates: config.extend_candidates,
1027 keep_pruned_connections: config.keep_pruned_connections,
1028 metric: config.metric,
1029 storage: config.storage,
1030 rerank_candidates: config.rerank_candidates,
1031 seed: config.seed,
1032 }
1033 }
1034}
1035
1036fn default_rerank_candidates() -> usize {
1037 100
1038}
1039
1040impl From<HNSWConfigWrapper> for crate::index::HNSWConfig {
1041 fn from(wrapper: HNSWConfigWrapper) -> Self {
1042 Self {
1043 metric: wrapper.metric,
1044 m: wrapper.m,
1045 m0: wrapper.m0,
1046 ef_construction: wrapper.ef_construction,
1047 ef_search: wrapper.ef_search,
1048 ml: wrapper.ml,
1049 use_heuristic: wrapper.use_heuristic,
1050 extend_candidates: wrapper.extend_candidates,
1051 keep_pruned_connections: wrapper.keep_pruned_connections,
1052 storage: wrapper.storage,
1053 rerank_candidates: wrapper.rerank_candidates,
1054 seed: wrapper.seed,
1055 build_strategy: crate::index::BuildStrategy::default(),
1062 }
1063 }
1064}
1065
1066impl HNSWIndexWrapper {
1067 pub fn from_index(index: &crate::index::HNSWIndex) -> Self {
1069 Self {
1070 embedding_dim: index.embedding_dim(),
1071 documents: index.get_all_documents(),
1072 config: HNSWConfigWrapper::from(index.config()),
1073 }
1074 }
1075
1076 pub fn to_index(&self) -> Result<crate::index::HNSWIndex> {
1078 let config: crate::index::HNSWConfig = self.config.clone().into();
1079 let quantized = config.storage != crate::index::Storage::F32;
1080 let mut index = crate::index::HNSWIndex::new(self.embedding_dim, config);
1081
1082 if quantized {
1093 let sample: Vec<Vec<f32>> =
1094 self.documents.iter().map(|d| d.embedding.clone()).collect();
1095 index.train(&sample)?;
1096 }
1097
1098 for doc in &self.documents {
1099 index.add(doc.clone())?;
1100 }
1101 Ok(index)
1102 }
1103}
1104
1105#[cfg(test)]
1106mod tests {
1107 use super::*;
1108 use std::sync::{Arc, Barrier};
1109 use std::thread;
1110 use tempfile::tempdir;
1111
1112 fn create_test_document(id: &str) -> Document {
1113 Document {
1114 id: id.to_string(),
1115 content: format!("Test content for {}", id),
1116 embedding: vec![0.1, 0.2, 0.3, 0.4, 0.5],
1117 metadata: Some(serde_json::json!({"test": true})),
1118 }
1119 }
1120
1121 fn create_test_flat_index() -> crate::index::FlatIndex {
1122 let mut index = crate::index::FlatIndex::new(5);
1123 index.add(create_test_document("doc1")).unwrap();
1124 index.add(create_test_document("doc2")).unwrap();
1125 index
1126 }
1127
1128 fn create_test_hnsw_index() -> crate::index::HNSWIndex {
1129 let mut index = crate::index::HNSWIndex::with_defaults(5);
1130 index.add(create_test_document("doc1")).unwrap();
1131 index.add(create_test_document("doc2")).unwrap();
1132 index
1133 }
1134
1135 #[test]
1136 fn test_new_storage() {
1137 let dir = tempdir().unwrap();
1138 let _storage = FileStorage::new(dir.path()).unwrap();
1139 assert!(dir.path().exists());
1140 assert!(dir.path().is_dir());
1141 }
1142
1143 #[test]
1144 fn test_new_storage_with_codec() {
1145 let dir = tempdir().unwrap();
1146 let _storage = FileStorage::with_codec(dir.path(), Codec::Gzip).unwrap();
1147 assert!(dir.path().exists());
1148 }
1149
1150 #[test]
1151 fn test_invalid_storage_path() {
1152 let dir = tempdir().unwrap();
1153 let file_path = dir.path().join("file.txt");
1154 std::fs::write(&file_path, b"test").unwrap();
1155
1156 let result = FileStorage::new(&file_path);
1157 assert!(result.is_err());
1158 }
1159
1160 #[test]
1161 fn test_document_save_load() {
1162 let dir = tempdir().unwrap();
1163 let storage = FileStorage::new(dir.path()).unwrap();
1164
1165 let doc = create_test_document("doc1");
1166 let stats = storage.save_document("doc1", &doc).unwrap();
1167
1168 assert!(stats.original_size > 0);
1169 assert_eq!(stats.codec, Codec::None);
1170
1171 let loaded = storage.load_document("doc1").unwrap();
1172 assert_eq!(loaded.id, doc.id);
1173 assert_eq!(loaded.content, doc.content);
1174 assert_eq!(loaded.embedding, doc.embedding);
1175 }
1176
1177 #[test]
1178 fn test_document_not_found() {
1179 let dir = tempdir().unwrap();
1180 let storage = FileStorage::new(dir.path()).unwrap();
1181
1182 let result = storage.load_document("nonexistent");
1183 assert!(result.is_err());
1184 }
1185
1186 #[test]
1187 fn test_flat_index_persistence() {
1188 let dir = tempdir().unwrap();
1189 let storage = FileStorage::new(dir.path()).unwrap();
1190
1191 let index = create_test_flat_index();
1192 let wrapper = FlatIndexWrapper::from_index(&index);
1193
1194 let stats = storage.save_flat_index("index1", &wrapper).unwrap();
1195 assert!(stats.original_size > 0);
1196
1197 let loaded_wrapper = storage.load_flat_index("index1").unwrap();
1198 let loaded_index = loaded_wrapper.to_index().unwrap();
1199
1200 assert_eq!(loaded_index.len(), index.len());
1201 assert_eq!(loaded_index.embedding_dim(), index.embedding_dim());
1202 }
1203
1204 #[test]
1205 fn test_hnsw_index_persistence() {
1206 let dir = tempdir().unwrap();
1207 let storage = FileStorage::new(dir.path()).unwrap();
1208
1209 let index = create_test_hnsw_index();
1210 let wrapper = HNSWIndexWrapper::from_index(&index);
1211
1212 let stats = storage.save_hnsw_index("index1", &wrapper).unwrap();
1213 assert!(stats.original_size > 0);
1214
1215 let loaded_wrapper = storage.load_hnsw_index("index1").unwrap();
1216 let loaded_index = loaded_wrapper.to_index().unwrap();
1217
1218 assert_eq!(loaded_index.len(), index.len());
1219 assert_eq!(loaded_index.embedding_dim(), index.embedding_dim());
1220 }
1221
1222 #[test]
1223 fn test_atomic_write() {
1224 let dir = tempdir().unwrap();
1225 let storage = FileStorage::new(dir.path()).unwrap();
1226
1227 let path = dir.path().join("test.data");
1228 let data = b"test data";
1229
1230 storage.write_atomic(&path, data).unwrap();
1231
1232 assert!(path.exists());
1233 let read_data = std::fs::read(&path).unwrap();
1234 assert_eq!(read_data, data);
1235
1236 let has_tmp = std::fs::read_dir(dir.path())
1238 .unwrap()
1239 .filter_map(|entry| entry.ok())
1240 .map(|entry| entry.file_name().to_string_lossy().to_string())
1241 .any(|name| name.ends_with(".tmp"));
1242 assert!(!has_tmp);
1243 }
1244
1245 #[test]
1246 fn concurrent_atomic_writes_to_sibling_paths_do_not_cross_contaminate() {
1247 let dir = tempdir().unwrap();
1248 let storage = Arc::new(FileStorage::new(dir.path()).unwrap());
1249 let data_path = dir.path().join("doc.data");
1250 let meta_path = dir.path().join("doc.meta");
1251
1252 for _ in 0..128 {
1253 let barrier = Arc::new(Barrier::new(3));
1254 let s1 = Arc::clone(&storage);
1255 let b1 = Arc::clone(&barrier);
1256 let data_path_1 = data_path.clone();
1257 let t1 = thread::spawn(move || {
1258 b1.wait();
1259 s1.write_atomic(&data_path_1, b"DATA").unwrap();
1260 });
1261
1262 let s2 = Arc::clone(&storage);
1263 let b2 = Arc::clone(&barrier);
1264 let meta_path_1 = meta_path.clone();
1265 let t2 = thread::spawn(move || {
1266 b2.wait();
1267 s2.write_atomic(&meta_path_1, b"META").unwrap();
1268 });
1269
1270 barrier.wait();
1271 t1.join().unwrap();
1272 t2.join().unwrap();
1273
1274 assert_eq!(std::fs::read(&data_path).unwrap(), b"DATA");
1275 assert_eq!(std::fs::read(&meta_path).unwrap(), b"META");
1276 }
1277 }
1278
1279 #[test]
1280 fn test_metadata() {
1281 let dir = tempdir().unwrap();
1282 let storage = FileStorage::new(dir.path()).unwrap();
1283
1284 let doc = create_test_document("doc1");
1285 storage.save_document("doc1", &doc).unwrap();
1286
1287 let metadata = storage.get_metadata("doc1").unwrap();
1288 assert_eq!(metadata.version, STORAGE_VERSION);
1289 assert_eq!(metadata.item_type, "document");
1290 assert!(metadata.created_at > 0);
1291 assert_eq!(metadata.created_at, metadata.updated_at);
1292 assert_eq!(metadata.compression, Codec::None);
1293 assert!(metadata.original_size > 0);
1294 }
1295
1296 #[test]
1297 fn test_metadata_update() {
1298 let dir = tempdir().unwrap();
1299 let storage = FileStorage::new(dir.path()).unwrap();
1300
1301 let doc = create_test_document("doc1");
1302 storage.save_document("doc1", &doc).unwrap();
1303
1304 let meta1 = storage.get_metadata("doc1").unwrap();
1305
1306 std::thread::sleep(std::time::Duration::from_millis(10));
1308
1309 storage.save_document("doc1", &doc).unwrap();
1311
1312 let meta2 = storage.get_metadata("doc1").unwrap();
1313 assert_eq!(meta2.created_at, meta1.created_at);
1314 assert!(meta2.updated_at >= meta1.updated_at);
1315 }
1316
1317 #[test]
1318 fn test_list_storage() {
1319 let dir = tempdir().unwrap();
1320 let storage = FileStorage::new(dir.path()).unwrap();
1321
1322 assert_eq!(storage.list().unwrap().len(), 0);
1323
1324 storage
1325 .save_document("doc1", &create_test_document("doc1"))
1326 .unwrap();
1327 storage
1328 .save_document("doc2", &create_test_document("doc2"))
1329 .unwrap();
1330 storage
1331 .save_document("doc3", &create_test_document("doc3"))
1332 .unwrap();
1333
1334 let items = storage.list().unwrap();
1335 assert_eq!(items.len(), 3);
1336 assert!(items.contains(&"doc1".to_string()));
1337 assert!(items.contains(&"doc2".to_string()));
1338 assert!(items.contains(&"doc3".to_string()));
1339 }
1340
1341 #[test]
1342 fn test_delete() {
1343 let dir = tempdir().unwrap();
1344 let storage = FileStorage::new(dir.path()).unwrap();
1345
1346 let doc = create_test_document("doc1");
1347 storage.save_document("doc1", &doc).unwrap();
1348
1349 assert!(storage.exists("doc1"));
1350 assert_eq!(storage.list().unwrap().len(), 1);
1351
1352 storage.delete("doc1").unwrap();
1353
1354 assert!(!storage.exists("doc1"));
1355 assert_eq!(storage.list().unwrap().len(), 0);
1356 }
1357
1358 #[test]
1359 fn test_delete_nonexistent() {
1360 let dir = tempdir().unwrap();
1361 let storage = FileStorage::new(dir.path()).unwrap();
1362
1363 let result = storage.delete("nonexistent");
1365 assert!(result.is_ok());
1366 }
1367
1368 #[test]
1369 fn test_clear() {
1370 let dir = tempdir().unwrap();
1371 let storage = FileStorage::new(dir.path()).unwrap();
1372
1373 storage
1374 .save_document("doc1", &create_test_document("doc1"))
1375 .unwrap();
1376 storage
1377 .save_document("doc2", &create_test_document("doc2"))
1378 .unwrap();
1379 storage
1380 .save_document("doc3", &create_test_document("doc3"))
1381 .unwrap();
1382
1383 assert_eq!(storage.list().unwrap().len(), 3);
1384
1385 storage.clear().unwrap();
1386
1387 assert_eq!(storage.list().unwrap().len(), 0);
1388 }
1389
1390 #[test]
1391 fn test_storage_size() {
1392 let dir = tempdir().unwrap();
1393 let storage = FileStorage::new(dir.path()).unwrap();
1394
1395 assert_eq!(storage.total_size().unwrap(), 0);
1396
1397 storage
1398 .save_document("doc1", &create_test_document("doc1"))
1399 .unwrap();
1400
1401 let size = storage.total_size().unwrap();
1402 assert!(size > 0);
1403
1404 storage
1405 .save_document("doc2", &create_test_document("doc2"))
1406 .unwrap();
1407
1408 let size2 = storage.total_size().unwrap();
1409 assert!(size2 > size);
1410 }
1411
1412 #[test]
1413 fn test_compression_codecs() {
1414 let dir = tempdir().unwrap();
1415
1416 #[allow(unused_mut)]
1418 let mut codecs = vec![Codec::None, Codec::Gzip];
1419
1420 #[cfg(feature = "zstd")]
1421 codecs.push(Codec::Zstd);
1422
1423 #[cfg(feature = "lz4")]
1424 codecs.push(Codec::Lz4);
1425
1426 for codec in codecs {
1427 let storage = FileStorage::with_codec(dir.path(), codec).unwrap();
1428 let doc = create_test_document("doc1");
1429
1430 let stats = storage.save_document("test", &doc).unwrap();
1431 assert!(stats.original_size > 0);
1432
1433 let loaded = storage.load_document("test").unwrap();
1434 assert_eq!(loaded.id, doc.id);
1435 assert_eq!(loaded.content, doc.content);
1436
1437 storage.delete("test").unwrap();
1438 }
1439 }
1440
1441 #[test]
1442 fn test_exists() {
1443 let dir = tempdir().unwrap();
1444 let storage = FileStorage::new(dir.path()).unwrap();
1445
1446 assert!(!storage.exists("doc1"));
1447
1448 storage
1449 .save_document("doc1", &create_test_document("doc1"))
1450 .unwrap();
1451
1452 assert!(storage.exists("doc1"));
1453 assert!(!storage.exists("doc2"));
1454 }
1455
1456 #[test]
1457 fn test_flat_index_wrapper_roundtrip() {
1458 let index = create_test_flat_index();
1459 let wrapper = FlatIndexWrapper::from_index(&index);
1460 let restored = wrapper.to_index().unwrap();
1461
1462 assert_eq!(restored.len(), index.len());
1463 assert_eq!(restored.embedding_dim(), index.embedding_dim());
1464
1465 let query = vec![0.1, 0.2, 0.3, 0.4, 0.5];
1467 let results = restored.search(&query, 2).unwrap();
1468 assert_eq!(results.len(), 2);
1469 }
1470
1471 #[test]
1472 fn test_hnsw_index_wrapper_roundtrip() {
1473 let index = create_test_hnsw_index();
1474 let wrapper = HNSWIndexWrapper::from_index(&index);
1475 let restored = wrapper.to_index().unwrap();
1476
1477 assert_eq!(restored.len(), index.len());
1478 assert_eq!(restored.embedding_dim(), index.embedding_dim());
1479
1480 let query = vec![0.1, 0.2, 0.3, 0.4, 0.5];
1482 let results = restored.search(&query, 2).unwrap();
1483 assert_eq!(results.len(), 2);
1484 }
1485
1486 #[test]
1487 fn quantized_index_survives_a_save_load_roundtrip() {
1488 use rand::{rngs::StdRng, RngExt, SeedableRng};
1507
1508 let dim = 16;
1509 let mut rng = StdRng::seed_from_u64(77);
1510 let centers: Vec<Vec<f32>> = (0..8)
1511 .map(|_| (0..dim).map(|_| rng.random::<f32>() * 10.0).collect())
1512 .collect();
1513 let base: Vec<Vec<f32>> = (0..200)
1514 .map(|i| {
1515 let c = ¢ers[i % 8];
1516 c.iter().map(|x| x + rng.random::<f32>() * 0.5).collect()
1517 })
1518 .collect();
1519 let queries: Vec<Vec<f32>> = (0..30)
1520 .map(|i| {
1521 let c = ¢ers[i % 8];
1522 c.iter().map(|x| x + rng.random::<f32>() * 0.5).collect()
1523 })
1524 .collect();
1525 let docs: Vec<Document> = base
1526 .iter()
1527 .enumerate()
1528 .map(|(i, v)| Document {
1529 id: i.to_string(),
1530 content: String::new(),
1531 embedding: v.clone(),
1532 metadata: None,
1533 })
1534 .collect();
1535
1536 let index = crate::index::HNSWIndex::build(
1537 docs.iter().map(|d| d.embedding.clone()).collect(),
1538 crate::index::HNSWConfig {
1539 storage: crate::index::Storage::SQ8,
1540 rerank_candidates: 100,
1541 metric: crate::index::DistanceMetric::L2,
1542 ..Default::default()
1543 },
1544 );
1545
1546 let wrapper = HNSWIndexWrapper::from_index(&index);
1547 let restored = wrapper
1548 .to_index()
1549 .expect("a quantized index must survive a save/load roundtrip");
1550
1551 assert_eq!(restored.len(), index.len());
1552 assert_eq!(
1553 restored.config().storage,
1554 crate::index::Storage::SQ8,
1555 "storage mode must persist — reloading SQ8 as F32 would reinterpret 8-bit codes \
1556 as f32 and return garbage"
1557 );
1558
1559 let k = 10;
1560 let mut total_recall = 0.0f32;
1561 for q in &queries {
1562 let mut exact: Vec<(f32, usize)> = base
1563 .iter()
1564 .enumerate()
1565 .map(|(i, v)| {
1566 let d: f32 = v.iter().zip(q).map(|(a, b)| (a - b) * (a - b)).sum();
1567 (d, i)
1568 })
1569 .collect();
1570 exact.sort_by(|a, b| a.0.total_cmp(&b.0));
1571 let truth: std::collections::HashSet<usize> =
1572 exact.iter().take(k).map(|(_, i)| *i).collect();
1573
1574 let got: std::collections::HashSet<usize> = restored
1575 .search(q, k)
1576 .expect("search")
1577 .into_iter()
1578 .filter_map(|r| r.id.parse::<usize>().ok())
1579 .collect();
1580
1581 total_recall += truth.intersection(&got).count() as f32 / k as f32;
1582 }
1583 let recall = total_recall / queries.len() as f32;
1584
1585 assert!(
1592 recall > 0.7,
1593 "restored SQ8 index recall@{k} against brute-force ground truth = {:.1}% — \
1594 the save/load roundtrip corrupted the codebook or the codes",
1595 recall * 100.0
1596 );
1597 }
1598
1599 #[test]
1600 fn test_concurrent_writes() {
1601 let dir = tempdir().unwrap();
1602 let storage = FileStorage::new(dir.path()).unwrap();
1603
1604 let doc = create_test_document("doc1");
1606
1607 for _ in 0..10 {
1608 storage.save_document("doc1", &doc).unwrap();
1609 let loaded = storage.load_document("doc1").unwrap();
1610 assert_eq!(loaded.id, doc.id);
1611 }
1612 }
1613
1614 #[test]
1615 fn test_large_document() {
1616 let dir = tempdir().unwrap();
1617 let storage = FileStorage::new(dir.path()).unwrap();
1618
1619 let mut large_doc = create_test_document("large");
1621 large_doc.embedding = vec![0.5; 10000];
1622 large_doc.content = "x".repeat(100000);
1623
1624 let stats = storage.save_document("large", &large_doc).unwrap();
1625 assert!(stats.original_size > 100000);
1626
1627 let loaded = storage.load_document("large").unwrap();
1628 assert_eq!(loaded.id, large_doc.id);
1629 assert_eq!(loaded.embedding.len(), 10000);
1630 assert_eq!(loaded.content.len(), 100000);
1631 }
1632
1633 #[test]
1634 fn test_rejects_path_traversal_item_names() {
1635 let dir = tempdir().unwrap();
1636 let storage = FileStorage::new(dir.path()).unwrap();
1637 let doc = create_test_document("doc1");
1638
1639 let result = storage.save_document("../outside", &doc);
1640 assert!(result.is_err(), "path traversal names should be rejected");
1641 }
1642
1643 #[test]
1644 fn test_is_invalid_item_name_comprehensive() {
1645 assert!(!FileStorage::is_invalid_item_name("hello"));
1647 assert!(!FileStorage::is_invalid_item_name("my_index"));
1648 assert!(!FileStorage::is_invalid_item_name("data-2024"));
1649 assert!(!FileStorage::is_invalid_item_name("file.txt"));
1650
1651 assert!(FileStorage::is_invalid_item_name(""));
1653
1654 assert!(FileStorage::is_invalid_item_name(".."));
1656 assert!(FileStorage::is_invalid_item_name("."));
1657 assert!(FileStorage::is_invalid_item_name("foo/bar"));
1658 assert!(FileStorage::is_invalid_item_name("foo\\bar"));
1659 assert!(FileStorage::is_invalid_item_name("../outside"));
1660
1661 assert!(FileStorage::is_invalid_item_name("/absolute"));
1663 #[cfg(target_os = "windows")]
1664 assert!(FileStorage::is_invalid_item_name("C:\\Windows\\System32"));
1665
1666 assert!(FileStorage::is_invalid_item_name("hello\0world"));
1668 assert!(FileStorage::is_invalid_item_name("\0"));
1669
1670 assert!(FileStorage::is_invalid_item_name("CON"));
1672 assert!(FileStorage::is_invalid_item_name("con"));
1673 assert!(FileStorage::is_invalid_item_name("Con"));
1674 assert!(FileStorage::is_invalid_item_name("PRN"));
1675 assert!(FileStorage::is_invalid_item_name("AUX"));
1676 assert!(FileStorage::is_invalid_item_name("NUL"));
1677 assert!(FileStorage::is_invalid_item_name("nul"));
1678 assert!(FileStorage::is_invalid_item_name("COM1"));
1679 assert!(FileStorage::is_invalid_item_name("com1"));
1680 assert!(FileStorage::is_invalid_item_name("COM9"));
1681 assert!(FileStorage::is_invalid_item_name("LPT1"));
1682 assert!(FileStorage::is_invalid_item_name("lpt1"));
1683 assert!(FileStorage::is_invalid_item_name("LPT9"));
1684
1685 assert!(FileStorage::is_invalid_item_name("CON.txt"));
1687 assert!(FileStorage::is_invalid_item_name("NUL.tar.gz"));
1688 assert!(FileStorage::is_invalid_item_name("com1.data"));
1689 assert!(FileStorage::is_invalid_item_name("lpt3.log"));
1690 }
1691}