Skip to main content

shuyi_cli/
pool.rs

1//! A pool of threads for the tools that play many slow tries (`playtest`,
2//! and later `autobalance`): every item of a work list run on one of `n`
3//! threads and timed, the answers given back in the list's order, so what
4//! is made of them never depends on the thread count. The clock is the
5//! tools'; the core never reads one.
6
7use std::sync::atomic::{AtomicUsize, Ordering};
8use std::thread;
9use std::time::{Duration, Instant};
10
11/// An item's place in the work list, and its timed answer.
12type Done<A, E> = (usize, Result<(A, Duration), E>);
13
14/// The threads to use when not told: one a core, or 1 if that's unknown.
15pub fn cores() -> usize {
16    thread::available_parallelism().map_or(1, std::num::NonZero::get)
17}
18
19/// Runs `run` on every item of `work` on `threads` threads (at least one,
20/// and no more than there are items): thread `w` takes items `w`,
21/// `w + threads`, `w + 2 × threads`, …, and times each. The answers come
22/// back in `work`'s order with their times, or the first error in that
23/// order, whatever the thread count: once an item fails, no thread starts
24/// a later one, but the earlier ones all run. A panic in `run` is resumed
25/// here, on the calling thread.
26pub 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    // The earliest item that failed yet.
35    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    /// Item `i` of `n` sleeps longer the earlier it is, so the threads
75    /// finish them out of order.
76    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            // Each item slept its time, and was timed doing it.
92            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        // On one thread, nothing after the failing item runs.
114        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        // On several threads too, though a later item may fail first.
122        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}