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

corium_store/
segment_cache.rs
