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://")
}
const MAX_EAGER_BYTES: u64 = 1 << 30;
const EAGER_SHARE: u64 = 16;
const UNKNOWN_MEMORY_LAZY_ABOVE_BYTES: u64 = 256 << 20;
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct LazyThreshold {
pub(crate) bytes: u64,
pub(crate) reason: String,
}
pub(crate) fn lazy_threshold(budget: Option<u64>) -> LazyThreshold {
let env = std::env::var("RETE_LOCAL_LAZY_ABOVE_MB")
.ok()
.and_then(|v| v.trim().parse::<u64>().ok());
let available = if env.is_some() {
None
} else {
available_memory()
};
lazy_threshold_from(env, available, budget)
}
pub(crate) fn lazy_threshold_from(
env_mb: Option<u64>,
available: Option<u64>,
budget: Option<u64>,
) -> LazyThreshold {
if let Some(mb) = env_mb {
return LazyThreshold {
bytes: mb.saturating_mul(1 << 20),
reason: format!("RETE_LOCAL_LAZY_ABOVE_MB={mb}"),
};
}
let mut bytes = MAX_EAGER_BYTES;
let mut reason = format!("cap {} MiB", MAX_EAGER_BYTES >> 20);
let (mem_limit, mem_reason) = match available {
Some(avail) => (
avail / EAGER_SHARE,
format!("1/{EAGER_SHARE} of {} MiB available", avail >> 20),
),
None => (
UNKNOWN_MEMORY_LAZY_ABOVE_BYTES,
"available memory unknown".to_string(),
),
};
if mem_limit < bytes {
bytes = mem_limit;
reason = mem_reason;
}
if let Some(b) = budget {
if b / 4 < bytes {
bytes = b / 4;
reason = format!("1/4 of the {} MiB memory budget", b >> 20);
}
}
LazyThreshold { bytes, reason }
}
pub(crate) fn available_memory() -> Option<u64> {
let meminfo = std::fs::read_to_string("/proc/meminfo")
.ok()
.and_then(|s| meminfo_available(&s));
let cgroup = cgroup_headroom(std::path::Path::new("/sys/fs/cgroup"));
match (meminfo, cgroup) {
(Some(a), Some(b)) => Some(a.min(b)),
(a, b) => a.or(b),
}
}
fn meminfo_available(meminfo: &str) -> Option<u64> {
meminfo.lines().find_map(|line| {
let rest = line.strip_prefix("MemAvailable:")?;
let kb: u64 = rest.trim().trim_end_matches("kB").trim().parse().ok()?;
Some(kb.saturating_mul(1024))
})
}
fn stat_field(stat: &str, key: &str) -> Option<u64> {
stat.lines().find_map(|line| {
let (k, v) = line.split_once(' ')?;
(k == key).then(|| v.trim().parse().ok())?
})
}
fn read_u64(path: &std::path::Path) -> Option<u64> {
std::fs::read_to_string(path).ok()?.trim().parse().ok()
}
fn cgroup_v2_headroom(dir: &std::path::Path) -> Option<u64> {
let max = read_u64(&dir.join("memory.max"))?;
let current = read_u64(&dir.join("memory.current")).unwrap_or(0);
let inactive = std::fs::read_to_string(dir.join("memory.stat"))
.ok()
.and_then(|s| stat_field(&s, "inactive_file"))
.unwrap_or(0);
Some(max.saturating_sub(current.saturating_sub(inactive)))
}
fn cgroup_v1_headroom(dir: &std::path::Path) -> Option<u64> {
let limit = read_u64(&dir.join("memory.limit_in_bytes")).filter(|&l| l < 1 << 60)?;
let usage = read_u64(&dir.join("memory.usage_in_bytes")).unwrap_or(0);
let inactive = std::fs::read_to_string(dir.join("memory.stat"))
.ok()
.and_then(|s| stat_field(&s, "total_inactive_file"))
.unwrap_or(0);
Some(limit.saturating_sub(usage.saturating_sub(inactive)))
}
fn cgroup_headroom(root: &std::path::Path) -> Option<u64> {
let own = std::fs::read_to_string("/proc/self/cgroup").ok();
let v2 = own
.as_deref()
.and_then(|s| s.lines().find_map(|l| l.strip_prefix("0::")))
.map(|p| p.trim().trim_start_matches('/').to_string());
let mut best: Option<u64> = None;
if let Some(rel) = v2 {
let mut dir = root.join(rel);
loop {
if let Some(h) = cgroup_v2_headroom(&dir) {
best = Some(best.map_or(h, |b| b.min(h)));
}
if dir.as_path() == root || !dir.pop() || !dir.starts_with(root) {
break;
}
}
}
best.or_else(|| cgroup_v1_headroom(&root.join("memory")))
}
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)),
}
}
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 = lazy_threshold(budget);
debug_open("open_local_budget", path, len, &threshold);
if len <= threshold.bytes {
let rete = Rete::open(&std::fs::read(path)?)?;
rete.set_spill_threshold(budget.map(rete_core::spill_threshold_for));
return Ok(rete);
}
open_local_lazy(path, len, budget)
}
fn debug_open(who: &str, path: &str, len: u64, threshold: &LazyThreshold) {
if std::env::var("RETE_OPEN_DEBUG").is_ok() {
eprintln!(
"[{who}] {path}: len={len} threshold={} ({}) -> {}",
threshold.bytes,
threshold.reason,
if len > threshold.bytes {
"LAZY"
} else {
"EAGER"
}
);
}
}
pub(crate) fn open_local(path: &str) -> anyhow::Result<Rete> {
let len = std::fs::metadata(path).map(|m| m.len()).unwrap_or(0);
let threshold = lazy_threshold(None);
debug_open("open_local", path, len, &threshold);
if len <= threshold.bytes {
return Ok(Rete::open(&std::fs::read(path)?)?);
}
open_local_lazy(path, len, None)
}
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)
}
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)?)?)
}
fn open_local_lazy(path: &str, len: u64, budget: Option<u64>) -> anyhow::Result<Rete> {
let reader = std::sync::Arc::new(LocalRangeReader::open(path)?);
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),
};
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))?
};
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),
}
}
fn read_many(&self, ranges: &[(u64, u64)]) -> io::Result<Vec<Vec<u8>>> {
match self {
Self::Local(r) => r.read_many(ranges),
Self::Http(r) => r.read_many(ranges),
}
}
fn concurrency(&self) -> usize {
match self {
Self::Local(r) => r.concurrency(),
Self::Http(r) => r.concurrency(),
}
}
}
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(())
}
#[cfg(test)]
mod tests {
use super::*;
const MIB: u64 = 1 << 20;
const GIB: u64 = 1 << 30;
#[test]
fn env_override_wins_over_everything() {
let t = lazy_threshold_from(Some(0), Some(64 * GIB), Some(4096 * MIB));
assert_eq!(t.bytes, 0, "0 forces lazy");
let t = lazy_threshold_from(Some(5000), Some(GIB), Some(MIB));
assert_eq!(t.bytes, 5000 * MIB, "the override is not clamped");
assert!(
t.reason.contains("RETE_LOCAL_LAZY_ABOVE_MB"),
"{}",
t.reason
);
}
#[test]
fn threshold_follows_available_memory() {
let avail = 3800 * MIB;
let t = lazy_threshold_from(None, Some(avail), None);
assert_eq!(t.bytes, avail / 16);
assert!(1_072_399_312 > t.bytes);
assert!(t.bytes > 250_000_000 / 2, "{t:?}");
let t = lazy_threshold_from(None, Some(48 * GIB), None);
assert_eq!(t.bytes, GIB);
assert!(t.reason.contains("cap"), "{}", t.reason);
}
#[test]
fn unknown_memory_falls_back_conservatively() {
let t = lazy_threshold_from(None, None, None);
assert_eq!(t.bytes, UNKNOWN_MEMORY_LAZY_ABOVE_BYTES);
assert!(t.reason.contains("unknown"), "{}", t.reason);
}
#[test]
fn budget_quarter_still_applies() {
let t = lazy_threshold_from(None, Some(48 * GIB), Some(2048 * MIB));
assert_eq!(t.bytes, 512 * MIB);
assert!(t.reason.contains("budget"), "{}", t.reason);
let t = lazy_threshold_from(None, Some(GIB), Some(4096 * MIB));
assert_eq!(t.bytes, GIB / 16);
}
#[test]
fn parses_meminfo() {
let s =
"MemTotal: 49283460 kB\nMemFree: 1234 kB\nMemAvailable: 40000000 kB\n";
assert_eq!(meminfo_available(s), Some(40_000_000 * 1024));
assert_eq!(meminfo_available("MemTotal: 1 kB\n"), None);
}
#[test]
fn cgroup_v2_headroom_takes_the_tightest_limit() {
let root = tempfile::tempdir().unwrap();
let leaf = root.path().join("a/b");
std::fs::create_dir_all(&leaf).unwrap();
std::fs::write(leaf.join("memory.max"), "max\n").unwrap();
std::fs::write(root.path().join("a/memory.max"), format!("{}\n", 4 * GIB)).unwrap();
std::fs::write(
root.path().join("a/memory.current"),
format!("{}\n", GIB + 256 * MIB),
)
.unwrap();
std::fs::write(
root.path().join("a/memory.stat"),
format!("anon 1\ninactive_file {}\nactive_file 9\n", 256 * MIB),
)
.unwrap();
assert_eq!(cgroup_v2_headroom(&leaf), None, "`max` is no limit");
assert_eq!(
cgroup_v2_headroom(&root.path().join("a")),
Some(3 * GIB),
"4 GiB limit, 1.25 GiB charged of which 256 MiB is page cache"
);
}
#[test]
fn cgroup_v1_unlimited_is_no_limit() {
let dir = tempfile::tempdir().unwrap();
std::fs::write(
dir.path().join("memory.limit_in_bytes"),
"9223372036854771712\n",
)
.unwrap();
assert_eq!(cgroup_v1_headroom(dir.path()), None);
std::fs::write(
dir.path().join("memory.limit_in_bytes"),
format!("{}", 2 * GIB),
)
.unwrap();
std::fs::write(dir.path().join("memory.usage_in_bytes"), format!("{GIB}")).unwrap();
assert_eq!(cgroup_v1_headroom(dir.path()), Some(GIB));
}
#[test]
fn available_memory_is_readable_on_linux() {
if cfg!(target_os = "linux") {
let a = available_memory().expect("Linux exposes MemAvailable");
assert!(a > 0);
}
}
#[test]
fn ranged_source_forwards_batches_and_concurrency() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("f.bin");
let data: Vec<u8> = (0..=255u8).cycle().take(4096).collect();
std::fs::write(&path, &data).unwrap();
let r = RangedSourceReader::open(path.to_str().unwrap()).unwrap();
assert_eq!(r.concurrency(), 1, "a local file reports no fan-out");
let got = r.read_many(&[(10, 5), (4000, 96), (0, 0)]).unwrap();
assert_eq!(got[0], &data[10..15]);
assert_eq!(got[1], &data[4000..4096]);
assert!(got[2].is_empty());
}
}