1use std::sync::atomic::{AtomicUsize, Ordering};
8use std::thread;
9use std::time::{Duration, Instant};
10
11type Done<A, E> = (usize, Result<(A, Duration), E>);
13
14pub fn cores() -> usize {
16 thread::available_parallelism().map_or(1, std::num::NonZero::get)
17}
18
19pub fn run<T, A, E, F>(work: &[T], threads: usize, run: &F) -> Result<Vec<(A, Duration)>, E>
27where
28 T: Sync,
29 A: Send,
30 E: Send,
31 F: Fn(&T) -> Result<A, E> + Sync,
32{
33 let threads = threads.clamp(1, work.len().max(1));
34 let failed = AtomicUsize::new(usize::MAX);
36 let failed = &failed;
37 let mut done: Vec<Done<A, E>> = thread::scope(|scope| {
38 let workers: Vec<_> = (0..threads)
39 .map(|worker| {
40 scope.spawn(move || {
41 let mut done = Vec::new();
42 for (i, item) in work.iter().enumerate().skip(worker).step_by(threads) {
43 if i > failed.load(Ordering::Relaxed) {
44 break;
45 }
46 let started = Instant::now();
47 let answer = run(item).map(|a| (a, started.elapsed()));
48 if answer.is_err() {
49 failed.fetch_min(i, Ordering::Relaxed);
50 }
51 done.push((i, answer));
52 }
53 done
54 })
55 })
56 .collect();
57 workers
58 .into_iter()
59 .flat_map(|worker| {
60 worker
61 .join()
62 .unwrap_or_else(|panic| std::panic::resume_unwind(panic))
63 })
64 .collect()
65 });
66 done.sort_by_key(|(i, _)| *i);
67 done.into_iter().map(|(_, answer)| answer).collect()
68}
69
70#[cfg(test)]
71mod tests {
72 use super::*;
73
74 fn slow(n: u64) -> impl Fn(&u64) -> Result<u64, String> + Sync {
77 move |&i| {
78 thread::sleep(Duration::from_millis(n - i));
79 Ok(i * 10)
80 }
81 }
82
83 #[test]
84 fn the_answers_come_back_in_the_works_order_whatever_the_thread_count() {
85 let work: Vec<u64> = (0..12).collect();
86 let want: Vec<u64> = work.iter().map(|i| i * 10).collect();
87 for threads in [0, 1, 3, 4, 12, 50] {
88 let done = run(&work, threads, &slow(12)).unwrap();
89 let answers: Vec<u64> = done.iter().map(|(a, _)| *a).collect();
90 assert_eq!(answers, want, "{threads} thread(s)");
91 assert!(done[0].1 >= Duration::from_millis(12), "{:?}", done[0].1);
93 }
94 let found = run(&[] as &[u64], 4, &slow(0)).unwrap();
95 assert!(found.is_empty(), "{found:?}");
96 assert_eq!(
97 cores(),
98 thread::available_parallelism().map_or(1, std::num::NonZero::get)
99 );
100 }
101
102 #[test]
103 fn the_first_error_in_the_works_order_stops_it() {
104 let work: Vec<u64> = (0..40).collect();
105 let fails = |&i: &u64| {
106 if i % 7 == 3 {
107 Err(format!("item {i}"))
108 } else {
109 Ok(i)
110 }
111 };
112 assert_eq!(run(&work, 1, &fails).unwrap_err(), "item 3");
113 let started = AtomicUsize::new(0);
115 let counted = |i: &u64| {
116 started.fetch_add(1, Ordering::Relaxed);
117 fails(i)
118 };
119 assert!(run(&work, 1, &counted).is_err());
120 assert_eq!(started.load(Ordering::Relaxed), 4);
121 for threads in [2, 4, 8] {
123 assert_eq!(run(&work, threads, &fails).unwrap_err(), "item 3");
124 }
125 }
126
127 #[test]
128 #[should_panic(expected = "a bug in a try")]
129 fn a_panic_on_a_thread_is_resumed_here() {
130 let work = [1, 2, 3];
131 let _ = run(&work, 3, &|&i: &i32| -> Result<i32, String> {
132 assert!(i != 2, "a bug in a try");
133 Ok(i)
134 });
135 }
136}