weavatrix_scan/parallel_multi/
visit.rs1use 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#[derive(Debug)]
10pub struct ParallelMultiWalkEvent<'a> {
11 pub root_index: usize,
12 pub root: &'a Path,
13 pub event: WalkEvent<'a>,
14}
15
16#[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 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 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}