corium_store/
segment_cache.rs

1//! Bounded memory and optional local-disk read-through segment cache.
2
3use std::collections::{BTreeMap, HashMap, VecDeque};
4use std::fs::{self, File, OpenOptions};
5use std::io::{self, Write};
6use std::path::{Path, PathBuf};
7use std::sync::atomic::{AtomicU64, Ordering};
8use std::sync::{Arc, Mutex, Weak};
9
10use fs2::FileExt;
11use tokio::sync::Mutex as AsyncMutex;
12
13use crate::{BlobId, BlobStore, StoreError, digest};
14use async_trait::async_trait;
15
16const ENTRY_METADATA_BYTES: u64 = 96;
17
18/// Minimal cache-neutral interface for immutable segment reads.
19#[async_trait]
20pub trait SegmentReader: Send + Sync {
21    /// Reads an immutable segment from authoritative storage.
22    async fn read_segment(&self, id: &BlobId) -> Result<Option<Vec<u8>>, StoreError>;
23}
24
25#[async_trait]
26impl<T: BlobStore + ?Sized> SegmentReader for T {
27    async fn read_segment(&self, id: &BlobId) -> Result<Option<Vec<u8>>, StoreError> {
28        self.get(id).await
29    }
30}
31
32/// Configuration for the peer's local segment cache.
33#[derive(Clone, Debug)]
34pub struct SegmentCacheConfig {
35    /// Dedicated cache directory, owned by one process.
36    pub directory: PathBuf,
37    /// Maximum accounted bytes in the SSD tier.
38    pub capacity_bytes: u64,
39    /// Maximum accounted bytes in the in-process tier.
40    pub memory_capacity_bytes: u64,
41}
42
43impl SegmentCacheConfig {
44    /// Default size of the in-process cache tier.
45    pub const DEFAULT_MEMORY_CAPACITY_BYTES: u64 = 64 * 1024 * 1024;
46
47    /// Validates the configured capacities.
48    ///
49    /// # Errors
50    /// Returns [`io::ErrorKind::InvalidInput`] for zero disk capacity or a
51    /// memory tier larger than the disk tier.
52    pub fn validate(&self) -> io::Result<()> {
53        if self.capacity_bytes == 0 {
54            return Err(io::Error::new(
55                io::ErrorKind::InvalidInput,
56                "cache capacity must be non-zero",
57            ));
58        }
59        if self.memory_capacity_bytes > self.capacity_bytes {
60            return Err(io::Error::new(
61                io::ErrorKind::InvalidInput,
62                "memory cache exceeds total capacity",
63            ));
64        }
65        Ok(())
66    }
67}
68
69/// Lock-free counters and gauges suitable for Prometheus adapters.
70#[derive(Default)]
71pub struct SegmentCacheMetrics {
72    /// Memory hits.
73    pub memory_hits: AtomicU64,
74    /// Memory misses.
75    pub memory_misses: AtomicU64,
76    /// Disk hits.
77    pub disk_hits: AtomicU64,
78    /// Disk misses.
79    pub disk_misses: AtomicU64,
80    /// Native found responses.
81    pub native_found: AtomicU64,
82    /// Native not-found responses.
83    pub native_not_found: AtomicU64,
84    /// Native errors.
85    pub native_errors: AtomicU64,
86    /// Coalesced callers.
87    pub coalesced_waiters: AtomicU64,
88    /// Successful admissions.
89    pub admissions: AtomicU64,
90    /// Oversize admission bypasses.
91    pub too_large: AtomicU64,
92    /// Admission I/O failures.
93    pub admission_errors: AtomicU64,
94    /// Capacity evictions.
95    pub evictions: AtomicU64,
96    /// Bytes removed by capacity eviction.
97    pub evicted_bytes: AtomicU64,
98    /// Invalid cache entries discarded.
99    pub corruptions: AtomicU64,
100    /// Current disk bytes.
101    pub disk_bytes: AtomicU64,
102    /// Current disk entries.
103    pub disk_entries: AtomicU64,
104    /// Current memory bytes.
105    pub memory_bytes: AtomicU64,
106    /// Current memory entries.
107    pub memory_entries: AtomicU64,
108    /// Bytes returned from memory.
109    pub bytes_memory: AtomicU64,
110    /// Bytes returned from disk.
111    pub bytes_disk: AtomicU64,
112    /// Bytes returned from native storage.
113    pub bytes_native: AtomicU64,
114}
115
116#[derive(Clone)]
117struct Entry {
118    len: u64,
119    generation: u64,
120}
121
122impl Entry {
123    fn accounted_bytes(&self) -> u64 {
124        self.len.saturating_add(ENTRY_METADATA_BYTES)
125    }
126}
127
128#[derive(Default)]
129struct MemoryTier {
130    entries: HashMap<BlobId, Arc<[u8]>>,
131    order: VecDeque<BlobId>,
132    bytes: u64,
133}
134
135struct DiskTier {
136    config: SegmentCacheConfig,
137    _lock: File,
138    entries: HashMap<BlobId, Entry>,
139    generation: u64,
140    bytes: u64,
141    readers: HashMap<BlobId, usize>,
142}
143
144#[derive(Clone)]
145enum FlightOutcome {
146    Found(Arc<[u8]>),
147    NotFound,
148}
149
150struct Flight {
151    lock: AsyncMutex<()>,
152    outcome: Mutex<Option<FlightOutcome>>,
153}
154
155/// Bounded read-through cache. Without a disk configuration it remains a
156/// bounded in-memory cache; its default capacity is 64 MiB.
157pub struct SegmentCache {
158    memory_capacity: u64,
159    memory: Mutex<MemoryTier>,
160    disk: Option<Mutex<DiskTier>>,
161    flights: Mutex<HashMap<BlobId, Weak<Flight>>>,
162    metrics: Arc<SegmentCacheMetrics>,
163}
164
165impl Default for SegmentCache {
166    fn default() -> Self {
167        Self::memory_only(SegmentCacheConfig::DEFAULT_MEMORY_CAPACITY_BYTES)
168    }
169}
170
171impl SegmentCache {
172    /// Creates a byte-bounded memory-only cache.
173    #[must_use]
174    pub fn memory_only(capacity_bytes: u64) -> Self {
175        Self {
176            memory_capacity: capacity_bytes,
177            memory: Mutex::default(),
178            disk: None,
179            flights: Mutex::default(),
180            metrics: Arc::default(),
181        }
182    }
183
184    /// Opens and reconciles an SSD cache, acquiring its exclusive ownership lock.
185    ///
186    /// # Errors
187    /// Returns an error for invalid configuration, inaccessible storage, or
188    /// when another process owns the directory.
189    pub fn open(config: &SegmentCacheConfig) -> io::Result<Self> {
190        config.validate()?;
191        fs::create_dir_all(config.directory.join("objects"))?;
192        fs::create_dir_all(config.directory.join("tmp"))?;
193        set_owner_only(&config.directory)?;
194        set_owner_only(&config.directory.join("objects"))?;
195        set_owner_only(&config.directory.join("tmp"))?;
196        let lock = OpenOptions::new()
197            .read(true)
198            .write(true)
199            .create(true)
200            .truncate(false)
201            .open(config.directory.join("LOCK"))?;
202        lock.try_lock_exclusive().map_err(|error| {
203            io::Error::new(
204                error.kind(),
205                format!("segment cache directory is already owned: {error}"),
206            )
207        })?;
208        for item in fs::read_dir(config.directory.join("tmp"))? {
209            let item = item?;
210            if item.file_type()?.is_file() {
211                fs::remove_file(item.path())?;
212            }
213        }
214        let mut tier = DiskTier {
215            config: config.clone(),
216            _lock: lock,
217            entries: HashMap::new(),
218            generation: 0,
219            bytes: 0,
220            readers: HashMap::new(),
221        };
222        tier.reconcile()?;
223        tier.evict_to_limit(config.capacity_bytes, None)?;
224        let metrics = Arc::new(SegmentCacheMetrics::default());
225        metrics.disk_bytes.store(tier.bytes, Ordering::Relaxed);
226        metrics
227            .disk_entries
228            .store(tier.entries.len() as u64, Ordering::Relaxed);
229        Ok(Self {
230            memory_capacity: config.memory_capacity_bytes,
231            memory: Mutex::default(),
232            disk: Some(Mutex::new(tier)),
233            flights: Mutex::default(),
234            metrics,
235        })
236    }
237
238    /// Returns the cache's metrics handle.
239    #[must_use]
240    pub fn metrics(&self) -> Arc<SegmentCacheMetrics> {
241        self.metrics.clone()
242    }
243
244    /// Returns cached bytes, loading and verifying the authoritative store on miss.
245    ///
246    /// # Errors
247    /// Returns errors from authoritative storage, including corrupt native bytes.
248    ///
249    /// # Panics
250    /// Panics if an internal cache mutex was poisoned by another panicking thread.
251    pub async fn get_or_load(
252        &self,
253        store: &dyn SegmentReader,
254        id: &BlobId,
255    ) -> Result<Option<Arc<[u8]>>, StoreError> {
256        if let Some(bytes) = self.memory_get(id) {
257            return Ok(Some(bytes));
258        }
259        if let Some(bytes) = self.disk_get(id).await {
260            self.memory_insert(id.clone(), bytes.clone());
261            return Ok(Some(bytes));
262        }
263        let (flight, joined) = {
264            let mut flights = self.flights.lock().expect("cache lock poisoned");
265            flights.retain(|_, flight| flight.strong_count() > 0);
266            if let Some(flight) = flights.get(id).and_then(Weak::upgrade) {
267                (flight, true)
268            } else {
269                let f = Arc::new(Flight {
270                    lock: AsyncMutex::new(()),
271                    outcome: Mutex::new(None),
272                });
273                flights.insert(id.clone(), Arc::downgrade(&f));
274                (f, false)
275            }
276        };
277        if joined {
278            self.metrics
279                .coalesced_waiters
280                .fetch_add(1, Ordering::Relaxed);
281        }
282        let _guard = flight.lock.lock().await;
283        if joined {
284            match flight.outcome.lock().expect("cache lock poisoned").clone() {
285                Some(FlightOutcome::Found(bytes)) => {
286                    self.metrics
287                        .bytes_native
288                        .fetch_add(bytes.len() as u64, Ordering::Relaxed);
289                    return Ok(Some(bytes));
290                }
291                Some(FlightOutcome::NotFound) => return Ok(None),
292                None => {}
293            }
294            if let Some(bytes) = self.memory_get(id) {
295                return Ok(Some(bytes));
296            }
297            if let Some(bytes) = self.disk_get(id).await {
298                return Ok(Some(bytes));
299            }
300        }
301        let loaded = store.read_segment(id).await;
302        let bytes = match loaded {
303            Ok(Some(bytes)) => {
304                self.metrics.native_found.fetch_add(1, Ordering::Relaxed);
305                bytes
306            }
307            Ok(None) => {
308                self.metrics
309                    .native_not_found
310                    .fetch_add(1, Ordering::Relaxed);
311                *flight.outcome.lock().expect("cache lock poisoned") =
312                    Some(FlightOutcome::NotFound);
313                return Ok(None);
314            }
315            Err(error) => {
316                self.metrics.native_errors.fetch_add(1, Ordering::Relaxed);
317                return Err(error);
318            }
319        };
320        if digest(&bytes) != *id {
321            return Err(StoreError::CorruptBlob(id.clone()));
322        }
323        self.metrics
324            .bytes_native
325            .fetch_add(bytes.len() as u64, Ordering::Relaxed);
326        let bytes: Arc<[u8]> = bytes.into();
327        self.disk_admit(id, &bytes);
328        self.memory_insert(id.clone(), bytes.clone());
329        *flight.outcome.lock().expect("cache lock poisoned") =
330            Some(FlightOutcome::Found(bytes.clone()));
331        Ok(Some(bytes))
332    }
333
334    fn memory_get(&self, id: &BlobId) -> Option<Arc<[u8]>> {
335        let mut memory = self.memory.lock().expect("cache lock poisoned");
336        let value = memory.entries.get(id).cloned();
337        if let Some(value) = value {
338            memory.order.retain(|key| key != id);
339            memory.order.push_back(id.clone());
340            self.metrics.memory_hits.fetch_add(1, Ordering::Relaxed);
341            self.metrics
342                .bytes_memory
343                .fetch_add(value.len() as u64, Ordering::Relaxed);
344            Some(value)
345        } else {
346            self.metrics.memory_misses.fetch_add(1, Ordering::Relaxed);
347            None
348        }
349    }
350
351    fn memory_insert(&self, id: BlobId, bytes: Arc<[u8]>) {
352        if self.memory_capacity == 0 || bytes.len() as u64 > self.memory_capacity {
353            return;
354        }
355        let mut memory = self.memory.lock().expect("cache lock poisoned");
356        if let Some(old) = memory.entries.remove(&id) {
357            memory.bytes -= old.len() as u64;
358            memory.order.retain(|key| key != &id);
359        }
360        memory.bytes += bytes.len() as u64;
361        memory.order.push_back(id.clone());
362        memory.entries.insert(id, bytes);
363        while memory.bytes > self.memory_capacity {
364            if let Some(old) = memory.order.pop_front()
365                && let Some(value) = memory.entries.remove(&old)
366            {
367                memory.bytes -= value.len() as u64;
368            }
369        }
370        self.metrics
371            .memory_bytes
372            .store(memory.bytes, Ordering::Relaxed);
373        self.metrics
374            .memory_entries
375            .store(memory.entries.len() as u64, Ordering::Relaxed);
376    }
377
378    async fn disk_get(&self, id: &BlobId) -> Option<Arc<[u8]>> {
379        let disk = self.disk.as_ref()?;
380        let (entry, path) = {
381            let mut tier = disk.lock().expect("cache lock poisoned");
382            let Some(entry) = tier.entries.get(id).cloned() else {
383                self.metrics.disk_misses.fetch_add(1, Ordering::Relaxed);
384                return None;
385            };
386            let path = tier.object_path(id);
387            *tier.readers.entry(id.clone()).or_default() += 1;
388            (entry, path)
389        };
390        let read = tokio::task::spawn_blocking(move || fs::read(path)).await;
391        let mut tier = disk.lock().expect("cache lock poisoned");
392        if let Some(readers) = tier.readers.get_mut(id) {
393            *readers -= 1;
394            if *readers == 0 {
395                tier.readers.remove(id);
396            }
397        }
398        match read {
399            Ok(Ok(bytes)) if bytes.len() as u64 == entry.len && digest(&bytes) == *id => {
400                tier.generation += 1;
401                let generation = tier.generation;
402                tier.entries.get_mut(id).expect("entry exists").generation = generation;
403                self.metrics.disk_hits.fetch_add(1, Ordering::Relaxed);
404                self.metrics
405                    .bytes_disk
406                    .fetch_add(entry.len, Ordering::Relaxed);
407                Some(bytes.into())
408            }
409            Ok(Ok(_) | Err(_)) | Err(_) => {
410                if let Err(error) = tier.remove_corrupt(id) {
411                    tracing::warn!(%error, segment_id = %id, "segment cache could not discard corrupt entry");
412                }
413                self.metrics.disk_misses.fetch_add(1, Ordering::Relaxed);
414                self.metrics.corruptions.fetch_add(1, Ordering::Relaxed);
415                self.metrics.disk_bytes.store(tier.bytes, Ordering::Relaxed);
416                self.metrics
417                    .disk_entries
418                    .store(tier.entries.len() as u64, Ordering::Relaxed);
419                None
420            }
421        }
422    }
423
424    fn disk_admit(&self, id: &BlobId, bytes: &[u8]) {
425        let Some(disk) = &self.disk else {
426            return;
427        };
428        let mut tier = disk.lock().expect("cache lock poisoned");
429        let accounted = (bytes.len() as u64).saturating_add(ENTRY_METADATA_BYTES);
430        if accounted > tier.config.capacity_bytes {
431            self.metrics.too_large.fetch_add(1, Ordering::Relaxed);
432            return;
433        }
434        let already_present = tier.entries.contains_key(id);
435        let evicted = if already_present {
436            Vec::new()
437        } else {
438            let limit = tier.config.capacity_bytes - accounted;
439            let before = tier
440                .entries
441                .iter()
442                .map(|(id, entry)| (id.clone(), entry.len))
443                .collect::<HashMap<_, _>>();
444            match tier.evict_to_limit(limit, None) {
445                Ok(evicted) => evicted,
446                Err(error) => {
447                    let evicted = before
448                        .into_iter()
449                        .filter(|(id, _)| !tier.entries.contains_key(id))
450                        .collect::<Vec<_>>();
451                    self.record_evictions(&evicted);
452                    self.metrics.disk_bytes.store(tier.bytes, Ordering::Relaxed);
453                    self.metrics
454                        .disk_entries
455                        .store(tier.entries.len() as u64, Ordering::Relaxed);
456                    self.metrics
457                        .admission_errors
458                        .fetch_add(1, Ordering::Relaxed);
459                    tracing::warn!(%error, segment_id = %id, "segment cache could not make admission space");
460                    return;
461                }
462            }
463        };
464        self.record_evictions(&evicted);
465        if let Err(error) = tier.admit(id, bytes) {
466            self.metrics
467                .admission_errors
468                .fetch_add(1, Ordering::Relaxed);
469            tracing::warn!(%error, segment_id = %id, "segment cache admission failed");
470        } else {
471            self.metrics.admissions.fetch_add(1, Ordering::Relaxed);
472        }
473        self.metrics.disk_bytes.store(tier.bytes, Ordering::Relaxed);
474        self.metrics
475            .disk_entries
476            .store(tier.entries.len() as u64, Ordering::Relaxed);
477    }
478
479    fn memory_remove(&self, id: &BlobId) {
480        let mut memory = self.memory.lock().expect("cache lock poisoned");
481        if let Some(value) = memory.entries.remove(id) {
482            memory.bytes -= value.len() as u64;
483            memory.order.retain(|key| key != id);
484            self.metrics
485                .memory_bytes
486                .store(memory.bytes, Ordering::Relaxed);
487            self.metrics
488                .memory_entries
489                .store(memory.entries.len() as u64, Ordering::Relaxed);
490        }
491    }
492
493    fn record_evictions(&self, evicted: &[(BlobId, u64)]) {
494        for (victim, _) in evicted {
495            self.memory_remove(victim);
496        }
497        self.metrics
498            .evictions
499            .fetch_add(evicted.len() as u64, Ordering::Relaxed);
500        self.metrics.evicted_bytes.fetch_add(
501            evicted.iter().map(|(_, len)| len).sum::<u64>(),
502            Ordering::Relaxed,
503        );
504    }
505}
506
507impl DiskTier {
508    fn object_path(&self, id: &BlobId) -> PathBuf {
509        self.config
510            .directory
511            .join("objects")
512            .join(&id.as_str()[..2])
513            .join(&id.as_str()[2..])
514    }
515    fn reconcile(&mut self) -> io::Result<()> {
516        let index_path = self.config.directory.join("index");
517        let persisted = match self.load_index() {
518            Ok(entries) => entries,
519            Err(error) => {
520                let quarantine = self.config.directory.join(format!(
521                    "index.corrupt-{}",
522                    std::time::SystemTime::now()
523                        .duration_since(std::time::UNIX_EPOCH)
524                        .unwrap_or_default()
525                        .as_nanos()
526                ));
527                fs::rename(&index_path, quarantine).map_err(|rename_error| {
528                    io::Error::new(
529                        rename_error.kind(),
530                        format!("cannot quarantine corrupt cache index ({error}): {rename_error}"),
531                    )
532                })?;
533                HashMap::new()
534            }
535        };
536        self.generation = persisted
537            .values()
538            .map(|entry| entry.generation)
539            .max()
540            .unwrap_or(0);
541        let root = self.config.directory.join("objects");
542        for fan in fs::read_dir(root)? {
543            let fan = fan?;
544            if !fan.file_type()?.is_dir() {
545                continue;
546            }
547            let prefix = fan.file_name().to_string_lossy().into_owned();
548            for file in fs::read_dir(fan.path())? {
549                let file = file?;
550                if !file.file_type()?.is_file() {
551                    continue;
552                }
553                let text = format!("{prefix}{}", file.file_name().to_string_lossy());
554                if let Some(id) = BlobId::from_hex(&text) {
555                    let len = file.metadata()?.len();
556                    let generation = persisted.get(&id).map_or_else(
557                        || {
558                            self.generation += 1;
559                            self.generation
560                        },
561                        |entry| entry.generation,
562                    );
563                    self.bytes += len.saturating_add(ENTRY_METADATA_BYTES);
564                    self.entries.insert(id, Entry { len, generation });
565                } else {
566                    fs::remove_file(file.path())?;
567                }
568            }
569        }
570        self.persist()
571    }
572
573    fn load_index(&self) -> io::Result<HashMap<BlobId, Entry>> {
574        let text = match fs::read_to_string(self.config.directory.join("index")) {
575            Ok(text) => text,
576            Err(error) if error.kind() == io::ErrorKind::NotFound => return Ok(HashMap::new()),
577            Err(error) => return Err(error),
578        };
579        let mut entries = HashMap::new();
580        for line in text.lines() {
581            let mut fields = line.split_whitespace();
582            let id = fields.next().and_then(BlobId::from_hex);
583            let len = fields.next().and_then(|field| field.parse::<u64>().ok());
584            let generation = fields.next().and_then(|field| field.parse::<u64>().ok());
585            if fields.next().is_some() || id.is_none() || len.is_none() || generation.is_none() {
586                return Err(io::Error::new(
587                    io::ErrorKind::InvalidData,
588                    "malformed segment cache index",
589                ));
590            }
591            entries.insert(
592                id.expect("checked"),
593                Entry {
594                    len: len.expect("checked"),
595                    generation: generation.expect("checked"),
596                },
597            );
598        }
599        Ok(entries)
600    }
601    fn admit(&mut self, id: &BlobId, bytes: &[u8]) -> io::Result<()> {
602        self.generation += 1;
603        if let Some(entry) = self.entries.get_mut(id) {
604            entry.generation = self.generation;
605            return self.persist();
606        }
607        let target = self.object_path(id);
608        let parent = target.parent().expect("object parent");
609        fs::create_dir_all(parent)?;
610        set_owner_only(parent)?;
611        let temp = self
612            .config
613            .directory
614            .join("tmp")
615            .join(format!("{}-{}", id, self.generation));
616        let mut file = OpenOptions::new()
617            .write(true)
618            .create_new(true)
619            .open(&temp)?;
620        set_file_owner_only(&temp)?;
621        file.write_all(bytes)?;
622        file.sync_all()?;
623        fs::rename(&temp, &target)?;
624        sync_directory(parent)?;
625        set_file_owner_only(&target)?;
626        let len = bytes.len() as u64;
627        self.entries.insert(
628            id.clone(),
629            Entry {
630                len,
631                generation: self.generation,
632            },
633        );
634        self.bytes += len.saturating_add(ENTRY_METADATA_BYTES);
635        self.persist()
636    }
637    fn evict_to_limit(
638        &mut self,
639        limit: u64,
640        protected: Option<&BlobId>,
641    ) -> io::Result<Vec<(BlobId, u64)>> {
642        let mut evicted = Vec::new();
643        while self.bytes > limit {
644            let victim = self
645                .entries
646                .iter()
647                .filter(|(id, _)| protected != Some(*id) && !self.readers.contains_key(*id))
648                .min_by_key(|(_, entry)| entry.generation)
649                .map(|(id, _)| id.clone());
650            let Some(victim) = victim else {
651                return Err(io::Error::new(
652                    io::ErrorKind::WouldBlock,
653                    "all cache eviction candidates are being read",
654                ));
655            };
656            let entry = self.entries.get(&victim).expect("victim exists").clone();
657            fs::remove_file(self.object_path(&victim))?;
658            self.entries.remove(&victim);
659            self.bytes -= entry.accounted_bytes();
660            evicted.push((victim, entry.len));
661        }
662        self.persist()?;
663        Ok(evicted)
664    }
665    fn remove_corrupt(&mut self, id: &BlobId) -> io::Result<()> {
666        match fs::remove_file(self.object_path(id)) {
667            Ok(()) => {}
668            Err(error) if error.kind() == io::ErrorKind::NotFound => {}
669            Err(error) => return Err(error),
670        }
671        if let Some(entry) = self.entries.remove(id) {
672            self.bytes -= entry.accounted_bytes();
673        }
674        self.persist()
675    }
676    fn persist(&self) -> io::Result<()> {
677        let mut sorted = BTreeMap::new();
678        for (id, entry) in &self.entries {
679            sorted.insert(id.as_str(), entry);
680        }
681        let temp = self.config.directory.join("index.tmp");
682        let mut file = File::create(&temp)?;
683        for (id, entry) in sorted {
684            writeln!(file, "{id} {} {}", entry.len, entry.generation)?;
685        }
686        file.sync_all()?;
687        fs::rename(temp, self.config.directory.join("index"))?;
688        sync_directory(&self.config.directory)
689    }
690}
691
692impl Drop for DiskTier {
693    fn drop(&mut self) {
694        if let Err(error) = self.persist() {
695            tracing::warn!(%error, "segment cache could not flush access generations during shutdown");
696        }
697    }
698}
699
700#[cfg(unix)]
701fn set_owner_only(path: &Path) -> io::Result<()> {
702    use std::os::unix::fs::PermissionsExt;
703    fs::set_permissions(path, fs::Permissions::from_mode(0o700))
704}
705#[cfg(not(unix))]
706fn set_owner_only(_path: &Path) -> io::Result<()> {
707    Ok(())
708}
709#[cfg(unix)]
710fn set_file_owner_only(path: &Path) -> io::Result<()> {
711    use std::os::unix::fs::PermissionsExt;
712    fs::set_permissions(path, fs::Permissions::from_mode(0o600))
713}
714#[cfg(not(unix))]
715fn set_file_owner_only(_path: &Path) -> io::Result<()> {
716    Ok(())
717}
718
719#[cfg(unix)]
720fn sync_directory(path: &Path) -> io::Result<()> {
721    File::open(path)?.sync_all()
722}
723#[cfg(not(unix))]
724fn sync_directory(_path: &Path) -> io::Result<()> {
725    Ok(())
726}
727
728#[cfg(test)]
729mod tests {
730    use super::*;
731    use crate::MemoryStore;
732    use std::sync::atomic::AtomicUsize;
733    use tokio::sync::Notify;
734
735    struct BlockingReader {
736        reads: AtomicUsize,
737        release: Notify,
738        bytes: Option<Vec<u8>>,
739    }
740
741    #[async_trait]
742    impl SegmentReader for BlockingReader {
743        async fn read_segment(&self, _id: &BlobId) -> Result<Option<Vec<u8>>, StoreError> {
744            self.reads.fetch_add(1, Ordering::Relaxed);
745            self.release.notified().await;
746            Ok(self.bytes.clone())
747        }
748    }
749
750    fn config(path: &Path, capacity: u64) -> SegmentCacheConfig {
751        SegmentCacheConfig {
752            directory: path.to_owned(),
753            capacity_bytes: capacity,
754            memory_capacity_bytes: 0,
755        }
756    }
757
758    #[tokio::test]
759    async fn disk_hit_survives_reopen_without_native_read() {
760        let directory = tempfile::tempdir().expect("tempdir");
761        let store = MemoryStore::default();
762        let id = store.put(b"cached on disk").await.expect("put");
763        {
764            let cache = SegmentCache::open(&config(directory.path(), 1024)).expect("open cache");
765            cache.get_or_load(&store, &id).await.expect("cold read");
766        }
767        let empty = MemoryStore::default();
768        let cache = SegmentCache::open(&config(directory.path(), 1024)).expect("reopen cache");
769        assert_eq!(
770            cache
771                .get_or_load(&empty, &id)
772                .await
773                .expect("warm read")
774                .as_deref(),
775            Some(b"cached on disk".as_slice())
776        );
777        assert_eq!(cache.metrics.disk_hits.load(Ordering::Relaxed), 1);
778    }
779
780    #[tokio::test]
781    async fn disk_capacity_evicts_least_recently_used_entry() {
782        let directory = tempfile::tempdir().expect("tempdir");
783        let store = MemoryStore::default();
784        let first = store.put(b"one").await.expect("first");
785        let second = store.put(b"two").await.expect("second");
786        let cache = SegmentCache::open(&config(directory.path(), 99)).expect("cache");
787        cache.get_or_load(&store, &first).await.expect("load first");
788        cache
789            .get_or_load(&store, &second)
790            .await
791            .expect("load second");
792        assert_eq!(cache.metrics.disk_bytes.load(Ordering::Relaxed), 99);
793        assert_eq!(cache.metrics.disk_entries.load(Ordering::Relaxed), 1);
794        assert_eq!(cache.metrics.evictions.load(Ordering::Relaxed), 1);
795        assert!(
796            !cache
797                .disk
798                .as_ref()
799                .expect("disk")
800                .lock()
801                .expect("lock")
802                .entries
803                .contains_key(&first)
804        );
805    }
806
807    #[tokio::test]
808    async fn persisted_recency_controls_eviction_after_restart() {
809        let directory = tempfile::tempdir().expect("tempdir");
810        let store = MemoryStore::default();
811        let first = store.put(b"one").await.expect("first");
812        let second = store.put(b"two").await.expect("second");
813        let third = store.put(b"tri").await.expect("third");
814        {
815            let cache = SegmentCache::open(&config(directory.path(), 198)).expect("cache");
816            cache.get_or_load(&store, &first).await.expect("first read");
817            cache
818                .get_or_load(&store, &second)
819                .await
820                .expect("second read");
821            cache
822                .get_or_load(&store, &first)
823                .await
824                .expect("refresh first");
825        }
826        let cache = SegmentCache::open(&config(directory.path(), 198)).expect("reopen");
827        cache.get_or_load(&store, &third).await.expect("third read");
828        let tier = cache.disk.as_ref().expect("disk").lock().expect("lock");
829        assert!(tier.entries.contains_key(&first));
830        assert!(!tier.entries.contains_key(&second));
831        assert!(tier.entries.contains_key(&third));
832    }
833
834    #[tokio::test]
835    async fn corrupt_disk_entry_self_heals_from_native_storage() {
836        let directory = tempfile::tempdir().expect("tempdir");
837        let store = MemoryStore::default();
838        let id = store.put(b"healthy").await.expect("put");
839        let cache = SegmentCache::open(&config(directory.path(), 1024)).expect("cache");
840        cache.get_or_load(&store, &id).await.expect("cold read");
841        let path = cache
842            .disk
843            .as_ref()
844            .expect("disk")
845            .lock()
846            .expect("lock")
847            .object_path(&id);
848        fs::write(path, b"broken").expect("corrupt file");
849        assert_eq!(
850            cache
851                .get_or_load(&store, &id)
852                .await
853                .expect("heal")
854                .as_deref(),
855            Some(b"healthy".as_slice())
856        );
857        assert_eq!(cache.metrics.corruptions.load(Ordering::Relaxed), 1);
858        assert_eq!(cache.metrics.native_found.load(Ordering::Relaxed), 2);
859    }
860
861    #[tokio::test]
862    async fn concurrent_native_misses_are_coalesced_including_not_found() {
863        let reader = Arc::new(BlockingReader {
864            reads: AtomicUsize::new(0),
865            release: Notify::new(),
866            bytes: None,
867        });
868        let cache = Arc::new(SegmentCache::memory_only(0));
869        let id = digest(b"missing");
870        let mut tasks = Vec::new();
871        for _ in 0..8 {
872            let cache = Arc::clone(&cache);
873            let reader = Arc::clone(&reader);
874            let id = id.clone();
875            tasks.push(tokio::spawn(async move {
876                cache.get_or_load(reader.as_ref(), &id).await
877            }));
878        }
879        while cache.metrics.coalesced_waiters.load(Ordering::Relaxed) < 7 {
880            tokio::task::yield_now().await;
881        }
882        assert_eq!(reader.reads.load(Ordering::Relaxed), 1);
883        reader.release.notify_waiters();
884        for task in tasks {
885            assert!(task.await.expect("join").expect("read").is_none());
886        }
887        assert_eq!(reader.reads.load(Ordering::Relaxed), 1);
888        assert_eq!(cache.metrics.coalesced_waiters.load(Ordering::Relaxed), 7);
889    }
890
891    #[tokio::test]
892    async fn object_larger_than_accounted_capacity_bypasses_disk() {
893        let directory = tempfile::tempdir().expect("tempdir");
894        let store = MemoryStore::default();
895        let id = store.put(b"large").await.expect("put");
896        let cache = SegmentCache::open(&config(directory.path(), 100)).expect("cache");
897        assert_eq!(
898            cache
899                .get_or_load(&store, &id)
900                .await
901                .expect("read")
902                .as_deref(),
903            Some(b"large".as_slice())
904        );
905        assert_eq!(cache.metrics.too_large.load(Ordering::Relaxed), 1);
906        assert_eq!(cache.metrics.disk_bytes.load(Ordering::Relaxed), 0);
907        assert_eq!(cache.metrics.disk_entries.load(Ordering::Relaxed), 0);
908    }
909
910    #[test]
911    fn rejects_second_process_and_invalid_capacity() {
912        let directory = tempfile::tempdir().expect("tempdir");
913        let cache = SegmentCache::open(&config(directory.path(), 10)).expect("first owner");
914        assert!(SegmentCache::open(&config(directory.path(), 10)).is_err());
915        drop(cache);
916        let mut invalid = config(directory.path(), 10);
917        invalid.memory_capacity_bytes = 11;
918        assert!(SegmentCache::open(&invalid).is_err());
919    }
920
921    #[tokio::test]
922    async fn malformed_index_is_quarantined_and_rebuilt() {
923        let directory = tempfile::tempdir().expect("tempdir");
924        let store = MemoryStore::default();
925        let id = store.put(b"survivor").await.expect("put");
926        {
927            let cache = SegmentCache::open(&config(directory.path(), 1024)).expect("cache");
928            cache.get_or_load(&store, &id).await.expect("read");
929        }
930        fs::write(directory.path().join("index"), "not an index\n").expect("break index");
931        let cache = SegmentCache::open(&config(directory.path(), 1024)).expect("rebuild");
932        assert!(
933            cache
934                .disk
935                .as_ref()
936                .expect("disk")
937                .lock()
938                .expect("lock")
939                .entries
940                .contains_key(&id)
941        );
942        assert!(
943            fs::read_dir(directory.path())
944                .expect("read directory")
945                .any(|item| {
946                    item.expect("entry")
947                        .file_name()
948                        .to_string_lossy()
949                        .starts_with("index.corrupt-")
950                })
951        );
952    }
953}
corium_store/segment_cache.rs

corium_store/
segment_cache.rs
