Skip to main content

dua/
walk.rs

1//! Parallel filesystem traversal backed by a work-stealing worker pool.
2//!
3//! [`walk`] yields the root first, then workers read directories and distribute newly discovered
4//! subdirectories among themselves. [`Order::ParentFirst`] publishes each directory's entries
5//! before scheduling its children, while [`Order::Completion`] allows descendant batches to arrive
6//! first when their reads finish sooner. Sibling order is unspecified in both modes.
7//!
8//! The `descend` predicate controls which directories are traversed; rejected directories are
9//! still yielded (but not traversed).
10//! Symbolic links are reported but never followed, and filesystem errors are
11//! returned as iterator items. Dropping the iterator stops and joins its workers.
12//!
13//! # Scheduling
14//!
15//! The root directory starts in a shared injector queue. Directory reads enqueue small metadata
16//! batches, and metadata batches enqueue accepted child directories. Every worker can run either
17//! kind of job from its local LIFO queue or steal from a peer. Each successful thief wakes another
18//! idle worker, ramping up only while work remains stealable. A worker parks when no queue has work
19//! and is unparked when new work arrives or the walk stops. The last completed job emits the
20//! finished event; dropping the iterator stops all workers, unparks them, and joins their threads.
21
22use 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
39/// Decides whether to traverse an entry's children for a given root index.
40/// Returning `false` prunes descendants but still emits the entry itself.
41type Descend = dyn Fn(usize, &Entry) -> bool + Send + Sync;
42/// Entries obtained from one directory read.
43/// The outer error means `fs::read_dir` could not open the directory; inner errors come from
44/// reading or converting individual directory entries.
45type Batch = io::Result<Vec<io::Result<Entry>>>;
46/// Number of directory entries grouped into each stealable metadata job.
47/// Small chunks expose parallel work while amortizing queueing overhead across several entries.
48const STAT_CHUNK_SIZE: usize = 4;
49
50/// Controls when entries are yielded relative to their descendants.
51#[derive(Clone, Copy)]
52pub enum Order {
53    /// Yield entries as their parent-directory reads complete.
54    Completion,
55    /// Yield every parent before its descendants.
56    ParentFirst,
57}
58
59/// A filesystem entry produced by [`walk`].
60pub struct Entry {
61    /// Distance from the walk root: `0` for the root, `1` for its children, and so on.
62    pub depth: usize,
63    /// File name relative to `parent_path`.
64    pub file_name: OsString,
65    /// Filesystem entry type without following symbolic links.
66    pub file_type: FileType,
67    /// Entry metadata, or the error encountered while reading it.
68    pub metadata: io::Result<Metadata>,
69    /// Path containing this entry.
70    pub parent_path: Arc<Path>,
71}
72
73enum Job {
74    /// Read a directory and schedule processing of its entries.
75    ReadDir {
76        root_idx: usize,
77        path: Arc<Path>,
78        /// Depth to be assigned to entries read from `path`; always at least `1`.
79        /// The directory at `path` is one level shallower.
80        entry_depth: usize,
81    },
82    /// Fetch metadata for a chunk of entries from a completed directory read.
83    StatCompletion {
84        root_idx: usize,
85        path: Arc<Path>,
86        /// Depth assigned to every entry in this chunk; always at least `1`, i.e. a file in a directory.
87        entry_depth: usize,
88        entries: Vec<fs::DirEntry>,
89    },
90}
91
92impl Job {
93    /// Return the index of the root path that this job belongs to.
94    fn root_idx(&self) -> usize {
95        match self {
96            Job::ReadDir { root_idx, .. } | Job::StatCompletion { root_idx, .. } => *root_idx,
97        }
98    }
99}
100
101/// Internal worker-channel events, including batches, per-root completion, and pool completion.
102enum Event {
103    Batch {
104        root_idx: usize,
105        batch: Batch,
106    },
107    /// All work for this root is complete; emitted after all of its batches.
108    /// Completion events for different roots may occur in any order.
109    RootFinished {
110        root_idx: usize,
111    },
112    /// Emitted once after all roots have emitted `RootFinished`; this is the final event.
113    Finished,
114}
115
116/// Per-root events exposed by [`RootWalk`].
117/// Unlike [`Event`], batches are flattened into entries and pool-wide completion ends the iterator
118/// instead of being yielded; `Finished` therefore means only that the associated root completed.
119/// [`RootWalk`] yields `(root_idx, event)`, separating root routing from event meaning. [`Event`]
120/// cannot do this uniformly because its `Finished` variant is pool-wide and has no root index.
121pub(crate) enum RootEvent {
122    Entry(io::Result<Entry>),
123    Finished,
124}
125
126struct PoolShared {
127    /// Global queue that makes the initial root job available to whichever worker starts first.
128    injector: Injector<Job>,
129    stealers: Vec<Stealer<Job>>,
130    stop: AtomicBool,
131    descend: Arc<Descend>,
132    events: SyncSender<Event>,
133    /// Number of roots with queued or running jobs.
134    active_roots: AtomicUsize,
135    /// Number of queued or running jobs for each root index.
136    /// A counter reaching zero emits that root's [`Event::RootFinished`].
137    jobs_per_root: Vec<AtomicUsize>,
138    order: Order,
139    /// Handles used to wake workers, indexed by worker number.
140    unparkers: Vec<Unparker>,
141    /// Whether each worker has announced that it is idle, indexed like `unparkers`.
142    /// `wake_worker` atomically claims one idle worker before unparking it.
143    idle: Vec<AtomicBool>,
144    /// A round-robin cursor for the first idle worker to inspect.
145    next_wake: AtomicUsize,
146}
147
148struct Pool {
149    shared: Arc<PoolShared>,
150    events: Receiver<Event>,
151    handles: Vec<thread::JoinHandle<()>>,
152}
153
154/// A multi-root iterator yielding each root index with entry and per-root completion events.
155/// Unlike [`Walk`], it preserves root identity and exposes when each root finishes.
156pub(crate) struct RootWalk {
157    /// Entries buffered for delivery, by root index.
158    next: Vec<(usize, RootEvent)>,
159    /// See [`Walk::pool`].
160    pool: Option<Pool>,
161}
162
163/// A single-root directory iterator whose directory reads happen in parallel.
164/// Unlike `RootWalk`, it yields entries directly and hides root identity and completion events.
165pub struct Walk {
166    /// Entries buffered for delivery.
167    ///
168    /// This vector is used as a stack: it starts with the root, and received batches are inserted
169    /// in reverse so popping preserves their original order.
170    ///
171    /// If consumption isn't as fast as its production, threads will block.
172    next: Vec<io::Result<Entry>>,
173    /// Owns the worker threads for as long as traversal is active.
174    ///
175    /// Clearing or dropping it requests shutdown, unparks every worker, and joins their threads.
176    pool: Option<Pool>,
177}
178
179/// Walk `root` without following symlinks.
180/// Unlike `walk_roots`, this yields entries directly for a single root and hides
181/// completion events.
182pub 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
245/// Walk multiple indexed roots without following symlinks.
246/// Unlike [`walk`], this preserves each root index and yields its completion as a [`RootEvent`].
247pub(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    /// Wake one worker that has announced it is idle.
314    fn wake_worker(&self) {
315        let len = self.idle.len();
316        // This cursor only distributes scan starting points, so relaxed races affect fairness, not
317        // correctness; the compare-exchange below exclusively claims the worker to wake.
318        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    /// Wake all threads unconditionally.
332    fn wake_workers(&self) {
333        for unparker in &self.unparkers {
334            unparker.unpark();
335        }
336    }
337}
338
339impl Entry {
340    /// Return the full path to this entry.
341    #[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
412/// Prepare initial root events and directory jobs.
413/// Returns events in stack order for [`RootWalk::next`] to pop, plus jobs requiring a worker pool.
414fn 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
444/// Seed an idle pool with one initial job per active root.
445/// Initializes per-root completion accounting, queues the jobs, and wakes workers to process them.
446fn 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            // A successful steal proves peer work is available; wake one more worker so
493            // concurrency ramps up only while work remains stealable.
494            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
510/// Find work in order of increasing synchronization cost.
511///
512/// The worker checks its own LIFO queue first, favoring locality and avoiding
513/// shared-queue contention. It next takes a batch from the injector, keeping one job and moving
514/// the rest into its local queue. Only then does it inspect other workers, because stealing from a
515/// peer is the most contentious path. Consequently, a worker with local jobs keeps processing
516/// them before helping elsewhere, and injector jobs take priority over peer jobs.
517///
518/// Returns the selected job and whether it was stolen from another worker; the caller uses a
519/// successful steal to wake another idle worker. Returns `None` when a full scan finds no work.
520fn 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
568/// Read a directory for completion-order traversal.
569/// Successful directory entries are split into stealable metadata jobs, while enumeration errors
570/// are emitted directly; the directory-read job completes after all chunks are queued.
571/// This adds parallelism within wide directories when metadata calls dominate. Both traversal
572/// orders already process separate directories concurrently, so typical trees may see no speedup.
573fn 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
683/// Read a directory for parent-first traversal.
684/// Entries are converted inline rather than scheduled as `StatCompletion` jobs, producing the
685/// complete parent batch and its child-directory jobs together. This lets `finish_directory` send
686/// the parent batch before making any child job available, preserving parent-before-descendant
687/// order. Metadata within one directory is serial, although separate directories still run in
688/// parallel; this often matches completion-order performance unless wide-directory metadata is the
689/// bottleneck.
690fn 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
723/// Publish a completed directory read and schedule its accepted child-directory jobs.
724/// `ParentFirst` sends the batch before exposing child jobs; `Completion` exposes child jobs first.
725/// Child jobs are counted before either action, and the current job is marked complete afterward.
726fn 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
767/// Mark one job complete for `root`.
768/// The last job emits `RootFinished`; if this was also the last active root, `Finished` follows.
769fn 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}