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