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` through a shared counter to the calling thread and every
4//! worker, and returns when all of them are done. Between jobs a worker spins for a while, since
5//! the next op of a plan is usually microseconds away, and then sleeps on a condition variable so
6//! an idle server costs nothing. A job is a borrowed closure, so running one allocates nothing.
7//!
8//! Which thread runs a task depends on timing, but what a task computes does not, and that is what
9//! keeps results independent of the thread count.
10
11use std::panic::{AssertUnwindSafe, catch_unwind};
12use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
13use std::sync::{Arc, Condvar, Mutex};
14use std::thread::JoinHandle;
15
16/// Spins before a worker sleeps. About 50 microseconds on a current core.
17const SPIN: u32 = 1 << 14;
18
19type Job = dyn Fn(usize, usize) + Sync;
20
21struct Inner {
22    /// Bumped once per job. A worker runs a job when it sees a value it has not seen.
23    generation: AtomicU64,
24    /// The current job. Written by the caller before it bumps `generation`, and read by workers
25    /// after they see the bump.
26    job: std::cell::UnsafeCell<*const Job>,
27    tasks: AtomicUsize,
28    next: AtomicUsize,
29    /// Workers still inside the current job.
30    busy: AtomicUsize,
31    sleepers: AtomicUsize,
32    panicked: AtomicBool,
33    stop: AtomicBool,
34    lock: Mutex<()>,
35    wake: Condvar,
36}
37
38// SAFETY: `job` is written only by the thread inside `Pool::run`, which holds the pool's run lock,
39// while every worker is outside a job (busy is 0), and read by workers only between seeing a new
40// generation and decrementing busy. The Release store of generation orders the write before the
41// reads. Everything else is atomics and a mutex.
42unsafe impl Sync for Inner {}
43// SAFETY: as above, the raw pointer is only dereferenced under that protocol.
44unsafe impl Send for Inner {}
45
46impl Inner {
47    fn work(&self, worker: usize) {
48        // SAFETY: see the Sync impl. The job outlives this call because the caller waits for busy
49        // to reach 0, and for its own part returns only after this returns.
50        let job = unsafe { &**self.job.get() };
51        let n = self.tasks.load(Ordering::Relaxed);
52        loop {
53            let i = self.next.fetch_add(1, Ordering::Relaxed);
54            if i >= n {
55                break;
56            }
57            job(i, worker);
58        }
59    }
60}
61
62/// Worker threads, created once.
63pub struct Pool {
64    inner: Arc<Inner>,
65    handles: Vec<JoinHandle<()>>,
66    run: Mutex<()>,
67}
68
69impl std::fmt::Debug for Pool {
70    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
71        f.debug_struct("Pool").field("threads", &self.threads()).finish()
72    }
73}
74
75impl Pool {
76    /// A pool of `threads` threads counting the caller, so `threads - 1` are spawned.
77    ///
78    /// # Panics
79    ///
80    /// If the OS will not start a thread.
81    #[must_use]
82    pub fn new(threads: usize) -> Self {
83        let noop: &'static Job = &|_, _| {};
84        let inner = Arc::new(Inner {
85            generation: AtomicU64::new(0),
86            job: std::cell::UnsafeCell::new(noop as *const Job),
87            tasks: AtomicUsize::new(0),
88            next: AtomicUsize::new(0),
89            busy: AtomicUsize::new(0),
90            sleepers: AtomicUsize::new(0),
91            panicked: AtomicBool::new(false),
92            stop: AtomicBool::new(false),
93            lock: Mutex::new(()),
94            wake: Condvar::new(),
95        });
96        let handles = (1..threads.max(1))
97            .map(|worker| {
98                let inner = Arc::clone(&inner);
99                std::thread::Builder::new()
100                    .name(format!("kime-cpu-{worker}"))
101                    .spawn(move || worker_loop(&inner, worker))
102                    .expect("spawn a worker thread")
103            })
104            .collect();
105        Self { inner, handles, run: Mutex::new(()) }
106    }
107
108    /// Threads that run tasks, the caller included.
109    #[must_use]
110    pub fn threads(&self) -> usize {
111        self.handles.len() + 1
112    }
113
114    /// Runs `f(task, worker)` for every task in `0..n`, where `worker` is below
115    /// [`Pool::threads`] and no two tasks run on the same worker at once. The caller is worker 0.
116    ///
117    /// # Panics
118    ///
119    /// If `f` panics on any thread.
120    pub fn run(&self, n: usize, f: &(dyn Fn(usize, usize) + Sync)) {
121        if n == 0 {
122            return;
123        }
124        if n == 1 || self.handles.is_empty() {
125            (0..n).for_each(|i| f(i, 0));
126            return;
127        }
128        let guard = self.run.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
129        let inner = &*self.inner;
130        // SAFETY: no worker is inside a job (the last run waited for busy to reach 0) and the run
131        // lock keeps other callers out, so nothing reads `job` now. The lifetime is erased, and
132        // the wait below keeps `f` borrowed until every worker is done with it.
133        unsafe {
134            *inner.job.get() =
135                std::mem::transmute::<&(dyn Fn(usize, usize) + Sync + '_), &'static Job>(f)
136                    as *const Job;
137        }
138        inner.tasks.store(n, Ordering::Relaxed);
139        inner.next.store(0, Ordering::Relaxed);
140        inner.busy.store(self.handles.len(), Ordering::Relaxed);
141        inner.generation.fetch_add(1, Ordering::SeqCst);
142        if inner.sleepers.load(Ordering::SeqCst) > 0 {
143            let _l = inner.lock.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
144            inner.wake.notify_all();
145        }
146        let mine = catch_unwind(AssertUnwindSafe(|| inner.work(0)));
147        while inner.busy.load(Ordering::Acquire) > 0 {
148            std::hint::spin_loop();
149        }
150        drop(guard);
151        if let Err(e) = mine {
152            std::panic::resume_unwind(e);
153        }
154        assert!(!inner.panicked.swap(false, Ordering::Relaxed), "a kime-cpu worker panicked");
155    }
156}
157
158fn worker_loop(inner: &Inner, worker: usize) {
159    let mut seen = 0;
160    loop {
161        let mut spins = 0u32;
162        loop {
163            let g = inner.generation.load(Ordering::Acquire);
164            if g != seen {
165                seen = g;
166                break;
167            }
168            if inner.stop.load(Ordering::Relaxed) {
169                return;
170            }
171            if spins < SPIN {
172                spins += 1;
173                std::hint::spin_loop();
174                continue;
175            }
176            let l = inner.lock.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
177            inner.sleepers.fetch_add(1, Ordering::SeqCst);
178            let l = if inner.generation.load(Ordering::SeqCst) == seen
179                && !inner.stop.load(Ordering::SeqCst)
180            {
181                inner.wake.wait(l).unwrap_or_else(std::sync::PoisonError::into_inner)
182            } else {
183                l
184            };
185            inner.sleepers.fetch_sub(1, Ordering::SeqCst);
186            drop(l);
187            spins = 0;
188        }
189        if catch_unwind(AssertUnwindSafe(|| inner.work(worker))).is_err() {
190            inner.panicked.store(true, Ordering::Relaxed);
191        }
192        inner.busy.fetch_sub(1, Ordering::Release);
193    }
194}
195
196impl Drop for Pool {
197    fn drop(&mut self) {
198        {
199            let _l = self.inner.lock.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
200            self.inner.stop.store(true, Ordering::SeqCst);
201            self.inner.wake.notify_all();
202        }
203        for h in self.handles.drain(..) {
204            let _ = h.join();
205        }
206    }
207}
208
209#[cfg(test)]
210mod tests {
211    use super::*;
212    use std::sync::atomic::AtomicU32;
213
214    #[test]
215    fn every_task_once_on_a_valid_worker() {
216        for threads in [1, 2, 5] {
217            let pool = Pool::new(threads);
218            for n in [0, 1, 3, 1000] {
219                let hits: Vec<AtomicU32> = (0..n).map(|_| AtomicU32::new(0)).collect();
220                pool.run(n, &|i, w| {
221                    assert!(w < threads);
222                    hits[i].fetch_add(1, Ordering::Relaxed);
223                });
224                assert!(hits.iter().all(|h| h.load(Ordering::Relaxed) == 1));
225            }
226        }
227    }
228
229    #[test]
230    fn wakes_after_sleeping() {
231        let pool = Pool::new(3);
232        let count = AtomicU32::new(0);
233        for _ in 0..3 {
234            std::thread::sleep(std::time::Duration::from_millis(30));
235            pool.run(64, &|_, _| {
236                count.fetch_add(1, Ordering::Relaxed);
237            });
238        }
239        assert_eq!(count.load(Ordering::Relaxed), 192);
240    }
241
242    #[test]
243    fn a_panic_reaches_the_caller_and_the_pool_survives() {
244        let pool = Pool::new(4);
245        let r = catch_unwind(AssertUnwindSafe(|| {
246            pool.run(100, &|i, _| assert!(i != 57, "task 57"));
247        }));
248        assert!(r.is_err());
249        let count = AtomicU32::new(0);
250        pool.run(10, &|_, _| {
251            count.fetch_add(1, Ordering::Relaxed);
252        });
253        assert_eq!(count.load(Ordering::Relaxed), 10);
254    }
255}
kime_cpu/pool.rs

kime_cpu/
pool.rs
