rete-cli 0.3.7

//! Shared local-or-HTTP range reader for commands that preview or execute
//! bounded reads over either a path or an HTTP(S) `.rete` URL.

use std::fs::File;
use std::io;

use rete_core::{auto_block, BlockCacheReader, RangeReader, Rete};

use crate::http::HttpRangeReader;

pub(crate) fn is_url(source: &str) -> bool {
    source.starts_with("http://") || source.starts_with("https://")
}

/// File size above which a local `.rete` is opened lazily instead of read whole.
/// Overridable with `RETE_LOCAL_LAZY_ABOVE_MB` (0 forces lazy for everything).
const DEFAULT_LAZY_ABOVE_BYTES: u64 = 1 << 30; // 1 GiB

fn lazy_threshold_bytes() -> u64 {
    std::env::var("RETE_LOCAL_LAZY_ABOVE_MB")
        .ok()
        .and_then(|v| v.parse::<u64>().ok())
        .map(|mb| mb.saturating_mul(1 << 20))
        .unwrap_or(DEFAULT_LAZY_ABOVE_BYTES)
}

/// The byte budget a query command runs under, from its `--memory-budget-mb`:
/// omitted = [`rete_core::DEFAULT_QUERY_BUDGET_MB`], `0` = unlimited (`None`).
pub(crate) fn query_budget(mb: Option<u64>) -> Option<u64> {
    match mb.unwrap_or(rete_core::DEFAULT_QUERY_BUDGET_MB) {
        0 => None,
        mb => Some(mb.saturating_mul(1 << 20)),
    }
}

/// [`open_local`] under an explicit memory budget (`None` = unlimited).
///
/// An eager open holds the file image, its decoded dictionary and every index
/// permutation — about three to four times the file — and none of that can be
/// evicted. So a file goes lazy once a quarter of the budget no longer covers
/// it (1 GiB at the 4096 MiB default, the long-standing threshold), and the lazy
/// open's caches are capped to the budget. `RETE_LOCAL_LAZY_ABOVE_MB` still
/// overrides the threshold.
pub(crate) fn open_local_budget(path: &str, budget: Option<u64>) -> anyhow::Result<Rete> {
    let len = std::fs::metadata(path).map(|m| m.len()).unwrap_or(0);
    let threshold = match std::env::var("RETE_LOCAL_LAZY_ABOVE_MB") {
        Ok(_) => lazy_threshold_bytes(),
        Err(_) => budget.map_or(DEFAULT_LAZY_ABOVE_BYTES, |b| b / 4),
    };
    if std::env::var("RETE_OPEN_DEBUG").is_ok() {
        eprintln!(
            "[open_local_budget] {path}: len={len} threshold={threshold} -> {}",
            if len > threshold { "LAZY" } else { "EAGER" }
        );
    }
    if len <= threshold {
        let rete = Rete::open(&std::fs::read(path)?)?;
        // A resident open has no evictable caches to cap, but a query's
        // `COUNT(DISTINCT …)` still spills past the budget's share.
        rete.set_spill_threshold(budget.map(rete_core::spill_threshold_for));
        return Ok(rete);
    }
    open_local_lazy(path, len, budget)
}

/// Open a local `.rete`, reading it whole only when that is actually cheaper.
///
/// `Rete::open` takes a full file image and decodes every dictionary chunk up
/// front, so peak memory is the file plus its *decompressed* dictionary. That is
/// fine for ordinary graphs and hopeless for one carrying embedded media: a
/// 23.5 GB file whose dictionary is 23.4 GB of base64 images needs ~55 GB to
/// open, and dies, even though answering a query touches a handful of chunks.
///
/// Above the threshold we therefore go through the same lazy path the HTTP
/// reader uses — `Rete::open_ranged_lazy` over a positional-read file handle —
/// which reads the header, the section directories and the index, then faults
/// dictionary chunks in on demand. Small files keep the read-it-all path, which
/// stays faster when the whole graph is going to be touched anyway.
///
/// The `BlockCacheReader` is not optional here. The engine re-reads the same
/// spans (a chunk directory during binary search, the header) many times over,
/// and a bare file handle re-does every one of them: a point lookup on a 23.5 GB
/// file read 15.9 GB and still had not finished. Block-aligned caching makes the
/// repeats free, which is what `sparql-url` has always done for the remote path.
pub(crate) fn open_local(path: &str) -> anyhow::Result<Rete> {
    let len = std::fs::metadata(path).map(|m| m.len()).unwrap_or(0);
    if std::env::var("RETE_OPEN_DEBUG").is_ok() {
        eprintln!(
            "[open_local] {path}: len={len} threshold={} -> {}",
            lazy_threshold_bytes(),
            if len > lazy_threshold_bytes() {
                "LAZY"
            } else {
                "EAGER"
            }
        );
    }
    if len <= lazy_threshold_bytes() {
        return Ok(Rete::open(&std::fs::read(path)?)?);
    }
    open_local_lazy(path, len, None)
}

/// Open a local `.rete` through the lazy RANGED reader **regardless of size** —
/// the header, section directories and index up front, then dictionary chunks
/// and tiles faulted in on demand over a positional-read file handle.
///
/// This is what `rete export` uses by default. The motivation is peak RSS: the
/// eager path (`Rete::open` over the whole file image, see [`open_local`]) holds
/// the entire file resident plus its decompressed dictionary and all six index
/// permutations, so export RSS scaled ~3 GB per GB of file — 6 GB for a 1.53 GB
/// graph, and a 52 GB crossref/datacite would need ~150 GB and OOM the 50 GB
/// Docker VM. The ranged reader never holds the raw file image or the unused
/// permutations, so a filtered export (one predicate/graph) is bounded to the
/// tiles it touches. A *full* dump still faults every dictionary chunk and those
/// stay resident (the decompressed dictionary is the floor either way), so the
/// win is largest for slices; but even a full dump avoids the raw-image and
/// spare-permutation overhead the eager open pays on top of that floor.
///
/// `--in-memory` (→ [`open_local_eager`]) forces the old whole-file load back,
/// which is faster for small files that fit in RAM.
///
/// `budget` is the total `--memory-budget-mb` in bytes (`None` = unlimited): it
/// caps the block cache proportionally here and the dictionary/tile caches on
/// the opened `Rete`, bounding peak RSS to roughly the budget (plan §5.2).
pub(crate) fn open_local_ranged(path: &str, budget: Option<u64>) -> anyhow::Result<Rete> {
    let len = std::fs::metadata(path)?.len();
    if std::env::var("RETE_OPEN_DEBUG").is_ok() {
        eprintln!("[open_local_ranged] {path}: len={len} -> LAZY");
    }
    open_local_lazy(path, len, budget)
}

/// Force the whole file resident: read the entire image and decode every
/// dictionary chunk and index permutation up front (`Rete::open`). Needs RAM
/// roughly equal to the file size; the opt-in `rete export --in-memory` path.
pub(crate) fn open_local_eager(path: &str) -> anyhow::Result<Rete> {
    if std::env::var("RETE_OPEN_DEBUG").is_ok() {
        let len = std::fs::metadata(path).map(|m| m.len()).unwrap_or(0);
        eprintln!("[open_local_eager] {path}: len={len} -> EAGER (whole file)");
    }
    Ok(Rete::open(&std::fs::read(path)?)?)
}

/// The lazy open shared by [`open_local`] (above its size threshold) and
/// [`open_local_ranged`] (always): a positional-read file handle behind a
/// block-aligned cache, then `Rete::open_ranged_lazy`.
fn open_local_lazy(path: &str, len: u64, budget: Option<u64>) -> anyhow::Result<Rete> {
    let reader = std::sync::Arc::new(LocalRangeReader::open(path)?);
    // `RETE_BLOCK_KB` wins (0 disables), else auto-tune by length — same knob
    // and heuristic as the URL commands, so local and remote behave alike.
    let block: u64 = match std::env::var("RETE_BLOCK_KB")
        .ok()
        .and_then(|v| v.parse::<u64>().ok())
    {
        Some(kb) => kb * 1024,
        None => auto_block(len),
    };
    // A finite budget caps the raw-block LRU too (its share of the split); an
    // unlimited budget keeps the reader's default cap.
    let block_cap = rete_core::split_memory_budget(budget.unwrap_or(u64::MAX)).block_cache;
    let rete = if block == 0 {
        Rete::open_ranged_lazy(reader)?
    } else {
        let mut bc = BlockCacheReader::new(reader, block);
        if block_cap != u64::MAX {
            bc = bc.with_cache_cap(block_cap);
        }
        Rete::open_ranged_lazy(std::sync::Arc::new(bc))?
    };
    // Cap the dictionary chunk cache and the index tile cache to the rest of the
    // budget (no-op when unlimited).
    rete.set_memory_budget(budget);
    Ok(rete)
}

pub(crate) enum RangedSourceReader {
    Local(LocalRangeReader),
    Http(HttpRangeReader),
}

impl RangedSourceReader {
    pub(crate) fn open(source: &str) -> anyhow::Result<Self> {
        if is_url(source) {
            Ok(Self::Http(HttpRangeReader::open(source)?))
        } else {
            Ok(Self::Local(LocalRangeReader::open(source)?))
        }
    }
}

impl RangeReader for RangedSourceReader {
    fn len(&self) -> u64 {
        match self {
            Self::Local(r) => r.len(),
            Self::Http(r) => r.len(),
        }
    }

    fn read_at(&self, offset: u64, len: u64) -> io::Result<Vec<u8>> {
        match self {
            Self::Local(r) => r.read_at(offset, len),
            Self::Http(r) => r.read_at(offset, len),
        }
    }
}

pub(crate) struct LocalRangeReader {
    file: File,
    len: u64,
}

impl LocalRangeReader {
    pub(crate) fn open(path: &str) -> anyhow::Result<Self> {
        let file = File::open(path)?;
        let len = file.metadata()?.len();
        Ok(Self { file, len })
    }
}

impl RangeReader for LocalRangeReader {
    fn len(&self) -> u64 {
        self.len
    }

    fn read_at(&self, offset: u64, len: u64) -> io::Result<Vec<u8>> {
        let end = offset
            .checked_add(len)
            .filter(|end| *end <= self.len)
            .ok_or_else(|| io::Error::new(io::ErrorKind::UnexpectedEof, "range out of bounds"))?;
        let size = usize::try_from(len).map_err(|_| {
            io::Error::new(io::ErrorKind::InvalidInput, "range too large for memory")
        })?;
        let mut buf = vec![0u8; size];
        read_exact_at(&self.file, &mut buf, offset)?;
        debug_assert_eq!(end - offset, len);
        Ok(buf)
    }
}

#[cfg(unix)]
fn read_some_at(file: &File, buf: &mut [u8], offset: u64) -> io::Result<usize> {
    use std::os::unix::fs::FileExt;
    file.read_at(buf, offset)
}

#[cfg(windows)]
fn read_some_at(file: &File, buf: &mut [u8], offset: u64) -> io::Result<usize> {
    use std::os::windows::fs::FileExt;
    file.seek_read(buf, offset)
}

fn read_exact_at(file: &File, mut buf: &mut [u8], mut offset: u64) -> io::Result<()> {
    while !buf.is_empty() {
        let n = read_some_at(file, buf, offset)?;
        if n == 0 {
            return Err(io::Error::new(
                io::ErrorKind::UnexpectedEof,
                "range out of bounds",
            ));
        }
        offset += n as u64;
        let tmp = buf;
        buf = &mut tmp[n..];
    }
    Ok(())
}