Skip to main content

weavatrix_scan/parallel_multi/
visit.rs

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