weavatrix_scan/parallel/
pull.rs

1use super::ParallelWalker;
2use super::dynamic;
3use crate::control::CancellationToken;
4use crate::walker::{WalkEntry, WalkError};
5use std::io;
6use std::sync::mpsc::{Receiver, SyncSender, sync_channel};
7use std::sync::{
8    Arc,
9    atomic::{AtomicBool, Ordering},
10};
11use std::thread::JoinHandle;
12
13pub(super) type PullItem = Result<WalkEntry, WalkError>;
14pub(super) type PullBatch = Vec<PullItem>;
15
16/// Bounded pull iterator backed by parallel traversal workers.
17///
18/// Entry order is intentionally unspecified. Dropping the iterator
19/// cooperatively cancels traversal and joins its coordinator thread.
20pub struct ParallelWalkIter {
21    receiver: Option<Receiver<PullBatch>>,
22    current: std::vec::IntoIter<PullItem>,
23    cancellation: CancellationToken,
24    coordinator: Option<JoinHandle<()>>,
25}
26
27impl ParallelWalker {
28    /// Starts parallel traversal with a bounded pull buffer.
29    ///
30    /// A capacity of zero is normalized to one. Workers stop producing when
31    /// the buffer is full until the consumer requests another item.
32    ///
33    /// # Panics
34    ///
35    /// Panics if the coordinator thread cannot be created. Use
36    /// [`Self::try_into_iter_bounded`] for fallible startup.
37    #[must_use]
38    pub fn into_iter_bounded(self, capacity: usize) -> ParallelWalkIter {
39        self.try_into_iter_bounded(capacity)
40            .expect("parallel pull coordinator thread can be created")
41    }
42
43    /// Fallible form of [`Self::into_iter_bounded`].
44    ///
45    /// This reports an operating-system thread creation failure instead of
46    /// panicking before traversal starts.
47    ///
48    /// # Errors
49    ///
50    /// Returns the coordinator thread spawn error.
51    pub fn try_into_iter_bounded(self, capacity: usize) -> io::Result<ParallelWalkIter> {
52        let capacity = capacity.max(1);
53        let batch_size = capacity.min(64);
54        let queued_batches = capacity.saturating_sub(batch_size) / batch_size;
55        let cancellation = CancellationToken::new();
56        let worker_cancellation = cancellation.clone();
57        let (sender, receiver) = sync_channel(queued_batches);
58        let coordinator = std::thread::Builder::new()
59            .name("weavatrix-scan-pull".to_owned())
60            .spawn(move || {
61                let emitted_error = Arc::new(AtomicBool::new(false));
62                let options = self.options.normalized();
63                let event_sender = sender.clone();
64                let visitor_emitted_error = Arc::clone(&emitted_error);
65                let result = dynamic::stream_batched(
66                    &self.root,
67                    options,
68                    self.parallelism,
69                    &self.runtime,
70                    &worker_cancellation,
71                    move |mut entries, errors| {
72                        if !errors.is_empty() {
73                            visitor_emitted_error.store(true, Ordering::Relaxed);
74                        }
75                        if self.skip_stdout.is_some() {
76                            entries.retain(|entry| !super::matches_stdout(entry, self.skip_stdout));
77                        }
78                        send_batches(&event_sender, batch_size, entries, errors)
79                    },
80                );
81                if let Err(error) = result
82                    && !emitted_error.load(Ordering::Relaxed)
83                {
84                    let _ = sender.send(vec![Err(error)]);
85                }
86            })?;
87        Ok(ParallelWalkIter {
88            receiver: Some(receiver),
89            current: Vec::new().into_iter(),
90            cancellation,
91            coordinator: Some(coordinator),
92        })
93    }
94}
95
96impl Iterator for ParallelWalkIter {
97    type Item = PullItem;
98
99    fn next(&mut self) -> Option<Self::Item> {
100        loop {
101            if let Some(item) = self.current.next() {
102                return Some(item);
103            }
104            if let Ok(batch) = self.receiver.as_ref()?.recv() {
105                self.current = batch.into_iter();
106            } else {
107                self.receiver.take();
108                self.join_coordinator();
109                return None;
110            }
111        }
112    }
113}
114
115impl Drop for ParallelWalkIter {
116    fn drop(&mut self) {
117        self.receiver.take();
118        self.cancellation.cancel();
119        self.join_coordinator();
120    }
121}
122
123impl ParallelWalkIter {
124    pub(super) fn from_coordinator(
125        receiver: Receiver<PullBatch>,
126        cancellation: CancellationToken,
127        coordinator: JoinHandle<()>,
128    ) -> Self {
129        Self {
130            receiver: Some(receiver),
131            current: Vec::new().into_iter(),
132            cancellation,
133            coordinator: Some(coordinator),
134        }
135    }
136
137    fn join_coordinator(&mut self) {
138        if let Some(coordinator) = self.coordinator.take() {
139            coordinator
140                .join()
141                .expect("parallel pull coordinator panicked");
142        }
143    }
144}
145
146fn send_batches(
147    sender: &SyncSender<PullBatch>,
148    batch_size: usize,
149    entries: Vec<WalkEntry>,
150    errors: &[WalkError],
151) -> bool {
152    let mut batch = Vec::with_capacity(batch_size.min(entries.len() + errors.len()));
153    for item in entries
154        .into_iter()
155        .map(Ok)
156        .chain(errors.iter().map(|error| Err(copy_walk_error(error))))
157    {
158        batch.push(item);
159        if batch.len() == batch_size && sender.send(std::mem::take(&mut batch)).is_err() {
160            return false;
161        }
162    }
163    batch.is_empty() || sender.send(batch).is_ok()
164}
165
166pub(super) fn copy_walk_error(error: &WalkError) -> WalkError {
167    WalkError::new(
168        error.path().to_path_buf(),
169        error.depth(),
170        error.operation(),
171        io::Error::new(error.io_error().kind(), error.io_error().to_string()),
172    )
173}
weavatrix_scan/parallel/pull.rs

weavatrix_scan/parallel/
pull.rs
