use crate::{LockExt, ProcessRegistry};
use std::sync::Mutex;
use std::time::{Duration, Instant};
const BUSY_WINDOW: Duration = Duration::from_secs(5);
const THREADS_EVERY: Duration = Duration::from_millis(100);
const OVERSUBSCRIBE: usize = 4;
const TURN_POLL: Duration = Duration::from_millis(10);
pub fn auto_threads() -> usize {
static LAST: Mutex<Option<(Instant, usize)>> = Mutex::new(None);
let mut last = LAST.lock_or_default();
if let Some((at, threads)) = *last
&& at.elapsed() < THREADS_EVERY
{
return threads;
}
let others = ProcessRegistry::default_dir().map(|dir| ProcessRegistry::others_in(&dir)).unwrap_or_default();
let busy = others.iter().filter(|other| other.busy.unwrap_or(other.idle < BUSY_WINDOW)).count();
let threads = auto_threads_for(cpus(), busy);
*last = Some((Instant::now(), threads));
threads
}
pub fn cpus() -> usize {
std::thread::available_parallelism().map_or(8, |n| n.get())
}
pub fn auto_threads_for(cpus: usize, busy_others: usize) -> usize {
(cpus * OVERSUBSCRIBE / (busy_others + 1)).clamp(cpus.min(2), cpus)
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum Threads {
#[default]
Auto,
Fixed(usize),
}
impl Threads {
pub fn count(self) -> usize {
match self {
Threads::Auto => auto_threads(),
Threads::Fixed(n) => n.max(1),
}
}
pub fn most(self) -> usize {
match self {
Threads::Auto => cpus(),
Threads::Fixed(n) => n.max(1),
}
}
pub fn wait_turn(self, k: usize, stop: impl Fn() -> bool) {
if self == Threads::Auto {
while k >= auto_threads() && !stop() {
std::thread::sleep(TURN_POLL);
}
}
}
}
impl From<usize> for Threads {
fn from(n: usize) -> Threads {
Threads::Fixed(n)
}
}
impl From<Option<usize>> for Threads {
fn from(n: Option<usize>) -> Threads {
n.map_or(Threads::Auto, Threads::Fixed)
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicUsize, Ordering};
#[test]
fn threads_split_between_busy_processes() {
assert_eq!(auto_threads_for(32, 0), 32);
assert_eq!(auto_threads_for(32, 3), 32, "4 busy processes keep every CPU");
assert_eq!(auto_threads_for(32, 7), 16);
assert_eq!(auto_threads_for(32, 15), 8);
assert_eq!(auto_threads_for(32, 255), 2, "never below 2");
assert_eq!(auto_threads_for(1, 3), 1, "nor above the CPUs");
assert!(auto_threads() >= 1);
}
#[test]
fn workers_wait_for_their_turn() {
Threads::Auto.wait_turn(0, || panic!("the first worker always runs"));
Threads::Fixed(1).wait_turn(5, || panic!("fixed counts never hold a worker back"));
let asked = AtomicUsize::new(0);
Threads::Auto.wait_turn(Threads::Auto.most(), || asked.fetch_add(1, Ordering::Relaxed) >= 2);
assert_eq!(asked.load(Ordering::Relaxed), 3);
assert_eq!((Threads::from(Some(3)), Threads::from(None)), (Threads::Fixed(3), Threads::Auto));
}
}