kime_cpu/
pool.rs

1//! A pool of worker threads that lives as long as the backend, from spec/10-cpu.md.
2//!
3//! [`Pool::run`] hands out tasks `0..n` to the calling thread and every worker, and returns when
4//! all of them are done. It waits for the tasks, not for the workers: a step of four tasks on a
5//! pool of ten returns once the four are done, even when the other workers have not woken up yet.
6//! That matters most on a busy machine, where a worker the OS has not scheduled would otherwise
7//! hold up every small step. Between jobs a worker spins for a while, since the next op of a plan
8//! is usually microseconds away, and then sleeps on a condition variable so an idle server costs
9//! nothing. A job is a borrowed closure, so running one allocates nothing.
10//!
11//! Tasks are claimed from one word that holds the job's generation, its task count and the next
12//! task, so a worker that wakes up late cannot claim a task of a job that has already finished.
13//!
14//! Which thread runs a task depends on timing, but what a task computes does not, and that is what
15//! keeps results independent of the thread count.
16
17use std::any::Any;
18use std::cell::UnsafeCell;
19use std::panic::{AssertUnwindSafe, catch_unwind};
20use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
21use std::sync::{Arc, Condvar, Mutex};
22use std::thread::JoinHandle;
23
24/// Spins before a worker sleeps. About 50 microseconds on a current core.
25const SPIN: u32 = 1 << 14;
26
27/// Bits of the claim word for the next task and for the task count. The generation takes the
28/// other 24.
29const BITS: u32 = 20;
30/// Tasks in one round of a job. Larger jobs run in rounds.
31const ROUND: usize = 1 << BITS;
32const MASK: u64 = (1 << BITS) - 1;
33
34type Job = dyn Fn(usize, usize) + Sync;
35
36/// The generation, the task count and the next task in a claim word.
37fn unpack(s: u64) -> (u64, usize, usize) {
38    (s >> (2 * BITS), ((s >> BITS) & MASK) as usize, (s & MASK) as usize)
39}
40
41struct Inner {
42    /// `generation << 40 | tasks << 20 | next`. A task is claimed by moving `next` up by one with a
43    /// compare and swap, which fails if the job changed in between.
44    state: AtomicU64,
45    /// The job of each generation, in the slot of its parity. Written by the caller before it
46    /// publishes the generation in `state`.
47    jobs: [UnsafeCell<*const Job>; 2],
48    /// Tasks of the current job that have finished.
49    done: AtomicUsize,
50    sleepers: AtomicUsize,
51    /// The first panic of the current job.
52    panic: Mutex<Option<Box<dyn Any + Send>>>,
53    stop: AtomicBool,
54    lock: Mutex<()>,
55    wake: Condvar,
56}
57
58// SAFETY: a job slot is written only by the thread inside `Pool::run`, which holds the pool's run
59// lock, and only for a generation whose slot was last used two jobs ago, which finished before the
60// last one started. A worker reads a slot only after it claimed a task of that slot's generation,
61// and the caller of that generation waits for the task to be done before it returns, so the slot
62// is not written while it is read. The Release store of `state` orders the write before the reads.
63// Everything else is atomics and mutexes.
64unsafe impl Sync for Inner {}
65// SAFETY: as above, the raw pointers are only dereferenced under that protocol.
66unsafe impl Send for Inner {}
67
68impl Inner {
69    /// Claims and runs tasks of the current job until none is left.
70    fn work(&self, worker: usize) {
71        let mut s = self.state.load(Ordering::Acquire);
72        loop {
73            let (generation, tasks, next) = unpack(s);
74            if next >= tasks {
75                return;
76            }
77            if let Err(now) =
78                self.state.compare_exchange_weak(s, s + 1, Ordering::Acquire, Ordering::Acquire)
79            {
80                s = now;
81                continue;
82            }
83            // SAFETY: see the Sync impl. This task belongs to `generation`, whose caller is still
84            // waiting for it.
85            let job = unsafe { &**self.jobs[(generation & 1) as usize].get() };
86            if let Err(e) = catch_unwind(AssertUnwindSafe(|| job(next, worker))) {
87                self.panic
88                    .lock()
89                    .unwrap_or_else(std::sync::PoisonError::into_inner)
90                    .get_or_insert(e);
91            }
92            self.done.fetch_add(1, Ordering::Release);
93            s = self.state.load(Ordering::Acquire);
94        }
95    }
96
97    /// Whether `state` has a task left to claim.
98    fn pending(&self, order: Ordering) -> bool {
99        let (_, tasks, next) = unpack(self.state.load(order));
100        next < tasks
101    }
102}
103
104/// Worker threads, created once.
105pub struct Pool {
106    inner: Arc<Inner>,
107    handles: Vec<JoinHandle<()>>,
108    run: Mutex<()>,
109}
110
111impl std::fmt::Debug for Pool {
112    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
113        f.debug_struct("Pool").field("threads", &self.threads()).finish()
114    }
115}
116
117impl Pool {
118    /// A pool of `threads` threads counting the caller, so `threads - 1` are spawned.
119    ///
120    /// # Panics
121    ///
122    /// If the OS will not start a thread.
123    #[must_use]
124    pub fn new(threads: usize) -> Self {
125        let noop: &'static Job = &|_, _| {};
126        let inner = Arc::new(Inner {
127            state: AtomicU64::new(0),
128            jobs: [UnsafeCell::new(noop as *const Job), UnsafeCell::new(noop as *const Job)],
129            done: AtomicUsize::new(0),
130            sleepers: AtomicUsize::new(0),
131            panic: Mutex::new(None),
132            stop: AtomicBool::new(false),
133            lock: Mutex::new(()),
134            wake: Condvar::new(),
135        });
136        let handles = (1..threads.max(1))
137            .map(|worker| {
138                let inner = Arc::clone(&inner);
139                std::thread::Builder::new()
140                    .name(format!("kime-cpu-{worker}"))
141                    .spawn(move || worker_loop(&inner, worker))
142                    .expect("spawn a worker thread")
143            })
144            .collect();
145        Self { inner, handles, run: Mutex::new(()) }
146    }
147
148    /// Threads that run tasks, the caller included.
149    #[must_use]
150    pub fn threads(&self) -> usize {
151        self.handles.len() + 1
152    }
153
154    /// Runs `f(task, worker)` for every task in `0..n`, where `worker` is below
155    /// [`Pool::threads`] and no two tasks run on the same worker at once. The caller is worker 0.
156    ///
157    /// # Panics
158    ///
159    /// If `f` panics on any thread.
160    pub fn run(&self, n: usize, f: &(dyn Fn(usize, usize) + Sync)) {
161        if n == 0 {
162            return;
163        }
164        if n == 1 || self.handles.is_empty() {
165            (0..n).for_each(|i| f(i, 0));
166            return;
167        }
168        let guard = self.run.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
169        for base in (0..n).step_by(ROUND) {
170            let tasks = ROUND.min(n - base);
171            self.round(tasks, &|i, worker| f(base + i, worker));
172        }
173        drop(guard);
174    }
175
176    /// One job of at most [`ROUND`] tasks. The caller holds the run lock.
177    fn round(&self, tasks: usize, f: &(dyn Fn(usize, usize) + Sync)) {
178        let inner = &*self.inner;
179        let (last, _, _) = unpack(inner.state.load(Ordering::Relaxed));
180        // A generation wraps after 16 million jobs, and a claim only goes wrong if a worker stalls
181        // between reading the word and swapping it for exactly that many jobs.
182        let generation = (last + 1) & ((1 << (64 - 2 * BITS)) - 1);
183        // SAFETY: see the Sync impl. The lifetime is erased, and the wait below keeps `f`
184        // borrowed until every task is done.
185        unsafe {
186            *inner.jobs[(generation & 1) as usize].get() =
187                std::mem::transmute::<&(dyn Fn(usize, usize) + Sync + '_), &'static Job>(f)
188                    as *const Job;
189        }
190        inner.done.store(0, Ordering::Relaxed);
191        inner.state.store(generation << (2 * BITS) | (tasks as u64) << BITS, Ordering::SeqCst);
192        if inner.sleepers.load(Ordering::SeqCst) > 0 {
193            let _l = inner.lock.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
194            inner.wake.notify_all();
195        }
196        inner.work(0);
197        while inner.done.load(Ordering::Acquire) < tasks {
198            std::hint::spin_loop();
199        }
200        let panic = inner.panic.lock().unwrap_or_else(std::sync::PoisonError::into_inner).take();
201        if let Some(e) = panic {
202            std::panic::resume_unwind(e);
203        }
204    }
205}
206
207fn worker_loop(inner: &Inner, worker: usize) {
208    loop {
209        let mut spins = 0u32;
210        while !inner.pending(Ordering::Acquire) {
211            if inner.stop.load(Ordering::Relaxed) {
212                return;
213            }
214            if spins < SPIN {
215                spins += 1;
216                std::hint::spin_loop();
217                continue;
218            }
219            let l = inner.lock.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
220            inner.sleepers.fetch_add(1, Ordering::SeqCst);
221            let l = if !inner.pending(Ordering::SeqCst) && !inner.stop.load(Ordering::SeqCst) {
222                inner.wake.wait(l).unwrap_or_else(std::sync::PoisonError::into_inner)
223            } else {
224                l
225            };
226            inner.sleepers.fetch_sub(1, Ordering::SeqCst);
227            drop(l);
228            spins = 0;
229        }
230        inner.work(worker);
231    }
232}
233
234impl Drop for Pool {
235    fn drop(&mut self) {
236        {
237            let _l = self.inner.lock.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
238            self.inner.stop.store(true, Ordering::SeqCst);
239            self.inner.wake.notify_all();
240        }
241        for h in self.handles.drain(..) {
242            let _ = h.join();
243        }
244    }
245}
246
247#[cfg(test)]
248mod tests {
249    use super::*;
250    use std::sync::atomic::AtomicU32;
251
252    #[test]
253    fn every_task_once_on_a_valid_worker() {
254        for threads in [1, 2, 5] {
255            let pool = Pool::new(threads);
256            for n in [0, 1, 3, 1000] {
257                let hits: Vec<AtomicU32> = (0..n).map(|_| AtomicU32::new(0)).collect();
258                pool.run(n, &|i, w| {
259                    assert!(w < threads);
260                    hits[i].fetch_add(1, Ordering::Relaxed);
261                });
262                assert!(hits.iter().all(|h| h.load(Ordering::Relaxed) == 1));
263            }
264        }
265    }
266
267    #[test]
268    fn wakes_after_sleeping() {
269        let pool = Pool::new(3);
270        let count = AtomicU32::new(0);
271        for _ in 0..3 {
272            std::thread::sleep(std::time::Duration::from_millis(30));
273            pool.run(64, &|_, _| {
274                count.fetch_add(1, Ordering::Relaxed);
275            });
276        }
277        assert_eq!(count.load(Ordering::Relaxed), 192);
278    }
279
280    #[test]
281    fn no_worker_runs_two_tasks_at_once_over_many_small_jobs() {
282        let pool = Pool::new(6);
283        let busy: Vec<AtomicBool> = (0..6).map(|_| AtomicBool::new(false)).collect();
284        for job in 0..20_000usize {
285            let n = 1 + job % 9;
286            let hits: Vec<AtomicU32> = (0..n).map(|_| AtomicU32::new(0)).collect();
287            pool.run(n, &|i, w| {
288                assert!(!busy[w].swap(true, Ordering::AcqRel), "worker {w} ran two tasks at once");
289                hits[i].fetch_add(1, Ordering::Relaxed);
290                busy[w].store(false, Ordering::Release);
291            });
292            assert!(hits.iter().all(|h| h.load(Ordering::Relaxed) == 1), "job {job}");
293        }
294    }
295
296    #[test]
297    fn a_panic_reaches_the_caller_and_the_pool_survives() {
298        let pool = Pool::new(4);
299        let r = catch_unwind(AssertUnwindSafe(|| {
300            pool.run(100, &|i, _| assert!(i != 57, "task 57"));
301        }));
302        assert!(r.is_err());
303        let count = AtomicU32::new(0);
304        pool.run(10, &|_, _| {
305            count.fetch_add(1, Ordering::Relaxed);
306        });
307        assert_eq!(count.load(Ordering::Relaxed), 10);
308    }
309}
kime_cpu/pool.rs

kime_cpu/
pool.rs
