use std::collections::{HashMap, HashSet};
use std::path::{Path, PathBuf};
use std::time::{Duration, Instant, SystemTime};
pub const DEFAULT_MAX_AGE: Duration = Duration::from_secs(60);
pub const DEFAULT_NODE_MAX_AGE: Duration = Duration::from_secs(900);
pub const DEFAULT_MAX_DEPTH: usize = 64;
pub const DEFAULT_WALK_BUDGET: Duration = Duration::from_secs(30);
const BUDGET_CHECK_INTERVAL: usize = 256;
const CLOCK_EVERY_DIRECTORY_BELOW: Duration = Duration::from_secs(1);
#[derive(Debug, Clone)]
pub struct IndexPolicy {
pub max_age: Duration,
pub node_max_age: Duration,
pub max_depth: usize,
pub walk_budget: Duration,
pub forbidden_roots: Vec<PathBuf>,
}
impl Default for IndexPolicy {
fn default() -> Self {
Self {
max_age: DEFAULT_MAX_AGE,
node_max_age: DEFAULT_NODE_MAX_AGE,
max_depth: DEFAULT_MAX_DEPTH,
walk_budget: DEFAULT_WALK_BUDGET,
forbidden_roots: default_forbidden_roots(),
}
}
}
fn default_forbidden_roots() -> Vec<PathBuf> {
let mut roots = vec![PathBuf::from("/")];
if let Some(home) = std::env::var_os("HOME") {
let home = PathBuf::from(home);
if !home.as_os_str().is_empty() {
roots.push(home);
}
}
roots
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
#[non_exhaustive]
pub enum SizeIndexError {
#[error(
"refusing to index `{path}`: it is, or contains, a root this index never walks by default"
)]
ForbiddenRoot {
path: PathBuf,
},
#[error("refusing to index `{path}`: not a directory")]
NotADirectory {
path: PathBuf,
},
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub struct DirSize {
pub path: PathBuf,
pub bytes: u64,
pub measured_at: SystemTime,
pub truncated: bool,
pub unreadable: Vec<PathBuf>,
pub from_cache: bool,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
#[non_exhaustive]
pub struct IndexStats {
pub directories_read: u64,
pub directories_revalidated: u64,
pub cache_hits: u64,
pub refreshes: u64,
}
#[derive(Debug)]
struct DirNode {
mtime: Option<SystemTime>,
own_bytes: u64,
children: Vec<PathBuf>,
read_at: Instant,
}
#[derive(Debug)]
struct RootRecord {
bytes: u64,
measured_at: SystemTime,
read_at: Instant,
truncated: bool,
unreadable: Vec<PathBuf>,
budget_exhausted: bool,
}
#[derive(Debug)]
pub struct DirSizeIndex {
policy: IndexPolicy,
nodes: HashMap<PathBuf, DirNode>,
roots: HashMap<PathBuf, RootRecord>,
stats: IndexStats,
}
impl Default for DirSizeIndex {
fn default() -> Self {
Self::with_policy(IndexPolicy::default())
}
}
impl DirSizeIndex {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn with_policy(mut policy: IndexPolicy) -> Self {
policy.forbidden_roots = policy
.forbidden_roots
.iter()
.map(|root| resolve_for_guard(root))
.collect();
Self {
policy,
nodes: HashMap::new(),
roots: HashMap::new(),
stats: IndexStats::default(),
}
}
#[must_use]
pub fn stats(&self) -> IndexStats {
self.stats
}
#[must_use]
pub fn policy(&self) -> &IndexPolicy {
&self.policy
}
pub fn invalidate(&mut self, root: &Path) {
self.roots.remove(root);
self.nodes.retain(|path, _| !path.starts_with(root));
}
pub fn measure(&mut self, root: &Path) -> Result<DirSize, SizeIndexError> {
self.measure_within(root, None)
}
pub fn measure_within(
&mut self,
root: &Path,
budget: Option<Duration>,
) -> Result<DirSize, SizeIndexError> {
self.guard_root(root)?;
let root = root.to_path_buf();
if let Some(record) = self.roots.get(&root)
&& record.read_at.elapsed() < self.policy.max_age
{
self.stats.cache_hits += 1;
return Ok(record.to_dir_size(root, true));
}
let ceiling = self.policy.walk_budget;
let record = self.refresh(&root, budget.map_or(ceiling, |asked| asked.min(ceiling)));
let size = record.to_dir_size(root.clone(), false);
if !record.budget_exhausted {
self.roots.insert(root, record);
}
Ok(size)
}
fn guard_root(&self, root: &Path) -> Result<(), SizeIndexError> {
let is_dir = std::fs::symlink_metadata(root).is_ok_and(|md| md.is_dir());
if !is_dir {
return Err(SizeIndexError::NotADirectory {
path: root.to_path_buf(),
});
}
let resolved = resolve_for_guard(root);
if self
.policy
.forbidden_roots
.iter()
.any(|forbidden| forbidden.starts_with(&resolved))
{
return Err(SizeIndexError::ForbiddenRoot {
path: root.to_path_buf(),
});
}
Ok(())
}
fn refresh(&mut self, root: &Path, budget: Duration) -> RootRecord {
self.stats.refreshes += 1;
let started = Instant::now();
let mut truncated = false;
let mut budget_exhausted = false;
let mut unreadable: Vec<PathBuf> = Vec::new();
let mut visited: Vec<PathBuf> = Vec::new();
let mut seen: HashSet<PathBuf> = HashSet::new();
let mut stack: Vec<(PathBuf, usize)> = vec![(root.to_path_buf(), 0)];
let mut opened = 0usize;
let check_every = if budget < CLOCK_EVERY_DIRECTORY_BELOW {
1
} else {
BUDGET_CHECK_INTERVAL
};
while let Some((dir, depth)) = stack.pop() {
opened += 1;
if opened.is_multiple_of(check_every) && started.elapsed() >= budget {
truncated = true;
budget_exhausted = true;
break;
}
if !seen.insert(dir.clone()) {
continue;
}
visited.push(dir.clone());
let mtime = std::fs::symlink_metadata(&dir)
.ok()
.and_then(|md| md.modified().ok());
if self.can_reuse(&dir, mtime) {
self.stats.directories_revalidated += 1;
truncated |= self.push_children_of(&dir, depth, &mut stack);
continue;
}
let Ok(entries) = std::fs::read_dir(&dir) else {
unreadable.push(dir.clone());
self.nodes.remove(&dir);
continue;
};
self.stats.directories_read += 1;
let (own_bytes, children) = sum_entries(entries);
self.nodes.insert(
dir.clone(),
DirNode {
mtime,
own_bytes,
children,
read_at: Instant::now(),
},
);
truncated |= self.push_children_of(&dir, depth, &mut stack);
}
if !truncated {
self.nodes
.retain(|path, _| !path.starts_with(root) || seen.contains(path));
}
RootRecord {
bytes: self.fold_totals(&visited, root),
measured_at: SystemTime::now(),
read_at: Instant::now(),
truncated,
unreadable,
budget_exhausted,
}
}
fn can_reuse(&self, dir: &Path, mtime: Option<SystemTime>) -> bool {
self.nodes.get(dir).is_some_and(|node| {
node.mtime.is_some()
&& node.mtime == mtime
&& node.read_at.elapsed() < self.policy.node_max_age
})
}
fn push_children_of(
&self,
dir: &Path,
depth: usize,
stack: &mut Vec<(PathBuf, usize)>,
) -> bool {
let Some(node) = self.nodes.get(dir) else {
return false;
};
if depth >= self.policy.max_depth {
return !node.children.is_empty();
}
stack.extend(node.children.iter().map(|c| (c.clone(), depth + 1)));
false
}
fn fold_totals(&self, visited: &[PathBuf], root: &Path) -> u64 {
let mut subtotal: HashMap<&Path, u64> = HashMap::with_capacity(visited.len());
for path in visited.iter().rev() {
let total = self.nodes.get(path.as_path()).map_or(0, |node| {
node.children
.iter()
.filter_map(|child| subtotal.get(child.as_path()))
.fold(node.own_bytes, |acc, bytes| acc.saturating_add(*bytes))
});
subtotal.insert(path.as_path(), total);
}
subtotal.get(root).copied().unwrap_or(0)
}
}
impl RootRecord {
fn to_dir_size(&self, path: PathBuf, from_cache: bool) -> DirSize {
DirSize {
path,
bytes: self.bytes,
measured_at: self.measured_at,
truncated: self.truncated,
unreadable: self.unreadable.clone(),
from_cache,
}
}
}
fn resolve_for_guard(path: &Path) -> PathBuf {
std::fs::canonicalize(path).unwrap_or_else(|_| lexical_normalize(path))
}
fn lexical_normalize(path: &Path) -> PathBuf {
let mut out = PathBuf::new();
for component in path.components() {
match component {
std::path::Component::CurDir => {}
std::path::Component::ParentDir => {
out.pop();
}
other => out.push(other.as_os_str()),
}
}
out
}
fn sum_entries(entries: std::fs::ReadDir) -> (u64, Vec<PathBuf>) {
let mut own_bytes = 0u64;
let mut children = Vec::new();
for entry in entries.flatten() {
let Ok(file_type) = entry.file_type() else {
continue;
};
if file_type.is_symlink() {
continue;
}
if file_type.is_dir() {
children.push(entry.path());
} else if file_type.is_file()
&& let Ok(md) = entry.metadata()
{
own_bytes = own_bytes.saturating_add(md.len());
}
}
(own_bytes, children)
}
#[cfg(test)]
#[path = "size_index_tests.rs"]
mod size_index_tests;