use std::collections::{HashSet, VecDeque};
use std::io::{self, BufRead};
use std::num::{NonZeroU16, NonZeroUsize};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use dua_core::Order;
use crate::Cancellation;
use crate::mincore::PageMap;
use crate::mode::DisplayMode;
use crate::ops::{FileRange, Op, Stats};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum Threads {
#[default]
All,
Exact(NonZeroU16),
}
impl Threads {
pub fn get(self) -> usize {
match self {
Self::All => std::thread::available_parallelism().map_or(1, NonZeroUsize::get),
Self::Exact(threads) => usize::from(threads.get()),
}
}
fn effective(self) -> usize {
self.get()
.min(std::thread::available_parallelism().map_or(1, NonZeroUsize::get))
}
}
impl From<u16> for Threads {
fn from(threads: u16) -> Self {
NonZeroU16::new(threads).map_or(Self::All, Self::Exact)
}
}
impl std::fmt::Display for Threads {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::All => f.write_str("0"),
Self::Exact(threads) => write!(f, "{threads}"),
}
}
}
impl std::str::FromStr for Threads {
type Err = std::num::ParseIntError;
fn from_str(value: &str) -> Result<Self, Self::Err> {
value.parse::<u16>().map(Self::from)
}
}
#[derive(Debug, Clone, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct CrawlConfig {
pub follow_symlinks: bool,
pub single_filesystem: bool,
pub count_hardlinks: bool,
pub ignore_patterns: Vec<String>,
pub filter_patterns: Vec<String>,
pub max_file_size: Option<u64>,
pub batch: Option<PathBuf>,
pub nul_delim: bool,
pub threads: Threads,
}
pub fn crawl_and_process<O: Op, PM: PageMap + Send + Sync, D: DisplayMode<PM>>(
paths: &[PathBuf],
crawl_config: &CrawlConfig,
op: &O,
range: &FileRange,
stats: &Stats,
display: &D,
cancellation: &Cancellation,
) -> crate::Result<Vec<O::Output>> {
cancellation.check()?;
tracing::info!("starting {} on {} path(s)", O::LABEL, paths.len());
validate_patterns(crawl_config)?;
let mut seen_inodes = HashSet::new();
let mut files = Vec::new();
let collection = collect_paths(
paths,
crawl_config,
&mut seen_inodes,
stats,
cancellation,
|path| {
files.push(path);
Ok(())
},
);
if let Err(error) = collection {
display.finish();
return Err(error);
}
let next = AtomicUsize::new(0);
let stop = AtomicBool::new(false);
let outputs = Mutex::new(Vec::with_capacity(files.len()));
let first_error = Mutex::new(None);
let workers = crawl_config.threads.effective().min(files.len());
let worker_result: crate::Result<()> = std::thread::scope(|scope| {
let mut handles = Vec::with_capacity(workers);
for _ in 0..workers {
handles.push(
std::thread::Builder::new()
.spawn_scoped(scope, || {
while !stop.load(Ordering::Relaxed) {
let index = next.fetch_add(1, Ordering::Relaxed);
let Some(path) = files.get(index) else {
break;
};
match display.process_one::<O>(op, path, range, stats, cancellation) {
Ok(Some(output)) => {
outputs.lock().unwrap().push((index, output));
}
Ok(None) => {}
Err(error) => {
stop.store(true, Ordering::Relaxed);
let mut first = first_error.lock().unwrap();
if first.is_none() {
*first = Some(error);
}
}
}
}
})
.map_err(|error| crate::Error::io("spawn file worker", error))?,
);
}
for handle in handles {
handle.join().map_err(|_| crate::Error::WorkerPanic)?;
}
Ok(())
});
display.finish();
worker_result?;
if let Some(error) = first_error.into_inner().unwrap() {
return Err(error);
}
cancellation.check()?;
op.finish()?;
tracing::info!(
"done: {} files, {} pages",
stats.total_files.load(Ordering::Relaxed),
stats.total_pages.load(Ordering::Relaxed),
);
let mut outputs = outputs.into_inner().unwrap();
outputs.sort_unstable_by_key(|(index, _)| *index);
Ok(outputs.into_iter().map(|(_, output)| output).collect())
}
pub fn validate_patterns(config: &CrawlConfig) -> crate::Result<()> {
build_overrides(Path::new("."), config).map(drop)
}
fn collect_paths(
paths: &[PathBuf],
crawl_config: &CrawlConfig,
seen_inodes: &mut HashSet<(u64, u64)>,
stats: &Stats,
cancellation: &Cancellation,
mut emit: impl FnMut(PathBuf) -> crate::Result<()>,
) -> crate::Result<()> {
cancellation.check()?;
let mut all_paths: Vec<PathBuf> = paths.to_vec();
if let Some(batch_path) = &crawl_config.batch {
let batch_paths = read_batch_paths(batch_path, crawl_config.nul_delim)
.map_err(|error| crate::Error::io(batch_path.display().to_string(), error))?;
all_paths.extend(batch_paths);
}
let needs_meta = crawl_config.max_file_size.is_some()
|| !crawl_config.count_hardlinks
|| crawl_config.single_filesystem;
for path in &all_paths {
if cancellation.is_cancelled() {
break;
}
let metadata = fs_err::metadata(path)
.map_err(|error| crate::Error::io(path.display().to_string(), error))?;
if metadata.is_dir() {
tracing::info!("crawling directory {}", path.display());
stats.total_dirs.fetch_add(1, Ordering::Relaxed);
walk_dir_entries(
path,
crawl_config,
needs_meta,
seen_inodes,
stats,
cancellation,
&mut emit,
)?;
} else if metadata.is_file() {
if explicit_file_allowed(path, &metadata, crawl_config, needs_meta, seen_inodes) {
emit(path.clone())?;
}
} else {
return Err(crate::Error::io(
path.display().to_string(),
io::Error::new(io::ErrorKind::InvalidInput, "not a file or directory"),
));
}
}
Ok(())
}
struct TraversalRoot {
physical: PathBuf,
logical: PathBuf,
device: Option<u64>,
overrides: Option<Arc<ignore::overrides::Override>>,
ignored: Option<Arc<ignore::overrides::Override>>,
}
impl TraversalRoot {
fn new(
physical: PathBuf,
logical: PathBuf,
device: Option<u64>,
config: &CrawlConfig,
) -> crate::Result<Self> {
Ok(Self {
physical,
overrides: build_overrides(&logical, config)?.map(Arc::new),
ignored: build_ignore_overrides(&logical, config)?.map(Arc::new),
logical,
device,
})
}
fn logical_path(&self, physical: &Path) -> PathBuf {
physical.strip_prefix(&self.physical).map_or_else(
|_| physical.to_owned(),
|relative| self.logical.join(relative),
)
}
}
fn walk_dir_entries(
root: &Path,
config: &CrawlConfig,
needs_meta: bool,
seen_inodes: &mut HashSet<(u64, u64)>,
stats: &Stats,
cancellation: &Cancellation,
mut emit: impl FnMut(PathBuf) -> crate::Result<()>,
) -> crate::Result<()> {
use std::os::unix::fs::MetadataExt as _;
let root_metadata = fs_err::metadata(root)
.map_err(|error| crate::Error::io(root.display().to_string(), error))?;
let root_device = config.single_filesystem.then(|| root_metadata.dev());
let physical_root = if config.follow_symlinks && root.is_symlink() {
fs_err::canonicalize(root)
.map_err(|error| crate::Error::io(root.display().to_string(), error))?
} else {
root.to_owned()
};
let mut pending = VecDeque::from([TraversalRoot::new(
physical_root.clone(),
root.to_owned(),
root_device,
config,
)?]);
let mut visited = HashSet::new();
if config.follow_symlinks {
visited.insert(fs_err::canonicalize(&physical_root).unwrap_or(physical_root));
}
while let Some(root) = pending.pop_front() {
if cancellation.is_cancelled() {
return Err(crate::Error::Cancelled);
}
let root = Arc::new(root);
let descend_root = Arc::clone(&root);
let descend_cancellation = cancellation.clone();
let walk_stop = Arc::new(AtomicBool::new(false));
let descend_stop = Arc::clone(&walk_stop);
let entries = dua_core::walk(
&root.physical,
config.threads.effective(),
Order::Completion,
move |entry| {
if descend_cancellation.is_cancelled() || descend_stop.load(Ordering::Relaxed) {
return false;
}
if entry.depth == 0 {
return true;
}
if descend_root.device.is_some_and(|device| {
entry
.metadata
.as_ref()
.map_or(true, |metadata| metadata.dev() != device)
}) {
return false;
}
let path = descend_root.logical_path(&entry.path());
!descend_root
.ignored
.as_ref()
.is_some_and(|overrides| overrides.matched(&path, true).is_ignore())
},
);
let mut walk_error = None;
for entry_result in entries {
if walk_error.is_some() {
continue;
}
if cancellation.is_cancelled() {
walk_stop.store(true, Ordering::Relaxed);
walk_error = Some(crate::Error::Cancelled);
continue;
}
let result = (|| -> crate::Result<()> {
let entry = entry_result
.map_err(|error| crate::Error::io(root.logical.display().to_string(), error))?;
if entry.depth == 0 && entry.file_type.is_dir() {
return Ok(());
}
let physical_path = entry.path();
let logical_path = root.logical_path(&physical_path);
if entry.file_type.is_symlink() {
if !config.follow_symlinks {
return Ok(());
}
let metadata = fs_err::metadata(&physical_path).map_err(|error| {
crate::Error::io(logical_path.display().to_string(), error)
})?;
if root.device.is_some_and(|device| metadata.dev() != device) {
return Ok(());
}
if metadata.is_dir() {
stats.total_dirs.fetch_add(1, Ordering::Relaxed);
let target = fs_err::canonicalize(&physical_path).map_err(|error| {
crate::Error::io(logical_path.display().to_string(), error)
})?;
if visited.insert(target.clone()) {
pending.push_back(TraversalRoot::new(
target,
logical_path,
root.device,
config,
)?);
}
} else if metadata.is_file()
&& path_allowed(&logical_path, false, config, root.overrides.as_deref())
&& file_allowed(Some(&metadata), config, needs_meta, seen_inodes)
{
emit(logical_path)?;
}
return Ok(());
}
let metadata = match entry.metadata.as_ref() {
Ok(metadata) => Some(metadata),
Err(error) if needs_meta => {
return Err(crate::Error::io(
logical_path.display().to_string(),
io::Error::new(error.kind(), error.to_string()),
));
}
Err(_) => None,
};
if entry.file_type.is_dir() {
stats.total_dirs.fetch_add(1, Ordering::Relaxed);
return Ok(());
}
if !entry.file_type.is_file()
|| !path_allowed(&logical_path, false, config, root.overrides.as_deref())
|| !file_allowed(metadata, config, needs_meta, seen_inodes)
{
return Ok(());
}
emit(logical_path)?;
Ok(())
})();
if let Err(error) = result {
walk_stop.store(true, Ordering::Relaxed);
walk_error = Some(error);
}
}
if let Some(error) = walk_error {
return Err(error);
}
}
Ok(())
}
fn explicit_file_allowed(
path: &Path,
metadata: &std::fs::Metadata,
config: &CrawlConfig,
needs_meta: bool,
seen_inodes: &mut HashSet<(u64, u64)>,
) -> bool {
let root = path
.parent()
.filter(|parent| !parent.as_os_str().is_empty())
.unwrap_or_else(|| Path::new("."));
let overrides =
build_overrides(root, config).expect("path patterns were validated before traversal");
if !path_allowed(path, false, config, overrides.as_ref()) {
return false;
}
file_allowed(Some(metadata), config, needs_meta, seen_inodes)
}
fn path_allowed(
path: &Path,
is_dir: bool,
config: &CrawlConfig,
overrides: Option<&ignore::overrides::Override>,
) -> bool {
overrides.is_none_or(|overrides| {
let matched = overrides.matched(path, is_dir);
!matched.is_ignore() && (config.filter_patterns.is_empty() || matched.is_whitelist())
})
}
fn file_allowed(
metadata: Option<&std::fs::Metadata>,
config: &CrawlConfig,
needs_meta: bool,
seen_inodes: &mut HashSet<(u64, u64)>,
) -> bool {
if needs_meta && metadata.is_none() {
return false;
}
if let Some(max_size) = config.max_file_size
&& let Some(metadata) = metadata
&& metadata.len() > max_size
{
return false;
}
if !config.count_hardlinks {
#[cfg(unix)]
{
use std::os::unix::fs::MetadataExt;
if let Some(metadata) = metadata
&& metadata.nlink() > 1
&& !seen_inodes.insert((metadata.dev(), metadata.ino()))
{
return false;
}
}
}
true
}
fn build_overrides(
root: &Path,
config: &CrawlConfig,
) -> crate::Result<Option<ignore::overrides::Override>> {
if config.ignore_patterns.is_empty() && config.filter_patterns.is_empty() {
return Ok(None);
}
let mut overrides = ignore::overrides::OverrideBuilder::new(root);
for pattern in &config.ignore_patterns {
overrides.add(&format!("!{pattern}"))?;
}
for pattern in &config.filter_patterns {
overrides.add(pattern)?;
}
Ok(Some(overrides.build()?))
}
fn build_ignore_overrides(
root: &Path,
config: &CrawlConfig,
) -> crate::Result<Option<ignore::overrides::Override>> {
if config.ignore_patterns.is_empty() {
return Ok(None);
}
let mut overrides = ignore::overrides::OverrideBuilder::new(root);
for pattern in &config.ignore_patterns {
overrides.add(&format!("!{pattern}"))?;
}
Ok(Some(overrides.build()?))
}
pub fn read_batch_paths(path: &Path, nul_delim: bool) -> io::Result<Vec<PathBuf>> {
use std::os::unix::ffi::OsStrExt;
let reader: Box<dyn BufRead> = if path == Path::new("-") {
Box::new(io::stdin().lock())
} else {
Box::new(io::BufReader::new(fs_err::File::open(path)?))
};
let delim = if nul_delim { b'\0' } else { b'\n' };
reader
.split(delim)
.filter_map(|r| match r {
Ok(buf) if !buf.is_empty() => {
Some(Ok(PathBuf::from(std::ffi::OsStr::from_bytes(&buf))))
}
Ok(_) => None,
Err(e) => Some(Err(e)),
})
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::Cancellation;
use crate::mode::Cli;
use crate::ops::{FileContext, Op, ResidencyEffect};
use std::sync::atomic::AtomicUsize;
#[test]
fn threads_parse_and_display() {
assert_eq!("0".parse::<Threads>().unwrap(), Threads::All);
assert_eq!("4".parse::<Threads>().unwrap().to_string(), "4");
assert_eq!(Threads::from(1).get(), 1);
assert!(Threads::All.get() > 0);
assert_eq!(
Threads::Exact(NonZeroU16::new(u16::MAX).unwrap()).effective(),
Threads::All.get()
);
}
#[test]
fn cancelled_collection_emits_no_paths() {
let file = tempfile::NamedTempFile::new().unwrap();
let config = CrawlConfig {
follow_symlinks: false,
single_filesystem: false,
count_hardlinks: true,
ignore_patterns: Vec::new(),
filter_patterns: Vec::new(),
max_file_size: None,
batch: None,
nul_delim: false,
threads: Threads::default(),
};
let mut seen_inodes = HashSet::new();
let stats = Stats::new();
let cancellation = Cancellation::new();
cancellation.cancel();
let mut emitted = Vec::new();
let result = collect_paths(
&[file.path().to_owned()],
&config,
&mut seen_inodes,
&stats,
&cancellation,
|path| {
emitted.push(path);
Ok(())
},
);
assert!(emitted.is_empty());
assert!(matches!(result, Err(crate::Error::Cancelled)));
}
#[test]
fn file_operations_use_requested_parallelism() {
struct ConcurrentOp {
active: AtomicUsize,
max_active: AtomicUsize,
}
impl Op for ConcurrentOp {
const LABEL: &'static str = "test";
const EFFECT: ResidencyEffect = ResidencyEffect::Preserve;
type Output = ();
fn execute<PM: PageMap + Sync>(&self, _ctx: &FileContext<'_, PM>) -> crate::Result<()> {
let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
self.max_active.fetch_max(active, Ordering::SeqCst);
std::thread::sleep(std::time::Duration::from_millis(20));
self.active.fetch_sub(1, Ordering::SeqCst);
Ok(())
}
}
let dir = tempfile::tempdir().unwrap();
for index in 0..4 {
fs_err::write(dir.path().join(index.to_string()), vec![0u8; 4096]).unwrap();
}
let config = CrawlConfig {
follow_symlinks: false,
single_filesystem: false,
count_hardlinks: true,
ignore_patterns: Vec::new(),
filter_patterns: Vec::new(),
max_file_size: None,
batch: None,
nul_delim: false,
threads: Threads::Exact(NonZeroU16::new(2).unwrap()),
};
let op = ConcurrentOp {
active: AtomicUsize::new(0),
max_active: AtomicUsize::new(0),
};
crawl_and_process::<_, crate::mincore::DefaultPageMap, _>(
&[dir.path().to_owned()],
&config,
&op,
&FileRange::full(),
&Stats::new(),
&Cli,
&Cancellation::new(),
)
.unwrap();
assert_eq!(op.max_active.load(Ordering::SeqCst), 2);
}
#[test]
fn cancellation_drains_wide_traversal() {
let dir = tempfile::tempdir().unwrap();
for index in 0..10_000 {
fs_err::write(dir.path().join(index.to_string()), []).unwrap();
}
let config = CrawlConfig {
follow_symlinks: false,
single_filesystem: false,
count_hardlinks: true,
ignore_patterns: Vec::new(),
filter_patterns: Vec::new(),
max_file_size: None,
batch: None,
nul_delim: false,
threads: Threads::Exact(NonZeroU16::new(4).unwrap()),
};
let cancellation = Cancellation::new();
let worker_cancellation = cancellation.clone();
let root = dir.path().to_owned();
let (started_sender, started_receiver) = std::sync::mpsc::channel();
let (resume_sender, resume_receiver) = std::sync::mpsc::channel();
let (result_sender, result_receiver) = std::sync::mpsc::channel();
std::thread::spawn(move || {
let mut seen = HashSet::new();
let mut started = false;
let result = walk_dir_entries(
&root,
&config,
false,
&mut seen,
&Stats::new(),
&worker_cancellation,
|_| {
if !started {
started = true;
started_sender.send(()).unwrap();
resume_receiver.recv().unwrap();
}
Ok(())
},
);
let _ = result_sender.send(result);
});
started_receiver
.recv_timeout(std::time::Duration::from_secs(5))
.expect("traversal did not start");
cancellation.cancel();
resume_sender.send(()).unwrap();
let result = result_receiver
.recv_timeout(std::time::Duration::from_secs(5))
.expect("cancelled traversal deadlocked");
assert!(matches!(result, Err(crate::Error::Cancelled)));
}
}