weavatrix_scan/stateful_walk/
parallel.rs

1use super::{DirectoryProcessor, StatefulWalkBuilder, StatefulWalkEntry};
2use crate::control::CancellationToken;
3use crate::walk_platform::{DirectoryIdentity, FileSystemId};
4use crate::{
5    ErrorPolicy, ParallelRuntime, WalkError, WalkOperation, WalkOptions, WalkSkipReason, Walker,
6};
7use std::any::Any;
8use std::collections::{HashMap, HashSet, VecDeque};
9use std::path::PathBuf;
10use std::sync::Arc;
11use std::sync::mpsc::{self, Receiver, SyncSender, sync_channel};
12use std::thread::JoinHandle;
13
14struct DirectoryTask<R> {
15    id: u64,
16    path: PathBuf,
17    depth: usize,
18    identity: Option<DirectoryIdentity>,
19    ancestors: Arc<HashSet<DirectoryIdentity>>,
20    read_state: R,
21}
22
23struct WorkerResult<R, E> {
24    id: u64,
25    outcome: Result<DirectoryBatch<R, E>, Box<dyn Any + Send>>,
26}
27
28struct DirectoryBatch<R, E> {
29    entries: Vec<Result<StatefulWalkEntry<E>, WalkError>>,
30    child_state: R,
31    ancestors: Arc<HashSet<DirectoryIdentity>>,
32}
33
34struct PreparedItem<E> {
35    item: Result<StatefulWalkEntry<E>, WalkError>,
36    child: Option<u64>,
37}
38
39struct DirectoryFrame<E> {
40    items: std::vec::IntoIter<PreparedItem<E>>,
41}
42
43struct OrderedStatefulScheduler<R, E> {
44    root: Arc<PathBuf>,
45    root_file_system: Option<FileSystemId>,
46    options: WalkOptions,
47    processor: Option<DirectoryProcessor<R, E>>,
48    cancellation: CancellationToken,
49    runtime: ParallelRuntime,
50    limit: usize,
51    next_id: u64,
52    queued: VecDeque<DirectoryTask<R>>,
53    outstanding: usize,
54    ready: HashMap<u64, DirectoryBatch<R, E>>,
55    result_sender: mpsc::Sender<WorkerResult<R, E>>,
56    result_receiver: mpsc::Receiver<WorkerResult<R, E>>,
57    schedule_error: Option<WalkError>,
58}
59
60/// Bounded stateful pull iterator with parallel directory processing and
61/// strict deterministic depth-first output.
62pub struct ParallelStatefulWalker<E> {
63    receiver: Option<Receiver<Result<StatefulWalkEntry<E>, WalkError>>>,
64    cancellation: CancellationToken,
65    coordinator: Option<JoinHandle<()>>,
66}
67
68impl<E> ParallelStatefulWalker<E> {
69    pub(super) fn start<R>(
70        builder: StatefulWalkBuilder<R, E>,
71        capacity: usize,
72    ) -> Result<Self, WalkError>
73    where
74        R: Clone + Send + 'static,
75        E: Default + Send + 'static,
76    {
77        let root = builder.root.clone();
78        let use_serial = builder.runtime.is_worker_thread();
79        let (sender, receiver) = sync_channel(capacity.max(1));
80        let cancellation = CancellationToken::new();
81        let coordinator_cancellation = cancellation.clone();
82        let coordinator = std::thread::Builder::new()
83            .name("weavatrix-scan-stateful".to_owned())
84            .spawn(move || {
85                if use_serial {
86                    run_serial(builder, &coordinator_cancellation, &sender);
87                } else {
88                    run_parallel(builder, &coordinator_cancellation, &sender);
89                }
90            })
91            .map_err(|source| WalkError::new(root, 0, WalkOperation::ScheduleWorker, source))?;
92        Ok(Self {
93            receiver: Some(receiver),
94            cancellation,
95            coordinator: Some(coordinator),
96        })
97    }
98
99    fn join_coordinator(&mut self) {
100        if let Some(coordinator) = self.coordinator.take() {
101            coordinator
102                .join()
103                .expect("parallel stateful coordinator panicked");
104        }
105    }
106}
107
108impl<E> Iterator for ParallelStatefulWalker<E> {
109    type Item = Result<StatefulWalkEntry<E>, WalkError>;
110
111    fn next(&mut self) -> Option<Self::Item> {
112        if let Ok(item) = self.receiver.as_ref()?.recv() {
113            Some(item)
114        } else {
115            self.receiver.take();
116            self.join_coordinator();
117            None
118        }
119    }
120}
121
122impl<E> Drop for ParallelStatefulWalker<E> {
123    fn drop(&mut self) {
124        self.receiver.take();
125        self.cancellation.cancel();
126        self.join_coordinator();
127    }
128}
129
130fn run_serial<R, E>(
131    builder: StatefulWalkBuilder<R, E>,
132    cancellation: &CancellationToken,
133    sender: &SyncSender<Result<StatefulWalkEntry<E>, WalkError>>,
134) where
135    R: Clone + Send + 'static,
136    E: Default + Send + 'static,
137{
138    let walker = match builder.build() {
139        Ok(walker) => walker,
140        Err(error) => {
141            let _ = sender.send(Err(error));
142            return;
143        }
144    };
145    for item in walker {
146        if cancellation.is_cancelled() || sender.send(item).is_err() {
147            break;
148        }
149    }
150}
151
152#[allow(clippy::too_many_lines)]
153fn run_parallel<R, E>(
154    builder: StatefulWalkBuilder<R, E>,
155    cancellation: &CancellationToken,
156    sender: &SyncSender<Result<StatefulWalkEntry<E>, WalkError>>,
157) where
158    R: Clone + Send + 'static,
159    E: Default + Send + 'static,
160{
161    let options = builder.options.normalized();
162    let mut root_options = options;
163    root_options.min_depth = 0;
164    root_options.error_policy = ErrorPolicy::Continue;
165    let mut root_walker = match Walker::with_options(&builder.root, root_options) {
166        Ok(walker) => walker,
167        Err(error) => {
168            let _ = sender.send(Err(error));
169            return;
170        }
171    };
172    let root_file_system = root_walker.root_file_system;
173    let root = Arc::clone(&root_walker.root);
174    let root_entry = match root_walker
175        .next()
176        .expect("a validated root yields one entry")
177    {
178        Ok(entry) => entry,
179        Err(error) => {
180            let _ = sender.send(Err(error));
181            return;
182        }
183    };
184    let root_identity = root_entry.directory_identity();
185    let can_descend = root_entry.is_dir() && root_entry.skip_reason().is_none();
186    if root_entry.depth() >= options.min_depth
187        && sender
188            .send(Ok(StatefulWalkEntry {
189                read_children: can_descend,
190                entry: root_entry,
191                state: E::default(),
192            }))
193            .is_err()
194    {
195        return;
196    }
197    if !can_descend || cancellation.is_cancelled() {
198        return;
199    }
200
201    let (result_sender, result_receiver) = mpsc::channel();
202    let limit = requested_workers(&builder.runtime, builder.parallelism, options.max_open);
203    let mut scheduler = OrderedStatefulScheduler {
204        root: Arc::clone(&root),
205        root_file_system,
206        options,
207        processor: builder.processor,
208        cancellation: cancellation.clone(),
209        runtime: builder.runtime,
210        limit,
211        next_id: 1,
212        queued: VecDeque::new(),
213        outstanding: 0,
214        ready: HashMap::new(),
215        result_sender,
216        result_receiver,
217        schedule_error: None,
218    };
219    let mut ancestors = HashSet::new();
220    if let Some(identity) = root_identity {
221        ancestors.insert(identity);
222    }
223    scheduler.queued.push_back(DirectoryTask {
224        id: 0,
225        path: root.as_ref().clone(),
226        depth: 0,
227        identity: root_identity,
228        ancestors: Arc::new(ancestors),
229        read_state: builder.root_read_dir_state,
230    });
231    scheduler.refill();
232    let root_batch = match scheduler.wait_for(0) {
233        Ok(Some(batch)) => batch,
234        Ok(None) => return,
235        Err(error) => {
236            let _ = sender.send(Err(error));
237            return;
238        }
239    };
240    let mut frames = vec![scheduler.prepare_frame(root_batch)];
241
242    while !cancellation.is_cancelled() {
243        let Some(frame) = frames.last_mut() else {
244            break;
245        };
246        let Some(prepared) = frame.items.next() else {
247            frames.pop();
248            continue;
249        };
250        let child = prepared.child;
251        let visible = prepared
252            .item
253            .as_ref()
254            .map_or(true, |entry| entry.depth() >= options.min_depth);
255        let abort = options.error_policy == ErrorPolicy::Abort && prepared.item.is_err();
256        if visible && sender.send(prepared.item).is_err() {
257            scheduler.cancel_and_drain();
258            return;
259        }
260        if abort {
261            scheduler.cancel_and_drain();
262            return;
263        }
264        if let Some(child) = child {
265            let batch = match scheduler.wait_for(child) {
266                Ok(Some(batch)) => batch,
267                Ok(None) => return,
268                Err(error) => {
269                    let _ = sender.send(Err(error));
270                    return;
271                }
272            };
273            frames.push(scheduler.prepare_frame(batch));
274        }
275    }
276    scheduler.cancel_and_drain();
277}
278
279impl<R, E> OrderedStatefulScheduler<R, E>
280where
281    R: Clone + Send + 'static,
282    E: Default + Send + 'static,
283{
284    fn refill(&mut self) {
285        while !self.cancellation.is_cancelled() && self.outstanding < self.limit {
286            let Some(task) = self.queued.pop_front() else {
287                break;
288            };
289            let root = Arc::clone(&self.root);
290            let result_sender = self.result_sender.clone();
291            let cancellation = self.cancellation.clone();
292            let root_file_system = self.root_file_system;
293            let options = self.options;
294            let processor = self.processor.clone();
295            let scheduled = self.runtime.try_execute(move || {
296                let id = task.id;
297                let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
298                    read_directory(
299                        &root,
300                        root_file_system,
301                        options,
302                        &cancellation,
303                        processor.as_ref(),
304                        task,
305                    )
306                }));
307                let _ = result_sender.send(WorkerResult { id, outcome });
308            });
309            match scheduled {
310                Ok(()) => self.outstanding += 1,
311                Err(source) => {
312                    self.cancellation.cancel();
313                    self.queued.clear();
314                    self.schedule_error = Some(WalkError::new(
315                        self.root.as_ref(),
316                        0,
317                        WalkOperation::ScheduleWorker,
318                        source,
319                    ));
320                    break;
321                }
322            }
323        }
324    }
325
326    fn wait_for(&mut self, id: u64) -> Result<Option<DirectoryBatch<R, E>>, WalkError> {
327        if let Some(batch) = self.ready.remove(&id) {
328            return Ok(Some(batch));
329        }
330        if let Some(error) = self.schedule_error.take() {
331            self.cancel_and_drain();
332            return Err(error);
333        }
334        loop {
335            let Ok(result) = self.result_receiver.recv() else {
336                if let Some(error) = self.schedule_error.take() {
337                    return Err(error);
338                }
339                return Ok(None);
340            };
341            self.outstanding = self.outstanding.saturating_sub(1);
342            match result.outcome {
343                Ok(batch) if result.id == id => {
344                    self.refill();
345                    if let Some(error) = self.schedule_error.take() {
346                        self.cancel_and_drain();
347                        return Err(error);
348                    }
349                    return Ok(Some(batch));
350                }
351                Ok(batch) => {
352                    self.ready.insert(result.id, batch);
353                    self.refill();
354                }
355                Err(payload) => {
356                    self.cancel_and_drain();
357                    std::panic::resume_unwind(payload);
358                }
359            }
360            if self.cancellation.is_cancelled() {
361                self.cancel_and_drain();
362                if let Some(error) = self.schedule_error.take() {
363                    return Err(error);
364                }
365                return Ok(None);
366            }
367        }
368    }
369
370    fn prepare_frame(&mut self, batch: DirectoryBatch<R, E>) -> DirectoryFrame<E> {
371        let mut items = Vec::with_capacity(batch.entries.len());
372        let mut children = Vec::new();
373        for item in batch.entries {
374            let child = item.as_ref().ok().and_then(|entry| {
375                entry.read_children.then(|| {
376                    let id = self.next_id;
377                    self.next_id = self.next_id.saturating_add(1);
378                    let identity = entry.entry.directory_identity();
379                    let ancestors = identity.map_or_else(
380                        || Arc::clone(&batch.ancestors),
381                        |identity| {
382                            let mut child = batch.ancestors.as_ref().clone();
383                            child.insert(identity);
384                            Arc::new(child)
385                        },
386                    );
387                    children.push(DirectoryTask {
388                        id,
389                        path: entry.path().to_path_buf(),
390                        depth: entry.depth(),
391                        identity,
392                        ancestors,
393                        read_state: batch.child_state.clone(),
394                    });
395                    id
396                })
397            });
398            items.push(PreparedItem { item, child });
399        }
400        for child in children.into_iter().rev() {
401            self.queued.push_front(child);
402        }
403        self.refill();
404        DirectoryFrame {
405            items: items.into_iter(),
406        }
407    }
408
409    fn cancel_and_drain(&mut self) {
410        self.cancellation.cancel();
411        self.queued.clear();
412        while self.outstanding > 0 {
413            if self.result_receiver.recv().is_err() {
414                break;
415            }
416            self.outstanding -= 1;
417        }
418    }
419}
420
421fn read_directory<R, E>(
422    root: &Arc<PathBuf>,
423    root_file_system: Option<FileSystemId>,
424    options: WalkOptions,
425    cancellation: &CancellationToken,
426    processor: Option<&DirectoryProcessor<R, E>>,
427    mut task: DirectoryTask<R>,
428) -> DirectoryBatch<R, E>
429where
430    R: Clone + Send + 'static,
431    E: Default + Send + 'static,
432{
433    let ancestors = Arc::clone(&task.ancestors);
434    let mut worker_options = options;
435    worker_options.error_policy = ErrorPolicy::Continue;
436    worker_options.min_depth = 0;
437    worker_options.max_open = 1;
438    worker_options.max_depth = Some(
439        options
440            .max_depth
441            .unwrap_or(task.depth.saturating_add(1))
442            .min(task.depth.saturating_add(1)),
443    );
444    let mut walker = Walker::from_known_directory_with_ancestry(
445        root,
446        task.path.clone(),
447        task.depth,
448        worker_options,
449        root_file_system,
450        task.identity,
451        task.ancestors.as_ref().clone(),
452    );
453    let mut entries = Vec::new();
454    while !cancellation.is_cancelled() {
455        let Some(item) = walker.next() else {
456            break;
457        };
458        match item {
459            Ok(mut entry) => {
460                if entry.is_dir()
461                    && entry.skip_reason() == Some(WalkSkipReason::MaxDepth)
462                    && options
463                        .max_depth
464                        .is_none_or(|maximum| entry.depth() < maximum)
465                {
466                    entry.clear_depth_skip();
467                }
468                if entry.is_dir() {
469                    walker.skip_current_dir();
470                }
471                entries.push(Ok(StatefulWalkEntry {
472                    read_children: entry.is_dir() && entry.skip_reason().is_none(),
473                    entry,
474                    state: E::default(),
475                }));
476            }
477            Err(error) => entries.push(Err(error)),
478        }
479    }
480    if let Some(processor) = processor {
481        processor(task.depth, &task.path, &mut task.read_state, &mut entries);
482    }
483    DirectoryBatch {
484        entries,
485        child_state: task.read_state,
486        ancestors,
487    }
488}
489
490fn requested_workers(runtime: &ParallelRuntime, parallelism: usize, max_open: usize) -> usize {
491    let available = runtime.parallelism();
492    let requested = if parallelism == 0 {
493        available.min(if cfg!(windows) { 16 } else { 8 })
494    } else {
495        parallelism.min(available)
496    };
497    requested.min(max_open.max(1)).max(1)
498}
weavatrix_scan/stateful_walk/parallel.rs

weavatrix_scan/stateful_walk/
parallel.rs
