use crate::control::CancellationToken;
use crate::report::FileIdentity;
use crate::runtime::ParallelRuntime;
use crate::walker::{ErrorPolicy, WalkEntry, WalkError, WalkOperation, WalkOptions};
use std::path::PathBuf;
use std::sync::{Arc, Mutex, mpsc};
mod collect;
pub(crate) mod dynamic;
mod ordered_pull;
mod pull;
mod visit;
mod visit_worker;
use collect::{DirectoryTask, collect_lane, collect_serial, collect_shallow, expand_frontier};
pub use pull::ParallelWalkIter;
pub use visit::{ParallelVisitReport, WalkControl, WalkEvent};
#[derive(Debug)]
pub struct ParallelWalkReport {
pub entries: Vec<WalkEntry>,
pub errors: Vec<WalkError>,
}
pub struct ParallelWalker {
pub(super) root: PathBuf,
pub(super) options: WalkOptions,
pub(super) parallelism: usize,
pub(super) skip_stdout: Option<FileIdentity>,
pub(super) runtime: ParallelRuntime,
}
impl ParallelWalker {
#[must_use]
pub fn new(root: impl Into<PathBuf>) -> Self {
Self {
root: root.into(),
options: WalkOptions::default(),
parallelism: 0,
skip_stdout: None,
runtime: ParallelRuntime::global(),
}
}
#[must_use]
pub const fn options(mut self, options: WalkOptions) -> Self {
self.options = options;
self
}
#[must_use]
pub const fn with_parallelism(mut self, parallelism: usize) -> Self {
self.parallelism = parallelism;
self
}
#[must_use]
pub fn runtime(mut self, runtime: ParallelRuntime) -> Self {
self.runtime = runtime;
self
}
#[must_use]
pub fn skip_stdout(mut self, enabled: bool) -> Self {
self.skip_stdout = enabled.then(crate::stdout::identity).flatten();
self
}
pub fn walk(mut self) -> Result<ParallelWalkReport, WalkError> {
self.options = self.options.normalized();
let skip_stdout = self.skip_stdout;
if self.runtime.is_worker_thread() {
return collect_serial(&self.root, self.options)
.map(|report| without_stdout(report, skip_stdout));
}
if self.options.follow_links {
return self
.walk_dynamic()
.map(|report| without_stdout(report, skip_stdout));
}
let mut shallow = collect_shallow(&self.root, self.options)?;
if self.options.error_policy == ErrorPolicy::Abort && !shallow.errors.is_empty() {
return Err(shallow.errors.into_iter().next().expect("error exists"));
}
if !self.options.same_file_system {
let target = requested_workers(&self.runtime, self.parallelism, self.options.max_open)
.saturating_mul(FRONTIER_TASKS_PER_WORKER);
if shallow.tasks.len() < target {
shallow = expand_frontier(shallow, self.options, target);
}
}
if shallow.tasks.is_empty() {
return Ok(without_stdout(
ParallelWalkReport {
entries: shallow.entries,
errors: shallow.errors,
},
skip_stdout,
));
}
self.walk_lanes(shallow)
.map(|report| without_stdout(report, skip_stdout))
}
fn walk_dynamic(self) -> Result<ParallelWalkReport, WalkError> {
let report = Arc::new(Mutex::new(ParallelWalkReport {
entries: Vec::new(),
errors: Vec::new(),
}));
let visitor_report = Arc::clone(&report);
dynamic::stream_batched(
&self.root,
self.options,
self.parallelism,
&self.runtime,
&CancellationToken::new(),
move |entries, errors| {
let mut report = visitor_report
.lock()
.expect("parallel walk report is not poisoned");
report.entries.extend(entries);
report.errors.extend(errors.iter().map(copy_walk_error));
true
},
)?;
let mut report = Arc::try_unwrap(report)
.expect("parallel walk visitor released report")
.into_inner()
.expect("parallel walk report is not poisoned");
report
.entries
.sort_unstable_by(|left, right| left.path().cmp(right.path()));
report.errors.sort_unstable_by(|left, right| {
left.path()
.cmp(right.path())
.then_with(|| left.depth().cmp(&right.depth()))
});
Ok(report)
}
fn walk_lanes(self, shallow: collect::ShallowWalk) -> Result<ParallelWalkReport, WalkError> {
let mut entries = shallow.entries;
let mut errors = shallow.errors;
let task_count = shallow.tasks.len();
let worker_count = parallel_worker_count(
&self.runtime,
self.parallelism,
self.options.max_open,
task_count,
);
let mut lanes = (0..worker_count)
.map(|_| Vec::<DirectoryTask>::new())
.collect::<Vec<_>>();
for (index, task) in shallow.tasks.into_iter().enumerate() {
lanes[index % worker_count].push(task);
}
let (sender, receiver) = mpsc::channel();
for (index, lane) in lanes.into_iter().enumerate() {
let sender = sender.clone();
let root = Arc::clone(&shallow.root);
let options = self.options;
self.runtime
.try_execute(move || {
let report = collect_lane(lane, options, &root);
let _ = sender.send((index, report));
})
.map_err(|source| schedule_error(&self.root, source))?;
}
drop(sender);
let mut completed = (0..worker_count).map(|_| None).collect::<Vec<_>>();
for (index, report) in receiver {
completed[index] = Some(report);
}
let mut lanes = completed
.into_iter()
.map(|lane| {
let lane = lane.expect("every parallel lane reports completion");
(
lane.report.entries.into_iter(),
lane.report.errors.into_iter(),
lane.segments,
)
})
.collect::<Vec<_>>();
for task_index in 0..task_count {
let (lane_entries, lane_errors, segments) = &mut lanes[task_index % worker_count];
let segment = segments
.pop_front()
.expect("every directory task has an output segment");
entries.extend(lane_entries.by_ref().take(segment.entries));
errors.extend(lane_errors.by_ref().take(segment.errors));
}
if self.options.error_policy == ErrorPolicy::Abort && !errors.is_empty() {
return Err(errors.remove(0));
}
Ok(ParallelWalkReport { entries, errors })
}
}
fn parallel_worker_count(
runtime: &ParallelRuntime,
parallelism: usize,
max_open: usize,
tasks: usize,
) -> usize {
requested_workers(runtime, parallelism, max_open)
.min(tasks)
.max(1)
}
fn requested_workers(runtime: &ParallelRuntime, parallelism: usize, max_open: usize) -> usize {
let available = runtime.parallelism();
if parallelism == 0 {
available.min(default_traversal_workers())
} else {
parallelism.min(available)
}
.min(max_open.max(1))
.max(1)
}
fn schedule_error(root: &std::path::Path, source: std::io::Error) -> WalkError {
WalkError::new(root, 0, WalkOperation::ScheduleWorker, source)
}
const fn default_traversal_workers() -> usize {
8
}
const FRONTIER_TASKS_PER_WORKER: usize = 2;
fn copy_walk_error(error: &WalkError) -> WalkError {
WalkError::new(
error.path().to_path_buf(),
error.depth(),
error.operation(),
std::io::Error::new(error.io_error().kind(), error.io_error().to_string()),
)
}
pub(super) fn matches_stdout(entry: &WalkEntry, identity: Option<FileIdentity>) -> bool {
identity.is_some_and(|identity| {
entry.is_file() && crate::stdout::path_matches(entry.path(), identity).unwrap_or(false)
})
}
fn without_stdout(
mut report: ParallelWalkReport,
identity: Option<FileIdentity>,
) -> ParallelWalkReport {
if identity.is_some() {
report
.entries
.retain(|entry| !matches_stdout(entry, identity));
}
report
}