newton-core 0.7.4

//! Local disk tier for content-addressed objects.
//!
//! Objects are immutable and named by CID, so a local copy never goes stale: the
//! backend (S3 in prod) stays the source of truth and is read only on a miss. Hits are
//! served from the filesystem (and the kernel page cache), not process memory.
//!
//! Layout: `<dir>/objects/<cid>` for objects, `<dir>/.staging/` for writes in progress.
//! Only CID-named files under `objects/` are ever indexed or deleted, so pointing `dir`
//! at a shared directory cannot evict unrelated files.
//!
//! Each file starts with the SHA2-256 multihash the resolver verified the object
//! against when it was filled. Files are written to staging, flushed, then renamed into
//! place, so a reader sees either nothing or a whole object. A crash can still leave a
//! truncated file behind, so every object found on disk at open is re-hashed against
//! its header on its first read; a mismatch is deleted and read from the backend again.

use std::{
    collections::{BTreeMap, HashMap},
    io::Write as _,
    path::{Path, PathBuf},
    sync::{
        atomic::{AtomicU64, Ordering},
        Arc, Mutex, OnceLock, Weak,
    },
};

use cid::{multihash::Multihash, Cid};
use eyre::{Result, WrapErr};
use tracing::{debug, warn};

use crate::cid::{sha2_256_multihash, RAW_CODEC};

/// Objects live here, inside the cache directory.
const OBJECTS_DIR: &str = "objects";
/// Staging directory for in-progress writes, inside the cache directory so the final
/// rename never crosses a filesystem.
const STAGING_DIR: &str = ".staging";
/// File header magic, followed by a one-byte multihash length and the multihash.
const MAGIC: &[u8; 4] = b"NOC1";

/// What a fill produced, shared with every task that waited on the same CID so a
/// failure or absence is not retried once per waiter.
pub type FillOutcome = std::result::Result<Option<Arc<[u8]>>, Arc<str>>;

/// Disk cache of verified objects keyed by CID, bounded by `max_bytes` with
/// least-recently-used eviction.
///
/// Every resolver in a process built for the same directory shares one instance (see
/// [`LocalObjectCache::shared`]), so concurrent misses for a CID make one backend read
/// and the size accounting covers the whole directory.
#[derive(Debug)]
pub struct LocalObjectCache {
    dir: PathBuf,
    max_bytes: u64,
    index: Mutex<Index>,
    /// Per-CID fills. Weak, so a fill lives only while a task holds or waits on it.
    fills: Mutex<HashMap<String, Weak<tokio::sync::Mutex<Option<FillOutcome>>>>>,
    staging_seq: AtomicU64,
}

/// What is on disk: size, recency and verification state per CID. Never bytes.
#[derive(Debug, Default)]
struct Index {
    entries: HashMap<String, Entry>,
    /// `last_used` -> CID, oldest first, so eviction never sorts.
    by_age: BTreeMap<u64, String>,
    total_bytes: u64,
    clock: u64,
}

#[derive(Debug, Clone, Copy)]
struct Entry {
    size: u64,
    last_used: u64,
    /// Hashed against its header since this process opened the cache.
    verified: bool,
}

impl Index {
    /// Record a use of `cid`, adding it if new.
    fn touch(&mut self, cid: &str, size: u64, verified: bool) {
        self.clock += 1;
        let clock = self.clock;
        match self.entries.get_mut(cid) {
            Some(entry) => {
                self.by_age.remove(&entry.last_used);
                entry.last_used = clock;
                entry.verified |= verified;
            }
            None => {
                self.entries.insert(
                    cid.to_owned(),
                    Entry {
                        size,
                        last_used: clock,
                        verified,
                    },
                );
                self.total_bytes += size;
            }
        }
        self.by_age.insert(clock, cid.to_owned());
    }

    fn remove(&mut self, cid: &str) {
        if let Some(entry) = self.entries.remove(cid) {
            self.by_age.remove(&entry.last_used);
            self.total_bytes = self.total_bytes.saturating_sub(entry.size);
        }
    }

    /// Drop least-recently-used entries, sparing `keep`, until `total_bytes <= max_bytes`.
    /// Returns the CIDs whose files should be deleted.
    fn evict(&mut self, max_bytes: u64, keep: &str) -> Vec<String> {
        let mut victims = Vec::new();
        let mut spared = None;
        while self.total_bytes > max_bytes {
            let Some((&clock, _)) = self.by_age.iter().next() else {
                break;
            };
            let cid = self.by_age.remove(&clock).expect("present");
            if cid == keep {
                spared = Some((clock, cid));
                continue;
            }
            if let Some(entry) = self.entries.remove(&cid) {
                self.total_bytes = self.total_bytes.saturating_sub(entry.size);
            }
            victims.push(cid);
        }
        if let Some((clock, cid)) = spared {
            self.by_age.insert(clock, cid);
        }
        victims
    }
}

impl LocalObjectCache {
    /// The process-wide cache for `dir`, opening it on first use. The directory scan
    /// runs on a blocking thread.
    ///
    /// # Errors
    ///
    /// Fails if the directory cannot be created or listed.
    pub async fn shared(dir: &Path, max_bytes: u64) -> Result<Arc<Self>> {
        static CACHES: OnceLock<tokio::sync::Mutex<HashMap<PathBuf, Arc<LocalObjectCache>>>> = OnceLock::new();
        let mut caches = CACHES.get_or_init(Default::default).lock().await;
        if let Some(cache) = caches.get(dir) {
            return Ok(cache.clone());
        }
        let owned = dir.to_path_buf();
        let cache = Arc::new(
            tokio::task::spawn_blocking(move || Self::open(&owned, max_bytes))
                .await
                .wrap_err("object cache open task panicked")??,
        );
        caches.insert(dir.to_path_buf(), cache.clone());
        Ok(cache)
    }

    /// Open `dir`: create it, clear interrupted writes, and index the objects already
    /// there, oldest-modified first so eviction starts with objects not refreshed
    /// recently. Blocking.
    ///
    /// # Errors
    ///
    /// Fails if the directory cannot be created or listed.
    pub fn open(dir: &Path, max_bytes: u64) -> Result<Self> {
        let staging = dir.join(STAGING_DIR);
        if staging.exists() {
            std::fs::remove_dir_all(&staging)
                .wrap_err_with(|| format!("clear object cache staging dir {}", staging.display()))?;
        }
        std::fs::create_dir_all(&staging).wrap_err_with(|| format!("create object cache dir {}", staging.display()))?;
        let objects = dir.join(OBJECTS_DIR);
        std::fs::create_dir_all(&objects).wrap_err_with(|| format!("create object cache dir {}", objects.display()))?;

        let mut found = Vec::new();
        for entry in
            std::fs::read_dir(&objects).wrap_err_with(|| format!("list object cache dir {}", objects.display()))?
        {
            // An entry that vanishes or cannot be stat'ed mid-scan is simply not indexed.
            let Ok(entry) = entry else { continue };
            let Ok(metadata) = entry.metadata() else { continue };
            if !metadata.is_file() {
                continue;
            }
            let Ok(cid) = entry.file_name().into_string() else {
                continue;
            };
            if Cid::try_from(cid.as_str()).is_err() {
                continue;
            }
            found.push((metadata.modified().ok(), cid, metadata.len()));
        }
        found.sort();

        let mut index = Index::default();
        for (_, cid, size) in found {
            index.touch(&cid, size, false);
        }
        debug!(
            dir = %dir.display(),
            objects = index.entries.len(),
            bytes = index.total_bytes,
            max_bytes,
            "opened local object cache"
        );

        let cache = Self {
            dir: dir.to_path_buf(),
            max_bytes,
            index: Mutex::new(index),
            fills: Mutex::new(HashMap::new()),
            staging_seq: AtomicU64::new(0),
        };
        let victims = cache.lock_index().evict(max_bytes, "");
        remove_files(&objects, victims);
        Ok(cache)
    }

    /// The cached bytes for `cid`, or `None` on a miss. An object not yet verified by
    /// this process is hashed against its header first; a corrupt one is deleted and
    /// reported as a miss. Read errors other than "not found" are logged and treated as
    /// a miss, so the caller falls through to the backend.
    pub async fn read(&self, cid: &str) -> Option<Vec<u8>> {
        let path = self.object_path(cid);
        let file = match tokio::fs::read(&path).await {
            Ok(file) => file,
            Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
                // Evicted or removed out from under the index.
                self.lock_index().remove(cid);
                return None;
            }
            Err(err) => {
                warn!(cid, error = %err, "local object cache read failed; reading backend");
                return None;
            }
        };
        let size = file.len() as u64;
        let verified = self.lock_index().entries.get(cid).is_some_and(|entry| entry.verified);
        let Some((expected, data_start)) = parse_header(&file) else {
            self.discard(cid, &path, "unreadable header").await;
            return None;
        };
        if !verified {
            let ok = expected_for(cid, &expected) && sha2_256_multihash(&file[data_start..]) == expected;
            if !ok {
                self.discard(cid, &path, "content does not match its checksum").await;
                return None;
            }
        }
        self.lock_index().touch(cid, size, true);
        let mut file = file;
        file.drain(..data_start);
        Some(file)
    }

    /// Hold while filling `cid`, so concurrent misses for one CID make one backend read.
    /// If the guard already holds an outcome, another task did the fill: use it instead
    /// of reading the backend. Otherwise re-check [`Self::read`], fill, and store the
    /// outcome for the tasks still waiting.
    pub async fn fill_lock(&self, cid: &str) -> tokio::sync::OwnedMutexGuard<Option<FillOutcome>> {
        let lock = {
            let mut fills = self.fills.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
            match fills.get(cid).and_then(Weak::upgrade) {
                Some(lock) => lock,
                None => {
                    fills.retain(|_, weak| weak.strong_count() > 0);
                    let lock = Arc::new(tokio::sync::Mutex::new(None));
                    fills.insert(cid.to_owned(), Arc::downgrade(&lock));
                    lock
                }
            }
        };
        lock.lock_owned().await
    }

    /// Store `bytes` for `cid`, already verified against the SHA2-256 multihash
    /// `expected`. Best effort: a failure is logged and the object is simply read from
    /// the backend again next time. An object larger than the whole budget is not
    /// stored: keeping it would leave the cache over `max_bytes` and evict everything
    /// else first.
    pub async fn write(&self, cid: &str, expected: &Multihash<64>, bytes: &[u8]) {
        let size = (MAGIC.len() + 1 + expected.encoded_len() + bytes.len()) as u64;
        if size > self.max_bytes {
            warn!(
                cid,
                size,
                max_bytes = self.max_bytes,
                "object is larger than the local object cache; not caching it"
            );
            return;
        }
        let seq = self.staging_seq.fetch_add(1, Ordering::Relaxed);
        let staging_dir = self.dir.join(STAGING_DIR);
        let staging = staging_dir.join(format!("{cid}.{}.{seq}", std::process::id()));
        let target = self.object_path(cid);
        let mut contents = Vec::with_capacity(MAGIC.len() + 1 + expected.encoded_len() + bytes.len());
        contents.extend_from_slice(MAGIC);
        let header = expected.to_bytes();
        contents.push(header.len() as u8);
        contents.extend_from_slice(&header);
        contents.extend_from_slice(bytes);

        let objects_dir = self.dir.join(OBJECTS_DIR);
        let result = tokio::task::spawn_blocking(move || {
            // Recreate the directories if they were removed while the process runs.
            if !staging_dir.is_dir() {
                std::fs::create_dir_all(&staging_dir)?;
            }
            if !objects_dir.is_dir() {
                std::fs::create_dir_all(&objects_dir)?;
            }
            let written = (|| {
                let mut file = std::fs::File::create(&staging)?;
                file.write_all(&contents)?;
                file.sync_all()?;
                std::fs::rename(&staging, &target)
            })();
            if written.is_err() {
                let _ = std::fs::remove_file(&staging);
            }
            written
        })
        .await;
        match result {
            Ok(Ok(())) => {}
            Ok(Err(err)) => {
                warn!(cid, error = %err, "local object cache write failed");
                return;
            }
            Err(err) => {
                warn!(cid, error = %err, "local object cache write task failed");
                return;
            }
        }
        let victims = {
            let mut index = self.lock_index();
            index.touch(cid, size, true);
            index.evict(self.max_bytes, cid)
        };
        if !victims.is_empty() {
            let objects = self.dir.join(OBJECTS_DIR);
            // A reader holding a file open keeps its bytes; later reads miss and refill.
            let _ = tokio::task::spawn_blocking(move || remove_files(&objects, victims)).await;
        }
    }

    /// Drop a corrupt object from disk and the index.
    async fn discard(&self, cid: &str, path: &Path, reason: &str) {
        warn!(
            cid,
            reason, "local object cache entry is corrupt; discarding and reading backend"
        );
        self.lock_index().remove(cid);
        if let Err(err) = tokio::fs::remove_file(path).await {
            if err.kind() != std::io::ErrorKind::NotFound {
                warn!(cid, error = %err, "local object cache discard failed");
            }
        }
    }

    fn object_path(&self, cid: &str) -> PathBuf {
        self.dir.join(OBJECTS_DIR).join(cid)
    }

    fn lock_index(&self) -> std::sync::MutexGuard<'_, Index> {
        self.index.lock().unwrap_or_else(std::sync::PoisonError::into_inner)
    }

    /// Bytes currently indexed on disk.
    #[must_use]
    pub fn total_bytes(&self) -> u64 {
        self.lock_index().total_bytes
    }
}

/// The expected multihash and the offset where object bytes start, if the header is
/// well formed and names a SHA2-256 digest.
fn parse_header(file: &[u8]) -> Option<(Multihash<64>, usize)> {
    let rest = file.strip_prefix(MAGIC.as_slice())?;
    let (&len, rest) = rest.split_first()?;
    let len = usize::from(len);
    let expected = Multihash::<64>::from_bytes(rest.get(..len)?).ok()?;
    if expected.code() != sha2_256_multihash(&[]).code() {
        return None;
    }
    Some((expected, MAGIC.len() + 1 + len))
}

/// Whether `expected` may stand for `cid`: a raw CID must carry exactly that digest. A
/// DAG-PB CID's digest is over the DAG node, so its header holds the checksum sidecar
/// verified at fill time instead.
fn expected_for(cid: &str, expected: &Multihash<64>) -> bool {
    match Cid::try_from(cid) {
        Ok(parsed) if parsed.codec() == RAW_CODEC => parsed.hash() == expected,
        Ok(_) => true,
        Err(_) => false,
    }
}

/// Delete evicted object files. Blocking.
fn remove_files(objects: &Path, victims: Vec<String>) {
    for cid in victims {
        if let Err(err) = std::fs::remove_file(objects.join(&cid)) {
            if err.kind() != std::io::ErrorKind::NotFound {
                warn!(cid, error = %err, "local object cache eviction failed");
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::cid::cid_v1_raw;

    fn object(data: &[u8]) -> (String, Multihash<64>) {
        (cid_v1_raw(data), sha2_256_multihash(data))
    }

    #[tokio::test]
    async fn write_then_read_round_trips_and_misses_are_none() {
        let tmp = tempfile::tempdir().unwrap();
        let cache = LocalObjectCache::open(tmp.path(), 1 << 20).unwrap();
        let (cid, mh) = object(b"hello");
        assert!(cache.read(&cid).await.is_none());
        cache.write(&cid, &mh, b"hello").await;
        assert_eq!(cache.read(&cid).await.unwrap(), b"hello");
        assert!(tmp.path().join(OBJECTS_DIR).join(&cid).is_file());
    }

    #[tokio::test]
    async fn evicts_least_recently_used_over_budget() {
        let tmp = tempfile::tempdir().unwrap();
        let (a, mh_a) = object(b"aaaa");
        let (b, mh_b) = object(b"bbbb");
        let (c, mh_c) = object(b"cccc");
        // Room for two objects with their headers, not three.
        let one = {
            let probe = tempfile::tempdir().unwrap();
            let cache = LocalObjectCache::open(probe.path(), u64::MAX).unwrap();
            cache.write(&a, &mh_a, b"aaaa").await;
            cache.total_bytes()
        };
        let cache = LocalObjectCache::open(tmp.path(), one * 2).unwrap();
        cache.write(&a, &mh_a, b"aaaa").await;
        cache.write(&b, &mh_b, b"bbbb").await;
        // Use `a`, so `b` is the least recently used.
        assert!(cache.read(&a).await.is_some());
        cache.write(&c, &mh_c, b"cccc").await;
        assert!(cache.read(&b).await.is_none());
        assert!(cache.read(&a).await.is_some());
        assert!(cache.read(&c).await.is_some());
        assert!(cache.total_bytes() <= one * 2);
    }

    #[tokio::test]
    async fn object_larger_than_the_budget_is_not_stored() {
        let tmp = tempfile::tempdir().unwrap();
        let (small, mh_small) = object(b"small");
        let one = {
            let probe = tempfile::tempdir().unwrap();
            let cache = LocalObjectCache::open(probe.path(), u64::MAX).unwrap();
            cache.write(&small, &mh_small, b"small").await;
            cache.total_bytes()
        };
        let cache = LocalObjectCache::open(tmp.path(), one * 2).unwrap();
        cache.write(&small, &mh_small, b"small").await;

        let big_bytes = vec![7u8; usize::try_from(one * 2).unwrap()];
        let (big, mh_big) = object(&big_bytes);
        cache.write(&big, &mh_big, &big_bytes).await;

        assert!(cache.read(&big).await.is_none());
        assert!(!tmp.path().join(OBJECTS_DIR).join(&big).exists());
        // Nothing already cached was evicted to make room for it.
        assert_eq!(cache.read(&small).await.unwrap(), b"small");
        assert_eq!(cache.total_bytes(), one);
    }

    #[tokio::test]
    async fn reopen_indexes_existing_objects_and_clears_staging() {
        let tmp = tempfile::tempdir().unwrap();
        let (cid, mh) = object(b"abc");
        {
            let cache = LocalObjectCache::open(tmp.path(), 1 << 20).unwrap();
            cache.write(&cid, &mh, b"abc").await;
        }
        std::fs::write(tmp.path().join(STAGING_DIR).join("partial"), b"x").unwrap();
        let cache = LocalObjectCache::open(tmp.path(), 1 << 20).unwrap();
        assert!(cache.total_bytes() > 0);
        assert!(!tmp.path().join(STAGING_DIR).join("partial").exists());
        assert_eq!(cache.read(&cid).await.unwrap(), b"abc");
    }

    #[tokio::test]
    async fn truncated_object_is_discarded_on_first_read_after_reopen() {
        let tmp = tempfile::tempdir().unwrap();
        let (cid, mh) = object(b"policy source");
        {
            let cache = LocalObjectCache::open(tmp.path(), 1 << 20).unwrap();
            cache.write(&cid, &mh, b"policy source").await;
        }
        // Simulate a crash that left the data blocks short.
        let path = tmp.path().join(OBJECTS_DIR).join(&cid);
        let mut file = std::fs::read(&path).unwrap();
        file.truncate(file.len() - 3);
        std::fs::write(&path, file).unwrap();

        let cache = LocalObjectCache::open(tmp.path(), 1 << 20).unwrap();
        assert!(cache.read(&cid).await.is_none());
        assert!(!path.exists());
        assert_eq!(cache.total_bytes(), 0);
    }

    #[tokio::test]
    async fn only_cid_named_files_are_indexed_or_evicted() {
        let tmp = tempfile::tempdir().unwrap();
        std::fs::create_dir_all(tmp.path().join(OBJECTS_DIR)).unwrap();
        std::fs::write(tmp.path().join("state.redb"), vec![0u8; 4096]).unwrap();
        std::fs::write(tmp.path().join(OBJECTS_DIR).join("not-a-cid"), vec![0u8; 4096]).unwrap();
        let cache = LocalObjectCache::open(tmp.path(), 16).unwrap();
        assert_eq!(cache.total_bytes(), 0);
        assert!(tmp.path().join("state.redb").exists());
        assert!(tmp.path().join(OBJECTS_DIR).join("not-a-cid").exists());
    }

    #[tokio::test]
    async fn write_recreates_directories_removed_at_runtime() {
        let tmp = tempfile::tempdir().unwrap();
        let cache = LocalObjectCache::open(tmp.path(), 1 << 20).unwrap();
        std::fs::remove_dir_all(tmp.path().join(STAGING_DIR)).unwrap();
        std::fs::remove_dir_all(tmp.path().join(OBJECTS_DIR)).unwrap();
        let (cid, mh) = object(b"again");
        cache.write(&cid, &mh, b"again").await;
        assert_eq!(cache.read(&cid).await.unwrap(), b"again");
    }

    #[tokio::test]
    async fn shared_returns_one_instance_per_dir() {
        let tmp = tempfile::tempdir().unwrap();
        let a = LocalObjectCache::shared(tmp.path(), 1024).await.unwrap();
        let b = LocalObjectCache::shared(tmp.path(), 1024).await.unwrap();
        assert!(Arc::ptr_eq(&a, &b));
    }

    #[tokio::test]
    async fn fill_lock_serializes_one_cid_and_shares_the_outcome() {
        let tmp = tempfile::tempdir().unwrap();
        let cache = Arc::new(LocalObjectCache::open(tmp.path(), 1024).unwrap());
        let mut guard = cache.fill_lock("a").await;
        let waiter = {
            let cache = cache.clone();
            tokio::spawn(async move { cache.fill_lock("a").await.clone() })
        };
        tokio::time::sleep(std::time::Duration::from_millis(20)).await;
        assert!(!waiter.is_finished());
        *guard = Some(Err(Arc::from("backend down")));
        drop(guard);
        let seen = waiter.await.unwrap();
        assert!(matches!(seen, Some(Err(msg)) if &*msg == "backend down"));
    }
}