1use std::path::{Path, PathBuf};
2use std::sync::atomic::{AtomicU64, Ordering};
3use std::sync::Arc;
4use std::time::Instant;
5
6use parking_lot::{Mutex, RwLock};
7use serde::{Deserialize, Serialize};
8
9use crate::cursor::PartitionSelector;
10use linkrs_core::types::{CommitLsn, IndexGeneration, SnapshotTimestamp};
11use linkrs_core::{StorageError, StorageResult};
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
19pub enum GenerationState {
20 Building,
23 CatchingUp,
26 Publishing,
29 Active,
32}
33
34#[derive(Debug, Clone, Serialize, Deserialize)]
37pub struct GenerationBuildState {
38 pub generation: IndexGeneration,
40 pub snapshot_timestamp: SnapshotTimestamp,
42 pub start_lsn: CommitLsn,
44 pub barrier_lsn: Option<CommitLsn>,
47 pub state: GenerationState,
49}
50
51impl GenerationBuildState {
52 pub fn new(
53 generation: IndexGeneration,
54 snapshot_timestamp: SnapshotTimestamp,
55 start_lsn: CommitLsn,
56 ) -> Self {
57 Self {
58 generation,
59 snapshot_timestamp,
60 start_lsn,
61 barrier_lsn: None,
62 state: GenerationState::Building,
63 }
64 }
65
66 pub fn transition_to_catching_up(&mut self) -> StorageResult<()> {
67 self.require_state(GenerationState::Building)?;
68 self.state = GenerationState::CatchingUp;
69 Ok(())
70 }
71
72 pub fn transition_to_publishing(&mut self, barrier_lsn: CommitLsn) -> StorageResult<()> {
73 self.require_state(GenerationState::CatchingUp)?;
74 if barrier_lsn < self.start_lsn {
75 return Err(StorageError::invalid_operation(
76 "Generation barrier LSN precedes the snapshot LSN",
77 ));
78 }
79 self.barrier_lsn = Some(barrier_lsn);
80 self.state = GenerationState::Publishing;
81 Ok(())
82 }
83
84 pub fn transition_to_active(&mut self) -> StorageResult<()> {
85 self.require_state(GenerationState::Publishing)?;
86 self.state = GenerationState::Active;
87 Ok(())
88 }
89
90 pub fn is_active(&self) -> bool {
91 self.state == GenerationState::Active
92 }
93
94 fn require_state(&self, expected: GenerationState) -> StorageResult<()> {
95 if self.state == expected {
96 Ok(())
97 } else {
98 Err(StorageError::invalid_operation(format!(
99 "Invalid generation transition from {:?}; expected {:?}",
100 self.state, expected
101 )))
102 }
103 }
104}
105
106#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
108pub struct IndexShard {
109 pub shard_id: u32,
110 pub lower: Option<Vec<u8>>,
111 pub upper: Option<Vec<u8>>,
112 pub checkpoint_file: PathBuf,
113 #[serde(default)]
114 pub checksum: Option<u32>,
115}
116
117impl IndexShard {
118 pub fn contains(&self, key: &[u8]) -> bool {
119 self.lower.as_deref().is_none_or(|lower| key >= lower)
120 && self.upper.as_deref().is_none_or(|upper| key < upper)
121 }
122
123 pub fn intersects(&self, lower: Option<&[u8]>, upper: Option<&[u8]>) -> bool {
124 self.upper
125 .as_deref()
126 .zip(lower)
127 .is_none_or(|(shard_upper, query_lower)| shard_upper > query_lower)
128 && upper
129 .zip(self.lower.as_deref())
130 .is_none_or(|(query_upper, shard_lower)| query_upper > shard_lower)
131 }
132
133 fn validate(&self) -> StorageResult<()> {
134 if self
135 .lower
136 .as_ref()
137 .zip(self.upper.as_ref())
138 .is_some_and(|(lower, upper)| lower >= upper)
139 {
140 return Err(StorageError::db_error(format!(
141 "Index shard {} has an empty or inverted range",
142 self.shard_id
143 )));
144 }
145 if self.checkpoint_file.as_os_str().is_empty() {
146 return Err(StorageError::db_error(format!(
147 "Index shard {} has no checkpoint file",
148 self.shard_id
149 )));
150 }
151 Ok(())
152 }
153
154 pub fn compute_checksum(&self) -> StorageResult<Option<u32>> {
163 if self.checkpoint_file.as_os_str().is_empty() {
164 return Ok(None);
165 }
166 if self.checkpoint_file.is_dir() {
167 return Ok(None);
168 }
169 let data = match std::fs::read(&self.checkpoint_file) {
170 Ok(data) => data,
171 Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None),
172 Err(e) => {
173 return Err(StorageError::db_error(format!(
174 "Read checkpoint for checksum: {}: {e}",
175 self.checkpoint_file.display()
176 )));
177 }
178 };
179 Ok(Some(crc32fast::hash(&data)))
180 }
181
182 pub fn verify_checksum(&self) -> StorageResult<()> {
183 let Some(expected) = self.checksum else {
184 return Ok(());
185 };
186 let Some(actual) = self.compute_checksum()? else {
187 return Ok(());
188 };
189 if actual != expected {
190 return Err(StorageError::db_error(format!(
191 "Shard {} checksum mismatch: expected {expected:#010x}, got {actual:#010x}",
192 self.shard_id
193 )));
194 }
195 Ok(())
196 }
197}
198
199#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
201pub struct IndexManifest {
202 pub space_id: u64,
205 pub index_id: u64,
206 pub generation: IndexGeneration,
207 pub shards: Vec<IndexShard>,
208}
209
210impl IndexManifest {
211 pub fn new(
212 space_id: u64,
213 index_id: u64,
214 generation: IndexGeneration,
215 shards: Vec<IndexShard>,
216 ) -> StorageResult<Self> {
217 let manifest = Self {
218 space_id,
219 index_id,
220 generation,
221 shards,
222 };
223 manifest.validate()?;
224 Ok(manifest)
225 }
226
227 pub fn validate(&self) -> StorageResult<()> {
228 if self.shards.is_empty() {
229 return Err(StorageError::db_error(
230 "Index manifest must contain at least one shard",
231 ));
232 }
233 if self
234 .shards
235 .first()
236 .is_some_and(|shard| shard.lower.is_some())
237 {
238 return Err(StorageError::db_error(
239 "The first index shard must have an unbounded lower range",
240 ));
241 }
242 if self
243 .shards
244 .last()
245 .is_some_and(|shard| shard.upper.is_some())
246 {
247 return Err(StorageError::db_error(
248 "The last index shard must have an unbounded upper range",
249 ));
250 }
251 for shard in &self.shards {
252 shard.validate()?;
253 }
254 for pair in self.shards.windows(2) {
255 if pair[0].shard_id == pair[1].shard_id {
256 return Err(StorageError::db_error(format!(
257 "Duplicate index shard id {}",
258 pair[0].shard_id
259 )));
260 }
261 if pair[0].upper != pair[1].lower {
262 return Err(StorageError::db_error(format!(
263 "Index shards {} and {} are not contiguous",
264 pair[0].shard_id, pair[1].shard_id
265 )));
266 }
267 }
268 Ok(())
269 }
270
271 pub fn route_key(&self, key: &[u8]) -> Option<&IndexShard> {
274 let idx = self
275 .shards
276 .partition_point(|shard| shard.lower.as_deref().is_none_or(|lower| lower <= key));
277 let shard = self.shards.get(idx.wrapping_sub(1))?;
278 shard.contains(key).then_some(shard)
279 }
280
281 pub fn select_shards(&self, selector: &PartitionSelector) -> Vec<&IndexShard> {
282 match selector {
283 PartitionSelector::All => self.shards.iter().collect(),
284 PartitionSelector::Shards(ids) => self
285 .shards
286 .iter()
287 .filter(|shard| ids.contains(&shard.shard_id))
288 .collect(),
289 PartitionSelector::KeyRange { lower, upper } => self
290 .shards
291 .iter()
292 .filter(|shard| shard.intersects(lower.as_deref(), upper.as_deref()))
293 .collect(),
294 }
295 }
296
297 pub fn scan_ranges_with_shard(
298 &self,
299 selector: &PartitionSelector,
300 query_lower: &[u8],
301 query_upper: &[u8],
302 ) -> Vec<(u32, Vec<u8>, Vec<u8>)> {
303 self.select_shards(selector)
304 .into_iter()
305 .filter(|shard| shard.intersects(Some(query_lower), Some(query_upper)))
306 .filter_map(|shard| {
307 let lower = shard.lower.as_deref().map_or_else(
308 || query_lower.to_vec(),
309 |value| value.max(query_lower).to_vec(),
310 );
311 let upper = shard.upper.as_deref().map_or_else(
312 || query_upper.to_vec(),
313 |value| value.min(query_upper).to_vec(),
314 );
315 (lower < upper).then_some((shard.shard_id, lower, upper))
316 })
317 .collect()
318 }
319
320 pub fn store(&self, path: &Path) -> StorageResult<()> {
321 self.validate()?;
322 let with_checksums = self.with_checksums()?;
323 let bytes = postcard::to_allocvec(&with_checksums).map_err(|error| {
324 StorageError::db_error(format!("Serialize index manifest: {error}"))
325 })?;
326 let mut wrapped = Vec::new();
327 crate::persistence::write_versioned_payload(&mut wrapped, &bytes);
328 crate::persistence::write_file_atomic(path, &wrapped)
329 }
330
331 pub fn load(path: &Path) -> StorageResult<Self> {
332 let mut file = std::fs::File::open(path)?;
333 let payload = crate::persistence::read_versioned_payload(
334 &mut file,
335 path.file_name()
336 .and_then(|n| n.to_str())
337 .unwrap_or("manifest.bin"),
338 )?;
339 let manifest: Self = postcard::from_bytes(&payload)
340 .map_err(|error| StorageError::db_error(format!("Read index manifest: {error}")))?;
341 manifest.validate()?;
342 for shard in &manifest.shards {
343 shard.verify_checksum()?;
344 }
345 Ok(manifest)
346 }
347
348 pub fn with_checksums(&self) -> StorageResult<Self> {
350 let mut clone = self.clone();
351 for shard in &mut clone.shards {
352 shard.checksum = shard.compute_checksum()?;
353 }
354 Ok(clone)
355 }
356}
357
358#[derive(Debug, Clone)]
372pub struct ManifestHandle(Arc<IndexManifest>);
373
374impl ManifestHandle {
375 pub fn manifest(&self) -> &IndexManifest {
376 &self.0
377 }
378}
379
380#[derive(Debug)]
381struct RetiredManifest {
382 manifest: Arc<IndexManifest>,
383 retired_at: Instant,
384}
385
386#[derive(Debug)]
396pub struct ManifestCatalog {
397 active: RwLock<Arc<IndexManifest>>,
398 retired: Mutex<Vec<RetiredManifest>>,
399 published: AtomicU64,
400 reclaimed_files: AtomicU64,
401}
402
403#[derive(Debug, Clone, Copy, PartialEq, Eq)]
404pub struct ManifestCatalogStats {
405 pub active_generation: IndexGeneration,
406 pub active_readers: u64,
407 pub retired_generations: u64,
408 pub published_manifests: u64,
409 pub reclaimed_files: u64,
410 pub oldest_retired_age_secs: u64,
411}
412
413impl ManifestCatalog {
414 pub fn new(manifest: IndexManifest) -> StorageResult<Self> {
415 manifest.validate()?;
416 Ok(Self {
417 active: RwLock::new(Arc::new(manifest)),
418 retired: Mutex::new(Vec::new()),
419 published: AtomicU64::new(0),
420 reclaimed_files: AtomicU64::new(0),
421 })
422 }
423
424 pub fn acquire(&self) -> ManifestHandle {
426 let guard = self.active.read();
427 #[cfg(debug_assertions)]
428 {
429 let count = Arc::strong_count(&guard);
430 log::trace!("ManifestCatalog acquire: readers={}", count + 1);
431 }
432 ManifestHandle(Arc::clone(&guard))
433 }
434
435 pub fn acquire_generation(&self, generation: IndexGeneration) -> Option<ManifestHandle> {
439 {
440 let active = self.active.read();
441 if active.generation == generation {
442 return Some(ManifestHandle(Arc::clone(&active)));
443 }
444 }
445 let retired = self.retired.lock();
446 retired
447 .iter()
448 .find(|entry| entry.manifest.generation == generation)
449 .map(|entry| ManifestHandle(Arc::clone(&entry.manifest)))
450 }
451
452 pub fn publish(&self, manifest: IndexManifest) -> StorageResult<ManifestHandle> {
453 manifest.validate()?;
454 let mut active = self.active.write();
455 if manifest.index_id != active.index_id {
456 return Err(StorageError::invalid_operation(
457 "Cannot publish a manifest for another index",
458 ));
459 }
460 if manifest.generation <= active.generation {
461 return Err(StorageError::invalid_operation(
462 "Index generation must increase",
463 ));
464 }
465
466 let next = Arc::new(manifest);
467 let previous = std::mem::replace(&mut *active, Arc::clone(&next));
468 #[cfg(debug_assertions)]
469 {
470 log::trace!(
471 "ManifestCatalog publish: gen={}, active_readers={}",
472 next.generation,
473 Arc::strong_count(&next) - 1,
474 );
475 }
476
477 self.retired.lock().push(RetiredManifest {
483 manifest: previous,
484 retired_at: Instant::now(),
485 });
486 self.published.fetch_add(1, Ordering::Relaxed);
487 Ok(ManifestHandle(next))
488 }
489
490 pub fn take_reclaimable_files(&self) -> Vec<PathBuf> {
493 self.take_reclaimable_manifests()
494 .into_iter()
495 .flat_map(|manifest| {
496 manifest
497 .shards
498 .into_iter()
499 .map(|shard| shard.checkpoint_file)
500 })
501 .collect()
502 }
503
504 pub fn take_reclaimable_manifests(&self) -> Vec<IndexManifest> {
508 let mut retired = self.retired.lock();
509 let mut manifests = Vec::new();
510 retired.retain(|entry| {
511 if Arc::strong_count(&entry.manifest) == 1 {
512 manifests.push((*entry.manifest).clone());
513 false
514 } else {
515 true
516 }
517 });
518 let total = manifests.iter().map(|m| m.shards.len() as u64).sum();
519 self.reclaimed_files.fetch_add(total, Ordering::Relaxed);
520 manifests
521 }
522
523 pub fn retired_reclaimable<F>(&self, matches: F) -> Vec<IndexManifest>
529 where
530 F: Fn(&IndexManifest) -> bool,
531 {
532 let retired = self.retired.lock();
533 retired
534 .iter()
535 .filter(|entry| Arc::strong_count(&entry.manifest) == 1 && matches(&entry.manifest))
536 .map(|entry| (*entry.manifest).clone())
537 .collect()
538 }
539
540 pub fn remove_retired(&self, generation: IndexGeneration) -> bool {
542 let mut retired = self.retired.lock();
543 let before = retired.len();
544 retired.retain(|entry| entry.manifest.generation != generation);
545 retired.len() != before
546 }
547
548 pub fn stats(&self) -> ManifestCatalogStats {
549 let active = self.active.read();
550 let retired = self.retired.lock();
551 let now = Instant::now();
552 let oldest_age = retired
553 .iter()
554 .map(|e| now.duration_since(e.retired_at).as_secs())
555 .max()
556 .unwrap_or(0);
557 ManifestCatalogStats {
558 active_generation: active.generation,
559 active_readers: Arc::strong_count(&active).saturating_sub(1) as u64,
560 retired_generations: retired.len() as u64,
561 published_manifests: self.published.load(Ordering::Relaxed),
562 reclaimed_files: self.reclaimed_files.load(Ordering::Relaxed),
563 oldest_retired_age_secs: oldest_age,
564 }
565 }
566}
567
568#[cfg(test)]
569mod tests {
570 use std::path::PathBuf;
571
572 use super::{IndexManifest, IndexShard, ManifestCatalog};
573 use crate::cursor::PartitionSelector;
574 use linkrs_core::types::IndexGeneration;
575
576 fn shard(shard_id: u32, lower: Option<&[u8]>, upper: Option<&[u8]>) -> IndexShard {
577 IndexShard {
578 shard_id,
579 lower: lower.map(<[u8]>::to_vec),
580 upper: upper.map(<[u8]>::to_vec),
581 checkpoint_file: format!("{shard_id}.index").into(),
582 checksum: None,
583 }
584 }
585
586 fn manifest(generation: u64, shards: Vec<IndexShard>) -> IndexManifest {
587 IndexManifest::new(1, 1, IndexGeneration::new(generation), shards)
588 .expect("manifest should be valid")
589 }
590
591 #[test]
592 fn manifest_routes_half_open_ranges_and_prunes_queries() {
593 let manifest = manifest(
594 1,
595 vec![shard(0, None, Some(b"m")), shard(1, Some(b"m"), None)],
596 );
597 assert_eq!(
598 manifest.route_key(b"a").map(|shard| shard.shard_id),
599 Some(0)
600 );
601 assert_eq!(
602 manifest.route_key(b"m").map(|shard| shard.shard_id),
603 Some(1)
604 );
605
606 let selected = manifest.select_shards(&PartitionSelector::KeyRange {
607 lower: Some(b"b".to_vec()),
608 upper: Some(b"m".to_vec()),
609 });
610 assert_eq!(
611 selected
612 .iter()
613 .map(|shard| shard.shard_id)
614 .collect::<Vec<_>>(),
615 vec![0]
616 );
617 }
618
619 #[test]
620 fn manifest_rejects_range_gaps() {
621 let result = IndexManifest::new(
622 1,
623 1,
624 IndexGeneration::new(1),
625 vec![shard(0, None, Some(b"m")), shard(1, Some(b"n"), None)],
626 );
627 assert!(result.is_err());
628 }
629
630 #[test]
631 fn reader_pin_keeps_retired_manifest_tracked_until_last_handle_drops() {
632 let catalog = ManifestCatalog::new(manifest(1, vec![shard(0, None, None)]))
633 .expect("catalog should be valid");
634 let old_reader = catalog.acquire();
635 catalog
636 .publish(manifest(
637 2,
638 vec![shard(0, None, Some(b"m")), shard(1, Some(b"m"), None)],
639 ))
640 .expect("new manifest should publish");
641
642 assert!(catalog.take_reclaimable_files().is_empty());
645 assert_eq!(catalog.stats().retired_generations, 1);
646 let pin = catalog
647 .acquire_generation(IndexGeneration::new(1))
648 .expect("retired generation should be re-acquirable");
649 drop(pin);
650 assert!(catalog.take_reclaimable_files().is_empty());
651
652 drop(old_reader);
653 assert_eq!(
654 catalog.take_reclaimable_files(),
655 vec![PathBuf::from("0.index")]
656 );
657 assert_eq!(catalog.stats().retired_generations, 0);
658 assert_eq!(catalog.stats().reclaimed_files, 1);
659 }
660
661 #[test]
662 fn retired_reclaimable_filters_by_caller_predicate() {
663 let catalog = ManifestCatalog::new(manifest(1, vec![shard(0, None, None)]))
664 .expect("catalog should be valid");
665 let old_reader = catalog.acquire();
666 catalog
667 .publish(manifest(
668 2,
669 vec![shard(0, None, Some(b"m")), shard(1, Some(b"m"), None)],
670 ))
671 .expect("new manifest should publish");
672 catalog
673 .publish(manifest(3, vec![shard(0, None, None)]))
674 .expect("new manifest should publish");
675
676 let candidates = catalog.retired_reclaimable(|m| m.generation.get() == 2);
678 assert_eq!(candidates.len(), 1);
679 assert_eq!(candidates[0].generation.get(), 2);
680
681 assert_eq!(catalog.stats().retired_generations, 2);
683
684 assert!(catalog.remove_retired(IndexGeneration::new(2)));
685 assert!(!catalog.remove_retired(IndexGeneration::new(99)));
686 assert_eq!(catalog.stats().retired_generations, 1);
687
688 drop(old_reader);
689 catalog.take_reclaimable_manifests();
690 assert_eq!(catalog.stats().retired_generations, 0);
691 }
692
693 #[test]
694 fn published_handle_still_reads_old_manifest_until_released() {
695 let catalog = ManifestCatalog::new(manifest(1, vec![shard(0, None, None)]))
696 .expect("catalog should be valid");
697 let first = catalog.acquire();
698 let second = catalog
699 .publish(manifest(2, vec![shard(0, None, None)]))
700 .expect("new manifest should publish");
701
702 assert_eq!(second.manifest().generation.get(), 2);
705 assert_eq!(first.manifest().generation.get(), 1);
706 assert!(catalog.take_reclaimable_files().is_empty());
707
708 drop(first);
709 assert_eq!(
710 catalog.take_reclaimable_files(),
711 vec![PathBuf::from("0.index")]
712 );
713 }
714
715 #[test]
716 fn persisted_manifest_roundtrips() {
717 let directory = tempfile::tempdir().expect("temporary directory should be created");
718 let path = directory.path().join("manifest.bin");
719 let manifest = manifest(1, vec![shard(0, None, None)]);
720 manifest.store(&path).expect("manifest should persist");
721 assert_eq!(
722 IndexManifest::load(&path).expect("manifest should load"),
723 manifest
724 );
725
726 let unsupported = IndexManifest {
727 shards: Vec::new(),
728 ..manifest
729 };
730 assert!(unsupported.validate().is_err());
731 }
732
733 #[test]
734 fn manifest_checksum_verifies_on_load() {
735 let directory = tempfile::tempdir().expect("temporary directory should be created");
736 let checkpoint_path = directory.path().join("checkpoint.bin");
737 std::fs::write(&checkpoint_path, b"test checkpoint data").expect("write checkpoint");
738
739 let shard_obj = IndexShard {
740 shard_id: 0,
741 lower: None,
742 upper: None,
743 checkpoint_file: checkpoint_path.clone(),
744 checksum: None,
745 };
746 let manifest = IndexManifest::new(1, 1, IndexGeneration::new(1), vec![shard_obj])
747 .expect("manifest should be valid");
748
749 let manifest_path = directory.path().join("manifest.bin");
750 manifest
751 .store(&manifest_path)
752 .expect("store should succeed");
753
754 let loaded = IndexManifest::load(&manifest_path).expect("load should succeed");
755 assert!(loaded.shards[0].checksum.is_some());
756
757 std::fs::write(&checkpoint_path, b"corrupted data").expect("corrupt checkpoint");
758 assert!(IndexManifest::load(&manifest_path).is_err());
759 }
760}