1use crossbeam::{
23 deque::{Injector, Steal, Stealer, Worker},
24 sync::{Parker, Unparker},
25};
26use std::{
27 ffi::OsString,
28 fs::{self, FileType, Metadata},
29 io,
30 path::{Path, PathBuf},
31 sync::{
32 Arc,
33 atomic::{AtomicBool, AtomicUsize, Ordering as AtomicOrdering},
34 mpsc::{Receiver, SyncSender, sync_channel},
35 },
36 thread,
37};
38
39type Descend = dyn Fn(usize, &Entry) -> bool + Send + Sync;
42type Batch = io::Result<Vec<io::Result<Entry>>>;
46const STAT_CHUNK_SIZE: usize = 4;
49
50#[derive(Clone, Copy)]
52pub enum Order {
53 Completion,
55 ParentFirst,
57}
58
59pub struct Entry {
61 pub depth: usize,
63 pub file_name: OsString,
65 pub file_type: FileType,
67 pub metadata: io::Result<Metadata>,
69 pub parent_path: Arc<Path>,
71}
72
73enum Job {
74 ReadDir {
76 root_idx: usize,
77 path: Arc<Path>,
78 entry_depth: usize,
81 },
82 StatCompletion {
84 root_idx: usize,
85 path: Arc<Path>,
86 entry_depth: usize,
88 entries: Vec<fs::DirEntry>,
89 },
90}
91
92impl Job {
93 fn root_idx(&self) -> usize {
95 match self {
96 Job::ReadDir { root_idx, .. } | Job::StatCompletion { root_idx, .. } => *root_idx,
97 }
98 }
99}
100
101enum Event {
103 Batch {
104 root_idx: usize,
105 batch: Batch,
106 },
107 RootFinished {
110 root_idx: usize,
111 },
112 Finished,
114}
115
116pub(crate) enum RootEvent {
122 Entry(io::Result<Entry>),
123 Finished,
124}
125
126struct PoolShared {
127 injector: Injector<Job>,
129 stealers: Vec<Stealer<Job>>,
130 stop: AtomicBool,
131 descend: Arc<Descend>,
132 events: SyncSender<Event>,
133 active_roots: AtomicUsize,
135 jobs_per_root: Vec<AtomicUsize>,
138 order: Order,
139 unparkers: Vec<Unparker>,
141 idle: Vec<AtomicBool>,
144 next_wake: AtomicUsize,
146}
147
148struct Pool {
149 shared: Arc<PoolShared>,
150 events: Receiver<Event>,
151 handles: Vec<thread::JoinHandle<()>>,
152}
153
154pub(crate) struct RootWalk {
157 next: Vec<(usize, RootEvent)>,
159 pool: Option<Pool>,
161}
162
163pub struct Walk {
166 next: Vec<io::Result<Entry>>,
173 pool: Option<Pool>,
177}
178
179pub fn walk(
183 root: &Path,
184 threads: usize,
185 order: Order,
186 descend: impl Fn(&Entry) -> bool + Send + Sync + 'static,
187) -> Walk {
188 let root = Entry::from_path(root);
189 let pool = match &root {
190 Ok(entry) if entry.file_type.is_dir() && descend(entry) => {
191 let path = Arc::from(entry.path());
192 let pool = start_pool(
193 threads.max(1),
194 1,
195 order,
196 Arc::new(move |_, entry| descend(entry)),
197 );
198 start_jobs(
199 &pool,
200 vec![Job::ReadDir {
201 root_idx: 0,
202 path,
203 entry_depth: 1,
204 }],
205 );
206 Some(pool)
207 }
208 _ => None,
209 };
210 Walk {
211 next: vec![root],
212 pool,
213 }
214}
215
216impl Iterator for Walk {
217 type Item = io::Result<Entry>;
218
219 fn next(&mut self) -> Option<Self::Item> {
220 loop {
221 if let Some(entry) = self.next.pop() {
222 return Some(entry);
223 }
224
225 match self.pool.as_ref()?.events.recv() {
226 Ok(Event::Batch {
227 batch: Ok(entries), ..
228 }) => {
229 self.next.extend(entries.into_iter().rev());
230 }
231 Ok(Event::Batch {
232 batch: Err(err), ..
233 }) => return Some(Err(err)),
234 Ok(Event::RootFinished { .. }) => {}
235 Ok(Event::Finished) => {
236 self.pool = None;
237 return None;
238 }
239 Err(_) => return Some(Err(io::Error::other("directory worker stopped"))),
240 }
241 }
242 }
243}
244
245pub(crate) fn walk_roots(
248 roots: impl IntoIterator<Item = (usize, PathBuf)>,
249 threads: usize,
250 order: Order,
251 descend: impl Fn(usize, &Entry) -> bool + Send + Sync + 'static,
252) -> RootWalk {
253 let roots = roots.into_iter().collect::<Vec<_>>();
254 let root_count = roots
255 .iter()
256 .map(|(root_idx, _)| *root_idx)
257 .max()
258 .unwrap_or(0)
259 + 1;
260 let descend = Arc::new(descend);
261 let (next, root_jobs) = begin_walks(roots, descend.as_ref());
262 let pool = if root_jobs.is_empty() {
263 None
264 } else {
265 let pool = start_pool(threads.max(1), root_count, order, descend);
266 start_jobs(&pool, root_jobs);
267 Some(pool)
268 };
269 RootWalk { next, pool }
270}
271
272impl Iterator for RootWalk {
273 type Item = (usize, RootEvent);
274
275 fn next(&mut self) -> Option<Self::Item> {
276 loop {
277 if let Some(entry) = self.next.pop() {
278 return Some(entry);
279 }
280 match self.pool.as_ref()?.events.recv() {
281 Ok(Event::Batch {
282 root_idx,
283 batch: Ok(entries),
284 }) => self.next.extend(
285 entries
286 .into_iter()
287 .rev()
288 .map(|entry| (root_idx, RootEvent::Entry(entry))),
289 ),
290 Ok(Event::Batch {
291 root_idx,
292 batch: Err(err),
293 }) => return Some((root_idx, RootEvent::Entry(Err(err)))),
294 Ok(Event::RootFinished { root_idx }) => {
295 return Some((root_idx, RootEvent::Finished));
296 }
297 Ok(Event::Finished) => {
298 self.pool = None;
299 return None;
300 }
301 Err(_) => {
302 return Some((
303 0,
304 RootEvent::Entry(Err(io::Error::other("directory worker stopped"))),
305 ));
306 }
307 }
308 }
309 }
310}
311
312impl PoolShared {
313 fn wake_worker(&self) {
315 let len = self.idle.len();
316 let start = self.next_wake.fetch_add(1, AtomicOrdering::Relaxed) % len;
319 for offset in 0..len {
320 let idx = (start + offset) % len;
321 if self.idle[idx]
322 .compare_exchange(true, false, AtomicOrdering::AcqRel, AtomicOrdering::Relaxed)
323 .is_ok()
324 {
325 self.unparkers[idx].unpark();
326 break;
327 }
328 }
329 }
330
331 fn wake_workers(&self) {
333 for unparker in &self.unparkers {
334 unparker.unpark();
335 }
336 }
337}
338
339impl Entry {
340 #[must_use]
342 pub fn path(&self) -> PathBuf {
343 self.parent_path.join(&self.file_name)
344 }
345
346 fn from_path(path: &Path) -> io::Result<Self> {
347 let metadata = fs::symlink_metadata(path)?;
348 Ok(Self {
349 depth: 0,
350 file_name: path.file_name().unwrap_or(path.as_os_str()).to_owned(),
351 file_type: metadata.file_type(),
352 metadata: Ok(metadata),
353 parent_path: Arc::from(path.parent().unwrap_or(Path::new(""))),
354 })
355 }
356
357 fn from_dir_entry(
358 depth: usize,
359 parent_path: Arc<Path>,
360 entry: fs::DirEntry,
361 ) -> io::Result<Self> {
362 Ok(Self {
363 depth,
364 file_name: entry.file_name(),
365 file_type: entry.file_type()?,
366 metadata: entry.metadata(),
367 parent_path,
368 })
369 }
370}
371
372fn start_pool(threads: usize, root_count: usize, order: Order, descend: Arc<Descend>) -> Pool {
373 let workers: Vec<_> = (0..threads).map(|_| Worker::new_lifo()).collect();
374 let parkers: Vec<_> = (0..threads).map(|_| Parker::new()).collect();
375 let (event_tx, event_rx) = sync_channel(threads * 2);
376 let shared = Arc::new(PoolShared {
377 injector: Injector::new(),
378 stealers: workers.iter().map(Worker::stealer).collect(),
379 stop: AtomicBool::new(false),
380 descend,
381 events: event_tx,
382 active_roots: AtomicUsize::new(0),
383 jobs_per_root: (0..root_count).map(|_| AtomicUsize::new(0)).collect(),
384 order,
385 unparkers: parkers
386 .iter()
387 .map(|parker| parker.unparker().clone())
388 .collect(),
389 idle: (0..threads).map(|_| AtomicBool::new(false)).collect(),
390 next_wake: AtomicUsize::new(0),
391 });
392 let handles: Vec<_> = workers
393 .into_iter()
394 .zip(parkers)
395 .enumerate()
396 .map(|(idx, (worker, parker))| {
397 let shared = Arc::clone(&shared);
398 thread::Builder::new()
399 .name(format!("dua-fs-walk-{idx}"))
400 .spawn(move || worker_loop(idx, worker, parker, shared))
401 .expect("filesystem worker thread can be spawned")
402 })
403 .collect();
404
405 Pool {
406 shared,
407 events: event_rx,
408 handles,
409 }
410}
411
412fn begin_walks(
415 roots: impl IntoIterator<Item = (usize, PathBuf)>,
416 descend: &Descend,
417) -> (Vec<(usize, RootEvent)>, Vec<Job>) {
418 let mut next = Vec::new();
419 let mut jobs = Vec::new();
420 for (root_idx, path) in roots {
421 let entry = Entry::from_path(&path);
422 let has_job = if let Ok(entry) = &entry
423 && entry.file_type.is_dir()
424 && descend(root_idx, entry)
425 {
426 jobs.push(Job::ReadDir {
427 root_idx,
428 path: Arc::from(entry.path()),
429 entry_depth: 1,
430 });
431 true
432 } else {
433 false
434 };
435 next.push((root_idx, RootEvent::Entry(entry)));
436 if !has_job {
437 next.push((root_idx, RootEvent::Finished));
438 }
439 }
440 next.reverse();
441 (next, jobs)
442}
443
444fn start_jobs(pool: &Pool, root_jobs: Vec<Job>) {
447 let wake_all = root_jobs.len() > 1;
448 debug_assert_eq!(
449 pool.shared.active_roots.load(AtomicOrdering::Relaxed),
450 0,
451 "initial jobs must be started on an idle pool"
452 );
453 debug_assert!(
454 root_jobs.iter().all(|j| match j {
455 Job::ReadDir { entry_depth, .. } | Job::StatCompletion { entry_depth, .. } =>
456 *entry_depth,
457 } == 1),
458 "the first jobs should be root jobs, so active_root counts match"
459 );
460 pool.shared
461 .active_roots
462 .store(root_jobs.len(), AtomicOrdering::Relaxed);
463 for job in &root_jobs {
464 add_pending(job.root_idx(), 1, &pool.shared);
465 }
466 for job in root_jobs {
467 pool.shared.injector.push(job);
468 }
469 if wake_all {
470 pool.shared.wake_workers();
471 } else {
472 pool.shared.wake_worker();
473 }
474}
475
476fn worker_loop(idx: usize, worker: Worker<Job>, parker: Parker, shared: Arc<PoolShared>) {
477 while !shared.stop.load(AtomicOrdering::Relaxed) {
478 let found = if let Some(found) = find_job(&worker, &shared) {
479 found
480 } else {
481 shared.idle[idx].store(true, AtomicOrdering::Release);
482 let Some(found) = find_job(&worker, &shared) else {
483 parker.park();
484 shared.idle[idx].store(false, AtomicOrdering::Release);
485 continue;
486 };
487 shared.idle[idx].store(false, AtomicOrdering::Release);
488 found
489 };
490 let (job, stolen) = found;
491 if stolen {
492 shared.wake_worker();
495 }
496 run_job(job, &worker, &shared);
497 }
498}
499
500impl Drop for Pool {
501 fn drop(&mut self) {
502 self.shared.stop.store(true, AtomicOrdering::Relaxed);
503 self.shared.wake_workers();
504 for handle in self.handles.drain(..) {
505 handle.join().ok();
506 }
507 }
508}
509
510fn find_job(worker: &Worker<Job>, shared: &PoolShared) -> Option<(Job, bool)> {
521 loop {
522 if let Some(job) = worker.pop() {
523 return Some((job, false));
524 }
525
526 match shared.injector.steal_batch_and_pop(worker) {
527 Steal::Success(job) => return Some((job, false)),
528 Steal::Retry => continue,
529 Steal::Empty => {}
530 }
531
532 let mut retry = false;
533 for stealer in &shared.stealers {
534 match stealer.steal() {
535 Steal::Success(job) => return Some((job, true)),
536 Steal::Retry => retry = true,
537 Steal::Empty => {}
538 }
539 }
540 if !retry {
541 return None;
542 }
543 }
544}
545
546fn run_job(job: Job, worker: &Worker<Job>, shared: &PoolShared) {
547 match job {
548 Job::ReadDir {
549 root_idx: root,
550 path,
551 entry_depth,
552 } => {
553 if matches!(shared.order, Order::Completion) {
554 read_dir_completion(root, path, entry_depth, worker, shared);
555 } else {
556 read_dir_parent_first(root, path, entry_depth, worker, shared);
557 }
558 }
559 Job::StatCompletion {
560 root_idx: root,
561 path,
562 entry_depth,
563 entries,
564 } => stat_entries_completion(root, path, entry_depth, entries, worker, shared),
565 }
566}
567
568fn read_dir_completion(
574 root_idx: usize,
575 path: Arc<Path>,
576 entry_depth: usize,
577 worker: &Worker<Job>,
578 shared: &PoolShared,
579) {
580 let dir_entries = match fs::read_dir(&path) {
581 Ok(entries) => entries,
582 Err(err) => {
583 if shared
584 .events
585 .send(Event::Batch {
586 root_idx,
587 batch: Err(err),
588 })
589 .is_err()
590 {
591 shared.stop.store(true, AtomicOrdering::Relaxed);
592 }
593 finish_pending(root_idx, shared);
594 return;
595 }
596 };
597 let mut chunk = Vec::with_capacity(STAT_CHUNK_SIZE);
598 let mut errors = Vec::new();
599 let mut has_jobs = false;
600 for entry in dir_entries {
601 match entry {
602 Ok(entry) => {
603 chunk.push(entry);
604 if chunk.len() == STAT_CHUNK_SIZE {
605 add_pending(root_idx, 1, shared);
606 worker.push(Job::StatCompletion {
607 root_idx,
608 path: Arc::clone(&path),
609 entry_depth,
610 entries: std::mem::replace(&mut chunk, Vec::with_capacity(STAT_CHUNK_SIZE)),
611 });
612 has_jobs = true;
613 }
614 }
615 Err(err) => errors.push(Err(err)),
616 }
617 }
618 if !chunk.is_empty() {
619 add_pending(root_idx, 1, shared);
620 worker.push(Job::StatCompletion {
621 root_idx,
622 path,
623 entry_depth,
624 entries: chunk,
625 });
626 has_jobs = true;
627 }
628 if has_jobs {
629 shared.wake_worker();
630 }
631 if !errors.is_empty()
632 && shared
633 .events
634 .send(Event::Batch {
635 root_idx,
636 batch: Ok(errors),
637 })
638 .is_err()
639 {
640 shared.stop.store(true, AtomicOrdering::Relaxed);
641 }
642 finish_pending(root_idx, shared);
643}
644
645fn stat_entries_completion(
646 root_idx: usize,
647 path: Arc<Path>,
648 depth: usize,
649 entries: Vec<fs::DirEntry>,
650 worker: &Worker<Job>,
651 shared: &PoolShared,
652) {
653 let mut jobs = Vec::new();
654 let entries = entries
655 .into_iter()
656 .map(|entry| {
657 Entry::from_dir_entry(depth, Arc::clone(&path), entry).inspect(|entry| {
658 if entry.file_type.is_dir() && (shared.descend)(root_idx, entry) {
659 jobs.push(Job::ReadDir {
660 root_idx,
661 path: Arc::from(entry.path()),
662 entry_depth: entry.depth + 1,
663 });
664 }
665 })
666 })
667 .collect();
668 add_pending(root_idx, jobs.len(), shared);
669 schedule_jobs(jobs, worker, shared);
670 if shared
671 .events
672 .send(Event::Batch {
673 root_idx,
674 batch: Ok(entries),
675 })
676 .is_err()
677 {
678 shared.stop.store(true, AtomicOrdering::Relaxed);
679 }
680 finish_pending(root_idx, shared);
681}
682
683fn read_dir_parent_first(
691 root_idx: usize,
692 path: Arc<Path>,
693 depth: usize,
694 worker: &Worker<Job>,
695 shared: &PoolShared,
696) {
697 let dir_entries = match fs::read_dir(&path) {
698 Ok(entries) => entries,
699 Err(err) => {
700 finish_directory(root_idx, Err(err), Vec::new(), worker, shared);
701 return;
702 }
703 };
704 let mut jobs = Vec::new();
705 let entries = dir_entries
706 .map(|entry| {
707 entry
708 .and_then(|entry| Entry::from_dir_entry(depth, Arc::clone(&path), entry))
709 .inspect(|entry| {
710 if entry.file_type.is_dir() && (shared.descend)(root_idx, entry) {
711 jobs.push(Job::ReadDir {
712 root_idx,
713 path: Arc::from(entry.path()),
714 entry_depth: depth + 1,
715 });
716 }
717 })
718 })
719 .collect();
720 finish_directory(root_idx, Ok(entries), jobs, worker, shared);
721}
722
723fn finish_directory(
727 root_idx: usize,
728 batch: Batch,
729 jobs: Vec<Job>,
730 worker: &Worker<Job>,
731 shared: &PoolShared,
732) {
733 add_pending(root_idx, jobs.len(), shared);
734
735 match shared.order {
736 Order::ParentFirst => {
737 if shared
738 .events
739 .send(Event::Batch { root_idx, batch })
740 .is_err()
741 {
742 shared.stop.store(true, AtomicOrdering::Relaxed);
743 return;
744 }
745 schedule_jobs(jobs, worker, shared);
746 }
747 Order::Completion => {
748 schedule_jobs(jobs, worker, shared);
749 if shared
750 .events
751 .send(Event::Batch { root_idx, batch })
752 .is_err()
753 {
754 shared.stop.store(true, AtomicOrdering::Relaxed);
755 return;
756 }
757 }
758 }
759
760 finish_pending(root_idx, shared);
761}
762
763fn add_pending(root: usize, count: usize, shared: &PoolShared) {
764 shared.jobs_per_root[root].fetch_add(count, AtomicOrdering::Relaxed);
765}
766
767fn finish_pending(root_idx: usize, shared: &PoolShared) {
770 if shared.jobs_per_root[root_idx].fetch_sub(1, AtomicOrdering::Relaxed) == 1 {
771 shared.events.send(Event::RootFinished { root_idx }).ok();
772 if shared.active_roots.fetch_sub(1, AtomicOrdering::Relaxed) == 1 {
773 shared.events.send(Event::Finished).ok();
774 }
775 }
776}
777
778fn schedule_jobs(jobs: Vec<Job>, worker: &Worker<Job>, shared: &PoolShared) {
779 let has_jobs = !jobs.is_empty();
780 for job in jobs {
781 worker.push(job);
782 }
783 if has_jobs {
784 shared.wake_worker();
785 }
786}
787
788#[cfg(test)]
789mod tests {
790 use super::*;
791
792 #[test]
793 fn parallel_walk_is_parent_first_and_does_not_follow_symlinks() {
794 let dir = tempfile::tempdir().unwrap();
795 fs::create_dir_all(dir.path().join("b/child")).unwrap();
796 fs::create_dir(dir.path().join("a")).unwrap();
797 fs::write(dir.path().join("b/child/file"), b"x").unwrap();
798
799 #[cfg(unix)]
800 std::os::unix::fs::symlink(dir.path().join("b"), dir.path().join("link")).unwrap();
801
802 #[cfg(unix)]
803 let expected = ["", "a", "b", "b/child", "b/child/file", "link"];
804 #[cfg(not(unix))]
805 let expected = ["", "a", "b", "b/child", "b/child/file"];
806 let expected = expected.into_iter().map(PathBuf::from).collect::<Vec<_>>();
807
808 for threads in [1, 4] {
809 let paths = walk(dir.path(), threads, Order::ParentFirst, |_| true)
810 .map(|entry| {
811 entry
812 .unwrap()
813 .path()
814 .strip_prefix(dir.path())
815 .unwrap()
816 .to_owned()
817 })
818 .collect::<Vec<_>>();
819 let mut sorted_paths = paths.clone();
820 sorted_paths.sort();
821 assert_eq!(
822 sorted_paths, expected,
823 "walk with {threads} threads should visit every expected path exactly once"
824 );
825
826 for path in paths.iter().filter(|path| path.components().count() > 1) {
827 let parent = path.parent().unwrap();
828 assert!(
829 paths.iter().position(|path| path == parent)
830 < paths.iter().position(|candidate| candidate == path),
831 "parent {parent:?} should precede child {path:?} with {threads} threads; \
832 traversal order: {paths:?}"
833 );
834 }
835 }
836 }
837
838 #[test]
839 fn pruning_keeps_the_directory_and_missing_roots_are_errors() {
840 let dir = tempfile::tempdir().unwrap();
841 fs::create_dir_all(dir.path().join("skip/child")).unwrap();
842
843 let paths = walk(dir.path(), 2, Order::Completion, |entry| {
844 entry.file_name != "skip"
845 })
846 .map(|entry| entry.unwrap().file_name)
847 .collect::<Vec<_>>();
848 assert_eq!(
849 paths,
850 vec![
851 dir.path().file_name().unwrap().to_owned(),
852 OsString::from("skip")
853 ],
854 "a pruned directory should be yielded without traversing its children"
855 );
856
857 assert!(
858 walk(&dir.path().join("missing"), 2, Order::Completion, |_| true)
859 .next()
860 .unwrap()
861 .is_err(),
862 "a missing root should be yielded as an I/O error"
863 );
864 }
865
866 #[test]
867 fn concurrent_roots_keep_their_identity() {
868 let dir = tempfile::tempdir().unwrap();
869 let roots = [dir.path().join("a"), dir.path().join("b")];
870 for root in &roots {
871 fs::create_dir_all(root.join("child")).unwrap();
872 }
873
874 let events = walk_roots(
875 roots.iter().cloned().enumerate(),
876 2,
877 Order::Completion,
878 |_, _| true,
879 )
880 .collect::<Vec<_>>();
881 let mut paths = Vec::new();
882 let mut last_entry = [0; 2];
883 let mut finished = [None; 2];
884 for (position, (root_idx, event)) in events.into_iter().enumerate() {
885 match event {
886 RootEvent::Entry(entry) => {
887 last_entry[root_idx] = position;
888 paths.push((
889 root_idx,
890 entry
891 .unwrap()
892 .path()
893 .strip_prefix(&roots[root_idx])
894 .unwrap()
895 .to_owned(),
896 ));
897 }
898 RootEvent::Finished => finished[root_idx] = Some(position),
899 }
900 }
901 paths.sort();
902 assert_eq!(
903 paths,
904 [
905 (0, PathBuf::new()),
906 (0, PathBuf::from("child")),
907 (1, PathBuf::new()),
908 (1, PathBuf::from("child")),
909 ]
910 );
911 for root_idx in 0..roots.len() {
912 assert!(
913 last_entry[root_idx] < finished[root_idx].unwrap(),
914 "root {root_idx} must finish after its last entry",
915 );
916 }
917 }
918
919 #[test]
920 fn wide_walk_wakes_multiple_idle_workers() {
921 let dir = tempfile::tempdir().unwrap();
922 for idx in 0..32 {
923 fs::create_dir_all(dir.path().join(format!("{idx}/child"))).unwrap();
924 }
925
926 let worker_threads = Arc::new(std::sync::Mutex::new(std::collections::HashSet::new()));
927 let seen_threads = Arc::clone(&worker_threads);
928 walk(dir.path(), 8, Order::Completion, move |entry| {
929 if entry.depth == 1 {
930 thread::sleep(std::time::Duration::from_millis(1));
931 } else if entry.depth == 2 {
932 seen_threads.lock().unwrap().insert(thread::current().id());
933 thread::sleep(std::time::Duration::from_millis(10));
934 }
935 true
936 })
937 .for_each(drop);
938
939 assert!(
940 worker_threads.lock().unwrap().len() >= 4,
941 "a wide directory should engage more than the producer and one thief"
942 );
943 }
944}