primase_foundations/
parallel.rs

1//! Independent items worked on a stated number of threads ([`Threads`]), with the serial result.
2//!
3//! [`map`] and [`rows_mut`] hand items to scoped threads by an atomic work index (or, for rows, a shared queue), each
4//! thread one item at a time, so at most that many items' working buffers exist at once. Results come back in item
5//! order and each item's work sees only that item, so the output is the bytes the serial run makes; a caller that
6//! tallies across items collects per item and combines in index order. A panicking item becomes a [`Panicked`] error
7//! naming it: the other threads take no new item and the call returns once they finish the one they hold. No global
8//! state: nothing here outlives the call.
9
10use std::num::NonZeroUsize;
11use std::panic::{AssertUnwindSafe, catch_unwind};
12use std::sync::Mutex;
13use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
14
15/// How many threads a run's parallel work may use; one is the serial run.
16#[derive(Clone, Copy, Debug, PartialEq, Eq)]
17pub struct Threads(pub NonZeroUsize);
18
19impl Threads {
20    /// One thread: the serial run.
21    pub const ONE: Threads = Threads(NonZeroUsize::MIN);
22
23    /// `n` threads; `None` for zero.
24    pub fn new(n: usize) -> Option<Threads> {
25        NonZeroUsize::new(n).map(Threads)
26    }
27
28    /// The thread count.
29    pub fn get(self) -> usize {
30        self.0.get()
31    }
32}
33
34impl Default for Threads {
35    /// One: the serial run.
36    fn default() -> Threads {
37        Threads::ONE
38    }
39}
40
41/// An item whose work panicked: its index and the panic's message.
42#[derive(Clone, Debug, PartialEq, Eq)]
43pub struct Panicked {
44    /// The index of the item (or row) whose work panicked; the lowest, when several did.
45    pub item: usize,
46    /// The panic message, or a note that there was none.
47    pub message: String,
48}
49
50impl std::fmt::Display for Panicked {
51    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
52        write!(f, "item {} panicked: {}", self.item, self.message)
53    }
54}
55
56impl std::error::Error for Panicked {}
57
58/// Why [`rows_mut`] did not write every row.
59#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
60pub enum RowsError {
61    /// `row_len` was zero, which divides the data into no rows; a row has at least one element.
62    #[error("rows of no length: expected a row length of at least 1")]
63    NoRowLength,
64    /// A row's work panicked.
65    #[error(transparent)]
66    Panicked(#[from] Panicked),
67}
68
69/// `f(index, item)` for every item of `items`, on up to `threads` threads; the results in item order. One thread (or
70/// one item) runs on the calling thread.
71pub fn map<T: Sync, R: Send>(
72    threads: Threads,
73    items: &[T],
74    f: impl Fn(usize, &T) -> R + Sync,
75) -> Result<Vec<R>, Panicked> {
76    let run = |i: usize| {
77        let item = items.get(i).ok_or_else(|| Panicked {
78            item: i,
79            message: format!("no item {i}: expected one of {}", items.len()),
80        })?;
81        catch_unwind(AssertUnwindSafe(|| f(i, item))).map_err(|p| panicked(i, &*p))
82    };
83    let workers = threads.get().min(items.len());
84    if workers <= 1 {
85        return (0..items.len()).map(run).collect();
86    }
87    let next = AtomicUsize::new(0);
88    let failed = AtomicBool::new(false);
89    let mut done: Vec<Vec<(usize, Result<R, Panicked>)>> = std::thread::scope(|s| {
90        let handles: Vec<_> = (0..workers)
91            .map(|_| {
92                s.spawn(|| {
93                    let mut mine = Vec::new();
94                    while !failed.load(Ordering::Relaxed) {
95                        let i = next.fetch_add(1, Ordering::Relaxed);
96                        if i >= items.len() {
97                            break;
98                        }
99                        let r = run(i);
100                        if r.is_err() {
101                            failed.store(true, Ordering::Relaxed);
102                        }
103                        mine.push((i, r));
104                    }
105                    mine
106                })
107            })
108            .collect();
109        // a worker's own panics are caught per item, so its join cannot fail
110        handles
111            .into_iter()
112            .map(|h| h.join().unwrap_or_default())
113            .collect()
114    });
115    let mut out: Vec<Option<R>> = std::iter::repeat_with(|| None).take(items.len()).collect();
116    let mut first_error: Option<Panicked> = None;
117    for (i, r) in done.drain(..).flatten() {
118        match r {
119            Ok(v) => {
120                if let Some(slot) = out.get_mut(i) {
121                    *slot = Some(v);
122                }
123            }
124            Err(e) => {
125                if first_error.as_ref().is_none_or(|f| e.item < f.item) {
126                    first_error = Some(e);
127                }
128            }
129        }
130    }
131    if let Some(e) = first_error {
132        return Err(e);
133    }
134    // no item failed, so every slot was filled; a missing one would be a bug here, reported as the item's panic
135    out.into_iter()
136        .enumerate()
137        .map(|(item, v)| {
138            v.ok_or_else(|| Panicked {
139                item,
140                message: "the item was never worked".to_owned(),
141            })
142        })
143        .collect()
144}
145
146/// `f(row, slice)` for every `row_len`-long row of `data` (the last may be shorter), on up to `threads` threads, each
147/// row's slice written in place. One thread (or one row) runs on the calling thread. A `row_len` of zero is
148/// [`RowsError::NoRowLength`], whatever `data` holds.
149pub fn rows_mut<T: Send>(
150    threads: Threads,
151    data: &mut [T],
152    row_len: usize,
153    f: impl Fn(usize, &mut [T]) + Sync,
154) -> Result<(), RowsError> {
155    if row_len == 0 {
156        return Err(RowsError::NoRowLength);
157    }
158    let rows = data.len().div_ceil(row_len);
159    let run =
160        |i: usize, row: &mut [T]| catch_unwind(AssertUnwindSafe(|| f(i, row))).map_err(|p| panicked(i, &*p));
161    let workers = threads.get().min(rows);
162    if workers <= 1 {
163        for (i, row) in data.chunks_mut(row_len).enumerate() {
164            run(i, row)?;
165        }
166        return Ok(());
167    }
168    let queue = Mutex::new(data.chunks_mut(row_len).enumerate());
169    let failed = AtomicBool::new(false);
170    let errors: Vec<Panicked> = std::thread::scope(|s| {
171        let handles: Vec<_> = (0..workers)
172            .map(|_| {
173                s.spawn(|| {
174                    let mut errors = Vec::new();
175                    while !failed.load(Ordering::Relaxed) {
176                        // a poisoned queue only means another row panicked, which `failed` already says
177                        let taken = queue.lock().map_or(None, |mut q| q.next());
178                        let Some((i, row)) = taken else { break };
179                        if let Err(e) = run(i, row) {
180                            failed.store(true, Ordering::Relaxed);
181                            errors.push(e);
182                        }
183                    }
184                    errors
185                })
186            })
187            .collect();
188        handles
189            .into_iter()
190            .flat_map(|h| h.join().unwrap_or_default())
191            .collect()
192    });
193    match errors.into_iter().min_by_key(|e| e.item) {
194        Some(e) => Err(e.into()),
195        None => Ok(()),
196    }
197}
198
199pub(crate) fn panicked(item: usize, payload: &(dyn std::any::Any + Send)) -> Panicked {
200    let message = payload
201        .downcast_ref::<&str>()
202        .map(|s| (*s).to_owned())
203        .or_else(|| payload.downcast_ref::<String>().cloned())
204        .unwrap_or_else(|| "a panic without a message".to_owned());
205    Panicked { item, message }
206}
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211
212    fn threads(n: usize) -> Threads {
213        Threads::new(n).unwrap()
214    }
215
216    /// The results are in item order and equal the serial run's, whatever the thread count.
217    #[test]
218    fn map_keeps_the_item_order_and_the_serial_result() {
219        let items: Vec<u64> = (0..1000).collect();
220        let f = |i: usize, x: &u64| (i as u64) * 31 + x * x;
221        let serial = map(Threads::ONE, &items, f).unwrap();
222        for n in [2, 3, 8, 64] {
223            assert_eq!(map(threads(n), &items, f).unwrap(), serial, "{n} threads");
224        }
225        assert_eq!(map(threads(8), &[] as &[u64], f).unwrap(), Vec::<u64>::new());
226    }
227
228    /// At most `threads` items are worked at once.
229    #[test]
230    fn map_works_at_most_threads_items_at_once() {
231        let busy = AtomicUsize::new(0);
232        let most = AtomicUsize::new(0);
233        map(threads(3), &[(); 200], |_, _| {
234            let now = busy.fetch_add(1, Ordering::SeqCst) + 1;
235            most.fetch_max(now, Ordering::SeqCst);
236            std::thread::sleep(std::time::Duration::from_micros(200));
237            busy.fetch_sub(1, Ordering::SeqCst);
238        })
239        .unwrap();
240        assert!(most.load(Ordering::SeqCst) <= 3, "{most:?}");
241    }
242
243    /// A panicking item is an error naming it (the lowest that panicked), on one thread or many.
244    #[test]
245    fn a_panic_is_an_error_naming_the_item() {
246        for n in [1, 4] {
247            let e = map(threads(n), &(0..100).collect::<Vec<u32>>(), |_, &x| {
248                assert!(x != 37, "item thirty-seven");
249                x
250            })
251            .unwrap_err();
252            assert_eq!(e.item, 37, "{n} threads");
253            assert!(e.message.contains("thirty-seven"), "{e}");
254            let mut data = vec![0u8; 100];
255            let e = rows_mut(threads(n), &mut data, 10, |row, _| assert!(row != 4, "row four")).unwrap_err();
256            let RowsError::Panicked(e) = e else {
257                panic!("{e:?}")
258            };
259            assert_eq!((e.item, e.message.contains("row four")), (4, true), "{n} threads");
260        }
261    }
262
263    /// Rows are written in place, the last one short, the same on any thread count.
264    #[test]
265    fn rows_mut_writes_every_row_in_place() {
266        let fill = |n: usize| {
267            let mut data = vec![0u32; 1003];
268            rows_mut(threads(n), &mut data, 10, |row, slice| {
269                for (k, v) in slice.iter_mut().enumerate() {
270                    *v = (row * 1000 + k) as u32;
271                }
272            })
273            .unwrap();
274            data
275        };
276        let serial = fill(1);
277        assert_eq!(serial[1002], 100_002);
278        for n in [2, 7, 32] {
279            assert_eq!(fill(n), serial, "{n} threads");
280        }
281    }
282}
primase_foundations/parallel.rs

primase_foundations/
parallel.rs
