Skip to main content

weavatrix_scan/parallel_multi/
visit.rs

1use super::{ParallelMultiWalker, root_worker_count};
2use crate::control::CancellationToken;
3use crate::parallel::{ParallelVisitReport, ParallelWalker, WalkControl, WalkEvent};
4use crate::runtime::ParallelRuntime;
5use crate::walk_types::WalkError;
6use std::path::Path;
7use std::sync::Arc;
8
9/// A streaming event tagged with its root's insertion index and path.
10#[derive(Debug)]
11pub struct ParallelMultiWalkEvent<'a> {
12    pub root_index: usize,
13    pub root: &'a Path,
14    pub event: WalkEvent<'a>,
15}
16
17/// Per-root streaming reports in root insertion order.
18#[derive(Debug)]
19pub struct ParallelMultiVisitReport {
20    pub reports: Vec<ParallelVisitReport>,
21}
22
23impl ParallelMultiVisitReport {
24    #[must_use]
25    pub const fn len(&self) -> usize {
26        self.reports.len()
27    }
28
29    #[must_use]
30    pub const fn is_empty(&self) -> bool {
31        self.reports.is_empty()
32    }
33
34    #[must_use]
35    pub fn visited(&self) -> u64 {
36        self.reports
37            .iter()
38            .fold(0_u64, |total, report| total.saturating_add(report.visited))
39    }
40
41    #[must_use]
42    pub fn quit(&self) -> bool {
43        self.reports.iter().any(|report| report.quit)
44    }
45
46    #[must_use]
47    pub fn cancelled(&self) -> bool {
48        self.reports.iter().any(|report| report.cancelled)
49    }
50}
51
52impl ParallelMultiWalker {
53    /// Streams every configured root through one concurrent callback.
54    ///
55    /// Callback order is intentionally unspecified. Each event carries its
56    /// root insertion index, while the returned per-root reports preserve
57    /// insertion order. `WalkControl::Quit` cancels every active root.
58    ///
59    /// # Errors
60    ///
61    /// Returns the first root error in insertion order after all started root
62    /// workers have joined.
63    ///
64    /// # Panics
65    ///
66    /// Panics if the callback or an internal root worker panics.
67    pub fn visit<F>(self, visitor: F) -> Result<ParallelMultiVisitReport, WalkError>
68    where
69        F: for<'entry> Fn(ParallelMultiWalkEvent<'entry>) -> WalkControl + Send + Sync + 'static,
70    {
71        self.visit_with_cancellation(&CancellationToken::new(), visitor)
72    }
73
74    /// Multi-root streaming traversal with shared cooperative cancellation.
75    ///
76    /// Returning `WalkControl::Quit` cancels the supplied token so every
77    /// active and not-yet-started root observes the same stop request.
78    ///
79    /// # Errors
80    ///
81    /// Returns the first root error in insertion order after all started root
82    /// workers have joined.
83    ///
84    /// # Panics
85    ///
86    /// Panics if the callback or an internal root worker panics.
87    pub fn visit_with_cancellation<F>(
88        self,
89        cancellation: &CancellationToken,
90        visitor: F,
91    ) -> Result<ParallelMultiVisitReport, WalkError>
92    where
93        F: for<'entry> Fn(ParallelMultiWalkEvent<'entry>) -> WalkControl + Send + Sync + 'static,
94    {
95        let worker_count = if self.runtime.is_worker_thread() {
96            1
97        } else {
98            root_worker_count(self.root_parallelism, self.roots.len())
99        };
100        let visitor = Arc::new(visitor);
101        if worker_count <= 1 {
102            let reports = self
103                .roots
104                .into_iter()
105                .enumerate()
106                .map(|(root_index, root)| {
107                    visit_root(
108                        root_index,
109                        root,
110                        self.options,
111                        self.traversal_parallelism,
112                        self.skip_stdout,
113                        self.runtime.clone(),
114                        cancellation,
115                        Arc::clone(&visitor),
116                    )
117                })
118                .collect::<Result<Vec<_>, _>>()?;
119            return Ok(ParallelMultiVisitReport { reports });
120        }
121
122        let chunk_size = self.roots.len().div_ceil(worker_count);
123        let indexed = self.roots.into_iter().enumerate().collect::<Vec<_>>();
124        let mut visited = std::thread::scope(|scope| {
125            indexed
126                .chunks(chunk_size)
127                .map(|chunk| {
128                    let visitor = Arc::clone(&visitor);
129                    let runtime = self.runtime.clone();
130                    let cancellation = cancellation.clone();
131                    scope.spawn(move || {
132                        chunk
133                            .iter()
134                            .map(|(root_index, root)| {
135                                (
136                                    *root_index,
137                                    visit_root(
138                                        *root_index,
139                                        root.clone(),
140                                        self.options,
141                                        self.traversal_parallelism,
142                                        self.skip_stdout,
143                                        runtime.clone(),
144                                        &cancellation,
145                                        Arc::clone(&visitor),
146                                    ),
147                                )
148                            })
149                            .collect::<Vec<_>>()
150                    })
151                })
152                .collect::<Vec<_>>()
153                .into_iter()
154                .flat_map(|handle| handle.join().expect("multi-root streaming worker panicked"))
155                .collect::<Vec<_>>()
156        });
157        visited.sort_unstable_by_key(|(index, _)| *index);
158        let reports = visited
159            .into_iter()
160            .map(|(_, report)| report)
161            .collect::<Result<Vec<_>, _>>()?;
162        Ok(ParallelMultiVisitReport { reports })
163    }
164}
165
166#[allow(clippy::too_many_arguments)]
167fn visit_root<F>(
168    root_index: usize,
169    root: std::path::PathBuf,
170    options: crate::WalkOptions,
171    traversal_parallelism: usize,
172    skip_stdout: bool,
173    runtime: ParallelRuntime,
174    cancellation: &CancellationToken,
175    visitor: Arc<F>,
176) -> Result<ParallelVisitReport, WalkError>
177where
178    F: for<'entry> Fn(ParallelMultiWalkEvent<'entry>) -> WalkControl + Send + Sync + 'static,
179{
180    let event_root = root.clone();
181    let quit_cancellation = cancellation.clone();
182    let result = ParallelWalker::new(root)
183        .options(options)
184        .with_parallelism(traversal_parallelism)
185        .runtime(runtime)
186        .skip_stdout(skip_stdout)
187        .visit_with_cancellation(cancellation, move |event| {
188            let control = visitor(ParallelMultiWalkEvent {
189                root_index,
190                root: &event_root,
191                event,
192            });
193            if control == WalkControl::Quit {
194                quit_cancellation.cancel();
195            }
196            control
197        });
198    if result.is_err() {
199        cancellation.cancel();
200    }
201    result
202}