use crossbeam::{
deque::{Injector, Steal, Stealer, Worker},
sync::{Parker, Unparker},
};
use std::{
ffi::OsString,
fs::{self, FileType, Metadata},
io,
path::{Path, PathBuf},
sync::{
Arc,
atomic::{AtomicBool, AtomicUsize, Ordering as AtomicOrdering},
mpsc::{Receiver, SyncSender, sync_channel},
},
thread,
};
type Descend = dyn Fn(usize, &Entry) -> bool + Send + Sync;
type Batch = io::Result<Vec<io::Result<Entry>>>;
const STAT_CHUNK_SIZE: usize = 4;
#[derive(Clone, Copy)]
pub enum Order {
Completion,
ParentFirst,
}
pub struct Entry {
pub depth: usize,
pub file_name: OsString,
pub file_type: FileType,
pub metadata: io::Result<Metadata>,
pub parent_path: Arc<Path>,
}
enum Job {
ReadDir {
root_idx: usize,
path: Arc<Path>,
entry_depth: usize,
},
StatCompletion {
root_idx: usize,
path: Arc<Path>,
entry_depth: usize,
entries: Vec<fs::DirEntry>,
},
}
impl Job {
fn root_idx(&self) -> usize {
match self {
Job::ReadDir { root_idx, .. } | Job::StatCompletion { root_idx, .. } => *root_idx,
}
}
}
enum Event {
Batch {
root_idx: usize,
batch: Batch,
},
RootFinished {
root_idx: usize,
},
Finished,
}
pub(crate) enum RootEvent {
Entry(io::Result<Entry>),
Finished,
}
struct PoolShared {
injector: Injector<Job>,
stealers: Vec<Stealer<Job>>,
stop: AtomicBool,
descend: Arc<Descend>,
events: SyncSender<Event>,
active_roots: AtomicUsize,
jobs_per_root: Vec<AtomicUsize>,
order: Order,
unparkers: Vec<Unparker>,
idle: Vec<AtomicBool>,
next_wake: AtomicUsize,
}
struct Pool {
shared: Arc<PoolShared>,
events: Receiver<Event>,
handles: Vec<thread::JoinHandle<()>>,
}
pub(crate) struct RootWalk {
next: Vec<(usize, RootEvent)>,
pool: Option<Pool>,
}
pub struct Walk {
next: Vec<io::Result<Entry>>,
pool: Option<Pool>,
}
pub fn walk(
root: &Path,
threads: usize,
order: Order,
descend: impl Fn(&Entry) -> bool + Send + Sync + 'static,
) -> Walk {
let root = Entry::from_path(root);
let pool = match &root {
Ok(entry) if entry.file_type.is_dir() && descend(entry) => {
let path = Arc::from(entry.path());
let pool = start_pool(
threads.max(1),
1,
order,
Arc::new(move |_, entry| descend(entry)),
);
start_jobs(
&pool,
vec![Job::ReadDir {
root_idx: 0,
path,
entry_depth: 1,
}],
);
Some(pool)
}
_ => None,
};
Walk {
next: vec![root],
pool,
}
}
impl Iterator for Walk {
type Item = io::Result<Entry>;
fn next(&mut self) -> Option<Self::Item> {
loop {
if let Some(entry) = self.next.pop() {
return Some(entry);
}
match self.pool.as_ref()?.events.recv() {
Ok(Event::Batch {
batch: Ok(entries), ..
}) => {
self.next.extend(entries.into_iter().rev());
}
Ok(Event::Batch {
batch: Err(err), ..
}) => return Some(Err(err)),
Ok(Event::RootFinished { .. }) => {}
Ok(Event::Finished) => {
self.pool = None;
return None;
}
Err(_) => return Some(Err(io::Error::other("directory worker stopped"))),
}
}
}
}
pub(crate) fn walk_roots(
roots: impl IntoIterator<Item = (usize, PathBuf)>,
threads: usize,
order: Order,
descend: impl Fn(usize, &Entry) -> bool + Send + Sync + 'static,
) -> RootWalk {
let roots = roots.into_iter().collect::<Vec<_>>();
let root_count = roots
.iter()
.map(|(root_idx, _)| *root_idx)
.max()
.unwrap_or(0)
+ 1;
let descend = Arc::new(descend);
let (next, root_jobs) = begin_walks(roots, descend.as_ref());
let pool = if root_jobs.is_empty() {
None
} else {
let pool = start_pool(threads.max(1), root_count, order, descend);
start_jobs(&pool, root_jobs);
Some(pool)
};
RootWalk { next, pool }
}
impl Iterator for RootWalk {
type Item = (usize, RootEvent);
fn next(&mut self) -> Option<Self::Item> {
loop {
if let Some(entry) = self.next.pop() {
return Some(entry);
}
match self.pool.as_ref()?.events.recv() {
Ok(Event::Batch {
root_idx,
batch: Ok(entries),
}) => self.next.extend(
entries
.into_iter()
.rev()
.map(|entry| (root_idx, RootEvent::Entry(entry))),
),
Ok(Event::Batch {
root_idx,
batch: Err(err),
}) => return Some((root_idx, RootEvent::Entry(Err(err)))),
Ok(Event::RootFinished { root_idx }) => {
return Some((root_idx, RootEvent::Finished));
}
Ok(Event::Finished) => {
self.pool = None;
return None;
}
Err(_) => {
return Some((
0,
RootEvent::Entry(Err(io::Error::other("directory worker stopped"))),
));
}
}
}
}
}
impl PoolShared {
fn wake_worker(&self) {
let len = self.idle.len();
let start = self.next_wake.fetch_add(1, AtomicOrdering::Relaxed) % len;
for offset in 0..len {
let idx = (start + offset) % len;
if self.idle[idx]
.compare_exchange(true, false, AtomicOrdering::AcqRel, AtomicOrdering::Relaxed)
.is_ok()
{
self.unparkers[idx].unpark();
break;
}
}
}
fn wake_workers(&self) {
for unparker in &self.unparkers {
unparker.unpark();
}
}
}
impl Entry {
#[must_use]
pub fn path(&self) -> PathBuf {
self.parent_path.join(&self.file_name)
}
fn from_path(path: &Path) -> io::Result<Self> {
let metadata = fs::symlink_metadata(path)?;
Ok(Self {
depth: 0,
file_name: path.file_name().unwrap_or(path.as_os_str()).to_owned(),
file_type: metadata.file_type(),
metadata: Ok(metadata),
parent_path: Arc::from(path.parent().unwrap_or(Path::new(""))),
})
}
fn from_dir_entry(
depth: usize,
parent_path: Arc<Path>,
entry: fs::DirEntry,
) -> io::Result<Self> {
Ok(Self {
depth,
file_name: entry.file_name(),
file_type: entry.file_type()?,
metadata: entry.metadata(),
parent_path,
})
}
}
fn start_pool(threads: usize, root_count: usize, order: Order, descend: Arc<Descend>) -> Pool {
let workers: Vec<_> = (0..threads).map(|_| Worker::new_lifo()).collect();
let parkers: Vec<_> = (0..threads).map(|_| Parker::new()).collect();
let (event_tx, event_rx) = sync_channel(threads * 2);
let shared = Arc::new(PoolShared {
injector: Injector::new(),
stealers: workers.iter().map(Worker::stealer).collect(),
stop: AtomicBool::new(false),
descend,
events: event_tx,
active_roots: AtomicUsize::new(0),
jobs_per_root: (0..root_count).map(|_| AtomicUsize::new(0)).collect(),
order,
unparkers: parkers
.iter()
.map(|parker| parker.unparker().clone())
.collect(),
idle: (0..threads).map(|_| AtomicBool::new(false)).collect(),
next_wake: AtomicUsize::new(0),
});
let handles: Vec<_> = workers
.into_iter()
.zip(parkers)
.enumerate()
.map(|(idx, (worker, parker))| {
let shared = Arc::clone(&shared);
thread::Builder::new()
.name(format!("dua-fs-walk-{idx}"))
.spawn(move || worker_loop(idx, worker, parker, shared))
.expect("filesystem worker thread can be spawned")
})
.collect();
Pool {
shared,
events: event_rx,
handles,
}
}
fn begin_walks(
roots: impl IntoIterator<Item = (usize, PathBuf)>,
descend: &Descend,
) -> (Vec<(usize, RootEvent)>, Vec<Job>) {
let mut next = Vec::new();
let mut jobs = Vec::new();
for (root_idx, path) in roots {
let entry = Entry::from_path(&path);
let has_job = if let Ok(entry) = &entry
&& entry.file_type.is_dir()
&& descend(root_idx, entry)
{
jobs.push(Job::ReadDir {
root_idx,
path: Arc::from(entry.path()),
entry_depth: 1,
});
true
} else {
false
};
next.push((root_idx, RootEvent::Entry(entry)));
if !has_job {
next.push((root_idx, RootEvent::Finished));
}
}
next.reverse();
(next, jobs)
}
fn start_jobs(pool: &Pool, root_jobs: Vec<Job>) {
let wake_all = root_jobs.len() > 1;
debug_assert_eq!(
pool.shared.active_roots.load(AtomicOrdering::Relaxed),
0,
"initial jobs must be started on an idle pool"
);
debug_assert!(
root_jobs.iter().all(|j| match j {
Job::ReadDir { entry_depth, .. } | Job::StatCompletion { entry_depth, .. } =>
*entry_depth,
} == 1),
"the first jobs should be root jobs, so active_root counts match"
);
pool.shared
.active_roots
.store(root_jobs.len(), AtomicOrdering::Relaxed);
for job in &root_jobs {
add_pending(job.root_idx(), 1, &pool.shared);
}
for job in root_jobs {
pool.shared.injector.push(job);
}
if wake_all {
pool.shared.wake_workers();
} else {
pool.shared.wake_worker();
}
}
fn worker_loop(idx: usize, worker: Worker<Job>, parker: Parker, shared: Arc<PoolShared>) {
while !shared.stop.load(AtomicOrdering::Relaxed) {
let found = if let Some(found) = find_job(&worker, &shared) {
found
} else {
shared.idle[idx].store(true, AtomicOrdering::Release);
let Some(found) = find_job(&worker, &shared) else {
parker.park();
shared.idle[idx].store(false, AtomicOrdering::Release);
continue;
};
shared.idle[idx].store(false, AtomicOrdering::Release);
found
};
let (job, stolen) = found;
if stolen {
shared.wake_worker();
}
run_job(job, &worker, &shared);
}
}
impl Drop for Pool {
fn drop(&mut self) {
self.shared.stop.store(true, AtomicOrdering::Relaxed);
self.shared.wake_workers();
for handle in self.handles.drain(..) {
handle.join().ok();
}
}
}
fn find_job(worker: &Worker<Job>, shared: &PoolShared) -> Option<(Job, bool)> {
loop {
if let Some(job) = worker.pop() {
return Some((job, false));
}
match shared.injector.steal_batch_and_pop(worker) {
Steal::Success(job) => return Some((job, false)),
Steal::Retry => continue,
Steal::Empty => {}
}
let mut retry = false;
for stealer in &shared.stealers {
match stealer.steal() {
Steal::Success(job) => return Some((job, true)),
Steal::Retry => retry = true,
Steal::Empty => {}
}
}
if !retry {
return None;
}
}
}
fn run_job(job: Job, worker: &Worker<Job>, shared: &PoolShared) {
match job {
Job::ReadDir {
root_idx: root,
path,
entry_depth,
} => {
if matches!(shared.order, Order::Completion) {
read_dir_completion(root, path, entry_depth, worker, shared);
} else {
read_dir_parent_first(root, path, entry_depth, worker, shared);
}
}
Job::StatCompletion {
root_idx: root,
path,
entry_depth,
entries,
} => stat_entries_completion(root, path, entry_depth, entries, worker, shared),
}
}
fn read_dir_completion(
root_idx: usize,
path: Arc<Path>,
entry_depth: usize,
worker: &Worker<Job>,
shared: &PoolShared,
) {
let dir_entries = match fs::read_dir(&path) {
Ok(entries) => entries,
Err(err) => {
if shared
.events
.send(Event::Batch {
root_idx,
batch: Err(err),
})
.is_err()
{
shared.stop.store(true, AtomicOrdering::Relaxed);
}
finish_pending(root_idx, shared);
return;
}
};
let mut chunk = Vec::with_capacity(STAT_CHUNK_SIZE);
let mut errors = Vec::new();
let mut has_jobs = false;
for entry in dir_entries {
match entry {
Ok(entry) => {
chunk.push(entry);
if chunk.len() == STAT_CHUNK_SIZE {
add_pending(root_idx, 1, shared);
worker.push(Job::StatCompletion {
root_idx,
path: Arc::clone(&path),
entry_depth,
entries: std::mem::replace(&mut chunk, Vec::with_capacity(STAT_CHUNK_SIZE)),
});
has_jobs = true;
}
}
Err(err) => errors.push(Err(err)),
}
}
if !chunk.is_empty() {
add_pending(root_idx, 1, shared);
worker.push(Job::StatCompletion {
root_idx,
path,
entry_depth,
entries: chunk,
});
has_jobs = true;
}
if has_jobs {
shared.wake_worker();
}
if !errors.is_empty()
&& shared
.events
.send(Event::Batch {
root_idx,
batch: Ok(errors),
})
.is_err()
{
shared.stop.store(true, AtomicOrdering::Relaxed);
}
finish_pending(root_idx, shared);
}
fn stat_entries_completion(
root_idx: usize,
path: Arc<Path>,
depth: usize,
entries: Vec<fs::DirEntry>,
worker: &Worker<Job>,
shared: &PoolShared,
) {
let mut jobs = Vec::new();
let entries = entries
.into_iter()
.map(|entry| {
Entry::from_dir_entry(depth, Arc::clone(&path), entry).inspect(|entry| {
if entry.file_type.is_dir() && (shared.descend)(root_idx, entry) {
jobs.push(Job::ReadDir {
root_idx,
path: Arc::from(entry.path()),
entry_depth: entry.depth + 1,
});
}
})
})
.collect();
add_pending(root_idx, jobs.len(), shared);
schedule_jobs(jobs, worker, shared);
if shared
.events
.send(Event::Batch {
root_idx,
batch: Ok(entries),
})
.is_err()
{
shared.stop.store(true, AtomicOrdering::Relaxed);
}
finish_pending(root_idx, shared);
}
fn read_dir_parent_first(
root_idx: usize,
path: Arc<Path>,
depth: usize,
worker: &Worker<Job>,
shared: &PoolShared,
) {
let dir_entries = match fs::read_dir(&path) {
Ok(entries) => entries,
Err(err) => {
finish_directory(root_idx, Err(err), Vec::new(), worker, shared);
return;
}
};
let mut jobs = Vec::new();
let entries = dir_entries
.map(|entry| {
entry
.and_then(|entry| Entry::from_dir_entry(depth, Arc::clone(&path), entry))
.inspect(|entry| {
if entry.file_type.is_dir() && (shared.descend)(root_idx, entry) {
jobs.push(Job::ReadDir {
root_idx,
path: Arc::from(entry.path()),
entry_depth: depth + 1,
});
}
})
})
.collect();
finish_directory(root_idx, Ok(entries), jobs, worker, shared);
}
fn finish_directory(
root_idx: usize,
batch: Batch,
jobs: Vec<Job>,
worker: &Worker<Job>,
shared: &PoolShared,
) {
add_pending(root_idx, jobs.len(), shared);
match shared.order {
Order::ParentFirst => {
if shared
.events
.send(Event::Batch { root_idx, batch })
.is_err()
{
shared.stop.store(true, AtomicOrdering::Relaxed);
return;
}
schedule_jobs(jobs, worker, shared);
}
Order::Completion => {
schedule_jobs(jobs, worker, shared);
if shared
.events
.send(Event::Batch { root_idx, batch })
.is_err()
{
shared.stop.store(true, AtomicOrdering::Relaxed);
return;
}
}
}
finish_pending(root_idx, shared);
}
fn add_pending(root: usize, count: usize, shared: &PoolShared) {
shared.jobs_per_root[root].fetch_add(count, AtomicOrdering::Relaxed);
}
fn finish_pending(root_idx: usize, shared: &PoolShared) {
if shared.jobs_per_root[root_idx].fetch_sub(1, AtomicOrdering::Relaxed) == 1 {
shared.events.send(Event::RootFinished { root_idx }).ok();
if shared.active_roots.fetch_sub(1, AtomicOrdering::Relaxed) == 1 {
shared.events.send(Event::Finished).ok();
}
}
}
fn schedule_jobs(jobs: Vec<Job>, worker: &Worker<Job>, shared: &PoolShared) {
let has_jobs = !jobs.is_empty();
for job in jobs {
worker.push(job);
}
if has_jobs {
shared.wake_worker();
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parallel_walk_is_parent_first_and_does_not_follow_symlinks() {
let dir = tempfile::tempdir().unwrap();
fs::create_dir_all(dir.path().join("b/child")).unwrap();
fs::create_dir(dir.path().join("a")).unwrap();
fs::write(dir.path().join("b/child/file"), b"x").unwrap();
#[cfg(unix)]
std::os::unix::fs::symlink(dir.path().join("b"), dir.path().join("link")).unwrap();
#[cfg(unix)]
let expected = ["", "a", "b", "b/child", "b/child/file", "link"];
#[cfg(not(unix))]
let expected = ["", "a", "b", "b/child", "b/child/file"];
let expected = expected.into_iter().map(PathBuf::from).collect::<Vec<_>>();
for threads in [1, 4] {
let paths = walk(dir.path(), threads, Order::ParentFirst, |_| true)
.map(|entry| {
entry
.unwrap()
.path()
.strip_prefix(dir.path())
.unwrap()
.to_owned()
})
.collect::<Vec<_>>();
let mut sorted_paths = paths.clone();
sorted_paths.sort();
assert_eq!(
sorted_paths, expected,
"walk with {threads} threads should visit every expected path exactly once"
);
for path in paths.iter().filter(|path| path.components().count() > 1) {
let parent = path.parent().unwrap();
assert!(
paths.iter().position(|path| path == parent)
< paths.iter().position(|candidate| candidate == path),
"parent {parent:?} should precede child {path:?} with {threads} threads; \
traversal order: {paths:?}"
);
}
}
}
#[test]
fn pruning_keeps_the_directory_and_missing_roots_are_errors() {
let dir = tempfile::tempdir().unwrap();
fs::create_dir_all(dir.path().join("skip/child")).unwrap();
let paths = walk(dir.path(), 2, Order::Completion, |entry| {
entry.file_name != "skip"
})
.map(|entry| entry.unwrap().file_name)
.collect::<Vec<_>>();
assert_eq!(
paths,
vec![
dir.path().file_name().unwrap().to_owned(),
OsString::from("skip")
],
"a pruned directory should be yielded without traversing its children"
);
assert!(
walk(&dir.path().join("missing"), 2, Order::Completion, |_| true)
.next()
.unwrap()
.is_err(),
"a missing root should be yielded as an I/O error"
);
}
#[test]
fn concurrent_roots_keep_their_identity() {
let dir = tempfile::tempdir().unwrap();
let roots = [dir.path().join("a"), dir.path().join("b")];
for root in &roots {
fs::create_dir_all(root.join("child")).unwrap();
}
let events = walk_roots(
roots.iter().cloned().enumerate(),
2,
Order::Completion,
|_, _| true,
)
.collect::<Vec<_>>();
let mut paths = Vec::new();
let mut last_entry = [0; 2];
let mut finished = [None; 2];
for (position, (root_idx, event)) in events.into_iter().enumerate() {
match event {
RootEvent::Entry(entry) => {
last_entry[root_idx] = position;
paths.push((
root_idx,
entry
.unwrap()
.path()
.strip_prefix(&roots[root_idx])
.unwrap()
.to_owned(),
));
}
RootEvent::Finished => finished[root_idx] = Some(position),
}
}
paths.sort();
assert_eq!(
paths,
[
(0, PathBuf::new()),
(0, PathBuf::from("child")),
(1, PathBuf::new()),
(1, PathBuf::from("child")),
]
);
for root_idx in 0..roots.len() {
assert!(
last_entry[root_idx] < finished[root_idx].unwrap(),
"root {root_idx} must finish after its last entry",
);
}
}
#[test]
fn wide_walk_wakes_multiple_idle_workers() {
let dir = tempfile::tempdir().unwrap();
for idx in 0..32 {
fs::create_dir_all(dir.path().join(format!("{idx}/child"))).unwrap();
}
let worker_threads = Arc::new(std::sync::Mutex::new(std::collections::HashSet::new()));
let seen_threads = Arc::clone(&worker_threads);
walk(dir.path(), 8, Order::Completion, move |entry| {
if entry.depth == 1 {
thread::sleep(std::time::Duration::from_millis(1));
} else if entry.depth == 2 {
seen_threads.lock().unwrap().insert(thread::current().id());
thread::sleep(std::time::Duration::from_millis(10));
}
true
})
.for_each(drop);
assert!(
worker_threads.lock().unwrap().len() >= 4,
"a wide directory should engage more than the producer and one thief"
);
}
}